#!/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///{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()