"""Teacher-forced perplexity through the decode path (one token per forward), so decode-only kernels are exercised."""
import os, sys, time
from pathlib import Path
import torch
from transformers import AutoTokenizer, MetalConfig, Qwen4ExpForCausalLM
root = Path("/private/tmp/ssd-moe-local/qwen38"); tag = sys.argv[1]; N = int(sys.argv[2]) if len(sys.argv) > 2 else 512
tok = AutoTokenizer.from_pretrained(root / "dense")
model = Qwen4ExpForCausalLM.from_pretrained(root / "dense", dtype=torch.bfloat16, device_map="mps", attn_implementation="sdpa",
    quantization_config=MetalConfig(bits=8, group_size=64, modules_to_not_convert=["shared_expert_gate"]),
    experts_offload={"cache_bytes": 12 << 30, "quantized_experts": root / "experts-q3", "io_threads": 8,
                     "disk_embeddings": {"model.layers.1.ple.ple_embedding.ngram_embedding": root / "ngram"}})
text = open("/private/tmp/transformers-moe-ssd-streaming/docs/source/en/philosophy.md").read()
ids = tok(text, return_tensors="pt").input_ids[:, :N].to("mps")
nll, top1 = [], []
t0 = time.perf_counter()
with torch.inference_mode():
    out = model(ids[:, :1], use_cache=True); cache = out.past_key_values
    for i in range(1, N):
        logits = out.logits[0, -1].float()
        nll.append(torch.nn.functional.cross_entropy(logits[None], ids[0, i:i + 1])); top1.append(logits.argmax())
        out = model(ids[:, i:i + 1], past_key_values=cache, use_cache=True)
nll = torch.stack(nll); top1 = torch.stack(top1).cpu()
torch.save(top1, f"{os.path.dirname(os.path.abspath(__file__))}/decode_top1_{tag}.pt")
print(tag, "decode-path tokens", N, "ppl", round(nll.mean().exp().item(), 4), f"({time.perf_counter() - t0:.0f} s)")