From a272d1b8f3fdd1d1c5687fd82a67345d12bd88be Mon Sep 17 00:00:00 2001
From: Fangyin Cheng <staneyffer@gmail.com>
Date: Fri, 29 Mar 2024 17:25:30 +0800
Subject: [PATCH] feat(model): Support Starling-LM-7B-beta (#1350)

---
 README.md                         |  1 +
 README.zh.md                      |  2 ++
 dbgpt/configs/model_config.py     |  1 +
 dbgpt/model/adapter/hf_adapter.py | 54 +++++++++++++++++++++++++++++++
 setup.py                          |  1 +
 5 files changed, 59 insertions(+)

diff --git a/README.md b/README.md
index 9e11e9ab1..730d7270b 100644
--- a/README.md
+++ b/README.md
@@ -158,6 +158,7 @@ At present, we have introduced several key features to showcase our current capa
   We offer extensive model support, including dozens of large language models (LLMs) from both open-source and API agents, such as LLaMA/LLaMA2, Baichuan, ChatGLM, Wenxin, Tongyi, Zhipu, and many more. 
 
   - News
+    - 🔥🔥🔥  [Starling-LM-7B-beta](https://huggingface.co/Nexusflow/Starling-LM-7B-beta)
     - 🔥🔥🔥  [gemma-7b-it](https://huggingface.co/google/gemma-7b-it)
     - 🔥🔥🔥  [gemma-2b-it](https://huggingface.co/google/gemma-2b-it)
     - 🔥🔥🔥  [SOLAR-10.7B](https://huggingface.co/upstage/SOLAR-10.7B-Instruct-v1.0)
diff --git a/README.zh.md b/README.zh.md
index 98bd9f3c9..2131d7b12 100644
--- a/README.zh.md
+++ b/README.zh.md
@@ -152,6 +152,7 @@
   海量模型支持，包括开源、API代理等几十种大语言模型。如LLaMA/LLaMA2、Baichuan、ChatGLM、文心、通义、智谱等。当前已支持如下模型: 
 
   - 新增支持模型
+    - 🔥🔥🔥  [Starling-LM-7B-beta](https://huggingface.co/Nexusflow/Starling-LM-7B-beta)
     - 🔥🔥🔥  [gemma-7b-it](https://huggingface.co/google/gemma-7b-it)
     - 🔥🔥🔥  [gemma-2b-it](https://huggingface.co/google/gemma-2b-it)
     - 🔥🔥🔥  [SOLAR-10.7B](https://huggingface.co/upstage/SOLAR-10.7B-Instruct-v1.0)
@@ -161,6 +162,7 @@
   - [更多开源模型](https://www.yuque.com/eosphoros/dbgpt-docs/iqaaqwriwhp6zslc#qQktR)
 
   - 支持在线代理模型 
+    - [x] [零一万物.Yi](https://platform.lingyiwanwu.com/docs)
     - [x] [OpenAI·ChatGPT](https://api.openai.com/)
     - [x] [百川·Baichuan](https://platform.baichuan-ai.com/)
     - [x] [阿里·通义](https://www.aliyun.com/product/dashscope)
diff --git a/dbgpt/configs/model_config.py b/dbgpt/configs/model_config.py
index a3e0605a2..81ff0fcfb 100644
--- a/dbgpt/configs/model_config.py
+++ b/dbgpt/configs/model_config.py
@@ -156,6 +156,7 @@ def get_device() -> str:
     "gemma-7b-it": os.path.join(MODEL_PATH, "gemma-7b-it"),
     # https://huggingface.co/google/gemma-2b-it
     "gemma-2b-it": os.path.join(MODEL_PATH, "gemma-2b-it"),
+    "starling-lm-7b-beta": os.path.join(MODEL_PATH, "Starling-LM-7B-beta"),
 }
 
 EMBEDDING_MODEL_CONFIG = {
diff --git a/dbgpt/model/adapter/hf_adapter.py b/dbgpt/model/adapter/hf_adapter.py
index 05429af3b..9c4f85a52 100644
--- a/dbgpt/model/adapter/hf_adapter.py
+++ b/dbgpt/model/adapter/hf_adapter.py
@@ -193,7 +193,61 @@ def do_match(self, lower_model_name_or_path: Optional[str] = None):
         )
 
 
+class StarlingLMAdapter(NewHFChatModelAdapter):
+    """
+    https://huggingface.co/Nexusflow/Starling-LM-7B-beta
+    """
+
+    support_4bit: bool = True
+    support_8bit: bool = True
+    support_system_message: bool = False
+
+    def do_match(self, lower_model_name_or_path: Optional[str] = None):
+        return (
+            lower_model_name_or_path
+            and "starling-" in lower_model_name_or_path
+            and "lm" in lower_model_name_or_path
+        )
+
+    def get_str_prompt(
+        self,
+        params: Dict,
+        messages: List[ModelMessage],
+        tokenizer: Any,
+        prompt_template: str = None,
+        convert_to_compatible_format: bool = False,
+    ) -> Optional[str]:
+        str_prompt = super().get_str_prompt(
+            params,
+            messages,
+            tokenizer,
+            prompt_template,
+            convert_to_compatible_format,
+        )
+        chat_mode = None
+        if params and "context" in params and "chat_mode" in params["context"]:
+            chat_mode = params["context"].get("chat_mode")
+        if chat_mode in [
+            "chat_dashboard",
+            "chat_with_db_execute",
+            "excel_learning",
+            "chat_excel",
+        ]:
+            # Coding conversation, use code prompt
+            # This is a temporary solution, we should use a better way to distinguish the conversation type
+            # https://huggingface.co/Nexusflow/Starling-LM-7B-beta#code-examples
+            str_prompt = str_prompt.replace("GPT4 Correct User:", "Code User:").replace(
+                "GPT4 Correct Assistant:", "Code Assistant:"
+            )
+            logger.info(
+                f"Use code prompt for chat_mode: {chat_mode}, transform 'GPT4 Correct User:' to 'Code User:' "
+                "and 'GPT4 Correct Assistant:' to 'Code Assistant:'"
+            )
+        return str_prompt
+
+
 register_model_adapter(YiAdapter)
 register_model_adapter(Mixtral8x7BAdapter)
 register_model_adapter(SOLARAdapter)
 register_model_adapter(GemmaAdapter)
+register_model_adapter(StarlingLMAdapter)
diff --git a/setup.py b/setup.py
index 68c445ab8..66750dd8d 100644
--- a/setup.py
+++ b/setup.py
@@ -630,6 +630,7 @@ def default_requires():
         "zhipuai",
         "dashscope",
         "chardet",
+        "sentencepiece",
     ]
     setup_spec.extras["default"] += setup_spec.extras["framework"]
     setup_spec.extras["default"] += setup_spec.extras["rag"]