code/eval_streak_ckpt.py

"""Standalone streak eval: python3 eval_streak_ckpt.py <ckpt> [max_heads]"""
import sys, torch
from config import Config
from model import build_model
from data_pipeline import load_tokens
from train_medusa import eval_streak

ckpt_path = sys.argv[1]
max_heads = int(sys.argv[2]) if len(sys.argv) > 2 else None

cfg = Config.v5_500m()
model = build_model(cfg, "cuda")
ckpt = torch.load(ckpt_path, map_location="cuda")
missing, unexpected = model.load_state_dict(ckpt["model"], strict=False)
print(f"loaded {ckpt_path} (step {ckpt.get('step')})  missing={len(missing)}")
model.eval()

cfg.use_markov_head = hasattr(model, "spec_markov_emb")
data = load_tokens("mixture500m").to("cuda")

for mk in ([True, False] if cfg.use_markov_head else [False]):
    cfg.use_markov_head = mk
    s, a = eval_streak(model, cfg, data, seq=1024, device="cuda",
                       max_heads=max_heads)
    print(f"markov={'on ' if mk else 'off'}  streak ~{s:.2f}  "
          f"head-0 acc {a:.3f}  (max_heads={max_heads})")