39 lines
2.0 KiB
Python
39 lines
2.0 KiB
Python
"""Run the pinned MTPLX final-state exit matrix without loading model weights."""
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
|
|
root = Path(sys.argv[1]).resolve()
|
|
assert subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
|
|
assert os.environ["MTPLX_CONTEXT_COPY"] == "0"
|
|
sys.path.insert(0, str(root))
|
|
spec = importlib.util.spec_from_file_location("final_state", root / "tests/test_mtp_final_state_covers_committed_tokens.py")
|
|
m = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(m)
|
|
for label, stop_index, depth, omniscient, temperature, maximum in [
|
|
("bonus primary", 8, 2, False, 0.0, 16),
|
|
("deferred greedy correction", 8, 3, False, 0.0, 16),
|
|
("accepted draft stop", 8, 3, True, 0.0, 16),
|
|
("sampled correction", 8, 3, False, 0.7, 16),
|
|
("first primary", 6, 3, False, 0.0, 16),
|
|
("fresh primary depth 1", 7, 1, False, 0.0, 16),
|
|
("max tokens depth 3", 10000, 3, False, 0.0, 7),
|
|
("max tokens depth 2", 10000, 2, False, 0.0, 5),
|
|
]:
|
|
sampler = m.SamplerConfig(temperature=temperature, top_k=1 if temperature else 0)
|
|
out = m.generate_mtpk(m._runtime(m._ScriptedModel(stop_index, omniscient)), m.PROMPT,
|
|
speculative_depth=depth, max_tokens=maximum, sampler=sampler, **m.COMMON)
|
|
gap, origin, safe = m._gap(m.PROMPT, out)
|
|
assert gap == 0 and safe
|
|
assert out.final_state.final_committed_mtp_cache[0].offset == len(m.PROMPT) + len(out.tokens) - 1
|
|
if stop_index < 1000:
|
|
assert out.tokens[-1] == m.STOP
|
|
print(json.dumps(dict(case="mtp_exit", label=label, stop_index=stop_index, depth=depth,
|
|
omniscient=omniscient, temperature=temperature, max_tokens=maximum,
|
|
tokens=out.tokens, gap=gap, origin=origin, safe=safe)), flush=True)
|
|
m.test_ar_final_state_control()
|
|
print(json.dumps(dict(case="ar_control", passed=True)), flush=True)
|