From dbdd3b5e5ace989923a5abb549780564980bc11e Mon Sep 17 00:00:00 2001 From: youkaichao Date: Wed, 23 Oct 2024 09:14:44 -0700 Subject: [PATCH] [misc] comment to avoid future confusion about baichuan (#9620) Signed-off-by: youkaichao --- vllm/model_executor/models/baichuan.py | 8 ++++++-- vllm/model_executor/models/registry.py | 6 ++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py index 54ed548ba8bc..767230aeacc3 100644 --- a/vllm/model_executor/models/baichuan.py +++ b/vllm/model_executor/models/baichuan.py @@ -432,7 +432,9 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP): class BaichuanForCausalLM(BaiChuanBaseForCausalLM): - """Baichuan 13B and Baichuan2 7B/13B.""" + """Baichuan 13B and Baichuan2 7B/13B. + NOTE: the class name has a lower case 'c'. + """ def __init__( self, @@ -450,7 +452,9 @@ class BaichuanForCausalLM(BaiChuanBaseForCausalLM): class BaiChuanForCausalLM(BaiChuanBaseForCausalLM): - """Baichuan 7B.""" + """Baichuan 7B. + NOTE: the class name has an upper case 'C'. + """ def __init__( self, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 787c65743e89..db5841429907 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -26,8 +26,10 @@ _TEXT_GENERATION_MODELS = { "AquilaModel": ("llama", "LlamaForCausalLM"), "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2 "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"), - "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), # baichuan-7b - "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), # baichuan-13b + # baichuan-7b, upper case 'C' in the class name + "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), + # baichuan-13b, lower case 'c' in the class name + "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), "BloomForCausalLM": ("bloom", "BloomForCausalLM"), # ChatGLMModel supports multimodal "CohereForCausalLM": ("commandr", "CohereForCausalLM"),