#!/usr/bin/env python3
"""Held-out eval: mean per-token logprob of the CANONICAL Apache-2.0 text
(never in the tournament corpus) + first-words probes. Compares v1 vs v2 when
run against each checkpoint.

Usage: legume_apache_eval.py <checkpoint> <out_json>
"""
import json
import math
import os
import sys

EVAL_TEXT = "/mnt/vault/llm/glm52-franken/corpus_eval_apache2.txt"
PROMPTS = ("Licensed under the Apache License, Version 2.0 (the",
           "   Copyright [yyyy] [name of copyright owner]\n\n   Licensed under",
           "Redistribution and use in source and binary forms",
           "The capital of France is")


def main() -> None:
    ckpt, out_json = sys.argv[1], sys.argv[2]
    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
    from vllm import LLM, SamplingParams
    from vllm.inputs import TokensPrompt
    from transformers import AutoTokenizer

    # Sized to the workload (256-tok chunks, 40-tok generations) with GiBs of
    # margin: three marginal-fit OOMs on the 31.4 GiB card taught us not to
    # carry 8k-ctx buffers into a 1k-ctx job. Eager: prefill-dominated.
    llm = LLM(model=ckpt, kv_cache_dtype="fp8_ds_mla", max_model_len=1024,
              max_num_seqs=2, max_num_batched_tokens=1024,
              trust_remote_code=False, kv_cache_memory_bytes=256 << 20,
              enforce_eager=True)
    tok = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=False)
    ids = tok.encode(open(EVAL_TEXT, encoding="utf-8").read())

    probe = SamplingParams(temperature=0.0, max_tokens=1, prompt_logprobs=1)
    total_lp, n = 0.0, 0
    # 256-token chunks — offline prefill >256 dispatches to unimplemented
    # forward_mha on this tuple (probe-verified boundary).
    for i in range(0, len(ids), 256):
        chunk = ids[i:i + 256]
        if len(chunk) < 16:
            break
        out = llm.generate([TokensPrompt(prompt_token_ids=chunk)], probe,
                           use_tqdm=False)[0]
        for pos_idx, pos in enumerate(out.prompt_logprobs or []):
            if pos is None:
                continue
            lp = pos.get(chunk[pos_idx])
            if lp is not None:
                total_lp += lp.logprob
                n += 1
    mean_lp = total_lp / max(n, 1)

    words = {}
    sp = SamplingParams(temperature=0.0, max_tokens=40)
    for prompt in PROMPTS:
        out = llm.generate([prompt], sp, use_tqdm=False)
        words[prompt] = out[0].outputs[0].text
        print(f"[eval] {prompt[:44]!r} -> {words[prompt][:80]!r}", flush=True)

    json.dump({"checkpoint": ckpt, "eval_tokens": n,
               "mean_logprob": mean_lp, "ppl": math.exp(-mean_lp),
               "first_words": words}, open(out_json, "w"))
    print(f"EVAL-DONE {out_json} mean_logprob={mean_lp:.4f} "
          f"ppl={math.exp(-mean_lp):.1f} over {n} tokens", flush=True)


if __name__ == "__main__":
    main()