"""Compare greedy continuations of torch and the exported ONNX model.

Usage: check_onnx.py EXPORT_DIR PROMPT [CHECKPOINT]
Run it after any change to quantize.py to catch degenerate output.
"""
import json
import sys

import numpy as np
import onnx
import onnxruntime as ort
import torch
from transformers import AutoTokenizer, GPT2LMHeadModel

d, prompt = sys.argv[1], sys.argv[2]
tok = AutoTokenizer.from_pretrained(d)
ids = tok(prompt)["input_ids"]
N = 40

model = GPT2LMHeadModel.from_pretrained(sys.argv[3] if len(sys.argv) > 3 else d).eval()
with torch.no_grad():
    out = model.generate(torch.tensor([ids]), max_new_tokens=N, do_sample=False)
print("torch :", repr(tok.decode(out[0, len(ids):])))

heads = json.load(open(f"{d}/config.json"))["n_head"]
for f in ["onnx/model_q4.onnx"]:
    m = onnx.load(f"{d}/{f}", load_external_data=False)
    big = sorted(((np.prod(i.dims), i.name, i.data_type) for i in m.graph.initializer), reverse=True)[:3]
    so = ort.SessionOptions()
    so.add_session_config_entry("session.disable_quant_qdq", "1")  # as web/brain.ts: fold to fp32
    sess = ort.InferenceSession(f"{d}/{f}", so, providers=["CPUExecutionProvider"])
    names = [i.name for i in sess.get_inputs()]
    L = sum(1 for n in names if n.endswith(".key"))
    kvt = {i.name: (np.float16 if i.type == "tensor(float16)" else np.float32) for i in sess.get_inputs()}
    past = {f"past_key_values.{i}.{k}": np.zeros((1, heads, 0, 64), kvt[f"past_key_values.{i}.{k}"])
            for i in range(L) for k in ("key", "value")}
    cur, seq = list(ids), []
    feed_ids = cur
    for step in range(N):
        n_past = len(cur) - len(feed_ids)
        feed = {"input_ids": np.array([feed_ids], np.int64),
                "attention_mask": np.ones((1, len(cur)), np.int64), **past}
        if "position_ids" in names:
            feed["position_ids"] = np.arange(n_past, len(cur), dtype=np.int64)[None]
        outs = sess.run(None, feed)
        nxt = int(outs[0][0, -1].argmax())
        seq.append(nxt)
        cur.append(nxt)
        feed_ids = [nxt]
        past = {k: v for k, v in zip(past.keys(), outs[1:])}
    print(f"{f.split('/')[-1]:22}:", repr(tok.decode(seq)), "| biggest", [(int(s), dt) for s, _, dt in big])