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
This commit is contained in:
2026-07-18 14:10:37 +08:00
parent 9d3ccfdffc
commit 28886e4241
3 changed files with 13 additions and 13 deletions
+2 -4
View File
@@ -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
]
+9 -8
View File
@@ -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__":
+2 -1
View File
@@ -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,