[Minor] Fix code style for baichuan (#2003)

fe470ae5 · Woosuk Kwon · GitHub · 3a8c2381 · fe470ae5
Unverified Commit fe470ae5 authored Dec 09, 2023 by Woosuk Kwon Committed by GitHub Dec 09, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 3 deletions

vllm/model_executor/models/baichuan.py vllm/model_executor/models/baichuan.py +4 -3

No files found.
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -366,8 +366,8 @@ class BaiChuanBaseForCausalLM(nn.Module):
                weight_loader(param, loaded_weight)


-class BaichuanForCausalLM(BaiChuanBaseForCausalLM
-                          ):  # baichuan 13b, baichuan2 13b, baichuan2 7b
+class BaichuanForCausalLM(BaiChuanBaseForCausalLM):
+    """Baichuan 13B and Baichuan2 7B/13B."""

    def __init__(self,
                 config,
@@ -378,7 +378,8 @@ class BaichuanForCausalLM(BaiChuanBaseForCausalLM
            super().__init__(config, "ALIBI", linear_method)


-class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):  # baichuan 7b
+class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):
+    """Baichuan 7B."""

    def __init__(self,
                 config,