Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152"""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 jsonimport sys
import numpy as npimport onnximport onnxruntime as ortimport torchfrom 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])