import argparse
import sys
import torch
from modeling import load_model

VERIFICATION_TOP5 = [319, 516, 22441, 497, 4102]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default=".")
    ap.add_argument("--capture", action="store_true")
    args = ap.parse_args()
    torch.set_num_threads(1)
    m = load_model(args.weights, threads=1)
    ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + 100
    with torch.no_grad():
        logits = m.model(ids)
    assert logits.dim() == 3 and logits.size(-1) == m.cfg.vocab_size, (
        f"Bad output shape {tuple(logits.shape)} (expected vocab {m.cfg.vocab_size})"
    )
    assert torch.isfinite(logits).all(), "Non finite logits"
    print(f"Output shape ok {tuple(logits.shape)}")
    top5 = logits[0, -1].topk(5).indices.tolist()
    if args.capture:
        print(f"VERIFICATION_TOP5 = {top5}")
        return
    if top5 != VERIFICATION_TOP5:
        print(
            f"[WARNING] top 5 tokens {top5} differ from the reference {VERIFICATION_TOP5}; small CPU numeric differences can move ties across machines"
        )
    else:
        print("Verification value ok")
    r = m.complete("import numpy as ", "", max_tokens=4)
    print(f"Sample completion: 'import numpy as ' -> {r['middle']!r}")


if __name__ == "__main__":
    main()