"""Rescore existing unconditional samples with independent causal language models. This leaves the canonical GPT-2 Large metrics untouched. Run from the repository root with the dlb-eval environment and one free GPU. """ from __future__ import annotations import argparse from dataclasses import asdict import json import math from pathlib import Path from dlb.io import atomic_json_write, sha256_file, validate_samples from evaluation.generative_perplexity import compute_gen_ppl ROOT = Path(__file__).resolve().parents[1] MODELS = { "pythia-1b": ( "EleutherAI/pythia-1b", "f73d7dcc545c8bd326d8559c8ef84ffe92fea6b2", ["*.json", "model.safetensors"], ), "opt-1.3b": ( "facebook/opt-1.3b", "3f5c25d0bc631cb57ac65913f76e22c2dfb61d62", ["*.json", "pytorch_model.bin", "merges.txt", "vocab.json"], ), } DATASETS = ("lm1b", "owt") GENERATORS = ("fmlm", "duo_dcd", "duo_di4c", "mdlm_sdtt", "mdlm_di4c") STEPS = (1, 2, 4) SAMPLE_COUNT = 1024 CONTEXT_LENGTH = 1024 def scorer_class(): import torch import torch.nn.functional as functional class Float32LossScorer: def __init__(self, model): self.model = model def eval(self): self.model.eval() return self def score_batch(self, input_ids, attention_mask): ids = torch.as_tensor(input_ids, dtype=torch.long, device="cuda") mask = torch.as_tensor(attention_mask, dtype=torch.bool, device="cuda") with torch.inference_mode(): logits = self.model(input_ids=ids, attention_mask=mask).logits labels = ids[:, 1:].clone() valid = mask[:, 1:] labels.masked_fill_(~valid, -100) loss = functional.cross_entropy( logits[:, :-1, :].float().reshape(-1, logits.size(-1)), labels.reshape(-1), ignore_index=-100, reduction="sum", ) return float(loss.item()), int(valid.sum().item()) return Float32LossScorer def load_model(model_key: str): import torch from huggingface_hub import snapshot_download from transformers import AutoModelForCausalLM, AutoTokenizer repo_id, revision, patterns = MODELS[model_key] snapshot = Path( snapshot_download(repo_id=repo_id, revision=revision, allow_patterns=patterns) ) tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True) tokenizer.padding_side = "right" tokenizer.truncation_side = "right" if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( snapshot, local_files_only=True, dtype=torch.float32 ).to("cuda").eval() return scorer_class()(model), tokenizer, snapshot def sample_path(dataset: str, generator: str, steps: int) -> Path: return ROOT / "results" / "samples" / dataset / generator / f"steps_{steps}" / "samples.jsonl" def output_path(output_root: Path, model_key: str, dataset: str, generator: str, steps: int) -> Path: return output_root / model_key / dataset / generator / f"steps_{steps}.json" def valid_existing(path: Path, sample_sha: str, model_key: str) -> bool: if not path.is_file(): return False try: row = json.loads(path.read_text(encoding="utf-8")) ppl = row["generative_perplexity"] return ( row["samples_sha256"] == sample_sha and row["sample_count"] == SAMPLE_COUNT and ppl["model_id"] == MODELS[model_key][0] and ppl["model_revision"] == MODELS[model_key][1] and ppl["forward_dtype"] == "float32" and ppl["valid_token_count"] > 0 and math.isfinite(ppl["perplexity"]) ) except (KeyError, TypeError, ValueError, OSError): return False def main() -> None: import torch import transformers parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model", choices=(*MODELS, "all"), default="all") parser.add_argument("--output-root", type=Path, default=ROOT / "results" / "ppl_rescore_20260923") parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--smoke-samples", type=int, default=0) args = parser.parse_args() if args.batch_size < 1 or args.smoke_samples < 0: parser.error("batch-size must be positive and smoke-samples must be nonnegative") # Validate the entire input matrix before loading a scorer or writing results. inputs = {} for dataset in DATASETS: for generator in GENERATORS: for steps in STEPS: path = sample_path(dataset, generator, steps) validate_samples(path, expected=SAMPLE_COUNT) inputs[(dataset, generator, steps)] = (path, sha256_file(path)) print(f"Validated {len(inputs)} sample files x {SAMPLE_COUNT} samples", flush=True) model_keys = tuple(MODELS) if args.model == "all" else (args.model,) for model_key in model_keys: pending = [ (dataset, generator, steps) for dataset in DATASETS for generator in GENERATORS for steps in STEPS if args.smoke_samples or not valid_existing( output_path(args.output_root, model_key, dataset, generator, steps), inputs[(dataset, generator, steps)][1], model_key, ) ] if not pending: print(f"{model_key}: all {len(inputs)} outputs already validated", flush=True) continue scorer, tokenizer, snapshot = load_model(model_key) print(f"{model_key}: loaded {snapshot}; pending={len(pending)}", flush=True) for index, (dataset, generator, steps) in enumerate(pending, 1): path, sample_sha = inputs[(dataset, generator, steps)] with path.open(encoding="utf-8") as handle: texts = [json.loads(line)["text"] for line in handle] if args.smoke_samples: texts = texts[:args.smoke_samples] result = compute_gen_ppl( texts, scorer, tokenizer, batch_size=args.batch_size, max_length=CONTEXT_LENGTH, model_revision=MODELS[model_key][1], tokenizer_revision=MODELS[model_key][1], ) metric = { **asdict(result), "model_id": MODELS[model_key][0], "tokenizer_id": MODELS[model_key][0], "model_snapshot": str(snapshot), "tokenizer_snapshot": str(snapshot), "forward_dtype": "float32", "loss_dtype": "float32", "torch_version": torch.__version__, "transformers_version": transformers.__version__, "offline_scoring": True, } print( f"{model_key} {index}/{len(pending)} {dataset}/{generator}/{steps} " f"samples={len(texts)} tokens={result.valid_token_count} " f"ppl={result.perplexity:.6f}", flush=True, ) if args.smoke_samples: break atomic_json_write( output_path(args.output_root, model_key, dataset, generator, steps), { "schema_version": 1, "evaluation": "independent_unconditional_ppl_rescore", "samples": str(path), "samples_sha256": sample_sha, "sample_count": SAMPLE_COUNT, "dataset": dataset, "generator": generator, "steps": steps, "generative_perplexity": metric, }, ) if __name__ == "__main__": main()