From 28886e42417f912c872dd0d3fd4d458a56cfe398 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 18 Jul 2026 14:10:37 +0800
Subject: [PATCH] fix: make system prompt optional across scripts

- stream_chat: default empty system_prompt, single-turn mode
- generate_batch: drop hardcoded system role
- generate.py: preserve original fields in messages branch
  and use response_key for the output column name
---
 scripts/demo/generate_batch.py |  6 ++----
 scripts/demo/stream_chat.py    | 17 +++++++++--------
 scripts/tools/generate.py      |  3 ++-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/scripts/demo/generate_batch.py b/scripts/demo/generate_batch.py
index 39c4975..da54922 100644
--- a/scripts/demo/generate_batch.py
+++ b/scripts/demo/generate_batch.py
@@ -26,11 +26,9 @@ def batch_generate():
 
     prompts = [
         tokenizer.apply_chat_template(
-            [
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": q},
-            ],
+            [{"role": "user", "content": q}],
             tokenize=False,
+            add_generation_prompt=True,
         )
         for q in inputs
     ]
diff --git a/scripts/demo/stream_chat.py b/scripts/demo/stream_chat.py
index 7ebd395..3fdc563 100644
--- a/scripts/demo/stream_chat.py
+++ b/scripts/demo/stream_chat.py
@@ -58,8 +58,8 @@ def parse_args():
     parser.add_argument(
         "--system_prompt",
         type=str,
-        default="You are a helpful assistant.",
-        help="Optional system prompt",
+        default="",
+        help="Optional system prompt (default: empty, model not SFT-trained on system role)",
     )
     return parser.parse_args()
 
@@ -73,18 +73,20 @@ def chat():
     model.to(device="cuda", dtype=torch.bfloat16)
     engine = InferenceEngine(model=model, tokenizer=tokenizer)
 
-    messages = [{"role": "system", "content": args.system_prompt}]
-
     while True:
         query = input(">> ")
         if query == "!exit":
             break
 
-        messages.append({"role": "user", "content": query})
+        msgs = []
+        if args.system_prompt:
+            msgs.append({"role": "system", "content": args.system_prompt})
+        msgs.append({"role": "user", "content": query})
+        prompt = tokenizer.apply_chat_template(
+            msgs, tokenize=False, add_generation_prompt=True
+        )
 
         full_response = ""
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
-
         for token in engine.generate(
             prompt=prompt,
             stream=True,
@@ -99,7 +101,6 @@ def chat():
             full_response += token
 
         print()
-        messages.append({"role": "assistant", "content": full_response.strip()})
 
 
 if __name__ == "__main__":
diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py
index e1f15ba..45eaaf3 100644
--- a/scripts/tools/generate.py
+++ b/scripts/tools/generate.py
@@ -99,7 +99,8 @@ def processor(
 
             for i, prompt in enumerate(chunk):
                 if input_data and "messages" in input_data[0]:
-                    output_item = {"response": resp_chunk[i]}
+                    orig = input_data[chunk_start + i]
+                    output_item = {**orig, response_key: resp_chunk[i]}
                 else:
                     output_item = {
                         question_key: prompt,