#!/usr/bin/env python3 """Post-hoc subgroup audit for the frozen DRiLL graph OOF predictions. This script does not fit or select a model. It asks whether the observed graph gain is confined to the easy document-clustering proxy, and whether queries whose candidate pools expose VBPL-derived causal-labelled links behave differently from queries for which those links are unavailable. """ from __future__ import annotations import argparse import json import platform import sys from datetime import datetime, timezone from pathlib import Path from typing import Mapping, Sequence import numpy as np REPO_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(REPO_ROOT / "src")) from scripts.run_drill_graph_evaluation import ( # noqa: E402 _ranking_observations, _stable_candidate_rankings, ) from telen.drill_graph import ( # noqa: E402 DrillLegalGraph, build_drill_legal_graph, load_drill_graph_dataset, ) from telen.static_benchmark import ( # noqa: E402 paired_bootstrap_difference, sha256_file, ) CAUSAL_RELATIONS = ("AMENDS", "REPEALS", "REPLACES", "TERMINATES_VERSION") def causal_candidate_masks( graph: DrillLegalGraph, candidates: np.ndarray, relevant_indices: Sequence[frozenset[int]], ) -> tuple[np.ndarray, np.ndarray]: """Return query masks for exposed causal links and gold-incident links.""" if candidates.ndim != 2 or len(candidates) != len(relevant_indices): raise ValueError("candidate matrix and qrels must align") candidate_exposed = np.zeros(len(candidates), dtype=bool) gold_incident = np.zeros(len(candidates), dtype=bool) reverse: dict[str, dict[str, set[str]]] = { relation: {} for relation in CAUSAL_RELATIONS } for relation in CAUSAL_RELATIONS: for source, targets in graph.document_neighbors[relation].items(): for target in targets: reverse[relation].setdefault(target, set()).add(source) for query_index, (row, relevant) in enumerate( zip(candidates, relevant_indices, strict=True) ): candidate_documents = { graph.article_document_ids[int(aid)] for aid in row } gold_documents = { graph.article_document_ids[int(aid)] for aid in relevant } for relation in CAUSAL_RELATIONS: outgoing = graph.document_neighbors[relation] incoming = reverse[relation] for document in candidate_documents: linked_candidates = ( set(outgoing.get(document, ())) | incoming.get(document, set()) ) & candidate_documents if linked_candidates: candidate_exposed[query_index] = True if document in gold_documents or linked_candidates & gold_documents: gold_incident[query_index] = True if candidate_exposed[query_index] and gold_incident[query_index]: break return candidate_exposed, gold_incident def build_subgroup_masks( article_document_ids: Sequence[str], relevant_indices: Sequence[frozenset[int]], causal_exposed: np.ndarray, causal_gold_incident: np.ndarray, ) -> Mapping[str, np.ndarray]: """Create mutually interpretable, label-audited query strata.""" if not ( len(relevant_indices) == len(causal_exposed) == len(causal_gold_incident) ): raise ValueError("subgroup inputs must align") qrel_counts = np.asarray(list(map(len, relevant_indices)), dtype=np.int64) gold_document_counts = np.asarray( [ len({article_document_ids[int(aid)] for aid in relevant}) for relevant in relevant_indices ], dtype=np.int64, ) all_queries = np.ones(len(relevant_indices), dtype=bool) return { "all_queries": all_queries, "single_qrel": qrel_counts == 1, "multiple_qrels": qrel_counts > 1, "multiple_qrels_single_gold_document": (qrel_counts > 1) & (gold_document_counts == 1), "multiple_gold_documents": gold_document_counts > 1, "causal_link_exposed_in_candidate_pool": causal_exposed, "no_causal_link_exposed_in_candidate_pool": ~causal_exposed, "gold_document_incident_to_exposed_causal_link": causal_gold_incident, "gold_document_not_incident_to_exposed_causal_link": ~causal_gold_incident, } def summarize_subgroups( observations: Mapping[str, np.ndarray], masks: Mapping[str, np.ndarray], *, bootstrap_samples: int, bootstrap_seed: int, ) -> dict[str, object]: """Summarize fixed graph/text/sham nDCG observations by subgroup.""" required = {"graph", "text", "sham", "causal_ablation"} if set(observations) != required: raise ValueError(f"observations must contain {sorted(required)}") size = len(observations["graph"]) if any(len(values) != size for values in observations.values()): raise ValueError("observation vectors must align") output: dict[str, object] = {} for name, mask in masks.items(): indices = np.flatnonzero(mask) if not len(indices): output[name] = {"query_count": 0} continue graph = observations["graph"][indices] text = observations["text"][indices] sham = observations["sham"][indices] causal_ablation = observations["causal_ablation"][indices] output[name] = { "query_count": int(len(indices)), "mean_ndcg_at_10": { "learned_text_residual": float(text.mean()), "learned_graph_residual": float(graph.mean()), "learned_sham_graph_residual": float(sham.mean()), "graph_vbpl_causal_labeled_edges_removed": float( causal_ablation.mean() ), }, "graph_minus_text": paired_bootstrap_difference( text, graph, samples=bootstrap_samples, seed=bootstrap_seed, ), "graph_minus_sham": paired_bootstrap_difference( sham, graph, samples=bootstrap_samples, seed=bootstrap_seed, ), "graph_minus_causal_edge_ablation": paired_bootstrap_difference( causal_ablation, graph, samples=bootstrap_samples, seed=bootstrap_seed, ), } return output def run(args: argparse.Namespace) -> dict[str, object]: dataset = load_drill_graph_dataset(args.corpus, args.questions) graph = build_drill_legal_graph( dataset, metadata_path=args.metadata, canonical_relationships_path=args.relationships, ) predictions = np.load(args.predictions) candidates = np.asarray(predictions["candidates"], dtype=np.int32) causal_exposed, causal_gold_incident = causal_candidate_masks( graph, candidates, dataset.retrieval.relevant_indices ) masks = build_subgroup_masks( dataset.article_document_ids, dataset.retrieval.relevant_indices, causal_exposed, causal_gold_incident, ) system_keys = { "graph": "scores_learned_graph_residual", "text": "scores_learned_text_residual", "sham": "scores_learned_sham_graph_residual", "causal_ablation": "scores_graph_vbpl_causal_labeled_edges_removed", } ndcg_observations: dict[str, np.ndarray] = {} for role, key in system_keys.items(): rankings = _stable_candidate_rankings( candidates, np.asarray(predictions[key], dtype=np.float32), dataset.retrieval.corpus_ids, ) _summary, per_query = _ranking_observations( rankings, dataset.retrieval.relevant_indices ) ndcg_observations[role] = per_query["ndcg_at_10"] subgroups = summarize_subgroups( ndcg_observations, masks, bootstrap_samples=args.bootstrap_samples, bootstrap_seed=args.bootstrap_seed, ) receipt: dict[str, object] = { "schema_version": 1, "artifact_type": "drill_graph_fixed_prediction_posthoc_subgroup_audit", "generated_at": datetime.now(timezone.utc).isoformat(), "claim_status": "exploratory_subgroup_analysis_not_model_selection", "protocol": { "models_refit": False, "hyperparameters_selected": False, "predictions_frozen_before_subgroup_definition": True, "multiplicity_adjustment": None, "interpretation": ( "descriptive heterogeneity audit; subgroup intervals do not replace " "the predeclared all-query graph credit gate" ), }, "dataset_audit": dataset.retrieval.audit, "graph_audit": graph.audit, "subgroups": subgroups, "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), }, "predictions": { "path": str(args.predictions), "sha256": sha256_file(args.predictions), }, }, "limitations": [ "Subgroups were defined after observing the primary aggregate result and are exploratory.", "VBPL-derived causal labels and canonical directions come from a pinned community mirror and are not expert-verified event truth.", "Candidate exposure depends on the shared RRF top-100 pool; absence of an exposed edge is not absence from the full legal graph.", ], "environment": {"python": platform.python_version(), "numpy": np.__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( "--predictions", type=Path, default=REPO_ROOT / "tmp/telen/static-benchmarks/cache/drill/drill-rrf-graph-semantic-rank64-oof-predictions.npz", ) parser.add_argument("--bootstrap-samples", type=int, default=10_000) parser.add_argument("--bootstrap-seed", type=int, default=20260829) parser.add_argument( "--output", type=Path, default=REPO_ROOT / "research/results/telen_drill_graph_subgroup_audit_2026-08-29.json", ) return parser.parse_args() if __name__ == "__main__": result = run(parse_args()) print(json.dumps(result, ensure_ascii=False, indent=2))