FEAT: support qwen 7b (#294)

xorbitsai · Aug 3, 2023 · b21d927 · b21d927
1 parent fed2655
commit b21d927
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 0 deletions.
diff --git a/xinference/model/llm/llm_family.json b/xinference/model/llm/llm_family.json
@@ -739,5 +739,40 @@
         11
       ]
     }
+  },
+  {
+    "version": 1,
+    "model_name": "qwen-chat",
+    "model_lang": [
+      "en", "zh"
+    ],
+    "model_ability": [
+      "embed",
+      "chat"
+    ],
+    "model_specs": [
+      {
+        "model_format": "pytorch",
+        "model_size_in_billions": 7,
+        "quantizations": [
+          "4-bit",
+          "8-bit",
+          "none"
+        ],
+        "model_id": "Qwen/Qwen-7B-Chat"
+      }
+    ],
+    "prompt_style": {
+      "style_name": "QWEN",
+      "system_prompt": "You are a helpful assistant.",
+      "roles": [
+        "user",
+        "assistant"
+      ],
+      "intra_message_sep": "\n",
+      "stop_token_ids": [
+        151643
+      ]
+    }
   }
 ]
diff --git a/xinference/model/llm/tests/test_utils.py b/xinference/model/llm/tests/test_utils.py
@@ -215,3 +215,26 @@ def test_prompt_style_chatglm_v2():
     assert expected == ChatModelMixin.get_prompt(
         "Write a poem.", chat_history, prompt_style
     )
+
+
+def test_prompt_style_qwen():
+    prompt_style = PromptStyleV1(
+        style_name="QWEN",
+        system_prompt="You are a helpful assistant.",
+        roles=["user", "assistant"],
+        intra_message_sep="\n",
+    )
+    chat_history = [
+        ChatCompletionMessage(role=prompt_style.roles[0], content="Hi there."),
+        ChatCompletionMessage(
+            role=prompt_style.roles[1], content="Hello, how may I help you?"
+        ),
+    ]
+    expected = (
+        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nHi there."
+        "<|im_end|>\n<|im_start|>assistant\nHello, how may I help you?<|im_end|>\n<|im_start|>"
+        "user\nWrite a poem.<|im_end|>\n<|im_start|>assistant\n"
+    )
+    assert expected == ChatModelMixin.get_prompt(
+        "Write a poem.", chat_history, prompt_style
+    )
diff --git a/xinference/model/llm/utils.py b/xinference/model/llm/utils.py
@@ -121,6 +121,18 @@ def get_prompt(
                 else:
                     ret += role + "："
             return ret
+        elif prompt_style.style_name == "QWEN":
+            ret = f"<|im_start|>system\n{prompt_style.system_prompt}<|im_end|>"
+            for message in chat_history:
+                role = message["role"]
+                content = message["content"]
+
+                ret += prompt_style.intra_message_sep
+                if content:
+                    ret += f"<|im_start|>{role}\n{content}<|im_end|>"
+                else:
+                    ret += f"<|im_start|>{role}\n"
+            return ret
         else:
             raise ValueError(f"Invalid prompt style: {prompt_style.style_name}")