Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
12 kB · 207 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208"""Held-out masked loss for a real checkpoint and its candidate ONNX weights.
Usage: quant_eval.py CHECKPOINT EXPORT_DIR DATA.jsonl [N] [--dtype q4|q8|fp32]Uses eval/model_<dtype>_full.onnx. Canonical validation, strict offset masks,explicit diary EOS and the seeded holdout are shared with training. JSON outputincludes denominators, conditional bits per supervised UTF-8 byte and actualartifact bytes; --report writes the same report to a file."""import argparseimport hashlibimport jsonimport mathfrom pathlib import Pathimport platform
import numpy as npimport torchimport torch.nn.functional as Ffrom transformers import AutoModelForCausalLM, AutoTokenizer
from model_data import encode, heldout_split, load_rows, rows_digestfrom onnx_utils import CausalSession, DTYPES, artifact_files, config_dimensions, model_path, provenance, sha256
def supervised_bytes(tok, row, ids, labels): """Union character offsets: byte-level tokens can share a Unicode character.
EOS contributes a loss term but zero bytes. This is conditional sequence NLLper observed supervised text byte, not a claim to compress the whole dialogue.""" encoded = tok(row["text"], add_special_tokens=False, return_offsets_mapping=True) count = min(len(encoded["input_ids"]), len(ids)) if ids[:count] != encoded["input_ids"][:count]: raise ValueError("loss-mask encoding and byte-accounting token IDs disagree") coverage = bytearray(len(row["text"])) for i, (start, end) in enumerate(encoded["offset_mapping"][:count]): if i > 0 and labels[i] != -100: coverage[start:end] = b"\x01" * (end - start) return sum(len(char.encode("utf-8")) for char, trained in zip(row["text"], coverage) if trained)
def nll(logits, labels): if not isinstance(logits, torch.Tensor): logits = torch.from_numpy(logits) logits = logits.float() if logits.ndim != 2 or logits.shape[0] != len(labels): raise ValueError(f"full-logit sequence shape mismatch: {tuple(logits.shape)}, labels={len(labels)}") targets = torch.tensor(labels[1:], dtype=torch.long, device=logits.device) count = int((targets != -100).sum().item()) if not count: raise ValueError("held-out row has no predictable supervised tokens within the context budget") value = float(F.cross_entropy(logits[:-1], targets, ignore_index=-100, reduction="sum").item()) if not math.isfinite(value): raise ValueError("nonfinite held-out loss") return value, count
def _summary(total_nll, tokens, byte_count): loss = total_nll / tokens return {"nll_nats": total_nll, "supervised_tokens": tokens, "supervised_utf8_bytes": byte_count, "loss_nats_per_token": loss, "perplexity": math.exp(loss) if loss < 700 else None, "bits_per_supervised_byte": total_nll / math.log(2) / byte_count if byte_count else None}
def evaluate(args): if args.n < 1: raise ValueError("N must be positive") checkpoint = Path(args.checkpoint) training_path = checkpoint / "training.json" training = json.loads(training_path.read_text()) if training_path.is_file() else {} settings = training.get("configuration", training.get("arguments", {})) seed = args.seed if args.seed is not None else settings.get("seed", 0) pet_training = bool(training) and training.get("stage", "pet") == "pet" trained_split_seed = settings.get("split_seed", training.get("split", {}).get("seed", 0)) if pet_training else None split_seed = args.split_seed if args.split_seed is not None else (trained_split_seed if trained_split_seed is not None else 0) if trained_split_seed is not None and split_seed != trained_split_seed: raise ValueError(f"split seed {split_seed} differs from training holdout seed {trained_split_seed}") trained_data = training.get("data", {}) data_hash = sha256(args.data) expected_hash = trained_data.get("sha256") if pet_training and isinstance(trained_data, dict) else None if expected_hash and expected_hash != data_hash: raise ValueError("evaluation data hash differs from pet training data: it cannot be called the training holdout") max_len = args.max_len if args.max_len is not None else settings.get("max_len", training.get("max_len", 1024)) if not isinstance(max_len, int) or max_len < 2: raise ValueError("max-len must be an integer of at least two") drop_stale = bool(args.drop_stale if args.drop_stale is not None else settings.get("drop_stale", False)) rows, validation = load_rows(args.data, drop_stale=drop_stale) fit, holdout = heldout_split(rows, split_seed) holdout_hash = rows_digest(holdout) expected_holdout_hash = training.get("split", {}).get("eval_rows_sha256") if pet_training else None if expected_holdout_hash and holdout_hash != expected_holdout_hash: raise ValueError("validated holdout rows differ from the training receipt (split/filtering mismatch)") selected = holdout[:args.n] if not selected: raise ValueError("empty held-out set") if args.threads is not None: if args.threads < 1: raise ValueError("threads must be positive") torch.set_num_threads(args.threads) torch.manual_seed(seed) if args.device == "cuda" and not torch.cuda.is_available(): raise ValueError("CUDA requested but unavailable") device = "cuda" if args.device == "cuda" or args.device == "auto" and torch.cuda.is_available() else "cpu" tok = AutoTokenizer.from_pretrained(checkpoint, use_fast=True) export_tok = AutoTokenizer.from_pretrained(args.export, use_fast=True) with open(Path(args.export) / "config.json") as f: config = json.load(f) context = config.get("max_position_embeddings", config.get("n_positions")) if context and max_len > context: raise ValueError(f"max-len {max_len} exceeds the exported model context {context}") full = model_path(args.export, args.dtype, full=True) browser = model_path(args.export, args.dtype) session = CausalSession(full, config, args.threads, args.qdq, args.prepack) model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.float32, attn_implementation="eager").to(device).eval() if config_dimensions(model.config.to_dict()) != session.dimensions: raise ValueError("torch and ONNX model configurations disagree") totals = {"torch": [0.0, 0], "onnx": [0.0, 0]} byte_count, transcripts = 0, [] with torch.no_grad(): for row_index, row in enumerate(selected): ids, labels = encode(tok, row, max_len) export_ids, export_labels = encode(export_tok, row, max_len) if ids != export_ids or labels != export_labels: raise ValueError("checkpoint/export tokenizers or loss masks disagree") bytes_in_row = supervised_bytes(tok, row, ids, labels) byte_count += bytes_in_row input_ids = torch.tensor([ids], dtype=torch.long, device=device) output = model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids), use_cache=False) torch_nll, count = nll(output.logits[0], labels) logits, _ = session.run(np.asarray(ids, dtype=np.int64)) onnx_nll, onnx_count = nll(logits[0], labels) if count != onnx_count: raise ValueError("torch and ONNX loss denominators disagree") for name, value in (("torch", torch_nll), ("onnx", onnx_nll)): totals[name][0] += value totals[name][1] += count transcripts.append({"holdout_index": row_index, "kind": row["kind"], "text_sha256": hashlib.sha256(row["text"].encode("utf-8")).hexdigest(), "sequence_tokens": len(ids), "supervised_tokens": count, "supervised_utf8_bytes": bytes_in_row, "torch_nll_nats": torch_nll, "onnx_nll_nats": onnx_nll, "diary_eos_supervised": row["kind"] == "diary" and len(ids) > len(tok(row["text"], add_special_tokens=False)["input_ids"])}) torch_summary = _summary(*totals["torch"], byte_count) onnx_summary = _summary(*totals["onnx"], byte_count) baseline = torch_summary["loss_nats_per_token"] delta = onnx_summary["loss_nats_per_token"] - baseline browser_files, full_files = artifact_files(browser), artifact_files(full) return {"status": "complete", "dtype": args.dtype, "fp_loss": baseline, "quant_loss": onnx_summary["loss_nats_per_token"], "relative_increase": delta / baseline if baseline else None, "N": len(selected), "dataset_sha256": data_hash, "graph_sha256": sha256(browser), "checkpoint": provenance(checkpoint), "export": provenance(args.export), "data": {"path": str(Path(args.data).resolve()), "sha256": data_hash, "valid_rows": len(rows), "validation": validation, "fit_rows": len(fit), "heldout_rows": len(holdout), "fit_rows_sha256": rows_digest(fit), "heldout_rows_sha256": holdout_hash, "evaluated_rows": len(selected), "requested_rows": args.n, "holdout_relation": "matching_pet_training_holdout" if expected_hash else "seeded_data_holdout_without_pet_training_receipt"}, "configuration": {"seed": seed, "split_seed": split_seed, "max_len": max_len, "device": device, "threads": torch.get_num_threads(), "drop_stale": drop_stale, "qdq_fusion": args.qdq, "prepacking": args.prepack, "add_special_tokens": False, "diary_eos": True, "byte_denominator": "union of supervised predicted-token character offsets encoded as UTF-8; EOS has zero bytes", "bits_interpretation": "conditional supervised sequence NLL including diary EOS; not whole-dialogue compression"}, "hardware": {"platform": platform.platform(), "processor": platform.processor(), "cuda_device": torch.cuda.get_device_name() if device == "cuda" else None}, "model": {"unique_parameters": sum(p.numel() for p in model.parameters()), "graph": session.metadata()}, "artifacts": {"browser": {"files": browser_files, "bytes": sum(browser_files.values()), "sha256": sha256(browser)}, "full_logits": {"files": full_files, "bytes": sum(full_files.values()), "sha256": sha256(full)}, "export_directory_bytes": sum(p.stat().st_size for p in Path(args.export).rglob("*") if p.is_file())}, "torch": torch_summary, "onnx": onnx_summary, "quantization_delta": {"loss_nats_per_token": delta, "relative_loss": delta / baseline if baseline else None}, "rows": transcripts}
def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("checkpoint") parser.add_argument("export") parser.add_argument("data") parser.add_argument("n", nargs="?", type=int, default=40) parser.add_argument("--dtype", choices=DTYPES, default="q4") parser.add_argument("--report") parser.add_argument("--seed", type=int) parser.add_argument("--split-seed", type=int, help="defaults to the pet training receipt's holdout seed, otherwise zero") parser.add_argument("--max-len", type=int) parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="cpu") parser.add_argument("--threads", type=int) parser.add_argument("--drop-stale", action="store_true", default=None) parser.add_argument("--qdq", action="store_true", help="allow QDQ fusion rather than fp32 folding") parser.add_argument("--prepack", action="store_true") args = parser.parse_args() try: report = evaluate(args) except Exception as error: report = {"status": "error", "error": f"{type(error).__name__}: {error}", "dtype": args.dtype} if args.report: with open(args.report, "w") as f: json.dump(report, f, indent=2, allow_nan=False) f.write("\n") print(json.dumps(report, allow_nan=False)) raise SystemExit(0 if report["status"] == "complete" else 1)
if __name__ == "__main__": main()