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