Fix multi-turn thinking block corruption for Kimi models

This commit is contained in:
oobabooga 2026-03-17 11:31:55 -07:00
parent 3f36189fa0
commit 27a6cdeec1

View file

@ -235,6 +235,7 @@ def generate_chat_prompt(user_input, state, **kwargs):
tools_in_user_message=False,
add_generation_prompt=False,
enable_thinking=state['enable_thinking'],
thinking=state['enable_thinking'],
reasoning_effort=state['reasoning_effort'],
thinking_budget=-1 if state.get('enable_thinking', True) else 0,
bos_token=shared.bos_token,
@ -351,6 +352,27 @@ def generate_chat_prompt(user_input, state, **kwargs):
messages.insert(insert_pos, msg_dict)
# Handle <think> blocks (Kimi, DeepSeek, Qwen, etc.)
elif '<think>' in assistant_msg:
thinking_content = ""
final_content = assistant_msg
parts = assistant_msg.split('<think>', 1)
if len(parts) > 1:
potential_content = parts[1]
if '</think>' in potential_content:
thinking_content = potential_content.split('</think>', 1)[0].strip()
final_content = parts[0] + potential_content.split('</think>', 1)[1]
else:
thinking_content = potential_content.strip()
final_content = parts[0]
msg_dict = {"role": "assistant", "content": final_content.strip()}
if thinking_content:
msg_dict["reasoning_content"] = thinking_content
messages.insert(insert_pos, msg_dict)
else:
# Default case (used by all other models)
messages.insert(insert_pos, {"role": "assistant", "content": assistant_msg})