#!/usr/bin/env python3 """Audit and evaluate TELEN on the licensed TVPL retrieval benchmark.""" from __future__ import annotations import argparse import hashlib import json import platform import sys import time from datetime import datetime, timezone from pathlib import Path import numpy as np import torch import torch.nn.functional as F from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer REPO_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(REPO_ROOT / "src")) from telen.static_benchmark import ( # noqa: E402 SparseBM25, dense_top_k, evaluate_rankings, expected_random_recall_at_k, load_tvpl_retrieval, ordered_text_sha256, paired_bootstrap_difference, reorder_ranking_prefix, reciprocal_rank_fusion, sha256_file, ) from telen.legacy.config import TELENConfig # noqa: E402 from telen.legacy.model import create_model # noqa: E402 from telen.adapters import StableTELENTextEmbedder # noqa: E402 DATASET_ID = "GreenNode/TVPL-Retrieval-VN" DATASET_REVISION = "67c901b65bfeb4e4b6e96eb123d1419c68ab90e8" DATASET_LICENSE = "CC-BY-SA-4.0" RERANKER_NAME = "BAAI/bge-reranker-v2-m3" RERANKER_REVISION = "953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e" TELEN_CHECKPOINT = REPO_ROOT / "data/checkpoints/telen/telen_best.pt" def resolve_device(requested: str) -> torch.device: if requested == "auto": return torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device(requested) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA was requested but is not available") return device def encode_bge_m3( texts: tuple[str, ...], *, model_name: str, revision: str, device: torch.device, max_length: int, batch_size: int, ) -> np.ndarray: tokenizer = AutoTokenizer.from_pretrained( model_name, revision=revision, local_files_only=True ) dtype = torch.float16 if device.type == "cuda" else torch.float32 model = AutoModel.from_pretrained( model_name, revision=revision, local_files_only=True, dtype=dtype, ).to(device) model.eval() vectors: list[torch.Tensor] = [] with torch.inference_mode(): for start in range(0, len(texts), batch_size): encoded = tokenizer( texts[start : start + batch_size], padding=True, truncation=True, max_length=max_length, return_tensors="pt", ).to(device) representation = model(**encoded).last_hidden_state[:, 0] vectors.append(F.normalize(representation.float(), p=2, dim=1).cpu()) del model if device.type == "cuda": torch.cuda.empty_cache() return torch.cat(vectors).numpy() def embedding_cache_path( cache_dir: Path, *, role: str, text_hash: str, model_revision: str, max_length: int, ) -> Path: return cache_dir / ( f"bge-m3-{role}-{model_revision[:12]}-{text_hash[:12]}-ml{max_length}.npy" ) def load_or_encode( texts: tuple[str, ...], ids: tuple[str, ...], *, role: str, cache_dir: Path, model_name: str, revision: str, device: torch.device, max_length: int, batch_size: int, ) -> tuple[np.ndarray, Path]: text_hash = ordered_text_sha256(ids, texts) path = embedding_cache_path( cache_dir, role=role, text_hash=text_hash, model_revision=revision, max_length=max_length, ) if path.is_file(): return np.load(path), path path.parent.mkdir(parents=True, exist_ok=True) embeddings = encode_bge_m3( texts, model_name=model_name, revision=revision, device=device, max_length=max_length, batch_size=batch_size, ) np.save(path, embeddings) return embeddings, path def load_or_encode_telen_base( texts: tuple[str, ...], ids: tuple[str, ...], *, role: str, cache_dir: Path, checkpoint: Path, device: torch.device, batch_size: int, ) -> tuple[np.ndarray, Path, str, str]: if not checkpoint.is_file(): raise FileNotFoundError(f"missing TELEN checkpoint: {checkpoint}") checkpoint_hash = sha256_file(checkpoint) config = TELENConfig() text_hash = ordered_text_sha256(ids, texts) path = cache_dir / ( f"predecessor-dense-{role}-{checkpoint_hash[:12]}-" f"{config.backbone_revision[:12]}-{text_hash[:12]}-" f"ml{config.max_seq_length}.npy" ) if path.is_file(): return np.load(path), path, checkpoint_hash, config.backbone_revision model = create_model(config).to(device) checkpoint_payload = torch.load( checkpoint, map_location=device, weights_only=False ) model.base_projection.load_state_dict(checkpoint_payload["base_projection"]) if "projection_norm" in checkpoint_payload: model.proj_norm.load_state_dict(checkpoint_payload["projection_norm"]) model.eval() embedder = StableTELENTextEmbedder(model, batch_size=batch_size) embeddings = embedder.encode(texts).numpy() del embedder, model, checkpoint_payload if device.type == "cuda": torch.cuda.empty_cache() path.parent.mkdir(parents=True, exist_ok=True) np.save(path, embeddings) return embeddings, path, checkpoint_hash, config.backbone_revision def save_rankings(path: Path, rankings: np.ndarray) -> None: path.parent.mkdir(parents=True, exist_ok=True) np.save(path, rankings) def rerank_with_cross_encoder( dataset, initial_rankings: np.ndarray, *, model_name: str, revision: str, device: torch.device, top_n: int, max_length: int, batch_size: int, score_cache: Path | None = None, ) -> tuple[np.ndarray, float]: if not 0 < top_n <= initial_rankings.shape[1]: raise ValueError("reranker top_n must be within the initial ranking depth") if batch_size < 1: raise ValueError("reranker batch size must be positive") total_pairs = len(dataset.query_texts) * top_n completed = 0 progress_path = None if score_cache is not None: score_cache.parent.mkdir(parents=True, exist_ok=True) progress_path = score_cache.with_suffix(score_cache.suffix + ".progress") if score_cache.is_file(): scores = np.lib.format.open_memmap(score_cache, mode="r+") if scores.shape != (total_pairs,) or scores.dtype != np.float32: raise ValueError("reranker score cache has an incompatible shape") if progress_path.is_file(): completed = int(progress_path.read_text(encoding="utf-8").strip()) else: scores = np.lib.format.open_memmap( score_cache, mode="w+", dtype=np.float32, shape=(total_pairs,) ) progress_path.write_text("0\n", encoding="utf-8") else: scores = np.empty(total_pairs, dtype=np.float32) if not 0 <= completed <= total_pairs: raise ValueError("reranker progress marker is out of range") started = time.perf_counter() if completed < total_pairs: 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(completed, total_pairs, batch_size): stop = min(start + batch_size, total_pairs) pairs: list[tuple[str, str]] = [] for flat_index in range(start, stop): query_index, local_rank = divmod(flat_index, top_n) document_index = int(initial_rankings[query_index, local_rank]) pairs.append( ( dataset.query_texts[query_index], dataset.corpus_texts[document_index], ) ) encoded = tokenizer( pairs, padding=True, truncation=True, max_length=max_length, return_tensors="pt", ).to(device) scores[start:stop] = ( model(**encoded, return_dict=True) .logits.view(-1) .float() .cpu() .numpy() ) if progress_path is not None: scores.flush() progress_path.write_text(f"{stop}\n", encoding="utf-8") del model if device.type == "cuda": torch.cuda.empty_cache() elapsed = time.perf_counter() - started reranked = reorder_ranking_prefix( initial_rankings, scores.reshape(len(dataset.query_texts), top_n), dataset.corpus_ids, ) return reranked, elapsed def parse_systems(value: str) -> tuple[str, ...]: systems = tuple(item.strip() for item in value.split(",") if item.strip()) allowed = {"bm25", "bge_m3", "rrf", "bge_reranker", "predecessor_dense"} unknown = set(systems) - allowed if unknown: raise argparse.ArgumentTypeError(f"unknown systems: {sorted(unknown)}") if "rrf" in systems and not {"bm25", "bge_m3"}.issubset(systems): raise argparse.ArgumentTypeError("rrf requires bm25 and bge_m3") if "bge_reranker" in systems and "bge_m3" not in systems: raise argparse.ArgumentTypeError("bge_reranker requires bge_m3") return systems def run(args: argparse.Namespace) -> dict: dataset = load_tvpl_retrieval(args.dataset_root) if args.retrieval_depth < 100: raise ValueError("retrieval depth must be at least 100") if args.retrieval_depth > len(dataset.corpus_ids): raise ValueError("retrieval depth exceeds the corpus size") rankings: dict[str, np.ndarray] = {} cache_artifacts: dict[str, str] = {} runtime_seconds: dict[str, float] = {} telen_configuration = None if "bm25" in args.systems: bm25 = SparseBM25(k1=args.bm25_k1, b=args.bm25_b) bm25.fit(dataset.corpus_texts) rankings["bm25"] = bm25.rank( dataset.query_texts, dataset.corpus_ids, k=args.retrieval_depth, batch_size=args.score_batch_size, ) path = args.cache_dir / f"tvpl-bm25-top{args.retrieval_depth}.npy" save_rankings(path, rankings["bm25"]) cache_artifacts["bm25_rankings"] = str(path) device = None if "bge_m3" in args.systems: device = resolve_device(args.device) corpus_embeddings, corpus_cache = load_or_encode( dataset.corpus_texts, dataset.corpus_ids, role="corpus", 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( dataset.query_texts, dataset.query_ids, role="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, ) rankings["bge_m3"] = dense_top_k( query_embeddings, corpus_embeddings, dataset.corpus_ids, k=args.retrieval_depth, batch_size=args.score_batch_size, ) path = args.cache_dir / f"tvpl-bge-m3-top{args.retrieval_depth}.npy" save_rankings(path, rankings["bge_m3"]) cache_artifacts.update( { "corpus_embeddings": str(corpus_cache), "query_embeddings": str(query_cache), "bge_m3_rankings": str(path), } ) if "predecessor_dense" in args.systems: telen_device = resolve_device(args.device) telen_started = time.perf_counter() ( telen_corpus, telen_corpus_cache, checkpoint_hash, backbone_revision, ) = load_or_encode_telen_base( dataset.corpus_texts, dataset.corpus_ids, role="corpus", cache_dir=args.cache_dir, checkpoint=args.telen_checkpoint, device=telen_device, batch_size=args.telen_batch_size, ) ( telen_queries, telen_query_cache, query_checkpoint_hash, query_backbone_revision, ) = load_or_encode_telen_base( dataset.query_texts, dataset.query_ids, role="queries", cache_dir=args.cache_dir, checkpoint=args.telen_checkpoint, device=telen_device, batch_size=args.telen_batch_size, ) if (checkpoint_hash, backbone_revision) != ( query_checkpoint_hash, query_backbone_revision, ): raise RuntimeError("TELEN corpus/query encoders do not share one checkpoint") rankings["predecessor_dense"] = dense_top_k( telen_queries, telen_corpus, dataset.corpus_ids, k=args.retrieval_depth, batch_size=args.score_batch_size, ) path = args.cache_dir / f"tvpl-predecessor-dense-top{args.retrieval_depth}.npy" save_rankings(path, rankings["predecessor_dense"]) runtime_seconds["predecessor_dense_encode_and_rank"] = ( time.perf_counter() - telen_started ) cache_artifacts.update( { "predecessor_dense_corpus_embeddings": str(telen_corpus_cache), "predecessor_dense_query_embeddings": str(telen_query_cache), "predecessor_dense_rankings": str(path), } ) telen_configuration = { "checkpoint": str(args.telen_checkpoint), "checkpoint_sha256": checkpoint_hash, "backbone": TELENConfig().backbone, "backbone_revision": backbone_revision, "max_length": TELENConfig().max_seq_length, "projection": "base_projection_then_projection_norm_without_graph_or_hypernetwork", "device": str(telen_device), } if "rrf" in args.systems: rankings["rrf_bm25_bge_m3"] = reciprocal_rank_fusion( (rankings["bm25"], rankings["bge_m3"]), dataset.corpus_ids, k=args.retrieval_depth, constant=args.rrf_constant, ) path = args.cache_dir / f"tvpl-rrf-top{args.retrieval_depth}.npy" save_rankings(path, rankings["rrf_bm25_bge_m3"]) cache_artifacts["rrf_rankings"] = str(path) reranker_device = None if "bge_reranker" in args.systems: reranker_device = resolve_device(args.device) source = rankings["bge_m3"] source_hash = hashlib.sha256(source.tobytes()).hexdigest()[:12] path = args.cache_dir / ( f"tvpl-bge-reranker-v2-m3-{args.reranker_revision[:12]}-" f"source{source_hash}-top{args.reranker_top_n}-" f"ml{args.reranker_max_length}.npy" ) score_path = path.with_name(path.stem + "-scores.npy") if path.is_file(): rankings["bge_m3_reranker_v2_m3"] = np.load(path) runtime_seconds["bge_m3_reranker_v2_m3_inference_this_run"] = 0.0 runtime_seconds["bge_m3_reranker_ranking_cache_hit"] = True else: ( rankings["bge_m3_reranker_v2_m3"], runtime_seconds["bge_m3_reranker_v2_m3_inference_this_run"], ) = rerank_with_cross_encoder( dataset, source, model_name=args.reranker_name, revision=args.reranker_revision, device=reranker_device, top_n=args.reranker_top_n, max_length=args.reranker_max_length, batch_size=args.reranker_batch_size, score_cache=score_path, ) save_rankings(path, rankings["bge_m3_reranker_v2_m3"]) runtime_seconds["bge_m3_reranker_ranking_cache_hit"] = False cache_artifacts["bge_reranker_rankings"] = str(path) cache_artifacts["bge_reranker_scores"] = str(score_path) results: dict[str, dict] = {} observations: dict[str, dict[str, np.ndarray]] = {} for name, values in rankings.items(): results[name], observations[name] = evaluate_rankings( values, dataset.relevant_indices ) comparisons: dict[str, dict] = {} if "bm25" in observations: for treatment in ("bge_m3", "rrf_bm25_bge_m3", "predecessor_dense"): if treatment not in observations: continue comparisons[f"{treatment}_minus_bm25"] = { metric: paired_bootstrap_difference( observations["bm25"][metric], observations[treatment][metric], samples=args.bootstrap_samples, seed=args.seed, ) for metric in ( "ndcg_at_10", "average_precision_at_100", "recall_at_100", ) } if "bge_m3" in observations: for treatment in ( "rrf_bm25_bge_m3", "bge_m3_reranker_v2_m3", "predecessor_dense", ): if treatment not in observations: continue comparisons[f"{treatment}_minus_bge_m3"] = { metric: paired_bootstrap_difference( observations["bge_m3"][metric], observations[treatment][metric], samples=args.bootstrap_samples, seed=args.seed, ) for metric in ( "ndcg_at_10", "average_precision_at_100", "recall_at_100", ) } receipt = { "artifact_type": "static_human_relevance_retrieval_benchmark", "created_at": datetime.now(timezone.utc).isoformat(), "benchmark": { "dataset_id": DATASET_ID, "revision": DATASET_REVISION, "license": DATASET_LICENSE, "task_semantics": ( "natural Vietnamese legal questions to human-associated answer passages; " "static snapshot, positive-only binary qrels" ), "limitations": [ "The MTEB conversion exposes positive-only qrels, so unjudged documents are treated as nonrelevant.", "The benchmark is a static snapshot and cannot establish version-aware adaptation.", "The dataset card declares human annotation but does not expose annotator agreement or adjudication metadata.", "This licensed fallback is not the official DRiLL or ALQAC split.", ], "audit": dataset.audit, }, "configuration": { "systems": list(args.systems), "retrieval_depth": args.retrieval_depth, "bm25": {"k1": args.bm25_k1, "b": args.bm25_b}, "bge_m3": { "model": args.model_name, "revision": args.model_revision, "max_length": args.max_length, "pooling": "normalized_cls", "device": str(device) if device is not None else None, }, "rrf_constant": args.rrf_constant, "bge_reranker_v2_m3": { "model": args.reranker_name, "revision": args.reranker_revision, "candidate_source": "bge_m3", "top_n": args.reranker_top_n, "max_length": args.reranker_max_length, "device": str(reranker_device) if reranker_device is not None else None, }, "predecessor_dense": telen_configuration, "bootstrap_samples": args.bootstrap_samples, "seed": args.seed, }, "sanity_floors": { "expected_random_recall_at_20": expected_random_recall_at_k( dataset.relevant_indices, len(dataset.corpus_ids), 20 ), "expected_random_recall_at_100": expected_random_recall_at_k( dataset.relevant_indices, len(dataset.corpus_ids), 100 ), }, "results": results, "paired_bootstrap": comparisons, "cache_artifacts": cache_artifacts, "runtime_seconds": runtime_seconds, "environment": { "python": platform.python_version(), "numpy": np.__version__, "torch": torch.__version__, "cuda_available": torch.cuda.is_available(), "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, }, } 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 main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--dataset-root", type=Path, default=REPO_ROOT / "tmp/telen/static-benchmarks/tvpl-retrieval-vn", ) parser.add_argument( "--output", type=Path, default=REPO_ROOT / "research/results/telen_static_benchmark_2026-08-29.json", ) parser.add_argument( "--cache-dir", type=Path, default=REPO_ROOT / "tmp/telen/static-benchmarks/cache", ) parser.add_argument("--systems", type=parse_systems, default=("bm25",)) parser.add_argument("--retrieval-depth", type=int, default=1000) parser.add_argument("--bm25-k1", type=float, default=1.5) parser.add_argument("--bm25-b", type=float, default=0.75) parser.add_argument("--model-name", default="BAAI/bge-m3") parser.add_argument( "--model-revision", default="5617a9f61b028005a4858fdac845db406aefb181", ) parser.add_argument("--device", default="auto") parser.add_argument("--max-length", type=int, default=2048) parser.add_argument("--corpus-batch-size", type=int, default=4) parser.add_argument("--query-batch-size", type=int, default=32) parser.add_argument("--score-batch-size", type=int, default=256) parser.add_argument("--rrf-constant", type=int, default=60) parser.add_argument("--reranker-name", default=RERANKER_NAME) parser.add_argument("--reranker-revision", default=RERANKER_REVISION) parser.add_argument("--reranker-top-n", type=int, default=20) parser.add_argument("--reranker-max-length", type=int, default=512) parser.add_argument("--reranker-batch-size", type=int, default=24) parser.add_argument("--telen-checkpoint", type=Path, default=TELEN_CHECKPOINT) parser.add_argument("--telen-batch-size", type=int, default=32) parser.add_argument("--bootstrap-samples", type=int, default=10_000) parser.add_argument("--seed", type=int, default=42) args = parser.parse_args() result = run(args) print(json.dumps(result, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()