fix telechat2 template (#2775)

modelscope · Dec 26, 2024 · 629dca3 · 629dca3
1 parent 37cb3c6
commit 629dca3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/swift/llm/template/template/llm.py b/swift/llm/template/template/llm.py
@@ -136,7 +136,7 @@ class TeleChatTemplateMeta(TemplateMeta):
     prompt: Prompt = field(default_factory=lambda: ['<_user>{{QUERY}}<_bot>'])
     chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<_end>'])
     suffix: Prompt = field(default_factory=lambda: ['<_end>'])
-    system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<_system>{{SYSTEM}}'])
+    system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<_system>{{SYSTEM}}\n'])
 
 
 register_template(TeleChatTemplateMeta(LLMTemplateType.telechat))

diff --git a/tests/test_align/test_template/test_llm.py b/tests/test_align/test_template/test_llm.py
@@ -96,9 +96,11 @@ def test_codegeex4():
 
 def test_telechat():
     pt_engine = PtEngine('TeleAI/TeleChat2-7B', torch_dtype=torch.float16)
-    _infer_model(pt_engine, messages=[{'role': 'user', 'content': '你是谁？'}])
+    messages = [{'role': 'system', 'content': '你是一个乐于助人的智能助手，请使用用户提问的语言进行有帮助的问答'}, {'role': 'user', 'content': '你好'}]
+    response = _infer_model(pt_engine, messages=messages)
     pt_engine.default_template.template_backend = 'jinja'
-    _infer_model(pt_engine, messages=[{'role': 'user', 'content': '你是谁？'}])
+    response2 = _infer_model(pt_engine, messages=messages)
+    assert response == response2
 
 
 def test_glm_edge():