scripts/run_drill_graph_evaluation.py

#!/usr/bin/env python3
"""Law-family-disjoint evaluation of a learned graph residual on DRiLL.

The text candidate pool and frozen cross-encoder scores are shared by every
system.  Only the treatment receives query-conditioned messages propagated over
qrel-independent Vietnamese legal relationships.  Model selection uses a
separate development fold; the reported fold is never used for early stopping.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import json
import platform
import random
import sys
import time
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Mapping, Sequence

import numpy as np
import torch
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from scripts.run_static_benchmark import (  # noqa: E402
    RERANKER_NAME,
    RERANKER_REVISION,
    load_or_encode,
    resolve_device,
)
from telen.drill_graph import (  # noqa: E402
    VBPL_DERIVED_RELATION_TYPES,
    build_drill_legal_graph,
    load_drill_graph_dataset,
    query_conditioned_graph_messages,
)
from telen.static_benchmark import (  # noqa: E402
    dense_top_k,
    evaluate_rankings,
    paired_bootstrap_difference,
    reciprocal_rank_fusion,
    sha256_file,
    SparseBM25,
)


MODEL_NAME = "BAAI/bge-m3"
MODEL_REVISION = "5617a9f61b028005a4858fdac845db406aefb181"
DEFAULT_SEEDS = (17, 29, 43)


def _state_dict_sha256(state: Mapping[str, torch.Tensor]) -> str:
    digest = hashlib.sha256()
    for name, value in sorted(state.items()):
        digest.update(name.encode("utf-8"))
        digest.update(value.detach().cpu().contiguous().numpy().tobytes())
    return digest.hexdigest()


def _stable_candidate_rankings(
    candidate_indices: np.ndarray,
    scores: np.ndarray,
    corpus_ids: Sequence[str],
) -> np.ndarray:
    if candidate_indices.shape != scores.shape:
        raise ValueError("candidate indices and score matrix must align")
    identifiers = np.asarray(corpus_ids, dtype=str)
    output = np.empty_like(candidate_indices, dtype=np.int32)
    for query_index, (candidates, values) in enumerate(
        zip(candidate_indices, scores, strict=True)
    ):
        order = np.lexsort((identifiers[candidates], -values))
        output[query_index] = candidates[order]
    return output


def score_candidates_with_cross_encoder_reuse(
    dataset,
    candidates: np.ndarray,
    *,
    model_name: str,
    revision: str,
    device: torch.device,
    max_length: int,
    batch_size: int,
    score_cache: Path,
    reuse_candidates: np.ndarray | None = None,
    reuse_scores: np.ndarray | None = None,
) -> tuple[np.ndarray, float, int]:
    """Score candidate pairs and reuse already-audited query/article pairs."""

    if candidates.ndim != 2 or batch_size < 1:
        raise ValueError("candidate matrix and positive batch size are required")
    if (reuse_candidates is None) != (reuse_scores is None):
        raise ValueError("reuse candidates and scores must be supplied together")
    shape = candidates.shape
    score_cache.parent.mkdir(parents=True, exist_ok=True)
    if score_cache.is_file():
        scores = np.lib.format.open_memmap(score_cache, mode="r+")
        if scores.shape != shape or scores.dtype != np.float32:
            raise ValueError("cross-encoder score cache has an incompatible shape")
    else:
        scores = np.lib.format.open_memmap(
            score_cache, mode="w+", dtype=np.float32, shape=shape
        )
        scores.fill(np.nan)
        if reuse_candidates is not None and reuse_scores is not None:
            if reuse_candidates.shape != reuse_scores.shape:
                raise ValueError("reuse candidate and score matrices do not align")
            for query_index in range(shape[0]):
                known = {
                    int(aid): float(score)
                    for aid, score in zip(
                        reuse_candidates[query_index],
                        reuse_scores[query_index],
                        strict=True,
                    )
                }
                for local_rank, aid in enumerate(candidates[query_index]):
                    if int(aid) in known:
                        scores[query_index, local_rank] = known[int(aid)]
        scores.flush()
    missing = np.argwhere(~np.isfinite(scores))
    reused_count = int(scores.size - len(missing))
    started = time.perf_counter()
    if len(missing):
        tokenizer = AutoTokenizer.from_pretrained(
            model_name, revision=revision, local_files_only=True
        )
        dtype = torch.float16 if device.type == "cuda" else torch.float32
        model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            revision=revision,
            local_files_only=True,
            dtype=dtype,
        ).to(device)
        model.eval()
        with torch.inference_mode():
            for start in range(0, len(missing), batch_size):
                batch_positions = missing[start : start + batch_size]
                pairs = [
                    (
                        dataset.query_texts[int(query_index)],
                        dataset.corpus_texts[
                            int(candidates[int(query_index), int(local_rank)])
                        ],
                    )
                    for query_index, local_rank in batch_positions
                ]
                encoded = tokenizer(
                    pairs,
                    padding=True,
                    truncation=True,
                    max_length=max_length,
                    return_tensors="pt",
                ).to(device)
                values = (
                    model(**encoded, return_dict=True)
                    .logits.view(-1)
                    .float()
                    .cpu()
                    .numpy()
                )
                scores[
                    batch_positions[:, 0], batch_positions[:, 1]
                ] = values.astype(np.float32)
                scores.flush()
        del model
        if device.type == "cuda":
            torch.cuda.empty_cache()
    return np.asarray(scores, dtype=np.float32), time.perf_counter() - started, reused_count


def _per_query_standardize(values: np.ndarray) -> np.ndarray:
    mean = values.mean(axis=1, keepdims=True)
    standard_deviation = values.std(axis=1, keepdims=True)
    return ((values - mean) / np.maximum(standard_deviation, 1e-6)).astype(
        np.float32
    )


def _reciprocal_ranks(values: np.ndarray, candidates: np.ndarray) -> np.ndarray:
    result = np.empty_like(values, dtype=np.float32)
    for query_index, (row, aids) in enumerate(zip(values, candidates, strict=True)):
        order = np.lexsort((aids, -row))
        result[query_index, order] = 1.0 / np.arange(1, len(row) + 1)
    return result


def expand_candidates_with_local_graph(
    seed_candidates: np.ndarray,
    fallback_candidates: np.ndarray,
    graph,
    *,
    depth: int,
) -> np.ndarray:
    """Add citation/adjacency neighbors without labels, then deterministic fill."""

    if seed_candidates.ndim != 2 or fallback_candidates.ndim != 2:
        raise ValueError("seed and fallback candidates must be matrices")
    if seed_candidates.shape[0] != fallback_candidates.shape[0]:
        raise ValueError("seed and fallback query counts must match")
    if not seed_candidates.shape[1] <= depth <= len(graph.article_document_ids):
        raise ValueError("expanded depth must cover the seed pool")
    relation_priority = ("ARTICLE_CITES", "ARTICLE_CITED_BY", "ADJACENT")
    output = np.empty((len(seed_candidates), depth), dtype=np.int32)
    for query_index, seed_row in enumerate(seed_candidates):
        seed = [int(value) for value in seed_row]
        selected = set(seed)
        proposed: dict[int, tuple[int, int, int]] = {}
        for seed_rank, aid in enumerate(seed):
            for priority, relation in enumerate(relation_priority):
                for neighbor in graph.article_neighbors[relation][aid]:
                    if neighbor not in selected:
                        key = (seed_rank, priority, neighbor)
                        proposed[neighbor] = min(proposed.get(neighbor, key), key)
        expanded = seed + [
            aid
            for aid, _key in sorted(proposed.items(), key=lambda item: item[1])
        ]
        selected.update(expanded)
        if len(expanded) < depth:
            for aid_value in fallback_candidates[query_index]:
                aid = int(aid_value)
                if aid not in selected:
                    expanded.append(aid)
                    selected.add(aid)
                    if len(expanded) == depth:
                        break
        if len(expanded) < depth:
            for aid in range(len(graph.article_document_ids)):
                if aid not in selected:
                    expanded.append(aid)
                    if len(expanded) == depth:
                        break
        output[query_index] = np.asarray(expanded[:depth], dtype=np.int32)
    return output


def build_gold_document_component_folds(
    relevant_indices: Sequence[frozenset[int]],
    article_document_ids: Sequence[str],
    fold_count: int,
) -> tuple[np.ndarray, dict[str, object]]:
    """Assign connected gold-document components to balanced folds."""

    if fold_count < 3:
        raise ValueError("at least three folds are required for train/dev/test")
    parent: dict[str, str] = {}

    def find(value: str) -> str:
        parent.setdefault(value, value)
        if parent[value] != value:
            parent[value] = find(parent[value])
        return parent[value]

    def union(left: str, right: str) -> None:
        left_root, right_root = find(left), find(right)
        if left_root != right_root:
            parent[right_root] = left_root

    query_documents: list[tuple[str, ...]] = []
    for qrels in relevant_indices:
        documents = tuple(sorted({article_document_ids[aid] for aid in qrels}))
        if not documents:
            raise ValueError("every DRiLL question must have a gold document")
        for document in documents:
            find(document)
        for document in documents[1:]:
            union(documents[0], document)
        query_documents.append(documents)

    components: Counter[str] = Counter()
    query_roots: list[str] = []
    for documents in query_documents:
        roots = {find(document) for document in documents}
        if len(roots) != 1:
            raise RuntimeError("gold component construction is inconsistent")
        root = next(iter(roots))
        query_roots.append(root)
        components[root] += 1

    loads = [0] * fold_count
    component_fold: dict[str, int] = {}
    for root, count in sorted(components.items(), key=lambda item: (-item[1], item[0])):
        fold = min(range(fold_count), key=lambda value: (loads[value], value))
        component_fold[root] = fold
        loads[fold] += count
    folds = np.asarray([component_fold[root] for root in query_roots], dtype=np.int64)
    return folds, {
        "fold_count": fold_count,
        "component_count": len(components),
        "fold_query_counts": {str(key): value for key, value in enumerate(loads)},
        "largest_component_query_count": max(components.values()),
        "gold_document_overlap_across_folds": 0,
    }


def _candidate_labels(
    candidates: np.ndarray, relevant_indices: Sequence[frozenset[int]]
) -> np.ndarray:
    labels = np.zeros_like(candidates, dtype=np.float32)
    for query_index, qrels in enumerate(relevant_indices):
        labels[query_index] = np.fromiter(
            (int(aid) in qrels for aid in candidates[query_index]),
            dtype=np.float32,
            count=candidates.shape[1],
        )
    return labels


def _selection_metrics(
    selected: np.ndarray,
    labels: np.ndarray,
    total_relevant: np.ndarray,
) -> tuple[dict[str, float], dict[str, np.ndarray]]:
    """Official-style macro P/R/F2 for variable-size returned sets."""

    if selected.shape != labels.shape or len(total_relevant) != len(selected):
        raise ValueError("selection, labels, and total relevance do not align")
    returned = selected.sum(axis=1).astype(np.float64)
    hits = (selected * labels).sum(axis=1).astype(np.float64)
    precision = np.divide(
        hits, returned, out=np.zeros_like(hits), where=returned > 0
    )
    recall = hits / np.maximum(total_relevant.astype(np.float64), 1.0)
    query_denominator = 4.0 * precision + recall
    query_f2 = np.divide(
        5.0 * precision * recall,
        query_denominator,
        out=np.zeros_like(precision),
        where=query_denominator > 0,
    )
    macro_precision = float(precision.mean())
    macro_recall = float(recall.mean())
    denominator = 4.0 * macro_precision + macro_recall
    return (
        {
            "macro_precision": macro_precision,
            "macro_recall": macro_recall,
            "macro_f2": (
                5.0 * macro_precision * macro_recall / denominator
                if denominator
                else 0.0
            ),
            "mean_query_f2": float(query_f2.mean()),
            "mean_returned_articles": float(returned.mean()),
            "empty_return_rate": float(np.mean(returned == 0)),
        },
        {
            "precision": precision,
            "recall": recall,
            "query_f2": query_f2,
            "returned": returned,
        },
    )


def apply_score_threshold(
    scores: np.ndarray, threshold: float, *, maximum_results: int
) -> np.ndarray:
    """Apply a global score threshold with a deterministic per-query cap."""

    if scores.ndim != 2 or maximum_results < 1:
        raise ValueError("score matrix and positive result cap are required")
    selected = np.zeros_like(scores, dtype=bool)
    for query_index, row in enumerate(scores):
        order = np.argsort(-row, kind="stable")[:maximum_results]
        kept = order[row[order] >= threshold]
        selected[query_index, kept] = True
    return selected


def select_f2_threshold(
    scores: np.ndarray,
    labels: np.ndarray,
    total_relevant: np.ndarray,
    *,
    maximum_results: int,
) -> tuple[float, dict[str, float]]:
    """Exactly select a development threshold over capped candidate events."""

    if scores.shape != labels.shape:
        raise ValueError("score and label matrices must align")
    events: list[tuple[float, int, bool]] = []
    for query_index, row in enumerate(scores):
        order = np.argsort(-row, kind="stable")[:maximum_results]
        events.extend(
            (float(row[position]), query_index, bool(labels[query_index, position]))
            for position in order
        )
    events.sort(key=lambda value: -value[0])
    selected_count = np.zeros(len(scores), dtype=np.float64)
    hit_count = np.zeros(len(scores), dtype=np.float64)
    best_threshold = float("inf")
    best_summary = {
        "macro_precision": 0.0,
        "macro_recall": 0.0,
        "macro_f2": 0.0,
        "mean_query_f2": 0.0,
        "mean_returned_articles": 0.0,
        "empty_return_rate": 1.0,
    }
    index = 0
    while index < len(events):
        threshold = events[index][0]
        stop = index
        while stop < len(events) and events[stop][0] == threshold:
            _score, query_index, relevant = events[stop]
            selected_count[query_index] += 1.0
            hit_count[query_index] += float(relevant)
            stop += 1
        precision = np.divide(
            hit_count,
            selected_count,
            out=np.zeros_like(hit_count),
            where=selected_count > 0,
        )
        recall = hit_count / np.maximum(total_relevant.astype(np.float64), 1.0)
        macro_precision = float(precision.mean())
        macro_recall = float(recall.mean())
        denominator = 4.0 * macro_precision + macro_recall
        macro_f2 = (
            5.0 * macro_precision * macro_recall / denominator
            if denominator
            else 0.0
        )
        if macro_f2 > best_summary["macro_f2"] + 1e-12:
            selected = apply_score_threshold(
                scores, threshold, maximum_results=maximum_results
            )
            best_summary, _ = _selection_metrics(
                selected, labels, total_relevant
            )
            best_threshold = threshold
        index = stop
    return best_threshold, best_summary


def _ranking_observations(
    rankings: np.ndarray, relevant_indices: Sequence[frozenset[int]]
) -> tuple[dict[str, float], dict[str, np.ndarray]]:
    summary, observations = evaluate_rankings(rankings, relevant_indices)
    precision_at_2 = np.zeros(len(rankings), dtype=np.float64)
    recall_at_2 = np.zeros(len(rankings), dtype=np.float64)
    f2_at_2 = np.zeros(len(rankings), dtype=np.float64)
    for query_index, (ranking, relevant) in enumerate(
        zip(rankings, relevant_indices, strict=True)
    ):
        hits = sum(int(aid) in relevant for aid in ranking[:2])
        precision = hits / 2.0
        recall = hits / len(relevant)
        precision_at_2[query_index] = precision
        recall_at_2[query_index] = recall
        denominator = 4.0 * precision + recall
        f2_at_2[query_index] = (
            5.0 * precision * recall / denominator if denominator else 0.0
        )
    macro_precision = float(precision_at_2.mean())
    macro_recall = float(recall_at_2.mean())
    denominator = 4.0 * macro_precision + macro_recall
    summary.update(
        {
            "precision_at_2": macro_precision,
            "recall_at_2": macro_recall,
            "macro_f2_from_macro_precision_recall_at_2": (
                5.0 * macro_precision * macro_recall / denominator
                if denominator
                else 0.0
            ),
            "mean_query_f2_at_2": float(f2_at_2.mean()),
        }
    )
    observations.update(
        {
            "precision_at_2": precision_at_2,
            "recall_at_2": recall_at_2,
            "query_f2_at_2": f2_at_2,
        }
    )
    return summary, observations


class _ResidualMLP(torch.nn.Module):
    def __init__(self, width: int, hidden_size: int, residual_cap: float) -> None:
        super().__init__()
        self.residual_cap = float(residual_cap)
        self.network = torch.nn.Sequential(
            torch.nn.Linear(width, hidden_size),
            torch.nn.GELU(),
            torch.nn.LayerNorm(hidden_size),
            torch.nn.Linear(hidden_size, hidden_size // 2),
            torch.nn.GELU(),
            torch.nn.Linear(hidden_size // 2, 1),
        )

    def forward(
        self, features: torch.Tensor, base_scores: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        residual = self.residual_cap * torch.tanh(self.network(features).squeeze(-1))
        return base_scores + residual, residual


class _SemanticResidualMLP(torch.nn.Module):
    """Residual ranker with a learned low-rank frozen-embedding interaction."""

    def __init__(
        self,
        feature_width: int,
        embedding_width: int,
        semantic_rank: int,
        hidden_size: int,
        residual_cap: float,
    ) -> None:
        super().__init__()
        if semantic_rank < 1:
            raise ValueError("semantic rank must be positive")
        self.residual_cap = float(residual_cap)
        self.query_projection = torch.nn.Linear(
            embedding_width, semantic_rank, bias=False
        )
        self.article_projection = torch.nn.Linear(
            embedding_width, semantic_rank, bias=False
        )
        self.network = torch.nn.Sequential(
            torch.nn.Linear(feature_width + semantic_rank, hidden_size),
            torch.nn.GELU(),
            torch.nn.LayerNorm(hidden_size),
            torch.nn.Linear(hidden_size, hidden_size // 2),
            torch.nn.GELU(),
            torch.nn.Linear(hidden_size // 2, 1),
        )

    def forward(
        self,
        features: torch.Tensor,
        base_scores: torch.Tensor,
        query_embeddings: torch.Tensor,
        article_embeddings: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        query_factors = self.query_projection(query_embeddings).unsqueeze(1)
        article_factors = self.article_projection(article_embeddings)
        semantic_interaction = query_factors * article_factors
        inputs = torch.cat((features, semantic_interaction), dim=-1)
        residual = self.residual_cap * torch.tanh(self.network(inputs).squeeze(-1))
        return base_scores + residual, residual


def _ndcg_at_10_from_scores(scores: np.ndarray, labels: np.ndarray) -> float:
    values: list[float] = []
    for row, targets in zip(scores, labels, strict=True):
        relevant_count = int(targets.sum())
        if not relevant_count:
            continue
        order = np.argsort(-row, kind="stable")[:10]
        gains = targets[order] / np.log2(np.arange(2, 12))
        ideal = np.sum(1.0 / np.log2(np.arange(2, min(relevant_count, 10) + 2)))
        values.append(float(gains.sum() / ideal))
    return float(np.mean(values)) if values else 0.0


def _predict(
    model: _ResidualMLP,
    features: np.ndarray,
    base_scores: np.ndarray,
    mean: np.ndarray,
    scale: np.ndarray,
    device: torch.device,
) -> np.ndarray:
    model.eval()
    normalized = (features - mean) / scale
    with torch.inference_mode():
        logits, _ = model(
            torch.from_numpy(normalized.astype(np.float32)).to(device),
            torch.from_numpy(base_scores.astype(np.float32)).to(device),
        )
    return logits.cpu().numpy()


def train_residual_ranker(
    features: np.ndarray,
    base_scores: np.ndarray,
    labels: np.ndarray,
    *,
    train_indices: np.ndarray,
    development_indices: np.ndarray,
    seed: int,
    device: torch.device,
    hidden_size: int,
    residual_cap: float,
    learning_rate: float,
    epochs: int,
    patience: int,
    query_batch_size: int,
) -> tuple[_ResidualMLP, np.ndarray, np.ndarray, dict[str, object]]:
    """Train with listwise relevance mass and candidate-level calibration."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)
    eligible_train = train_indices[labels[train_indices].sum(axis=1) > 0]
    eligible_development = development_indices[
        labels[development_indices].sum(axis=1) > 0
    ]
    if not len(eligible_train) or not len(eligible_development):
        raise ValueError("train and development folds need candidate-positive queries")
    flat_train = features[eligible_train].reshape(-1, features.shape[-1])
    mean = flat_train.mean(axis=0, dtype=np.float64).astype(np.float32)
    scale = flat_train.std(axis=0, dtype=np.float64).astype(np.float32)
    scale = np.maximum(scale, 1e-5)
    model = _ResidualMLP(features.shape[-1], hidden_size, residual_cap).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=learning_rate, weight_decay=1e-4
    )
    best_state: dict[str, torch.Tensor] | None = None
    best_epoch = -1
    best_metric = -float("inf")
    stale = 0
    generator = np.random.default_rng(seed)
    mean_tensor = torch.from_numpy(mean).to(device)
    scale_tensor = torch.from_numpy(scale).to(device)

    for epoch in range(epochs):
        model.train()
        shuffled = generator.permutation(eligible_train)
        for start in range(0, len(shuffled), query_batch_size):
            query_indices = shuffled[start : start + query_batch_size]
            batch_features = torch.from_numpy(features[query_indices]).to(device)
            batch_features = (batch_features - mean_tensor) / scale_tensor
            batch_base = torch.from_numpy(base_scores[query_indices]).to(device)
            batch_labels = torch.from_numpy(labels[query_indices]).to(device)
            logits, residual = model(batch_features, batch_base)
            positive_logits = logits.masked_fill(batch_labels == 0, -torch.inf)
            listwise = (
                torch.logsumexp(logits, dim=1)
                - torch.logsumexp(positive_logits, dim=1)
            ).mean()
            positive_count = batch_labels.sum().clamp_min(1.0)
            negative_count = batch_labels.numel() - positive_count
            positive_weight = (negative_count / positive_count).clamp(max=25.0)
            calibration = F.binary_cross_entropy_with_logits(
                logits,
                batch_labels,
                pos_weight=positive_weight,
            )
            loss = listwise + 0.15 * calibration + 0.002 * residual.square().mean()
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()

        development_scores = _predict(
            model,
            features[eligible_development],
            base_scores[eligible_development],
            mean,
            scale,
            device,
        )
        metric = _ndcg_at_10_from_scores(
            development_scores, labels[eligible_development]
        )
        if metric > best_metric + 1e-6:
            best_metric = metric
            best_epoch = epoch
            best_state = copy.deepcopy(model.state_dict())
            stale = 0
        else:
            stale += 1
            if stale >= patience:
                break
    if best_state is None:
        raise RuntimeError("residual training did not produce a checkpoint")
    model.load_state_dict(best_state)
    return model, mean, scale, {
        "seed": seed,
        "best_epoch_zero_based": best_epoch,
        "development_ndcg_at_10": best_metric,
        "model_state_sha256": _state_dict_sha256(best_state),
        "eligible_training_queries": int(len(eligible_train)),
        "eligible_development_queries": int(len(eligible_development)),
    }


def _predict_semantic(
    model: _SemanticResidualMLP,
    features: np.ndarray,
    base_scores: np.ndarray,
    query_embeddings: np.ndarray,
    corpus_embeddings: np.ndarray,
    candidates: np.ndarray,
    mean: np.ndarray,
    scale: np.ndarray,
    device: torch.device,
    *,
    query_batch_size: int,
) -> np.ndarray:
    """Predict in bounded batches without materializing the candidate tensor."""

    model.eval()
    output = np.empty_like(base_scores, dtype=np.float32)
    with torch.inference_mode():
        for start in range(0, len(features), query_batch_size):
            stop = min(start + query_batch_size, len(features))
            normalized = (features[start:stop] - mean) / scale
            article_batch = np.ascontiguousarray(
                corpus_embeddings[candidates[start:stop]], dtype=np.float32
            )
            logits, _ = model(
                torch.from_numpy(normalized.astype(np.float32)).to(device),
                torch.from_numpy(base_scores[start:stop].astype(np.float32)).to(device),
                torch.from_numpy(
                    np.ascontiguousarray(query_embeddings[start:stop], dtype=np.float32)
                ).to(device),
                torch.from_numpy(article_batch).to(device),
            )
            output[start:stop] = logits.float().cpu().numpy()
    return output


def train_semantic_residual_ranker(
    features: np.ndarray,
    base_scores: np.ndarray,
    labels: np.ndarray,
    query_embeddings: np.ndarray,
    corpus_embeddings: np.ndarray,
    candidates: np.ndarray,
    *,
    train_indices: np.ndarray,
    development_indices: np.ndarray,
    seed: int,
    device: torch.device,
    hidden_size: int,
    semantic_rank: int,
    residual_cap: float,
    learning_rate: float,
    epochs: int,
    patience: int,
    query_batch_size: int,
) -> tuple[_SemanticResidualMLP, np.ndarray, np.ndarray, dict[str, object]]:
    """Train a frozen-embedding semantic adaptor under nested OOF selection."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)
    eligible_train = train_indices[labels[train_indices].sum(axis=1) > 0]
    eligible_development = development_indices[
        labels[development_indices].sum(axis=1) > 0
    ]
    if not len(eligible_train) or not len(eligible_development):
        raise ValueError("train and development folds need candidate-positive queries")
    flat_train = features[eligible_train].reshape(-1, features.shape[-1])
    mean = flat_train.mean(axis=0, dtype=np.float64).astype(np.float32)
    scale = flat_train.std(axis=0, dtype=np.float64).astype(np.float32)
    scale = np.maximum(scale, 1e-5)
    model = _SemanticResidualMLP(
        features.shape[-1],
        query_embeddings.shape[-1],
        semantic_rank,
        hidden_size,
        residual_cap,
    ).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=learning_rate, weight_decay=1e-4
    )
    best_state: dict[str, torch.Tensor] | None = None
    best_epoch = -1
    best_metric = -float("inf")
    stale = 0
    generator = np.random.default_rng(seed)
    mean_tensor = torch.from_numpy(mean).to(device)
    scale_tensor = torch.from_numpy(scale).to(device)

    for epoch in range(epochs):
        model.train()
        shuffled = generator.permutation(eligible_train)
        for start in range(0, len(shuffled), query_batch_size):
            query_indices = shuffled[start : start + query_batch_size]
            batch_features = torch.from_numpy(features[query_indices]).to(device)
            batch_features = (batch_features - mean_tensor) / scale_tensor
            batch_base = torch.from_numpy(base_scores[query_indices]).to(device)
            batch_labels = torch.from_numpy(labels[query_indices]).to(device)
            batch_queries = torch.from_numpy(
                np.ascontiguousarray(query_embeddings[query_indices], dtype=np.float32)
            ).to(device)
            batch_articles = torch.from_numpy(
                np.ascontiguousarray(
                    corpus_embeddings[candidates[query_indices]], dtype=np.float32
                )
            ).to(device)
            logits, residual = model(
                batch_features,
                batch_base,
                batch_queries,
                batch_articles,
            )
            positive_logits = logits.masked_fill(batch_labels == 0, -torch.inf)
            listwise = (
                torch.logsumexp(logits, dim=1)
                - torch.logsumexp(positive_logits, dim=1)
            ).mean()
            positive_count = batch_labels.sum().clamp_min(1.0)
            negative_count = batch_labels.numel() - positive_count
            positive_weight = (negative_count / positive_count).clamp(max=25.0)
            calibration = F.binary_cross_entropy_with_logits(
                logits,
                batch_labels,
                pos_weight=positive_weight,
            )
            loss = listwise + 0.15 * calibration + 0.002 * residual.square().mean()
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()

        development_scores = _predict_semantic(
            model,
            features[eligible_development],
            base_scores[eligible_development],
            query_embeddings[eligible_development],
            corpus_embeddings,
            candidates[eligible_development],
            mean,
            scale,
            device,
            query_batch_size=query_batch_size,
        )
        metric = _ndcg_at_10_from_scores(
            development_scores, labels[eligible_development]
        )
        if metric > best_metric + 1e-6:
            best_metric = metric
            best_epoch = epoch
            best_state = copy.deepcopy(model.state_dict())
            stale = 0
        else:
            stale += 1
            if stale >= patience:
                break
    if best_state is None:
        raise RuntimeError("semantic residual training did not produce a checkpoint")
    model.load_state_dict(best_state)
    return model, mean, scale, {
        "seed": seed,
        "best_epoch_zero_based": best_epoch,
        "development_ndcg_at_10": best_metric,
        "model_state_sha256": _state_dict_sha256(best_state),
        "eligible_training_queries": int(len(eligible_train)),
        "eligible_development_queries": int(len(eligible_development)),
        "semantic_projection_rank": semantic_rank,
        "trainable_parameter_count": sum(
            parameter.numel() for parameter in model.parameters()
        ),
    }


def continue_semantic_residual_ranker(
    model: _SemanticResidualMLP,
    features: np.ndarray,
    base_scores: np.ndarray,
    labels: np.ndarray,
    query_embeddings: np.ndarray,
    corpus_embeddings: np.ndarray,
    candidates: np.ndarray,
    mean: np.ndarray,
    scale: np.ndarray,
    *,
    train_indices: np.ndarray,
    development_indices: np.ndarray,
    seed: int,
    device: torch.device,
    learning_rate: float,
    epochs: int,
    patience: int,
    query_batch_size: int,
) -> tuple[_SemanticResidualMLP, dict[str, object]]:
    """Continue a pretrained adaptor with fixed feature normalization."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)
    eligible_train = train_indices[labels[train_indices].sum(axis=1) > 0]
    eligible_development = development_indices[
        labels[development_indices].sum(axis=1) > 0
    ]
    if not len(eligible_train) or not len(eligible_development):
        raise ValueError("train and development folds need candidate-positive queries")
    model = model.to(device)
    initial_sha = _state_dict_sha256(model.state_dict())
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=learning_rate, weight_decay=1e-4
    )
    best_state = copy.deepcopy(model.state_dict())
    best_epoch = -1
    best_metric = _ndcg_at_10_from_scores(
        _predict_semantic(
            model,
            features[eligible_development],
            base_scores[eligible_development],
            query_embeddings[eligible_development],
            corpus_embeddings,
            candidates[eligible_development],
            mean,
            scale,
            device,
            query_batch_size=query_batch_size,
        ),
        labels[eligible_development],
    )
    stale = 0
    generator = np.random.default_rng(seed + 1_000_003)
    mean_tensor = torch.from_numpy(mean).to(device)
    scale_tensor = torch.from_numpy(scale).to(device)

    for epoch in range(epochs):
        model.train()
        shuffled = generator.permutation(eligible_train)
        for start in range(0, len(shuffled), query_batch_size):
            query_indices = shuffled[start : start + query_batch_size]
            batch_features = torch.from_numpy(features[query_indices]).to(device)
            batch_features = (batch_features - mean_tensor) / scale_tensor
            batch_base = torch.from_numpy(base_scores[query_indices]).to(device)
            batch_labels = torch.from_numpy(labels[query_indices]).to(device)
            batch_queries = torch.from_numpy(
                np.ascontiguousarray(query_embeddings[query_indices], dtype=np.float32)
            ).to(device)
            batch_articles = torch.from_numpy(
                np.ascontiguousarray(
                    corpus_embeddings[candidates[query_indices]], dtype=np.float32
                )
            ).to(device)
            logits, residual = model(
                batch_features, batch_base, batch_queries, batch_articles
            )
            positive_logits = logits.masked_fill(batch_labels == 0, -torch.inf)
            listwise = (
                torch.logsumexp(logits, dim=1)
                - torch.logsumexp(positive_logits, dim=1)
            ).mean()
            positive_count = batch_labels.sum().clamp_min(1.0)
            negative_count = batch_labels.numel() - positive_count
            calibration = F.binary_cross_entropy_with_logits(
                logits,
                batch_labels,
                pos_weight=(negative_count / positive_count).clamp(max=25.0),
            )
            loss = listwise + 0.15 * calibration + 0.002 * residual.square().mean()
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()

        development_scores = _predict_semantic(
            model,
            features[eligible_development],
            base_scores[eligible_development],
            query_embeddings[eligible_development],
            corpus_embeddings,
            candidates[eligible_development],
            mean,
            scale,
            device,
            query_batch_size=query_batch_size,
        )
        metric = _ndcg_at_10_from_scores(
            development_scores, labels[eligible_development]
        )
        if metric > best_metric + 1e-6:
            best_metric = metric
            best_epoch = epoch
            best_state = copy.deepcopy(model.state_dict())
            stale = 0
        else:
            stale += 1
            if stale >= patience:
                break
    model.load_state_dict(best_state)
    return model, {
        "seed": seed,
        "initial_model_state_sha256": initial_sha,
        "final_model_state_sha256": _state_dict_sha256(best_state),
        "best_epoch_zero_based": best_epoch,
        "development_ndcg_at_10": best_metric,
        "eligible_training_queries": int(len(eligible_train)),
        "eligible_development_queries": int(len(eligible_development)),
        "feature_normalization": "frozen_from_drill_pretraining",
    }


def _relation_feature_columns(
    text_width: int,
    relation_slices: Mapping[str, slice],
    relations: Sequence[str],
) -> np.ndarray:
    columns: list[int] = []
    for relation in relations:
        section = relation_slices[relation]
        columns.extend(range(text_width + section.start, text_width + section.stop))
    return np.asarray(sorted(set(columns)), dtype=np.int64)


def run(args: argparse.Namespace) -> dict:
    dataset = load_drill_graph_dataset(args.corpus, args.questions)
    graph = build_drill_legal_graph(
        dataset,
        metadata_path=args.metadata,
        canonical_relationships_path=args.relationships,
    )
    retrieval = dataset.retrieval
    device = resolve_device(args.device)
    corpus_embeddings, corpus_cache = load_or_encode(
        retrieval.corpus_texts,
        retrieval.corpus_ids,
        role="drill-articles",
        cache_dir=args.cache_dir,
        model_name=args.model_name,
        revision=args.model_revision,
        device=device,
        max_length=args.max_length,
        batch_size=args.corpus_batch_size,
    )
    query_embeddings, query_cache = load_or_encode(
        retrieval.query_texts,
        retrieval.query_ids,
        role="drill-train-queries",
        cache_dir=args.cache_dir,
        model_name=args.model_name,
        revision=args.model_revision,
        device=device,
        max_length=args.max_length,
        batch_size=args.query_batch_size,
    )
    corpus_embeddings = np.asarray(corpus_embeddings, dtype=np.float32)
    query_embeddings = np.asarray(query_embeddings, dtype=np.float32)
    seed_depth = (
        args.seed_depth if args.candidate_mode == "rrf_graph" else args.candidate_depth
    )
    if seed_depth > args.candidate_depth:
        raise ValueError("seed depth cannot exceed final candidate depth")
    dense_candidates_path = args.cache_dir / f"drill-bge-top{seed_depth}.npy"
    if dense_candidates_path.is_file():
        dense_candidates = np.load(dense_candidates_path)
    else:
        dense_candidates = dense_top_k(
            query_embeddings,
            corpus_embeddings,
            retrieval.corpus_ids,
            k=seed_depth,
            batch_size=args.score_batch_size,
        )
        dense_candidates_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(dense_candidates_path, dense_candidates)
    bm25_candidates_path = args.cache_dir / f"drill-bm25-top{seed_depth}.npy"
    bm25_candidates: np.ndarray | None = None
    rrf_seed_candidates: np.ndarray | None = None
    if args.candidate_mode in {"rrf", "rrf_graph"}:
        if bm25_candidates_path.is_file():
            bm25_candidates = np.load(bm25_candidates_path)
        else:
            bm25 = SparseBM25(k1=args.bm25_k1, b=args.bm25_b)
            bm25.fit(retrieval.corpus_texts)
            bm25_candidates = bm25.rank(
                retrieval.query_texts,
                retrieval.corpus_ids,
                k=seed_depth,
                batch_size=args.score_batch_size,
            )
            np.save(bm25_candidates_path, bm25_candidates)
        rrf_seed_path = args.cache_dir / f"drill-rrf-top{seed_depth}.npy"
        if rrf_seed_path.is_file():
            rrf_seed_candidates = np.load(rrf_seed_path)
        else:
            rrf_seed_candidates = reciprocal_rank_fusion(
                (bm25_candidates, dense_candidates),
                retrieval.corpus_ids,
                k=seed_depth,
                constant=args.rrf_constant,
            )
            np.save(rrf_seed_path, rrf_seed_candidates)
        if args.candidate_mode == "rrf_graph":
            fallback_path = args.cache_dir / f"drill-bge-top{args.candidate_depth}.npy"
            if fallback_path.is_file():
                fallback_candidates = np.load(fallback_path)
            else:
                fallback_candidates = dense_top_k(
                    query_embeddings,
                    corpus_embeddings,
                    retrieval.corpus_ids,
                    k=args.candidate_depth,
                    batch_size=args.score_batch_size,
                )
                np.save(fallback_path, fallback_candidates)
            candidates_path = args.cache_dir / (
                f"drill-rrf-graph-expanded-seed{seed_depth}-"
                f"top{args.candidate_depth}.npy"
            )
            if candidates_path.is_file():
                candidates = np.load(candidates_path)
            else:
                candidates = expand_candidates_with_local_graph(
                    rrf_seed_candidates,
                    fallback_candidates,
                    graph,
                    depth=args.candidate_depth,
                )
                np.save(candidates_path, candidates)
        else:
            candidates_path = rrf_seed_path
            candidates = rrf_seed_candidates
    else:
        candidates_path = dense_candidates_path
        candidates = dense_candidates
    candidate_hash = hashlib.sha256(candidates.tobytes()).hexdigest()[:12]
    score_path = args.cache_dir / (
        f"drill-reranker-pairmatrix-{args.reranker_revision[:12]}-"
        f"source{candidate_hash}-top{args.candidate_depth}-"
        f"ml{args.reranker_max_length}-scores.npy"
    )
    dense_hash = hashlib.sha256(dense_candidates.tobytes()).hexdigest()[:12]
    legacy_dense_scores_path = args.cache_dir / (
        f"drill-reranker-{args.reranker_revision[:12]}-source{dense_hash}-"
        f"top{seed_depth}-ml{args.reranker_max_length}-scores.npy"
    )
    reuse_scores = None
    if legacy_dense_scores_path.is_file():
        reuse_scores = np.asarray(
            np.load(legacy_dense_scores_path, mmap_mode="r"), dtype=np.float32
        ).reshape(len(retrieval.query_ids), seed_depth)
    reuse_candidates = dense_candidates if reuse_scores is not None else None
    if args.candidate_mode == "rrf_graph" and rrf_seed_candidates is not None:
        rrf_seed_hash = hashlib.sha256(rrf_seed_candidates.tobytes()).hexdigest()[:12]
        rrf_seed_score_path = args.cache_dir / (
            f"drill-reranker-pairmatrix-{args.reranker_revision[:12]}-"
            f"source{rrf_seed_hash}-top{seed_depth}-"
            f"ml{args.reranker_max_length}-scores.npy"
        )
        if rrf_seed_score_path.is_file():
            candidate_reuse_scores = np.asarray(
                np.load(rrf_seed_score_path, mmap_mode="r"), dtype=np.float32
            )
            if np.isfinite(candidate_reuse_scores).all():
                reuse_candidates = rrf_seed_candidates
                reuse_scores = candidate_reuse_scores
    cross_scores, reranker_runtime, reused_pair_count = (
        score_candidates_with_cross_encoder_reuse(
        retrieval,
        candidates,
        model_name=args.reranker_name,
        revision=args.reranker_revision,
        device=device,
        max_length=args.reranker_max_length,
        batch_size=args.reranker_batch_size,
        score_cache=score_path,
        reuse_candidates=reuse_candidates,
        reuse_scores=reuse_scores,
        )
    )
    dense_scores = np.einsum(
        "qd,qkd->qk",
        np.asarray(query_embeddings, dtype=np.float32),
        np.asarray(corpus_embeddings, dtype=np.float32)[candidates],
        optimize=True,
    ).astype(np.float32)
    dense_z = _per_query_standardize(dense_scores)
    cross_z = _per_query_standardize(cross_scores)
    dense_rr = _reciprocal_ranks(dense_scores, candidates)
    cross_rr = _reciprocal_ranks(cross_scores, candidates)
    feature_channels = [dense_z, cross_z, dense_rr, cross_rr]
    text_feature_names_list = ["dense_z", "cross_z", "dense_rr", "cross_rr"]
    if bm25_candidates is not None:
        bm25_rr = np.zeros_like(dense_rr)
        for query_index in range(len(candidates)):
            ranks = {
                int(aid): 1.0 / (rank + 1)
                for rank, aid in enumerate(bm25_candidates[query_index])
            }
            bm25_rr[query_index] = np.fromiter(
                (ranks.get(int(aid), 0.0) for aid in candidates[query_index]),
                dtype=np.float32,
                count=args.candidate_depth,
            )
        feature_channels.append(bm25_rr)
        text_feature_names_list.append("bm25_rr")
    text_features = np.stack(feature_channels, axis=-1)
    text_feature_names = tuple(text_feature_names_list)
    graph_channels = np.stack((dense_z, cross_z), axis=-1)
    graph_messages, graph_feature_names, relation_slices = (
        query_conditioned_graph_messages(
            graph,
            candidates,
            graph_channels,
            channel_names=("dense_z", "cross_z"),
        )
    )
    graph_features = np.concatenate((text_features, graph_messages), axis=-1)
    sham_permutation = np.random.default_rng(args.sham_seed).permutation(
        len(retrieval.query_ids)
    )
    sham_graph_features = np.concatenate(
        (text_features, graph_messages[sham_permutation]), axis=-1
    )
    labels = _candidate_labels(candidates, retrieval.relevant_indices)
    total_gold = np.asarray(
        list(map(len, retrieval.relevant_indices)), dtype=np.float32
    )
    folds, fold_audit = build_gold_document_component_folds(
        retrieval.relevant_indices, dataset.article_document_ids, args.fold_count
    )
    base_scores = cross_z
    system_seed_scores: dict[str, list[np.ndarray]] = defaultdict(list)
    system_seed_decisions: dict[str, list[np.ndarray]] = defaultdict(list)
    training_receipts: list[dict[str, object]] = []
    all_graph_relations = tuple(relation_slices)
    local_relations = (
        "SAME_DOCUMENT",
        "ADJACENT",
        "ARTICLE_CITES",
        "ARTICLE_CITED_BY",
    )
    causal_relations = tuple(
        f"{relation}_{direction}"
        for relation in ("AMENDS", "REPEALS", "REPLACES", "TERMINATES_VERSION")
        for direction in ("OUT", "IN")
    )
    contextual_relations = tuple(
        f"{relation}_{direction}"
        for relation in ("CITES", "IMPLEMENTS")
        for direction in ("OUT", "IN")
    )
    intervention_relations = {
        "graph_all_edges_removed": all_graph_relations,
        "graph_local_edges_removed": local_relations,
        "graph_vbpl_causal_labeled_edges_removed": causal_relations,
        "graph_vbpl_context_edges_removed": contextual_relations,
    }

    for seed in args.seeds:
        seed_text = np.zeros_like(base_scores, dtype=np.float32)
        seed_graph = np.zeros_like(base_scores, dtype=np.float32)
        seed_sham = np.zeros_like(base_scores, dtype=np.float32)
        seed_decisions = {
            name: np.zeros_like(base_scores, dtype=bool)
            for name in (
                "frozen_bge_reranker",
                "learned_text_residual",
                "learned_graph_residual",
                "learned_sham_graph_residual",
            )
        }
        seed_interventions = {
            name: np.zeros_like(base_scores, dtype=np.float32)
            for name in intervention_relations
        }
        seed_shuffled = np.zeros_like(base_scores, dtype=np.float32)
        for test_fold in range(args.fold_count):
            development_fold = (test_fold + 1) % args.fold_count
            test_indices = np.flatnonzero(folds == test_fold)
            development_indices = np.flatnonzero(folds == development_fold)
            training_indices = np.flatnonzero(
                (folds != test_fold) & (folds != development_fold)
            )
            training_function = (
                train_semantic_residual_ranker
                if args.semantic_rank > 0
                else train_residual_ranker
            )
            semantic_arguments = (
                {
                    "query_embeddings": query_embeddings,
                    "corpus_embeddings": corpus_embeddings,
                    "candidates": candidates,
                    "semantic_rank": args.semantic_rank,
                }
                if args.semantic_rank > 0
                else {}
            )
            text_model, text_mean, text_scale, text_receipt = training_function(
                text_features,
                base_scores,
                labels,
                train_indices=training_indices,
                development_indices=development_indices,
                seed=seed,
                device=device,
                hidden_size=args.hidden_size,
                residual_cap=args.residual_cap,
                learning_rate=args.learning_rate,
                epochs=args.epochs,
                patience=args.patience,
                query_batch_size=args.training_query_batch_size,
                **semantic_arguments,
            )
            graph_model, graph_mean, graph_scale, graph_receipt = training_function(
                graph_features,
                base_scores,
                labels,
                train_indices=training_indices,
                development_indices=development_indices,
                seed=seed,
                device=device,
                hidden_size=args.hidden_size,
                residual_cap=args.residual_cap,
                learning_rate=args.learning_rate,
                epochs=args.epochs,
                patience=args.patience,
                query_batch_size=args.training_query_batch_size,
                **semantic_arguments,
            )
            sham_model, sham_mean, sham_scale, sham_receipt = training_function(
                sham_graph_features,
                base_scores,
                labels,
                train_indices=training_indices,
                development_indices=development_indices,
                seed=seed,
                device=device,
                hidden_size=args.hidden_size,
                residual_cap=args.residual_cap,
                learning_rate=args.learning_rate,
                epochs=args.epochs,
                patience=args.patience,
                query_batch_size=args.training_query_batch_size,
                **semantic_arguments,
            )
            def predict_subset(model, values, mean, scale, indices):
                if args.semantic_rank > 0:
                    return _predict_semantic(
                        model,
                        values[indices],
                        base_scores[indices],
                        query_embeddings[indices],
                        corpus_embeddings,
                        candidates[indices],
                        mean,
                        scale,
                        device,
                        query_batch_size=args.training_query_batch_size,
                    )
                return _predict(
                    model,
                    values[indices],
                    base_scores[indices],
                    mean,
                    scale,
                    device,
                )

            seed_text[test_indices] = predict_subset(
                text_model, text_features, text_mean, text_scale, test_indices
            )
            seed_graph[test_indices] = predict_subset(
                graph_model, graph_features, graph_mean, graph_scale, test_indices
            )
            seed_sham[test_indices] = predict_subset(
                sham_model,
                sham_graph_features,
                sham_mean,
                sham_scale,
                test_indices,
            )
            decision_receipts = {}
            decision_systems = (
                (
                    "frozen_bge_reranker",
                    base_scores[development_indices],
                    base_scores[test_indices],
                ),
                (
                    "learned_text_residual",
                    predict_subset(
                        text_model,
                        text_features,
                        text_mean,
                        text_scale,
                        development_indices,
                    ),
                    seed_text[test_indices],
                ),
                (
                    "learned_graph_residual",
                    predict_subset(
                        graph_model,
                        graph_features,
                        graph_mean,
                        graph_scale,
                        development_indices,
                    ),
                    seed_graph[test_indices],
                ),
                (
                    "learned_sham_graph_residual",
                    predict_subset(
                        sham_model,
                        sham_graph_features,
                        sham_mean,
                        sham_scale,
                        development_indices,
                    ),
                    seed_sham[test_indices],
                ),
            )
            for system_name, development_scores, test_scores in decision_systems:
                threshold, development_summary = select_f2_threshold(
                    development_scores,
                    labels[development_indices],
                    total_gold[development_indices],
                    maximum_results=args.maximum_returned_articles,
                )
                seed_decisions[system_name][test_indices] = apply_score_threshold(
                    test_scores,
                    threshold,
                    maximum_results=args.maximum_returned_articles,
                )
                decision_receipts[system_name] = {
                    "threshold": threshold,
                    "development": development_summary,
                }
            for name, relations in intervention_relations.items():
                corrupted = graph_features[test_indices].copy()
                columns = _relation_feature_columns(
                    len(text_feature_names), relation_slices, relations
                )
                corrupted[:, :, columns] = 0.0
                if args.semantic_rank > 0:
                    seed_interventions[name][test_indices] = _predict_semantic(
                        graph_model,
                        corrupted,
                        base_scores[test_indices],
                        query_embeddings[test_indices],
                        corpus_embeddings,
                        candidates[test_indices],
                        graph_mean,
                        graph_scale,
                        device,
                        query_batch_size=args.training_query_batch_size,
                    )
                else:
                    seed_interventions[name][test_indices] = _predict(
                        graph_model,
                        corrupted,
                        base_scores[test_indices],
                        graph_mean,
                        graph_scale,
                        device,
                    )
            shuffled = graph_features[test_indices].copy()
            graph_width = graph_messages.shape[-1]
            permutation = np.random.default_rng(seed + test_fold * 1009).permutation(
                len(test_indices)
            )
            shuffled[:, :, -graph_width:] = shuffled[permutation, :, -graph_width:]
            if args.semantic_rank > 0:
                seed_shuffled[test_indices] = _predict_semantic(
                    graph_model,
                    shuffled,
                    base_scores[test_indices],
                    query_embeddings[test_indices],
                    corpus_embeddings,
                    candidates[test_indices],
                    graph_mean,
                    graph_scale,
                    device,
                    query_batch_size=args.training_query_batch_size,
                )
            else:
                seed_shuffled[test_indices] = _predict(
                    graph_model,
                    shuffled,
                    base_scores[test_indices],
                    graph_mean,
                    graph_scale,
                    device,
                )
            training_receipts.append(
                {
                    "seed": seed,
                    "test_fold": test_fold,
                    "development_fold": development_fold,
                    "training_fold_count": args.fold_count - 2,
                    "test_query_count": int(len(test_indices)),
                    "text_control": text_receipt,
                    "graph_treatment": graph_receipt,
                    "matched_parameter_sham_graph": sham_receipt,
                    "development_only_f2_decisions": decision_receipts,
                }
            )
            del text_model, graph_model, sham_model
            if device.type == "cuda":
                torch.cuda.empty_cache()
        system_seed_scores["learned_text_residual"].append(seed_text)
        system_seed_scores["learned_graph_residual"].append(seed_graph)
        system_seed_scores["learned_sham_graph_residual"].append(seed_sham)
        for name, selected in seed_decisions.items():
            system_seed_decisions[name].append(selected)
        for name, scores in seed_interventions.items():
            system_seed_scores[name].append(scores)
        system_seed_scores["graph_messages_query_shuffled"].append(seed_shuffled)

    system_scores = {
        "frozen_bge_reranker": cross_scores,
        **{
            name: np.mean(np.stack(values), axis=0)
            for name, values in system_seed_scores.items()
        },
    }
    official_f2_decisions: dict[str, object] = {
        "maximum_returned_articles": args.maximum_returned_articles,
        "selection_rule": "one global threshold selected exactly on the nested development fold; deterministic score order cap",
        "systems": {},
    }
    selection_observations: dict[str, dict[str, np.ndarray]] = {}
    for name, seed_values in system_seed_decisions.items():
        seed_summaries = []
        for selected in seed_values:
            summary, _observations = _selection_metrics(
                selected, labels, total_gold
            )
            seed_summaries.append(summary)
        majority = np.mean(np.stack(seed_values), axis=0) >= 0.5
        majority_summary, majority_observations = _selection_metrics(
            majority, labels, total_gold
        )
        official_f2_decisions["systems"][name] = {
            "seed_results": seed_summaries,
            "majority_vote_ensemble": majority_summary,
        }
        selection_observations[name] = majority_observations
    official_f2_decisions["paired_bootstrap"] = {
        "learned_graph_residual_minus_learned_text_residual": {
            "query_f2": paired_bootstrap_difference(
                selection_observations["learned_text_residual"]["query_f2"],
                selection_observations["learned_graph_residual"]["query_f2"],
                samples=args.bootstrap_samples,
                seed=args.bootstrap_seed,
            )
        },
        "learned_graph_residual_minus_learned_sham_graph_residual": {
            "query_f2": paired_bootstrap_difference(
                selection_observations["learned_sham_graph_residual"]["query_f2"],
                selection_observations["learned_graph_residual"]["query_f2"],
                samples=args.bootstrap_samples,
                seed=args.bootstrap_seed,
            )
        },
    }
    seed_stability: dict[str, list[dict[str, float]]] = {}
    for name, seed_values in system_seed_scores.items():
        seed_stability[name] = []
        for values in seed_values:
            seed_rankings = _stable_candidate_rankings(
                candidates, values, retrieval.corpus_ids
            )
            seed_summary, _ = _ranking_observations(
                seed_rankings, retrieval.relevant_indices
            )
            seed_stability[name].append(seed_summary)
    results: dict[str, dict[str, float]] = {}
    observations: dict[str, dict[str, np.ndarray]] = {}
    rankings: dict[str, np.ndarray] = {}
    for name, scores in system_scores.items():
        rankings[name] = _stable_candidate_rankings(
            candidates, scores, retrieval.corpus_ids
        )
        results[name], observations[name] = _ranking_observations(
            rankings[name], retrieval.relevant_indices
        )
    comparisons: dict[str, dict[str, dict[str, float | int]]] = {}
    for baseline, treatment in (
        ("frozen_bge_reranker", "learned_text_residual"),
        ("learned_text_residual", "learned_graph_residual"),
        ("learned_sham_graph_residual", "learned_graph_residual"),
        ("frozen_bge_reranker", "learned_graph_residual"),
        ("graph_all_edges_removed", "learned_graph_residual"),
        ("graph_local_edges_removed", "learned_graph_residual"),
        ("graph_vbpl_causal_labeled_edges_removed", "learned_graph_residual"),
        ("graph_vbpl_context_edges_removed", "learned_graph_residual"),
        ("graph_messages_query_shuffled", "learned_graph_residual"),
    ):
        comparisons[f"{treatment}_minus_{baseline}"] = {
            metric: paired_bootstrap_difference(
                observations[baseline][metric],
                observations[treatment][metric],
                samples=args.bootstrap_samples,
                seed=args.bootstrap_seed,
            )
            for metric in ("ndcg_at_10", "recall_at_2", "query_f2_at_2")
        }

    gold_in_pool = labels.sum(axis=1)
    semantic_suffix = (
        f"-semantic-rank{args.semantic_rank}" if args.semantic_rank > 0 else ""
    )
    prediction_path = args.cache_dir / (
        f"drill-{args.candidate_mode}-graph{semantic_suffix}-oof-predictions.npz"
    )
    np.savez_compressed(
        prediction_path,
        candidates=candidates,
        folds=folds,
        labels=labels,
        **{f"scores_{name}": values for name, values in system_scores.items()},
        **{
            f"scores_{name}_seed_{seed}": values[index]
            for name, values in system_seed_scores.items()
            for index, seed in enumerate(args.seeds)
        },
        **{
            f"selected_{name}_seed_{seed}": values[index]
            for name, values in system_seed_decisions.items()
            for index, seed in enumerate(args.seeds)
        },
    )
    text_credit = comparisons[
        "learned_graph_residual_minus_learned_text_residual"
    ]["ndcg_at_10"]
    sham_credit = comparisons[
        "learned_graph_residual_minus_learned_sham_graph_residual"
    ]["ndcg_at_10"]
    causal_credit = comparisons[
        "learned_graph_residual_minus_graph_vbpl_causal_labeled_edges_removed"
    ]["ndcg_at_10"]
    causal_changes = int(
        np.count_nonzero(
            rankings["learned_graph_residual"]
            != rankings["graph_vbpl_causal_labeled_edges_removed"]
        )
    )
    graph_credit = bool(
        text_credit["observed_difference"] > 0
        and text_credit["confidence_low"] > 0
        and sham_credit["observed_difference"] > 0
        and sham_credit["confidence_low"] > 0
        and causal_credit["observed_difference"] > 0
        and causal_credit["confidence_low"] > 0
        and causal_changes > 0
    )
    receipt = {
        "schema_version": 1,
        "artifact_type": "drill_law_family_disjoint_learned_graph_evaluation",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "claim_status": (
            "learned_graph_incremental_utility_passed"
            if graph_credit
            else "learned_graph_incremental_utility_not_established"
        ),
        "benchmark": {
            "name": "VLSP 2025 DRiLL public training split",
            "semantics": "expert-refined real Vietnamese legal questions to entailing statutory articles",
            "redistribution": "raw organizer/mirror data stays untracked; receipt contains hashes and aggregates only",
            "dataset_audit": retrieval.audit,
        },
        "graph": graph.audit,
        "candidate_pool": {
            "mode": args.candidate_mode,
            "seed_depth": seed_depth,
            "depth": args.candidate_depth,
            "queries_with_at_least_one_gold": int(np.count_nonzero(gold_in_pool)),
            "queries_with_all_gold": int(np.count_nonzero(gold_in_pool == total_gold)),
            "micro_gold_recall": float(gold_in_pool.sum() / total_gold.sum()),
            "macro_gold_recall": float(np.mean(gold_in_pool / total_gold)),
        },
        "split_protocol": fold_audit,
        "configuration": {
            "bge_m3": {
                "model": args.model_name,
                "revision": args.model_revision,
                "max_length": args.max_length,
                "pooling": "normalized_cls",
            },
            "cross_encoder": {
                "model": args.reranker_name,
                "revision": args.reranker_revision,
                "max_length": args.reranker_max_length,
                "frozen": True,
                "pair_scores_reused_from_dense_run": reused_pair_count,
            },
            "candidate_generation": {
                "mode": args.candidate_mode,
                "seed_depth": seed_depth,
                "graph_expansion_relations": (
                    ["ARTICLE_CITES", "ARTICLE_CITED_BY", "ADJACENT"]
                    if args.candidate_mode == "rrf_graph"
                    else []
                ),
                "bm25_k1": args.bm25_k1 if bm25_candidates is not None else None,
                "bm25_b": args.bm25_b if bm25_candidates is not None else None,
                "rrf_constant": (
                    args.rrf_constant if bm25_candidates is not None else None
                ),
            },
            "residual": {
                "architecture": (
                    "frozen-embedding low-rank bilinear adaptor plus two-layer MLP residual over typed query-conditioned one-hop graph messages"
                    if args.semantic_rank > 0
                    else "two-layer MLP residual over typed query-conditioned one-hop graph messages"
                ),
                "semantic_projection_rank": args.semantic_rank,
                "semantic_backbone_frozen": True,
                "text_feature_names": list(text_feature_names),
                "graph_feature_count": len(graph_feature_names),
                "text_control_input_width": len(text_feature_names),
                "graph_and_sham_input_width": int(graph_features.shape[-1]),
                "sham_seed": args.sham_seed,
                "sham_definition": "same graph-message tensor and parameter count, deterministically permuted across queries before nested training",
                "hidden_size": args.hidden_size,
                "residual_cap": args.residual_cap,
                "learning_rate": args.learning_rate,
                "maximum_epochs": args.epochs,
                "early_stopping_patience": args.patience,
                "seeds": list(args.seeds),
                "no_hyperparameter_sweep": True,
            },
            "vbpl_derived_relation_types": list(VBPL_DERIVED_RELATION_TYPES),
            "reranker_runtime_this_run_seconds": reranker_runtime,
            "device": str(device),
        },
        "results": results,
        "official_f2_decisions": official_f2_decisions,
        "seed_stability": seed_stability,
        "fold_stability": {
            str(fold): {
                "query_count": int(np.count_nonzero(folds == fold)),
                "graph_minus_text_ndcg_at_10": float(
                    np.mean(
                        observations["learned_graph_residual"]["ndcg_at_10"][
                            folds == fold
                        ]
                        - observations["learned_text_residual"]["ndcg_at_10"][
                            folds == fold
                        ]
                    )
                ),
            }
            for fold in range(args.fold_count)
        },
        "paired_bootstrap": comparisons,
        "causal_diagnostics": {
            "vbpl_causal_labeled_edge_ablation_changed_ranking_positions": causal_changes,
            "all_edge_ablation_changed_ranking_positions": int(
                np.count_nonzero(
                    rankings["learned_graph_residual"]
                    != rankings["graph_all_edges_removed"]
                )
            ),
            "query_shuffled_messages_changed_ranking_positions": int(
                np.count_nonzero(
                    rankings["learned_graph_residual"]
                    != rankings["graph_messages_query_shuffled"]
                )
            ),
        },
        "credit_gate": {
            "passed": graph_credit,
            "requirements": [
                "learned graph beats the shared-backbone learned text residual on law-family-disjoint OOF nDCG@10",
                "learned graph beats a matched-parameter sham graph on law-family-disjoint OOF nDCG@10",
                "paired 95% bootstrap interval excludes zero in the positive direction",
                "removing VBPL-derived causal-labeled edges causes a positive nDCG loss with a paired interval excluding zero",
            ],
            "observed_graph_minus_text": text_credit,
            "observed_graph_minus_matched_sham": sham_credit,
            "observed_graph_minus_causal_edge_ablation": causal_credit,
        },
        "training_receipts": training_receipts,
        "artifacts": {
            "corpus": {"path": str(args.corpus), "sha256": sha256_file(args.corpus)},
            "questions": {
                "path": str(args.questions),
                "sha256": sha256_file(args.questions),
            },
            "metadata": {
                "path": str(args.metadata),
                "sha256": sha256_file(args.metadata),
            },
            "relationships": {
                "path": str(args.relationships),
                "sha256": sha256_file(args.relationships),
            },
            "corpus_embeddings": str(corpus_cache),
            "query_embeddings": str(query_cache),
            "candidates": str(candidates_path),
            "cross_encoder_scores": str(score_path),
            "oof_predictions": {
                "path": str(prediction_path),
                "sha256": sha256_file(prediction_path),
            },
        },
        "limitations": [
            "The official DRiLL public/private test labels are not present, so learned results are law-family-disjoint out-of-fold estimates on the public training questions.",
            "The relationship enrichment is a pinned VBPL-derived community mirror; canonical edge directions are heuristic discovery evidence, not expert-verified event truth, and are not admissible under the original DRiLL no-external-data rule.",
            "This static task tests expert article relevance; temporal update claims still require a separate version/event evaluation.",
            "Three fixed seeds quantify optimization variation but are not independent dataset replications.",
        ],
        "environment": {
            "python": platform.python_version(),
            "numpy": np.__version__,
            "torch": torch.__version__,
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return receipt


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--corpus",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/drill-public/legal-corpus/data/train-00000-of-00001.parquet",
    )
    parser.add_argument(
        "--questions",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/drill-public/train/data/train-00000-of-00001.parquet",
    )
    parser.add_argument(
        "--metadata",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/vietnamese-legal-documents-source/data/metadata.parquet",
    )
    parser.add_argument(
        "--relationships",
        type=Path,
        default=REPO_ROOT
        / "tmp/telen/bootstrap-audit/canonical_relationships.parquet",
    )
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=REPO_ROOT / "tmp/telen/static-benchmarks/cache/drill",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=REPO_ROOT
        / "research/results/telen_drill_graph_oof_evaluation_2026-08-29.json",
    )
    parser.add_argument("--model-name", default=MODEL_NAME)
    parser.add_argument("--model-revision", default=MODEL_REVISION)
    parser.add_argument("--reranker-name", default=RERANKER_NAME)
    parser.add_argument("--reranker-revision", default=RERANKER_REVISION)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--max-length", type=int, default=2048)
    parser.add_argument("--reranker-max-length", type=int, default=512)
    parser.add_argument("--corpus-batch-size", type=int, default=8)
    parser.add_argument("--query-batch-size", type=int, default=24)
    parser.add_argument("--reranker-batch-size", type=int, default=24)
    parser.add_argument("--score-batch-size", type=int, default=128)
    parser.add_argument("--candidate-depth", type=int, default=100)
    parser.add_argument(
        "--candidate-mode",
        choices=("dense", "rrf", "rrf_graph"),
        default="rrf",
    )
    parser.add_argument("--seed-depth", type=int, default=100)
    parser.add_argument("--bm25-k1", type=float, default=1.5)
    parser.add_argument("--bm25-b", type=float, default=0.75)
    parser.add_argument("--rrf-constant", type=int, default=60)
    parser.add_argument("--fold-count", type=int, default=5)
    parser.add_argument("--seeds", type=int, nargs="+", default=list(DEFAULT_SEEDS))
    parser.add_argument("--hidden-size", type=int, default=64)
    parser.add_argument(
        "--semantic-rank",
        type=int,
        default=0,
        help="rank of a learned bilinear adaptor over frozen BGE embeddings; zero disables it",
    )
    parser.add_argument("--residual-cap", type=float, default=3.0)
    parser.add_argument("--learning-rate", type=float, default=1e-3)
    parser.add_argument("--epochs", type=int, default=40)
    parser.add_argument("--patience", type=int, default=6)
    parser.add_argument("--training-query-batch-size", type=int, default=64)
    parser.add_argument("--maximum-returned-articles", type=int, default=10)
    parser.add_argument("--bootstrap-samples", type=int, default=10_000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260829)
    parser.add_argument("--sham-seed", type=int, default=314159)
    return parser.parse_args()


if __name__ == "__main__":
    payload = run(parse_args())
    print(json.dumps({"claim_status": payload["claim_status"], "results": payload["results"]}, indent=2))