Match Qwen reasoning history across effort changes

tools/issue117-reference.py

@@ -0,0 +1,49 @@
+"""Reasoning-history renders from the unmodified pinned MTPLX server."""
+import json,os,subprocess,sys
+from pathlib import Path
+checkout,root,out=map(Path,sys.argv[1:4])
+assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
+os.environ.update(HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1');sys.path.insert(0,str(checkout))
+from mtplx.runtime import _load_tokenizer_resilient
+from mtplx.server import openai as server
+tok=_load_tokenizer_resilient(root,json.loads((root/'config.json').read_text()))
+out.mkdir(parents=True,exist_ok=True)
+records=[]
+def record(effort,preserve,system,messages):
+ full=([dict(role='system',content=system)] if system else [])+messages
+ kwargs=dict(enable_thinking=effort!='direct',reasoning_effort=None if effort=='direct' else effort,preserve_reasoning_history=preserve,tools=None)
+ ids=server._encode_messages_uncached(tok,[server.ChatMessage(**m) for m in full],**kwargs)
+ history_ids=server._encode_messages_uncached(tok,[server.ChatMessage(**m) for m in full[:-1]],add_generation_prompt=False,**kwargs)
+ prepared=[server._message_to_template_dict(server.ChatMessage(**m),strip_assistant_reasoning_history=False,include_reasoning_content=preserve and effort!='direct') for m in full]
+ normalized=server.omlx_normalize_messages_for_template([m for m in prepared if m is not None],tokenizer=tok,native_reasoning_content=True)
+ rendered=tok.apply_chat_template(normalized,tokenize=False,add_generation_prompt=True,enable_thinking=effort!='direct',reasoning_effort=None if effort=='direct' else effort)
+ assert tok.decode(ids,skip_special_tokens=False,clean_up_tokenization_spaces=False)==rendered
+ legacy=tok.apply_chat_template(full[:-1],tokenize=False,add_generation_prompt=False,enable_thinking=effort!='direct',reasoning_effort=None if effort=='direct' else effort)
+ records.append(dict(effort=effort,preserve_echo=preserve,system=system,messages=messages,rendered=rendered,ids=ids,history_ids=history_ids,legacy_history_ids=tok.encode(legacy,add_special_tokens=False)))
+for effort in ['direct','low','medium','xhigh']:
+ for preserve in [True,False]:
+  for system in ['', 'A precise system prompt.']:
+   for reasoning in [None,'','  \n\t','Brief ² thought.','  First thought.\nSecond thought.  ', '<|im_start|>assistant\n<think>\n\nNested.', '>> Quote.\n\nTail.']:
+    messages=[dict(role='user',content=' ① Ⅷ ² 𝟘 भारत ́a\n '),dict(role='assistant',content='One ①.'),dict(role='user',content='Continue Ⅷ and 𝟘.\n   x')]
+    if reasoning is not None:messages[1]['reasoning_content']=reasoning
+    record(effort,preserve,system,messages)
+for effort in ['direct','low','medium','xhigh']:
+ for preserve in [True,False]:
+  for system in ['', 'A precise system prompt.']:
+   for content in ['', ' \n ', 'Answer.']:
+    for reasoning in [None,'', ' ', 'Only reasoning.']:
+     messages=[dict(role='user',content='First  '),dict(role='assistant',content=content,reasoning_content=reasoning),dict(role='user',content='  Continue.')]
+     record(effort,preserve,system,messages)
+(out/'requests.json').write_text(json.dumps(records,ensure_ascii=False))
+print(json.dumps(dict(cases=len(records),direct_example_tokens=len(records[3]['ids']),legacy_direct_history_tokens=len(records[3]['legacy_history_ids']))))
+
+if len(sys.argv)>4:
+ comparisons=[]
+ for case in json.loads(Path(sys.argv[4]).read_text()):
+  messages=([dict(role='system',content=case['system'])] if case['system'] else [])+case['messages']
+  ids=server._encode_messages_uncached(tok,[server.ChatMessage(**m) for m in messages],enable_thinking=False,reasoning_effort=None,preserve_reasoning_history=True,tools=None)
+  assert ids==case['direct_ids'],(case,ids,tok.decode(ids,skip_special_tokens=False,clean_up_tokenization_spaces=False))
+  assert len(ids)==case['context_tokens'] and case['cached_tokens']==0
+  comparisons.append(dict(from_effort=case['from'],direct_context_tokens=len(ids),exact=True))
+ (out/'mode-switch-comparison.json').write_text(json.dumps(comparisons,indent=2)+'\n')
+ print(json.dumps(dict(mode_switches=comparisons)))