From cfba5f43f28e6289326100b4f6b4ba24956426de Mon Sep 17 00:00:00 2001
From: Jingsong-Yan <75230787+Jingsong-Yan@users.noreply.github.com>
Date: Sun, 26 Nov 2023 16:05:02 +0800
Subject: [PATCH] Fix YiAdapter (#2730)

---
 fastchat/model/model_adapter.py  | 2 +-
 fastchat/model/model_registry.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/fastchat/model/model_adapter.py b/fastchat/model/model_adapter.py
index f94b5cf61..c5b3df832 100644
--- 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
index a8e603c72..ed1e3337c 100644
--- 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.",