validate_release.py

"""Validate published regression checkpoints against their fixed evaluation data.

Paths supplied on the command line are relative to the repository root. Output
directories must be new, so validation cannot overwrite archived evidence.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent
REGRESSION_BASE = "RNASeek-regression-pretrained"
DEFAULTS = {
    "efficiency": {
        "checkpoint": "efficiency_figure2/qwen_regression_ckpt/clean_cosine_restart_besthp_preview_fixed-wd-0.9_reproduce/checkpoint-304419",
        "data": "efficiency_figure2/evenBetterDataFolded-vl.json",
        "metric": "r2",
        "target": 0.389,
    },
    "stability": {
        "checkpoint": "regression_stability_functionalviral/rnaseek_abs20_tuning/hpcc_best_wave29_copyback/best_model",
        "data": "regression_stability_functionalviral/rnaseek_abs20_tuning/hpcc_best_wave29_copyback/valid_predictions.tsv",
        "metric": "spearman_r",
        "target": 0.342,
    },
}


def local_path(value):
    path = Path(value).expanduser()
    return path.resolve() if path.is_absolute() else (ROOT / path).resolve()


def load_regression_checkpoint(checkpoint, tokenizer_path=None, *, initialize_head=False):
    """Strictly load a predictor, or explicitly initialize a training-only head."""
    import torch
    from torch import nn
    from safetensors.torch import load_file
    from transformers import AutoConfig, AutoModel, AutoTokenizer
    checkpoint = local_path(checkpoint)
    config = AutoConfig.from_pretrained(checkpoint, local_files_only=True)
    tokenizer = AutoTokenizer.from_pretrained(
        local_path(tokenizer_path) if tokenizer_path else checkpoint,
        local_files_only=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # Strict tensor loading supports both archived layouts: a standalone
    # backbone plus head, or a single state dict prefixed with backbone.
    index = checkpoint / "model.safetensors.index.json"
    shards = sorted(set(json.loads(index.read_text())["weight_map"].values())) if index.exists() else ["model.safetensors"]
    state = {}
    for shard in shards:
        state.update(load_file(str(checkpoint / shard)))
    if any(key.startswith("backbone.") for key in state):
        head_state = {key.removeprefix("regression_head."): value for key, value in state.items() if key.startswith("regression_head.")}
        backbone_state = {key.removeprefix("backbone."): value for key, value in state.items() if key.startswith("backbone.")}
        unexpected = [key for key in state if not key.startswith(("backbone.", "regression_head."))]
        if unexpected:
            raise ValueError("Unexpected tensors in combined checkpoint")
    elif initialize_head and not (checkpoint / "regression_head.pt").exists() and any(
        name.endswith("ForCausalLM") for name in (config.architectures or [])
    ):
        unexpected = [key for key in state if not key.startswith("model.") and key != "lm_head.weight"]
        if unexpected:
            raise ValueError("Unexpected tensors in pretrained causal-language-model checkpoint")
        backbone_state = {key.removeprefix("model."): value for key, value in state.items() if key.startswith("model.")}
        head_state = None
    else:
        backbone_state = state
        if not (checkpoint / "regression_head.pt").exists():
            raise ValueError("This checkpoint has no trained regression head. Use train_regression.py to fine-tune the pretrained base first.")
        head_state = torch.load(checkpoint / "regression_head.pt", map_location="cpu", weights_only=True)
    linear_index = 2 if head_state is None or "net.2.weight" in head_state else 1
    head = nn.Module()
    layers = [nn.LayerNorm(config.hidden_size)]
    if linear_index == 2:
        layers.append(nn.Dropout(0.1))
    layers.append(nn.Linear(config.hidden_size, 1))
    head.net = nn.Sequential(*layers)
    if head_state is not None:
        head.load_state_dict(head_state, strict=True)
    backbone = AutoModel.from_config(config, attn_implementation="sdpa")
    backbone.load_state_dict(backbone_state, strict=True)
    del state, backbone_state, head_state
    backbone.config.use_cache = False

    return backbone, head, tokenizer


def evaluate(args):
    import numpy as np
    import pandas as pd
    import torch
    from safetensors.torch import load_file
    from scipy.stats import pearsonr, spearmanr
    from sklearn.metrics import r2_score
    from torch import nn
    from transformers import AutoConfig, AutoModel, AutoTokenizer

    spec = DEFAULTS[args.task]
    checkpoint = local_path(args.checkpoint or spec["checkpoint"])
    data = local_path(args.data or spec["data"])
    output = local_path(args.output)
    if output.exists():
        raise FileExistsError("Validation requires a new output directory.")
    backbone, head, tokenizer = load_regression_checkpoint(checkpoint, args.tokenizer)
    device = torch.device(args.device)
    dtype = torch.bfloat16 if args.precision == "bfloat16" else torch.float32
    backbone.to(device=device, dtype=dtype).eval()
    head.to(device=device, dtype=dtype).eval()
    if args.task == "efficiency":
        examples = json.loads(data.read_text())
        texts, labels = list(examples), list(examples.values())
        add_special_tokens = True
    else:
        frame = pd.read_csv(data, sep="\t")
        texts, labels = frame["text"].tolist(), frame["label"].tolist()
        add_special_tokens = False
    if args.limit:
        texts, labels = texts[:args.limit], labels[:args.limit]
    predictions = []
    with torch.inference_mode(), torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=args.precision == "autocast"):
        for start in range(0, len(texts), args.batch_size):
            batch = tokenizer(texts[start:start + args.batch_size], padding=True, pad_to_multiple_of=8, truncation=True, max_length=512, return_tensors="pt", add_special_tokens=add_special_tokens)
            batch = {key: batch[key].to(device) for key in ("input_ids", "attention_mask")}
            hidden = backbone(**batch).last_hidden_state
            positions = torch.arange(hidden.shape[1], device=device)[None, :].expand(hidden.shape[:2])
            last = positions.masked_fill(~batch["attention_mask"].bool(), -1).max(dim=1).values
            if (last < 0).any():
                raise ValueError("Empty tokenized sequence")
            pooled = hidden[torch.arange(hidden.shape[0], device=device), last]
            predictions.extend(head.net(pooled).flatten().float().cpu().tolist())
            if start % (args.batch_size * 20) == 0:
                print(f"{args.task}: {min(start + args.batch_size, len(texts))}/{len(texts)} evaluated", flush=True)
    labels, predictions = np.asarray(labels), np.asarray(predictions)
    if not np.isfinite(predictions).all():
        raise ValueError("Non-finite model predictions")
    metrics = {
        "task": args.task, "n": len(labels), "precision": args.precision,
        "r2": float(r2_score(labels, predictions)),
        "pearson_r": float(pearsonr(labels, predictions).statistic),
        "spearman_r": float(spearmanr(labels, predictions).statistic),
        "full_validation": not bool(args.limit),
        "target_metric": spec["metric"], "target": spec["target"],
    }
    metrics["allowed_decrease"] = args.allowed_decrease
    metrics["target_passed"] = not args.limit and metrics[spec["metric"]] >= spec["target"] - args.allowed_decrease
    output.mkdir(parents=True, exist_ok=False)
    pd.DataFrame({"label": labels, "prediction": predictions}).to_csv(output / "predictions.tsv", sep="\t", index=False)
    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
    print(json.dumps(metrics, indent=2), flush=True)
    if not args.limit and not metrics["target_passed"]:
        raise SystemExit("Metric fell below the reproduction tolerance; inspect results before publishing.")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("task", choices=DEFAULTS)
    parser.add_argument("--checkpoint")
    parser.add_argument("--tokenizer")
    parser.add_argument("--data")
    parser.add_argument("--output", required=True)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--precision", choices=["float32", "bfloat16", "autocast"], default="float32")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--allowed-decrease", type=float, default=0.01,
                        help="Maximum absolute decrease from the reference metric (default: 0.01).")
    args = parser.parse_args()
    if args.batch_size < 1 or args.limit < 0 or args.allowed_decrease < 0:
        parser.error("batch-size must be positive; limit and allowed-decrease must be nonnegative")
    evaluate(args)


if __name__ == "__main__":
    main()