"""GPU-free load-and-generate smoke test for the quantized tree.

§4.4 says test on a temp port, never on the live seat -- but GPU1 has 19.9 GB free
against 19.5 GB of weights, so a vLLM serve test cannot happen without displacing a
live seat, and that is not my call. This is what CAN be established without one:
that the checkpoint's tensor names map cleanly onto the architecture (no missing or
unexpected keys), that compressed-tensors can decompress it, and that it emits
plausible tokens rather than garbage.

What it does NOT establish: that vLLM's sm_120 NVFP4 kernels serve it, or anything
about long-context quality. Those still need the GPU. Saying so is part of the
result -- a smoke test whose limits go unstated gets read as more than it is.

Deliberately greedy and short. This is an "is it wired up" check, not an eval; n=1
proves nothing about quality and is not offered as if it did.
"""
import sys
import time

import torch
from transformers import AutoTokenizer

path = sys.argv[1]
print(f"tree: {path}", flush=True)

t0 = time.time()
tok = AutoTokenizer.from_pretrained(path)
print(f"tokenizer OK ({time.time()-t0:.1f}s)  truncation_side={tok.truncation_side}", flush=True)

t0 = time.time()
from transformers import AutoModelForImageTextToText as M
model = M.from_pretrained(path, dtype=torch.bfloat16, device_map=None)
print(f"model loaded on CPU ({time.time()-t0:.1f}s)  {type(model).__name__}", flush=True)

n = sum(p.numel() for p in model.parameters())
print(f"parameters: {n/1e9:.2f} B", flush=True)

# Any tensor still sitting on meta means a weight the checkpoint never supplied --
# from_pretrained does not always raise on that, it just leaves the hole.
meta = [k for k, v in model.state_dict().items() if v.is_meta]
print(f"tensors still on meta device: {len(meta)}"
      + (f"  *** {meta[:5]}" if meta else "  (none -- every weight was materialised)"), flush=True)

msgs = [{"role": "user", "content": "In one sentence, what is a lighthouse for?"}]
enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt")
# transformers 5.x hands back a BatchEncoding here, not a bare tensor.
ids = enc["input_ids"] if hasattr(enc, "keys") else enc
print(f"prompt tokens: {ids.shape[-1]}", flush=True)

t0 = time.time()
with torch.inference_mode():
    out = model.generate(ids, max_new_tokens=24, do_sample=False)
dt = time.time() - t0
text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True)
print(f"generated {out.shape[-1]-ids.shape[-1]} tokens in {dt:.1f}s "
      f"({dt/max(1,out.shape[-1]-ids.shape[-1]):.1f}s/tok, CPU)", flush=True)
print(f"OUTPUT: {text!r}", flush=True)