From 629dca31d757503161f5e72c605f2d94e5a3ae8b Mon Sep 17 00:00:00 2001 From: Jintao Date: Thu, 26 Dec 2024 16:21:03 +0800 Subject: [PATCH] fix telechat2 template (#2775) --- swift/llm/template/template/llm.py | 2 +- tests/test_align/test_template/test_llm.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/swift/llm/template/template/llm.py b/swift/llm/template/template/llm.py index 0cd125d8a..f88583eaf 100644 --- 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 index b1f8272be..b997ec4fb 100644 --- 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():