scripts/run_static_benchmark.py

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