Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980"""Teacher-forced loss on held-out pet data: torch fp32 vs the exported 4-bit model.
Usage: quant_eval.py CHECKPOINT EXPORT_DIR DATA.jsonl [N]Uses EXPORT_DIR/eval/model_q4_full.onnx, the published weights with logits forevery position (the published file only computes the last one)."""import jsonimport randomimport sys
import numpy as npimport onnxruntime as ortimport torchfrom transformers import AutoTokenizer, GPT2LMHeadModel
import gen_data as G
def fresh(row): try: G.recheck(row) return True except ValueError: return False
ckpt, exp, data = sys.argv[1:4]N = int(sys.argv[4]) if len(sys.argv) > 4 else 40tok = AutoTokenizer.from_pretrained(ckpt)rows = [json.loads(l) for l in open(data)]rows = [r for r in rows if fresh(r)] # as train.py --drop-stale; a no-op on clean datarandom.seed(0)random.shuffle(rows)rows = rows[:N] # train.py holds out the first rows of the same shuffle
def encode(row): enc = tok(row["text"], return_offsets_mapping=True) mark = bytearray(len(row["text"]) + 1) for a, b in row["spans"]: mark[a:b] = b"\x01" * (b - a) ids = enc["input_ids"][:1024] lab = [any(mark[s:e]) for s, e in enc["offset_mapping"][:1024]] return np.array(ids), np.array(lab)
def nll(logits, ids, lab): lp = torch.log_softmax(torch.tensor(logits[:-1], dtype=torch.float32), -1) tgt = torch.tensor(ids[1:]) l = -lp[torch.arange(len(tgt)), tgt] m = torch.tensor(lab[1:]) return l[m].sum().item(), int(m.sum())
model = GPT2LMHeadModel.from_pretrained(ckpt).eval()results = {"torch fp32": [0, 0]}so = ort.SessionOptions()so.add_session_config_entry("session.disable_quant_qdq", "1") # as web/brain.ts: fold to fp32sessions = {"q4 (published)": ort.InferenceSession(f"{exp}/eval/model_q4_full.onnx", so, providers=["CPUExecutionProvider"])}heads = model.config.n_headfor f in sessions: results[f] = [0, 0]for row in rows: ids, lab = encode(row) with torch.no_grad(): lg = model(torch.tensor(ids)[None]).logits[0].numpy() a, b = nll(lg, ids, lab) results["torch fp32"][0] += a; results["torch fp32"][1] += b for f, s in sessions.items(): ins = {i.name: i for i in s.get_inputs()} feed = {"input_ids": ids[None].astype(np.int64), "attention_mask": np.ones((1, len(ids)), np.int64)} if "position_ids" in ins: feed["position_ids"] = np.arange(len(ids), dtype=np.int64)[None] for k, v in ins.items(): if k.startswith("past"): feed[k] = np.zeros((1, heads, 0, 64), np.float32) a, b = nll(s.run(["logits"], feed)[0][0], ids, lab) results[f][0] += a; results[f][1] += bfor k, (a, b) in results.items(): print(f"{k:22} loss {a / b:.4f} ppl {np.exp(a / b):.3f} ({b} tokens)")