code/Q-Prefer/validation/run_local_benchmark.py

#!/usr/bin/env python3
"""Run Q-Prefer on a local pairwise benchmark without modifying source data.

The input JSONL rows must contain ``pair_id``, ``prompt``, ``label_quality``,
and ``label_alignment``. Media are resolved as
``MEDIA_ROOT/<task>/<md5(pair_id)[:16]>/{a.mp4,b.mp4,cond.jpg?}``.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import math
import time
from collections import defaultdict
from pathlib import Path

import torch

from qprefer_reward import QPreferConfig, QPreferScorer


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--media-root", type=Path, required=True)
    parser.add_argument("--task", choices=["t2v", "i2v", "near_pair", "mixed"], required=True)
    parser.add_argument("--adapter", required=True)
    parser.add_argument("--adapter-revision")
    parser.add_argument("--base-model", required=True)
    parser.add_argument("--base-revision")
    parser.add_argument(
        "--calc-accuracy",
        type=Path,
        default=Path(__file__).resolve().parent / "metrics.py",
        help="Metric module; defaults to the exact implementation bundled with this release.",
    )
    parser.add_argument("--expected-predictions", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--progress-every", type=int, default=10)
    parser.add_argument("--resume", action="store_true")
    return parser.parse_args()


def load_accuracy_module(path: Path):
    spec = importlib.util.spec_from_file_location("qprefer_calc_accuracy", path)
    if spec is None or spec.loader is None:
        raise ImportError(f"cannot import accuracy functions from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def pair_directory(media_root: Path, task: str, pair_id: str) -> Path:
    digest = hashlib.md5(pair_id.encode()).hexdigest()[:16]
    return media_root / task / digest


def row_task(row: dict, fallback: str) -> str:
    task = row.get("source_task") or row.get("task") or fallback
    if task == "human_valid":
        task = row.get("source_task")
    if task not in {"t2v", "i2v", "near_pair"}:
        raise ValueError(f"cannot resolve source task for {row.get('pair_id')}: {task!r}")
    return task


def load_jsonl(path: Path) -> list[dict]:
    with path.open() as handle:
        return [json.loads(line) for line in handle if line.strip()]


def load_expected(path: Path | None) -> dict[str, dict]:
    if path is None:
        return {}
    data = json.loads(path.read_text())
    if not isinstance(data, list):
        raise TypeError(f"expected a prediction list in {path}")
    return {row["pair_id"]: row for row in data if not row.get("error")}


def metric_cell(rows: list[dict], dimension: str, accuracy) -> dict:
    label_key = f"label_{dimension}"
    margin_key = "m_quality" if dimension == "quality" else "m_alignment"
    valid = [row for row in rows if not row.get("error") and row.get(label_key) in (-1, 0, 1)]
    labels = [row[label_key] for row in valid]
    margins = [row[margin_key] for row in valid]
    if not valid:
        return {"n": 0}
    return {
        "n": len(valid),
        "n_gt_tie": sum(label == 0 for label in labels),
        "acc_with_ties": round(accuracy.calc_accuracy_with_ties(labels, margins), 4),
        "acc_without_ties": round(accuracy.calc_accuracy_without_ties(labels, margins), 4),
        "acc_with_ties_fixed_epsilon0": round(
            accuracy.calc_accuracy_with_ties_fixed(labels, margins, 0), 4
        ),
        "stats_epsilon0": accuracy.calc_stats_fixed(labels, margins, 0),
    }


def compare_expected(rows: list[dict], expected: dict[str, dict]) -> dict:
    if not expected:
        return {"enabled": False}
    differences = []
    missing = []
    for row in rows:
        if row.get("error"):
            continue
        reference = expected.get(row["pair_id"])
        if reference is None:
            missing.append(row["pair_id"])
            continue
        differences.append(
            {
                "pair_id": row["pair_id"],
                "m_quality": abs(row["m_quality"] - reference["m_quality"]),
                "m_alignment": abs(row["m_alignment"] - reference["m_alignment"]),
            }
        )
    exact = [
        item for item in differences if item["m_quality"] == 0.0 and item["m_alignment"] == 0.0
    ]
    max_quality = max((item["m_quality"] for item in differences), default=math.nan)
    max_alignment = max((item["m_alignment"] for item in differences), default=math.nan)
    return {
        "enabled": True,
        "expected_rows": len(expected),
        "compared_rows": len(differences),
        "exact_rows": len(exact),
        "missing_expected_rows": missing,
        "max_abs_difference": {
            "m_quality": max_quality,
            "m_alignment": max_alignment,
        },
        "all_compared_rows_exact": len(exact) == len(differences),
    }


def atomic_write_json(path: Path, data: object) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n")
    temporary.replace(path)


def main() -> None:
    args = parse_args()
    args.output.parent.mkdir(parents=True, exist_ok=True)
    rows = load_jsonl(args.manifest)
    expected = load_expected(args.expected_predictions)
    accuracy = load_accuracy_module(args.calc_accuracy)

    prediction_path = args.output.with_name(args.output.stem + "_predictions.json")
    predictions: list[dict] = []
    if args.resume and prediction_path.is_file():
        loaded = json.loads(prediction_path.read_text())
        predictions = [row for row in loaded if not row.get("error")]
    completed = {row["pair_id"] for row in predictions}

    started = time.perf_counter()
    scorer = QPreferScorer(
        QPreferConfig(
            adapter=args.adapter,
            adapter_revision=args.adapter_revision,
            base_model=args.base_model,
            base_revision=args.base_revision,
            device=args.device,
            dtype=torch.bfloat16,
            merge_lora=True,
        )
    )
    load_seconds = time.perf_counter() - started

    inference_started = time.perf_counter()
    for index, row in enumerate(rows, start=1):
        if row["pair_id"] in completed:
            continue
        task = row_task(row, args.task)
        directory = pair_directory(args.media_root, task, row["pair_id"])
        video_a = directory / "a.mp4"
        video_b = directory / "b.mp4"
        reference = directory / "cond.jpg"
        references = [reference, reference] if reference.is_file() else None
        record = {
            "pair_id": row["pair_id"],
            "source_task": task,
            "label_quality": row.get("label_quality"),
            "label_alignment": row.get("label_alignment"),
        }
        try:
            scores = scorer.score_batch(
                [video_a, video_b],
                [row["prompt"], row["prompt"]],
                reference_images=references,
                batch_size=1,
            )
            visual = [float(value) for value in scores.visual_quality]
            alignment = [float(value) for value in scores.text_alignment]
            record.update(
                {
                    "video_a": {
                        "visual_quality": visual[0],
                        "text_alignment": alignment[0],
                    },
                    "video_b": {
                        "visual_quality": visual[1],
                        "text_alignment": alignment[1],
                    },
                    "m_quality": visual[0] - visual[1],
                    "m_alignment": alignment[0] - alignment[1],
                }
            )
        except Exception as error:
            record["error"] = repr(error)
        predictions.append(record)
        if index % args.progress_every == 0 or index == len(rows):
            atomic_write_json(prediction_path, predictions)
            errors = sum("error" in item for item in predictions)
            elapsed = time.perf_counter() - inference_started
            print(
                f"[{args.task}] {index}/{len(rows)} errors={errors} elapsed={elapsed:.1f}s",
                flush=True,
            )

    by_source: dict[str, list[dict]] = defaultdict(list)
    for prediction in predictions:
        by_source[prediction["source_task"]].append(prediction)
    cells = {}
    for source, source_rows in sorted(by_source.items()):
        for dimension in ("quality", "alignment"):
            cells[f"{source}/{dimension}"] = metric_cell(source_rows, dimension, accuracy)
    if len(by_source) > 1:
        for dimension in ("quality", "alignment"):
            cells[f"all/{dimension}"] = metric_cell(predictions, dimension, accuracy)

    primary_keys = [
        key for key in cells if not key.startswith("all/") and cells[key].get("n", 0) > 0
    ]
    summary = {
        "manifest": str(args.manifest.resolve()),
        "media_root": str(args.media_root.resolve()),
        "adapter": args.adapter,
        "base_model": args.base_model,
        "device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else args.device,
        "torch": torch.__version__,
        "n_manifest_rows": len(rows),
        "n_predictions": len(predictions),
        "n_errors": sum("error" in item for item in predictions),
        "load_seconds": load_seconds,
        "inference_seconds": time.perf_counter() - inference_started,
        "cells": cells,
        "macro_with_ties": round(
            sum(cells[key]["acc_with_ties"] for key in primary_keys) / len(primary_keys), 4
        ),
        "macro_without_ties": round(
            sum(cells[key]["acc_without_ties"] for key in primary_keys) / len(primary_keys), 4
        ),
        "legacy_parity": compare_expected(predictions, expected),
        "predictions": str(prediction_path.resolve()),
    }
    atomic_write_json(args.output, summary)
    print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()