script.py

#!/usr/bin/env python3
"""
BenchLabs universal evaluation script.

One script, every BenchLabs benchmark. Downloads the datasets straight from the
Hugging Face Hub, runs your model, and prints an in-depth report with
category / subcategory breakdowns -- in the exact shape the
BenchLabs-Leaderboard `models.json` expects.

Benchmarks covered
  bench-effortless-7-2026   dual-mode: generative + log-likelihood (tier 1, latest)
  bench-easy-7-2026         dual-mode: generative + log-likelihood (tier 2, latest)
  bench-mid-7-2026          dual-mode: generative + log-likelihood (tier 3, latest)
  bench-effortless-6-2026   generative, exact-match (tier 1, legacy)
  bench-easy-6-2026         generative, hybrid category-aware (tier 2, legacy)
  bench-mid-6-2026          multiple-choice, log-likelihood (tier 3, legacy)
  bench-AGI                 skipped -- scoring pipeline under maintenance

Dual-mode (7-2026, schema v2)
  Every item carries a gold answer + aliases AND target_scores choices, so each
  benchmark is scored BOTH ways in one run:
    generative      exact_match (alias-aware) + hybrid_score routed per item by
                    its `gen_scoring` field (strict / semantic / fuzzy)
    loglikelihood   lm-eval style over choices: acc, acc_norm, soft_score,
                    soft_score_norm, with per-choice log-probs recorded
  Headline metric stays tier-conventional: exact_match (effortless),
  hybrid_score (easy), soft_score_norm (mid). Per-item detail -- raw
  generation, extracted answer, per-choice log-probs raw and per-byte -- is
  written to samples_<id>.jsonl next to the usual CSV.

Install
  pip install torch transformers
  pip install sentence-transformers   # optional: better semantic scoring on Easy
  pip install accelerate               # optional: faster / multi-GPU loading

Run
  python script.py --model Qwen/Qwen2.5-0.5B
  python script.py --model Qwen/Qwen2.5-1.5B-Instruct --benchmarks easy,mid
  python script.py --model ./my-local-checkpoint --device cuda --batch-size 16
  python script.py --model Qwen/Qwen2.5-0.5B --limit 10          # quick smoke test
  python script.py --model Qwen/Qwen2.5-0.5B --leaderboard       # print models.json entry

Outputs (under --output-dir, default benchlabs_results/<model>/)
  results.json        full report: every benchmark, category, subcategory, sample counts
  samples_<id>.csv     per-sample predictions and scores for each benchmark
  leaderboard.json    ready-to-paste `models.json` entry for the leaderboard PR

Scoring conventions
  Effortless    exact match after normalization (strip, lowercase, drop punctuation).
  Easy          hybrid category-aware scoring, identical to the official
                benchmark.ipynb: strict categories are binary exact-match, soft
                categories get semantic similarity, hybrid categories get fuzzy
                string similarity. Plain exact-match is also reported.
  Mid           lm-eval style log-likelihood over the `target_scores` candidates:
                  acc = argmax raw log-likelihood is the 1.0 answer
                  acc_norm = argmax log-likelihood / byte-length of the answer
                  soft_score / soft_score_norm = target_scores value of the picked
                  answer (partial credit on distractors with non-zero scores)
                Headline score = soft_score_norm, matching the leaderboard.
  Multiple-choice prompt format: "Q: {input}\nA:" with candidates " {choice}".

Reasoning / CoT models
  <think>...</think> blocks are stripped before answer extraction: only the
  text after the final </think> is scored. Raise --max-new-tokens (2048+) so
  the model can finish thinking -- the default 32 is sized for direct-answer
  models. A generation cut off mid-think (unclosed <think>) scores as an
  empty answer. Mid is scored by log-likelihood over the answer choices with
  no generation at all, so thinking never happens there.
  Known limit: a model that reasons in plain prose with NO tags slips the
  strip -- its first prose line is what gets scored. The scorer trusts the
  tag convention; script_sha256 pins which scorer said so.

Reproducibility
  --revision pins the exact model commit to evaluate. Every run records
  model_revision (the snapshot actually loaded) + script_sha256 (the exact
  scorer bytes); re-running the pinned script with --revision <recorded sha>
  reproduces the run even if the model's main branch moved since.

Bench Labs - Simple, Reliable, Open sourced
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import os
import re
import sys
import urllib.request
from collections import defaultdict
from dataclasses import dataclass, field
from difflib import SequenceMatcher
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

# --------------------------------------------------------------------------- #
# Benchmark registry
# --------------------------------------------------------------------------- #

HUB_BASE = "https://huggingface.co/datasets/bench-labs/{id}/resolve/main/eval.jsonl"

BENCHMARKS: Dict[str, Dict[str, Any]] = {
    "effortless7": {
        "id": "bench-effortless-7-2026",
        "tier": 1,
        "kind": "dual",
        "metric": "exact_match",
        "generation": "7-2026",
        "description": "Sanity-layer QA, dual-mode (generative + log-likelihood).",
    },
    "easy7": {
        "id": "bench-easy-7-2026",
        "tier": 2,
        "kind": "dual",
        "metric": "hybrid_score",
        "generation": "7-2026",
        "description": "Easy-tier QA, dual-mode with per-item scorer routing.",
    },
    "mid7": {
        "id": "bench-mid-7-2026",
        "tier": 3,
        "kind": "dual",
        "metric": "soft_score_norm",
        "generation": "7-2026",
        "description": "Mid-tier QA, dual-mode (headline: log-likelihood soft_score_norm).",
    },
    "effortless": {
        "id": "bench-effortless-6-2026",
        "tier": 1,
        "kind": "generative",
        "metric": "exact_match",
        "generation": "6-2026",
        "description": "Sanity-layer QA: unambiguous single-answer questions.",
    },
    "easy": {
        "id": "bench-easy-6-2026",
        "tier": 2,
        "kind": "generative",
        "metric": "hybrid_score",
        "generation": "6-2026",
        "description": "Easy-tier QA with hybrid category-aware scoring.",
    },
    "mid": {
        "id": "bench-mid-6-2026",
        "tier": 3,
        "kind": "multiple_choice",
        "metric": "soft_score_norm",
        "generation": "6-2026",
        "description": "Mid-tier multiple-choice QA via log-likelihood.",
    },
    "agi": {
        "id": "bench-AGI",
        "tier": 4,
        "kind": "rank_order",
        "metric": "rank_order",
        "generation": "6-2026",
        "description": "Hard open-ended questions, panel-graded rank order.",
        "unavailable": "Scoring pipeline under maintenance -- see the dataset README.",
    },
}

DUAL_METRICS_GEN = ("exact_match", "hybrid_score")
DUAL_METRICS_LL = ("acc", "acc_norm", "soft_score", "soft_score_norm")

# Easy-tier category routing, identical to the official benchmark.ipynb.
STRICT_CATEGORIES = {
    "Math-arithmetic", "Math-pattern",
    "Logic-deduction", "Logic-pattern", "Logic-consistency",
    "Knowledge-basic", "Pattern-matching",
}
SOFT_CATEGORIES = {
    "Commonsense-simulation", "Commonsense-causality", "Commonsense-reasoning",
    "Language-comprehension", "Knowledge-definitions",
}
HYBRID_CATEGORIES = {
    "Language-structure", "Language-transformation",
}

SYSTEM_PROMPT = "You are a precise assistant. Give only the final answer, without explanation."
MC_PROMPT = "Q: {input}\nA:"

ANSWER_PREFIXES = re.compile(
    r"^(the answer is|answer\s*[:=]|final answer\s*[:=]?|it is|it's)\s*", re.IGNORECASE
)

# Reasoning-model tags. THINK_CLOSE also matches a bare closing tag: some chat
# templates open <think> inside the prompt, so the generation contains only
# the reasoning and a </think>.
THINK_CLOSE = re.compile(r"</think(?:ing)?>\s*", re.IGNORECASE)
THINK_OPEN = re.compile(r"<think(?:ing)?>.*", re.IGNORECASE | re.DOTALL)

# --------------------------------------------------------------------------- #
# Text normalization and scoring
# --------------------------------------------------------------------------- #

def normalize(text: str) -> str:
    text = str(text).strip().lower()
    text = re.sub(r"[\u201c\u201d\"'`]", "", text)
    text = text.replace("\u2019", "'")
    text = re.sub(r"[\.\,\!\?\:\;\(\)\[\]\{\}]", "", text)
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def extract_answer(text: str) -> str:
    """First line of the generation after any <think> block, minus boilerplate prefixes.

    Only text after the final </think> is scored. An unclosed <think> means the
    generation ran out of budget mid-reasoning, so there is no answer to extract.
    """
    text = str(text)
    parts = THINK_CLOSE.split(text)
    if len(parts) > 1:
        text = parts[-1]
    else:
        text = THINK_OPEN.sub("", text)
    text = text.strip()
    if "\n" in text:
        text = text.split("\n", 1)[0]
    text = ANSWER_PREFIXES.sub("", text.strip())
    return text.strip()


def strict_score(pred: str, gold: str) -> float:
    return 1.0 if normalize(pred) == normalize(gold) else 0.0


def fuzzy_score(pred: str, gold: str) -> float:
    p, g = normalize(pred), normalize(gold)
    if p == g:
        return 1.0
    return max(0.0, min(1.0, SequenceMatcher(None, p, g).ratio()))


class SemanticScorer:
    """Sentence-embedding similarity with a fuzzy-string fallback."""

    def __init__(self) -> None:
        self._embedder = None
        try:
            from sentence_transformers import SentenceTransformer  # type: ignore
            self._embedder = SentenceTransformer("all-MiniLM-L6-v2")
        except Exception:
            self._embedder = None

    @property
    def backend(self) -> str:
        return "sentence-transformers/all-MiniLM-L6-v2" if self._embedder else "difflib-fallback"

    def score(self, pred: str, gold: str) -> float:
        p, g = normalize(pred), normalize(gold)
        if p == g:
            return 1.0
        if self._embedder is not None:
            try:
                import numpy as np
                pv, gv = self._embedder.encode([p, g], normalize_embeddings=True)
                cos = float(np.dot(pv, gv))
                return max(0.0, min(1.0, (cos + 1.0) / 2.0))
            except Exception:
                pass
        return fuzzy_score(pred, gold)


def easy_hybrid_score(category: str, pred: str, gold: str, semantic: SemanticScorer) -> float:
    if category in STRICT_CATEGORIES:
        return strict_score(pred, gold)
    if category in SOFT_CATEGORIES:
        return semantic.score(pred, gold)
    if category in HYBRID_CATEGORIES:
        return fuzzy_score(pred, gold)
    return fuzzy_score(pred, gold)  # unknown categories: fuzzy, never hard-fail


# -- v2 (7-2026) scoring: alias-aware, routed per item by `gen_scoring` ------ #

def strict_score_multi(pred: str, golds: Sequence[str]) -> float:
    return 1.0 if any(normalize(pred) == normalize(g) for g in golds) else 0.0


def routed_score(gen_scoring: str, pred: str, golds: Sequence[str],
                 semantic: SemanticScorer) -> float:
    """v2 generation-mode score: routing comes from the item, not category tables."""
    if gen_scoring == "strict":
        return strict_score_multi(pred, golds)
    if gen_scoring == "semantic":
        return max(semantic.score(pred, g) for g in golds)
    return max(fuzzy_score(pred, g) for g in golds)  # "fuzzy"


# --------------------------------------------------------------------------- #
# Dataset loading (no `datasets` dependency -- each benchmark is one eval.jsonl)
# --------------------------------------------------------------------------- #

def cache_dir() -> Path:
    return Path(os.environ.get("BENCHLABS_CACHE", Path.home() / ".cache" / "benchlabs"))


def load_benchmark_rows(bench_id: str, refresh: bool = False) -> List[dict]:
    path = cache_dir() / f"{bench_id}.jsonl"
    if refresh or not path.exists():
        url = HUB_BASE.format(id=bench_id)
        print(f"  downloading {url}")
        path.parent.mkdir(parents=True, exist_ok=True)
        req = urllib.request.Request(url)
        token = os.environ.get("HF_TOKEN")
        if token:
            req.add_header("Authorization", f"Bearer {token}")
        with urllib.request.urlopen(req) as resp:
            path.write_bytes(resp.read())
    rows = []
    with path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def split_category(cat: str) -> Tuple[str, Optional[str]]:
    """'Commonsense-causality' -> ('Commonsense', 'causality'); 'Math' -> ('Math', None)."""
    if "-" in cat:
        top, sub = cat.split("-", 1)
        return top, sub
    return cat, None


# --------------------------------------------------------------------------- #
# Model backend (lazy torch/transformers import)
# --------------------------------------------------------------------------- #

class HFModel:
    """Thin wrapper: batched greedy generation + batched log-likelihood scoring."""

    def __init__(self, name: str, device: str, dtype: str, trust_remote_code: bool,
                 use_chat_template: bool, revision: Optional[str] = None) -> None:
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except ImportError as e:
            sys.exit(f"Missing dependency ({e.name}). Install with: pip install torch transformers")

        self.torch = torch
        if device == "auto":
            if torch.cuda.is_available():
                device = "cuda"
            elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"
        self.device = device

        if dtype == "auto":
            torch_dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() \
                else (torch.float16 if device in ("cuda", "mps") else torch.float32)
        else:
            torch_dtype = {"float16": torch.float16, "fp16": torch.float16,
                           "bfloat16": torch.bfloat16, "bf16": torch.bfloat16,
                           "float32": torch.float32, "fp32": torch.float32}[dtype.lower()]

        pin = f", revision={revision}" if revision else ""
        print(f"Loading model: {name} (device={device}, dtype={torch_dtype}{pin})")
        self.tokenizer = AutoTokenizer.from_pretrained(name, trust_remote_code=trust_remote_code,
                                                        revision=revision)
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.model = AutoModelForCausalLM.from_pretrained(
            name, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code,
            revision=revision,
        ).to(device)
        self.model.eval()
        self.use_chat = use_chat_template and self.tokenizer.chat_template is not None
        print(f"  chat template: {'yes' if self.use_chat else 'no (plain QA prompt)'}")

        # -- resolved model provenance ------------------------------------ #
        # The weights already carry their commit: from_pretrained records the
        # snapshot it actually loaded in config._commit_hash, no second Hub
        # lookup. Asking the Hub afterwards can pin a different commit if the
        # branch moved between load and lookup, so _commit_hash is primary.
        self.resolved_revision: Optional[str] = getattr(self.model.config, "_commit_hash", None)
        if self.resolved_revision is None and "/" in name and not Path(name).exists():
            # Fallback for transformers versions that don't record it. Local
            # checkpoints stay None, which is honest: they have no hub revision.
            try:
                from huggingface_hub import HfApi
                self.resolved_revision = HfApi().model_info(name).sha
            except Exception:
                pass

    # -- generation -------------------------------------------------------- #

    def _format_prompt(self, question: str) -> str:
        if self.use_chat:
            return self.tokenizer.apply_chat_template(
                [{"role": "system", "content": SYSTEM_PROMPT},
                 {"role": "user", "content": question}],
                tokenize=False, add_generation_prompt=True,
            )
        return f"Question: {question}\nAnswer:"

    def generate(self, questions: Sequence[str], batch_size: int, max_new_tokens: int,
                 progress: str = "") -> List[str]:
        torch = self.torch
        tok = self.tokenizer
        preds: List[str] = []
        old_side = tok.padding_side
        tok.padding_side = "left"
        try:
            with torch.no_grad():
                for start in range(0, len(questions), batch_size):
                    chunk = questions[start:start + batch_size]
                    prompts = [self._format_prompt(q) for q in chunk]
                    inputs = tok(prompts, return_tensors="pt", padding=True,
                                 truncation=True).to(self.device)
                    out = self.model.generate(
                        **inputs, max_new_tokens=max_new_tokens, do_sample=False,
                        pad_token_id=tok.pad_token_id,
                    )
                    gen = out[:, inputs["input_ids"].shape[1]:]
                    preds.extend(tok.decode(g, skip_special_tokens=True).strip() for g in gen)
                    _progress(progress, len(preds), len(questions))
        finally:
            tok.padding_side = old_side
        return preds

    # -- log-likelihood ---------------------------------------------------- #

    def loglikelihoods(self, pairs: Sequence[Tuple[str, str]], batch_size: int,
                        progress: str = "") -> List[float]:
        """Sum of log-probs of `continuation` given `context` for each pair."""
        torch = self.torch
        tok = self.tokenizer
        encoded = []
        for ctx, cont in pairs:
            ctx_ids = tok.encode(ctx)
            full_ids = tok.encode(ctx + cont)
            n_cont = len(full_ids) - len(ctx_ids)
            if n_cont <= 0:  # tokenizer merged across the boundary; re-split manually
                cont_ids = tok.encode(cont, add_special_tokens=False)
                full_ids = ctx_ids + cont_ids
                n_cont = len(cont_ids)
            encoded.append((full_ids, n_cont))

        results: List[float] = []
        with torch.no_grad():
            for start in range(0, len(encoded), batch_size):
                chunk = encoded[start:start + batch_size]
                maxlen = max(len(ids) for ids, _ in chunk)
                pad_id = tok.pad_token_id
                input_ids = torch.full((len(chunk), maxlen), pad_id, dtype=torch.long)
                attn = torch.zeros((len(chunk), maxlen), dtype=torch.long)
                for i, (ids, _) in enumerate(chunk):
                    input_ids[i, :len(ids)] = torch.tensor(ids)
                    attn[i, :len(ids)] = 1
                input_ids, attn = input_ids.to(self.device), attn.to(self.device)
                logits = self.model(input_ids=input_ids, attention_mask=attn).logits
                logprobs = torch.log_softmax(logits.float(), dim=-1)
                for i, (ids, n_cont) in enumerate(chunk):
                    total = 0.0
                    for pos in range(len(ids) - n_cont, len(ids)):
                        total += logprobs[i, pos - 1, ids[pos]].item()
                    results.append(total)
                _progress(progress, len(results), len(pairs))
        return results


_progress_t0: Dict[str, float] = {}


def _progress(label: str, done: int, total: int) -> None:
    if not label:
        return
    import time
    t0 = _progress_t0.setdefault(label, time.monotonic())
    elapsed = time.monotonic() - t0
    eta = ""
    if 0 < done < total and elapsed > 2:
        remain = elapsed / done * (total - done)
        eta = f" · {int(remain // 60)}m{int(remain % 60):02d}s left"
    print(f"\r  {label}: {done}/{total} ({100 * done // max(1, total)}%){eta}   ",
          end="", flush=True)
    if done >= total:
        _progress_t0.pop(label, None)
        print(f"\r  {label}: {total}/{total} done in {int(elapsed // 60)}m{int(elapsed % 60):02d}s")


# --------------------------------------------------------------------------- #
# Aggregation
# --------------------------------------------------------------------------- #

@dataclass
class Sample:
    idx: int
    category: str
    question: str
    gold: str
    pred: str
    scores: Dict[str, float] = field(default_factory=dict)
    detail: Optional[Dict[str, Any]] = None  # dual-mode per-item record (samples_<id>.jsonl)


def mean(xs: Sequence[float]) -> float:
    return sum(xs) / len(xs) if xs else 0.0


def stderr_of(xs: Sequence[float]) -> float:
    if len(xs) < 2:
        return 0.0
    m = mean(xs)
    var = sum((x - m) ** 2 for x in xs) / (len(xs) - 1)
    return math.sqrt(var / len(xs))


def aggregate(samples: List[Sample], metrics: Sequence[str]) -> Dict[str, Any]:
    """Overall + per-category + per-subcategory rollups for each metric."""
    by_cat: Dict[str, List[Sample]] = defaultdict(list)
    by_top: Dict[str, List[Sample]] = defaultdict(list)
    for s in samples:
        by_cat[s.category].append(s)
        by_top[split_category(s.category)[0]].append(s)

    def block(rows: List[Sample]) -> Dict[str, Any]:
        out: Dict[str, Any] = {"n": len(rows)}
        for m in metrics:
            vals = [s.scores[m] for s in rows]
            out[m] = round(mean(vals), 4)
        return out

    return {
        "overall": {**block(samples),
                    "stderr": round(stderr_of([s.scores[metrics[-1]] for s in samples]), 4)},
        "categories": {cat: block(rows) for cat, rows in sorted(by_cat.items())},
        "category_groups": {top: block(rows) for top, rows in sorted(by_top.items())},
        "macro_avg": {m: round(mean([mean([s.scores[m] for s in rows])
                                      for rows in by_cat.values()]), 4) for m in metrics},
    }


# --------------------------------------------------------------------------- #
# Benchmark runners
# --------------------------------------------------------------------------- #

def run_generative(key: str, rows: List[dict], model: HFModel, args) -> Tuple[List[Sample], Dict]:
    bench = BENCHMARKS[key]
    questions = [str(r["question"]) for r in rows]
    raw_preds = model.generate(questions, args.batch_size, args.max_new_tokens,
                                progress=f"{bench['id']} generate")

    semantic = SemanticScorer() if key == "easy" else None
    if semantic:
        print(f"  semantic scorer: {semantic.backend}")

    samples: List[Sample] = []
    for i, (row, raw) in enumerate(zip(rows, raw_preds)):
        pred = extract_answer(raw)
        gold = str(row["answer"])
        cat = str(row["category"])
        scores = {"exact_match": strict_score(pred, gold)}
        if key == "easy":
            scores["hybrid_score"] = easy_hybrid_score(cat, pred, gold, semantic)
        samples.append(Sample(i, cat, str(row["question"]), gold, pred, scores))

    metrics = ["exact_match"] + (["hybrid_score"] if key == "easy" else [])
    return samples, aggregate(samples, metrics)


def run_multiple_choice(key: str, rows: List[dict], model: HFModel, args) -> Tuple[List[Sample], Dict]:
    bench = BENCHMARKS[key]
    pairs: List[Tuple[str, str]] = []
    index: List[Tuple[int, List[str]]] = []
    for i, row in enumerate(rows):
        choices = list(row["target_scores"].keys())
        ctx = MC_PROMPT.format(input=row["input"])
        for c in choices:
            pairs.append((ctx, f" {c}"))
        index.append((i, choices))

    lls = model.loglikelihoods(pairs, args.batch_size, progress=f"{bench['id']} loglikelihood")

    samples: List[Sample] = []
    pos = 0
    for i, choices in index:
        row = rows[i]
        tgt = row["target_scores"]
        chunk = lls[pos:pos + len(choices)]
        pos += len(choices)
        norm = [ll / max(1, len(c.encode("utf-8"))) for ll, c in zip(chunk, choices)]
        pick_raw = choices[max(range(len(choices)), key=lambda j: chunk[j])]
        pick_norm = choices[max(range(len(choices)), key=lambda j: norm[j])]
        gold = max(tgt, key=tgt.get)
        samples.append(Sample(
            i, str(row["category"]), str(row["input"]), gold, pick_norm,
            scores={
                "acc": 1.0 if tgt.get(pick_raw) == 1 else 0.0,
                "acc_norm": 1.0 if tgt.get(pick_norm) == 1 else 0.0,
                "soft_score": float(tgt.get(pick_raw, 0.0)),
                "soft_score_norm": float(tgt.get(pick_norm, 0.0)),
            },
        ))
    return samples, aggregate(samples, ["acc", "acc_norm", "soft_score", "soft_score_norm"])


def run_dual(key: str, rows: List[dict], model: HFModel, args) -> Tuple[List[Sample], Dict]:
    """v2 (7-2026) benchmarks: run BOTH modes over every item.

    Generative: greedy generation, alias-aware exact match + hybrid_score
    routed per item by its `gen_scoring` field.
    Log-likelihood: lm-eval style over `target_scores` choices, per-choice
    log-probs (raw and per-byte) recorded in the sample detail.
    """
    bench = BENCHMARKS[key]

    # -- generative pass ---------------------------------------------------- #
    questions = [str(r["question"]) for r in rows]
    raw_preds = model.generate(questions, args.batch_size, args.max_new_tokens,
                                progress=f"{bench['id']} generate")
    semantic = SemanticScorer()
    print(f"  semantic scorer: {semantic.backend}")

    # -- log-likelihood pass ------------------------------------------------ #
    pairs: List[Tuple[str, str]] = []
    index: List[List[str]] = []
    for row in rows:
        choices = list(row["target_scores"].keys())
        ctx = MC_PROMPT.format(input=row["question"])
        for c in choices:
            pairs.append((ctx, f" {c}"))
        index.append(choices)
    lls = model.loglikelihoods(pairs, args.batch_size,
                                progress=f"{bench['id']} loglikelihood")

    samples: List[Sample] = []
    pos = 0
    for i, (row, raw) in enumerate(zip(rows, raw_preds)):
        gold = str(row["answer"])
        aliases = [str(a) for a in row.get("answer_aliases", [])]
        golds = [gold] + aliases
        gen_scoring = str(row.get("gen_scoring", "strict"))
        cat = str(row["category"])

        pred = extract_answer(raw)
        exact = strict_score_multi(pred, golds)
        hybrid = routed_score(gen_scoring, pred, golds, semantic)

        tgt = row["target_scores"]
        choices = index[i]
        chunk = lls[pos:pos + len(choices)]
        pos += len(choices)
        norm = [ll / max(1, len(c.encode("utf-8"))) for ll, c in zip(chunk, choices)]
        pick_raw = choices[max(range(len(choices)), key=lambda j: chunk[j])]
        pick_norm = choices[max(range(len(choices)), key=lambda j: norm[j])]

        scores = {
            "exact_match": exact,
            "hybrid_score": hybrid,
            "acc": 1.0 if tgt.get(pick_raw) == 1 else 0.0,
            "acc_norm": 1.0 if tgt.get(pick_norm) == 1 else 0.0,
            "soft_score": float(tgt.get(pick_raw, 0.0)),
            "soft_score_norm": float(tgt.get(pick_norm, 0.0)),
        }
        detail = {
            "id": row.get("id", i),
            "category": cat,
            "gold": gold,
            "answer_aliases": aliases,
            "gen_scoring": gen_scoring,
            "preferred_mode": row.get("preferred_mode"),
            "generative": {
                "raw": raw,
                "extracted": pred,
                "exact_match": exact,
                "hybrid_score": round(hybrid, 4),
            },
            "loglikelihood": {
                "choices": {
                    c: {"logprob": round(ll, 4), "logprob_per_byte": round(nb, 5)}
                    for c, ll, nb in zip(choices, chunk, norm)
                },
                "pick_raw": pick_raw,
                "pick_norm": pick_norm,
                **{m: scores[m] for m in DUAL_METRICS_LL},
            },
        }
        samples.append(Sample(i, cat, str(row["question"]), gold, pred, scores, detail))

    # stderr is computed on metrics[-1]; keep the headline metric last.
    metrics = [m for m in (*DUAL_METRICS_GEN, *DUAL_METRICS_LL) if m != bench["metric"]]
    metrics.append(bench["metric"])
    return samples, aggregate(samples, metrics)


# --------------------------------------------------------------------------- #
# Reporting
# --------------------------------------------------------------------------- #

def print_report(bench_key: str, agg: Dict[str, Any]) -> None:
    bench = BENCHMARKS[bench_key]
    headline = bench["metric"]
    overall = agg["overall"]
    print(f"\n=== {bench['id']} (tier {bench['tier']}) ===")
    print(f"  headline [{headline}]: {overall[headline]:.4f} "
          f"(n={overall['n']}, stderr={overall['stderr']:.4f})")
    others = [m for m in overall if m not in ("n", "stderr", headline)]
    if others:
        print("  also: " + " ".join(f"{m}={overall[m]:.4f}" for m in others))
    print(f"  macro avg [{headline}]: {agg['macro_avg'][headline]:.4f}")

    print(f"  {'category':<28}{'n':>4} {headline}")
    current_top = None
    for cat, stats in agg["categories"].items():
        top, sub = split_category(cat)
        if top != current_top:
            group = agg["category_groups"][top]
            print(f"  {top:<28}{group['n']:>4} {group[headline]:.3f}")
            current_top = top
        if sub is not None:
            print(f"    - {sub:<24}{stats['n']:>4} {stats[headline]:.3f}")


def leaderboard_entry(model_name: str, results: Dict[str, Any],
                       model_revision: Optional[str]) -> Dict[str, Any]:
    """A ready-to-paste entry for the leaderboard's models.json `models` array."""
    runs: Dict[str, Any] = {}
    for key, bench in BENCHMARKS.items():
        bid = bench["id"]
        if key not in results:
            runs[bid] = {"score": None, "n": None, "notes": "Not yet evaluated on this tier."}
            continue
        agg = results[key]["aggregate"]
        overall = agg["overall"]
        entry: Dict[str, Any] = {"score": overall[bench["metric"]], "n": overall["n"]}
        if bench["kind"] == "dual":
            # v2: uniform shape -- run-level metrics{} split by mode, and every
            # category block carries a generic "score" (the headline metric).
            entry["stderr"] = overall["stderr"]
            entry["metrics"] = {
                "generative": {m: overall[m] for m in DUAL_METRICS_GEN},
                "loglikelihood": {m: overall[m] for m in DUAL_METRICS_LL},
            }
            entry["categories"] = {
                cat: {"n": s["n"], "score": s[bench["metric"]],
                      "exact_match": s["exact_match"], "acc_norm": s["acc_norm"]}
                for cat, s in agg["categories"].items()
            }
        elif bench["kind"] == "multiple_choice":
            entry.update({m: overall[m] for m in ("acc", "acc_norm", "soft_score", "soft_score_norm")})
            entry["stderr"] = overall["stderr"]
            entry["categories"] = {
                cat: {"n": s["n"], "acc": s["acc"], "acc_norm": s["acc_norm"]}
                for cat, s in agg["categories"].items()
            }
        else:
            entry["notes"] = ("Exact-match, normalized." if bench["metric"] == "exact_match"
                              else "Hybrid category-aware scoring (strict / flexible / semantic).")
            entry["categories"] = {
                cat: {"n": s["n"], bench["metric"]: s[bench["metric"]]}
                for cat, s in agg["categories"].items()
            }
        runs[bid] = entry

    slug = re.sub(r"[^a-z0-9.]+", "-", model_name.lower()).strip("-")
    return {
        "id": slug.split("/")[-1] if "/" in slug else slug,
        "name": model_name,
        "org": model_name.split("/")[0] if "/" in model_name else "",
        "params_b": None,
        "license": None,
        "architecture": None,
        "url": f"https://huggingface.co/{model_name}" if "/" in model_name else None,
        "model_revision": model_revision,
        "script_sha256": script_sha256(),
        "runs": runs,
    }


MODELS_JSON_URL = ("https://huggingface.co/spaces/bench-labs/BenchLabs-Leaderboard/"
                   "resolve/main/models.json")


def merge_into_models_json(entry: Dict[str, Any], evaluated_bench_ids: List[str],
                            out_dir: Path) -> Optional[Path]:
    """Fetch the live models.json and merge this run's entry into it.

    The result is written to <out_dir>/models.json, ready to upload as-is --
    no hand-pasting. Merge rules:
      * matched by `id`: only the benchmarks evaluated THIS run are replaced;
        scores from other tiers and hand-curated metadata (params_b, license,
        architecture) are kept.
      * unmatched: the entry is appended.
    Returns the written path, or None if the live file could not be fetched.
    """
    try:
        req = urllib.request.Request(MODELS_JSON_URL)
        token = os.environ.get("HF_TOKEN")
        if token:
            req.add_header("Authorization", f"Bearer {token}")
        with urllib.request.urlopen(req, timeout=30) as resp:
            board = json.loads(resp.read().decode("utf-8"))
    except Exception as e:
        print(f"  could not fetch live models.json ({e}); skipping auto-merge")
        return None

    existing = next((m for m in board.get("models", []) if m.get("id") == entry["id"]), None)
    if existing is None:
        board.setdefault("models", []).append(entry)
    else:
        for bid in evaluated_bench_ids:
            existing.setdefault("runs", {})[bid] = entry["runs"][bid]
        existing["model_revision"] = entry["model_revision"]
        existing["script_sha256"] = entry["script_sha256"]
        for meta in ("name", "org", "url"):
            existing.setdefault(meta, entry[meta])

    import datetime as _dt
    board["updated"] = _dt.date.today().isoformat()

    path = out_dir / "models.json"
    path.write_text(json.dumps(board, indent=2, ensure_ascii=False), encoding="utf-8")
    return path


def script_sha256() -> str:
    """SHA-256 of this file's own bytes.

    Written for content, not label: it lets a maintainer re-run the pinned
    copy of this script and compare hashes, rather than trusting a static
    version string that an edited copy would still print unchanged.
    """
    return hashlib.sha256(Path(__file__).read_bytes()).hexdigest()


def save_outputs(out_dir: Path, model_name: str, results: Dict[str, Any], args,
                  model_revision: Optional[str]) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)

    report = {
        "model": model_name,
        "model_revision": model_revision,
        "script_sha256": script_sha256(),
        "config": {
            "device": args.device, "dtype": args.dtype, "batch_size": args.batch_size,
            "max_new_tokens": args.max_new_tokens, "limit": args.limit,
            "chat_template": not args.no_chat_template, "seed": "greedy/deterministic",
            "requested_revision": args.revision,
        },
        "benchmarks": {
            BENCHMARKS[k]["id"]: {"metric": BENCHMARKS[k]["metric"], **v["aggregate"]}
            for k, v in results.items()
        },
    }
    (out_dir / "results.json").write_text(json.dumps(report, indent=2, ensure_ascii=False),
                                           encoding="utf-8")

    for key, res in results.items():
        path = out_dir / f"samples_{BENCHMARKS[key]['id']}.csv"
        with path.open("w", newline="", encoding="utf-8") as f:
            w = csv.writer(f)
            metric_names = list(res["samples"][0].scores.keys()) if res["samples"] else []
            w.writerow(["idx", "category", *metric_names, "question", "gold", "pred"])
            for s in res["samples"]:
                w.writerow([s.idx, s.category, *[f"{s.scores[m]:.4f}" for m in metric_names],
                            s.question, s.gold, s.pred])

    # dual-mode benchmarks additionally get a rich per-item JSONL: raw
    # generation, extracted answer, per-choice log-probs, both metric families.
    for key, res in results.items():
        if BENCHMARKS[key]["kind"] != "dual":
            continue
        path = out_dir / f"samples_{BENCHMARKS[key]['id']}.jsonl"
        with path.open("w", encoding="utf-8") as f:
            for s in res["samples"]:
                if s.detail is not None:
                    f.write(json.dumps(s.detail, ensure_ascii=False) + "\n")

    entry = leaderboard_entry(model_name, results, model_revision)
    (out_dir / "leaderboard.json").write_text(json.dumps(entry, indent=2, ensure_ascii=False),
                                               encoding="utf-8")

    evaluated = [BENCHMARKS[k]["id"] for k in results]
    merged = merge_into_models_json(entry, evaluated, out_dir)

    print(f"\nSaved: {out_dir / 'results.json'}")
    print(f"Saved: {out_dir / 'leaderboard.json'} (single entry, for reference)")
    if merged:
        print(f"Saved: {merged} <- live leaderboard with this run merged in; "
              f"upload this file to the Space as-is")
    for key in results:
        print(f"Saved: {out_dir / ('samples_' + BENCHMARKS[key]['id'] + '.csv')}")


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #

def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Universal BenchLabs evaluator -- one script, every benchmark.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Example: python script.py --model Qwen/Qwen2.5-0.5B",
    )
    p.add_argument("--model", required=True, help="HF model id or local checkpoint path")
    p.add_argument("--revision", default=None,
                   help="pin an exact model commit (SHA / tag / branch). A maintainer "
                        "re-running with the recorded model_revision + the pinned "
                        "script copy reproduces the run byte-for-byte even if the "
                        "model's main branch has moved since")
    p.add_argument("--benchmarks", default="latest",
                   help="'latest' (default: the 7-2026 tiers), 'all' (both generations), "
                        "or comma-separated keys: effortless7,easy7,mid7 / "
                        "effortless,easy,mid (legacy 6-2026)")
    p.add_argument("--device", default="auto", help="auto | cuda | cpu | mps")
    p.add_argument("--dtype", default="auto", help="auto | float16 | bfloat16 | float32")
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--max-new-tokens", type=int, default=32,
                   help="generation budget; raise to 2048+ for reasoning models "
                        "that emit <think> blocks (default: 32)")
    p.add_argument("--limit", type=int, default=None, help="cap rows per benchmark (smoke test)")
    p.add_argument("--output-dir", default=None,
                   help="default: benchlabs_results/<model-name>")
    p.add_argument("--no-chat-template", action="store_true",
                   help="force plain 'Question:/Answer:' prompting even for instruct models")
    p.add_argument("--trust-remote-code", action="store_true")
    p.add_argument("--refresh-data", action="store_true", help="re-download datasets")
    p.add_argument("--leaderboard", action="store_true",
                   help="also print the models.json entry to stdout")
    return p.parse_args(argv)


LATEST_GENERATION = "7-2026"


def resolve_benchmarks(spec: str) -> List[str]:
    spec = spec.strip().lower()
    if spec == "latest":
        keys = [k for k, b in BENCHMARKS.items()
                if "unavailable" not in b and b.get("generation") == LATEST_GENERATION]
    elif spec == "all":
        keys = [k for k, b in BENCHMARKS.items() if "unavailable" not in b]
    else:
        keys = [s.strip().lower() for s in spec.split(",") if s.strip()]
    unknown = [k for k in keys if k not in BENCHMARKS]
    if unknown:
        sys.exit(f"Unknown benchmark(s): {unknown}. Choose from: {list(BENCHMARKS)}")
    for k in list(keys):
        if "unavailable" in BENCHMARKS[k]:
            print(f"Skipping {BENCHMARKS[k]['id']}: {BENCHMARKS[k]['unavailable']}")
            keys.remove(k)
    return keys


def main(argv: Optional[Sequence[str]] = None, model_factory=None) -> int:
    args = parse_args(argv)
    keys = resolve_benchmarks(args.benchmarks)
    if not keys:
        sys.exit("No runnable benchmarks selected.")

    print("Loading datasets...")
    data: Dict[str, List[dict]] = {}
    for k in keys:
        rows = load_benchmark_rows(BENCHMARKS[k]["id"], refresh=args.refresh_data)
        if args.limit:
            rows = rows[:args.limit]
        data[k] = rows
        print(f"  {BENCHMARKS[k]['id']}: {len(rows)} rows")

    factory = model_factory or (lambda: HFModel(
        args.model, args.device, args.dtype, args.trust_remote_code,
        use_chat_template=not args.no_chat_template, revision=args.revision))
    model = factory()

    results: Dict[str, Any] = {}
    for k in keys:
        bench = BENCHMARKS[k]
        print(f"\nRunning {bench['id']} ({bench['kind']}, {len(data[k])} rows)...")
        if bench["kind"] == "dual":
            samples, agg = run_dual(k, data[k], model, args)
        elif bench["kind"] == "generative":
            samples, agg = run_generative(k, data[k], model, args)
        else:
            samples, agg = run_multiple_choice(k, data[k], model, args)
        results[k] = {"samples": samples, "aggregate": agg}
        print_report(k, agg)

    out_dir = Path(args.output_dir) if args.output_dir else \
        Path("benchlabs_results") / re.sub(r"[^A-Za-z0-9._-]+", "_", args.model)
    save_outputs(out_dir, args.model, results, args, model.resolved_revision)

    if args.leaderboard:
        print("\n=== leaderboard entry (models.json) ===")
        print(json.dumps(leaderboard_entry(args.model, results, model.resolved_revision),
                          indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())