| |
| """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") |
| |
| |
| |
| 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() |
|
|
