fix hiyouga#2421

painkillernhat · Feb 4, 2024 · d63a3a7 · d63a3a7
1 parent 78e1a33
commit d63a3a7
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/llmtuner/data/template.py b/src/llmtuner/data/template.py
@@ -339,7 +339,9 @@ def get_template_and_fix_tokenizer(name: str, tokenizer: "PreTrainedTokenizer")
         slots=[{"token": "[gMASK]"}, {"token": "sop"}, {"token": "<|system|>"}, "\n", "{{content}}"]
     ),
     format_function=FunctionFormatter(slots=["{{name}}\n{{arguments}}"]),
-    format_observation=StringFormatter(slots=[{"token": "<|observation|>"}, "\n", "{{content}}"]),
+    format_observation=StringFormatter(
+        slots=[{"token": "<|observation|>"}, "\n", "{{content}}", {"token": "<|assistant|>"}]
+    ),
     default_system=(
         "You are ChatGLM3, a large language model trained by Zhipu.AI. "
         "Follow the user's instructions carefully. Respond using markdown."