Fix YiAdapter (#2730)

lm-sys · Nov 26, 2023 · cfba5f4 · cfba5f4
1 parent c199c8d
commit cfba5f4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fastchat/model/model_adapter.py b/fastchat/model/model_adapter.py
@@ -1916,7 +1916,7 @@ class YiAdapter(BaseModelAdapter):
     """The model adapter for Yi models"""
 
     def match(self, model_path: str):
-        return "yi-34b-chat" in model_path.lower()
+        return "yi-" in model_path.lower() and "chat" in model_path.lower()
 
     def get_default_conv_template(self, model_path: str) -> Conversation:
         return get_conv_template("Yi-34b-chat")

diff --git a/fastchat/model/model_registry.py b/fastchat/model/model_registry.py
@@ -399,7 +399,7 @@ def get_model_info(name: str) -> ModelInfo:
 )
 
 register_model_info(
-    ["Yi-34B-Chat"],
+    ["Yi-34B-Chat", "Yi-6B-Chat"],
     "Yi-Chat",
     "https://huggingface.co/01-ai",
     "A large language model by 01.AI.",