#!/usr/bin/env python3 """Evaluate public ELF OWT-prefix samples with the legacy 64+512 contract. The ELF checkpoint still samples its native 1024-position T5 canvas. This evaluator defines the comparable text-level contract: exactly 64 GPT-2 prompt tokens and the first 512 GPT-2 response tokens per completion. """ from __future__ import annotations import argparse import csv import json from pathlib import Path import statistics import sys from elf_common import ROOT, asset, manifest, offline, read_jsonl, require_server, sha256, write_json from evaluate_elf import validate_records def _load_prompt_rows(root: Path): rows = read_jsonl(root / "data/conditional/owt-c64/prompts.jsonl") prepared = read_jsonl(root / "data/elf/owt-prefix.jsonl") by_id = {row["prompt_id"]: row for row in rows} prepared_by_id = {row["id"]: row for row in prepared} expected_ids = list(range(1024)) if len(rows) != 1024 or sorted(by_id) != expected_ids: raise ValueError("canonical OWT C64 prompt file must contain 1024 unique prompt IDs") if len(prepared) != 1024 or sorted(prepared_by_id) != expected_ids: raise ValueError("prepared ELF OWT-prefix input must contain 1024 unique prompt IDs") for prompt_id, row in by_id.items(): if len(row["prefix_token_ids"]) != 64: raise ValueError(f"prompt_id={prompt_id} does not contain 64 canonical GPT-2 tokens") row["prepared_input"] = prepared_by_id[prompt_id]["input"] return by_id def _tokenize_contract(rows, prompts, tokenizer, response_tokens: int): """Return (prompt IDs, response IDs, short IDs) for every saved completion.""" prompt_ids, response_ids, short_ids = [], [], [] for row in rows: expected = prompts[row["prompt_id"]] if row["input"] != expected["prepared_input"]: raise ValueError(f"prompt_id={row['prompt_id']} input text differs from the prepared ELF input") # Use the canonical source IDs for GPT-2 scoring. GPT-2 byte-BPE # decode->encode is not guaranteed to preserve token boundaries, even # when the decoded input text is unchanged. prompt = list(expected["prefix_token_ids"]) response = tokenizer( row["generated"], add_special_tokens=False, truncation=True, max_length=response_tokens, )["input_ids"] prompt_ids.append(prompt) response_ids.append(response) if len(response) < response_tokens: short_ids.append(row["id"]) return prompt_ids, response_ids, short_ids def _score_ppl(prompt_ids, response_ids, model, tokenizer, batch_size, response_tokens): from evaluation.conditional_perplexity import PromptExcludedCausalLMScorer from evaluation.generative_perplexity import aggregate_nll scorer = PromptExcludedCausalLMScorer(model, device="cuda").eval() parts = [] for start in range(0, len(prompt_ids), batch_size): p_batch = prompt_ids[start:start + batch_size] r_batch = [row[:response_tokens] for row in response_ids[start:start + batch_size]] combined = [p + r for p, r in zip(p_batch, r_batch, strict=True)] width = max(map(len, combined)) ids = [row + [tokenizer.pad_token_id] * (width - len(row)) for row in combined] attention = [[1] * len(row) + [0] * (width - len(row)) for row in combined] targets = [] expected = 0 for p, r in zip(p_batch, r_batch, strict=True): mask = [0] * len(p) + [1] * len(r) mask += [0] * (width - len(mask)) targets.append(mask) expected += len(r) nll, count = scorer.score_prompt_excluded_batch(ids, attention, targets) if count != expected: raise ValueError(f"response target count mismatch: {count} != {expected}") parts.append((nll, count)) return aggregate_nll(parts) def main(argv=None): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--root", type=Path, default=ROOT) parser.add_argument("--input-root", type=Path, required=True, help="quality root containing owt-prefix/c64-text-t5/steps_N/... runs") parser.add_argument("--output", type=Path, required=True) parser.add_argument("--steps", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 1024]) parser.add_argument("--response-tokens", type=int, default=512) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--allow-short", action="store_true", help="score observed shorter responses; result is not strict 64+512") args = parser.parse_args(argv) if args.response_tokens < 2 or args.batch_size < 1: parser.error("response-tokens must be >=2 and batch-size must be positive") root = args.root.resolve() input_root = args.input_root.resolve() output = args.output.resolve() sys.path[:0] = [str(root), str(root / "src")] require_server() offline() from transformers import AutoModelForCausalLM, AutoTokenizer from evaluation.unigram_entropy import unigram_entropy prompts = _load_prompt_rows(root) assets = manifest(root)["assets"] gpt2 = AutoTokenizer.from_pretrained( str(asset(root, "gpt2")), local_files_only=True, clean_up_tokenization_spaces=False, ) gpt2.pad_token = gpt2.eos_token gpt2.padding_side = "right" model = AutoModelForCausalLM.from_pretrained( str(asset(root, "gpt2-large")), local_files_only=True ).cuda().eval() output.mkdir(parents=True, exist_ok=True) summary = [] for steps in args.steps: run = input_root / "quality" / "owt-prefix" / "c64-text-t5" / f"steps_{steps}" / "seed_42" / "compile_0" samples_path = run / "samples.jsonl" if not samples_path.exists(): raise FileNotFoundError(f"missing generated samples: {samples_path}") rows = read_jsonl(samples_path) validate_records(rows, {"sample_count": len(rows), "prompt_count": 1024, "diversity": True}) if len(rows) != 2048: raise ValueError(f"expected 2048 completions at {run}, found {len(rows)}") prompt_ids, response_ids, short_ids = _tokenize_contract(rows, prompts, gpt2, args.response_tokens) invalid = [] if args.allow_short else short_ids result = { "schema": "dlb-elf-owt-c64r512-v1", "protocol": "owt_c64_r512_observed_v1" if args.allow_short else "owt_c64_r512_strict_v1", "steps": steps, "sample_count": len(rows), "prompt_count": 1024, "completion_schedule": "1024 primary + 256 prompts x 4 extra completions", "prompt_tokenizer": "gpt2", "prompt_tokens": 64, "response_tokenizer": "gpt2", "response_tokens_requested": args.response_tokens, "response_short_count": len(short_ids), "short_sample_ids": short_ids, "source_run": str(run), "samples_sha256": sha256(samples_path), "generation_manifest_sha256": sha256(run / "generation.json"), "assets_manifest_sha256": sha256(root / "data/elf/assets.json"), "metrics": {}, "valid": not invalid, "policy": "all 2048 completions; first 512 GPT-2 response tokens; prompt loss excluded", "native_generation_note": "ELF samples a 1024-position T5 canvas before GPT-2 text retokenization", "length_diagnostics": { "response_gpt2_capped_mean": statistics.fmean(map(len, response_ids)), "response_gpt2_capped_min": min(map(len, response_ids)), "response_gpt2_capped_max": max(map(len, response_ids)), }, } if invalid: result["reason"] = "strict 64+512 contract rejected short responses; no samples were dropped or padded" else: effective = [row[:args.response_tokens] for row in response_ids] result["metrics"]["entropy_gpt2_nats"] = statistics.fmean( unigram_entropy(row) for row in effective ) ppl = _score_ppl(prompt_ids, response_ids, model, gpt2, args.batch_size, args.response_tokens) result["metrics"]["conditional_ppl"] = { "perplexity": ppl.perplexity, "total_nll": ppl.total_nll, "valid_token_count": ppl.valid_token_count, "sample_count": len(rows), "context_length": 64 + args.response_tokens, "model_id": "gpt2-large", "model_revision": assets["gpt2-large"]["revision"], "tokenizer_id": "gpt2", "tokenizer_revision": assets["gpt2"]["revision"], } destination = output / f"steps_{steps}" / "metrics.json" write_json(destination, result) summary.append({ "steps": steps, "sample_count": len(rows), "valid": result["valid"], "short_count": len(short_ids), "ppl": result["metrics"].get("conditional_ppl", {}).get("perplexity"), "entropy": result["metrics"].get("entropy_gpt2_nats"), "metrics_path": str(destination), }) print(json.dumps(summary[-1], ensure_ascii=False), flush=True) with (output / "summary.csv").open("w", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=list(summary[0])) writer.writeheader() writer.writerows(summary) print(f"SUMMARY: {output / 'summary.csv'}", flush=True) if any(not row["valid"] for row in summary): raise SystemExit(2) if __name__ == "__main__": main()