50 lines
4.0 KiB
Python
50 lines
4.0 KiB
Python
"""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)))
|