Add TinyLlama (lm-sys#2889)

silk55 · Jan 26, 2024 · 265bd2e · 265bd2e
1 parent 5142ead
commit 265bd2e
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 0 deletions.
diff --git a/docs/model_support.md b/docs/model_support.md
@@ -47,6 +47,7 @@
 - [project-baize/baize-v2-7b](https://huggingface.co/project-baize/baize-v2-7b)
 - [Qwen/Qwen-7B-Chat](https://huggingface.co/Qwen/Qwen-7B-Chat)
 - [rishiraj/CatPPT](https://huggingface.co/rishiraj/CatPPT)
+- [TinyLlama/TinyLlama-1.1B-Chat-v1.0](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0)
 - [Salesforce/codet5p-6b](https://huggingface.co/Salesforce/codet5p-6b)
 - [StabilityAI/stablelm-tuned-alpha-7b](https://huggingface.co/stabilityai/stablelm-tuned-alpha-7b)
 - [THUDM/chatglm-6b](https://huggingface.co/THUDM/chatglm-6b)

diff --git a/fastchat/conversation.py b/fastchat/conversation.py
@@ -1374,6 +1374,20 @@ def get_conv_template(name: str) -> Conversation:
     )
 )
 
+# TinyLlama template
+# reference: https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0
+register_conv_template(
+    Conversation(
+        name="TinyLlama",
+        system_template="<|system|>\n{system_message}",
+        roles=("<|user|>", "<|assistant|>"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="</s>",
+        stop_token_ids=[2],
+        stop_str="</s>",
+    )
+)
+
 # Orca-2 template
 # reference: https://huggingface.co/microsoft/Orca-2-7b
 register_conv_template(

diff --git a/fastchat/model/model_adapter.py b/fastchat/model/model_adapter.py
@@ -1993,6 +1993,16 @@ def get_default_conv_template(self, model_path: str) -> Conversation:
         return get_conv_template("catppt")
 
 
+class TinyLlamaAdapter(BaseModelAdapter):
+    """The model adapter for TinyLlama (e.g. TinyLlama/TinyLlama-1.1B-Chat-v1.0)"""
+
+    def match(self, model_path: str):
+        return "TinyLlama" in model_path.lower()
+
+    def get_default_conv_template(self, model_path: str) -> Conversation:
+        return get_conv_template("TinyLlama")
+
+
 class XwinLMAdapter(BaseModelAdapter):
     """The model adapter for Xwin-LM V0.1 and V0.2 series of models(e.g., Xwin-LM/Xwin-LM-70B-V0.1)"""
 
@@ -2185,6 +2195,7 @@ def get_default_conv_template(self, model_path: str) -> Conversation:
 register_model_adapter(Llama2ChangAdapter)
 register_model_adapter(ZephyrAdapter)
 register_model_adapter(CatPPTAdapter)
+register_model_adapter(TinyLlamaAdapter)
 register_model_adapter(XwinLMAdapter)
 register_model_adapter(LemurAdapter)
 register_model_adapter(PygmalionAdapter)

diff --git a/fastchat/model/model_registry.py b/fastchat/model/model_registry.py
@@ -175,6 +175,13 @@ def get_model_info(name: str) -> ModelInfo:
     "a chatbot fine-tuned from a SLERP merged model by Rishiraj Acharya",
 )
 
+register_model_info(
+    ["TinyLlama"],
+    "TinyLlama",
+    "https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "The TinyLlama project is an open endeavor to pretrain a 1.1B Llama model on 3 trillion tokens.",
+)
+
 register_model_info(
     ["qwen-14b-chat"],
     "Qwen",