Reproducible TEI ModernBERT benchmark on Natural Questions

benchmark_gte_modernbert_nq.py

#!/usr/bin/env python3
"""Benchmark a TEI ModernBERT server on Natural Questions."""

from __future__ import annotations

import argparse
import json
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import requests
from datasets import load_dataset


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="http://127.0.0.1:8080/embed")
    parser.add_argument("--dataset", default="sentence-transformers/natural-questions")
    parser.add_argument("--dataset-config", default="pair")
    parser.add_argument("--split", default="train[:20000]")
    parser.add_argument("--text-column", default="answer")
    parser.add_argument("--batch-size", type=int, default=1024)
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--warmup-batches", type=int, default=4)
    parser.add_argument("--timeout", type=float, default=600.0)
    parser.add_argument("--label", default="tei")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--save-embeddings", type=Path)
    parser.add_argument("--compare-to", type=Path)
    return parser.parse_args()


def post_batch(url: str, texts: list[str], timeout: float) -> list[list[float]]:
    response = requests.post(
        url,
        json={"inputs": texts, "normalize": True, "truncate": True},
        timeout=timeout,
    )
    response.raise_for_status()
    return response.json()


def compare_embeddings(reference_path: Path, actual: np.ndarray) -> dict[str, float]:
    reference = np.load(reference_path)
    if reference.shape != actual.shape:
        raise ValueError(f"embedding shape mismatch: {reference.shape} != {actual.shape}")

    difference = np.abs(reference - actual)
    reference_norm = np.linalg.norm(reference, axis=1)
    actual_norm = np.linalg.norm(actual, axis=1)
    cosine = np.sum(reference * actual, axis=1) / (reference_norm * actual_norm)
    return {
        "min_cosine_similarity": float(np.min(cosine)),
        "mean_cosine_similarity": float(np.mean(cosine)),
        "max_cosine_similarity": float(np.max(cosine)),
        "max_absolute_difference": float(np.max(difference)),
        "mean_absolute_difference": float(np.mean(difference)),
        "max_norm_error": float(np.max(np.abs(actual_norm - 1.0))),
    }


def main() -> None:
    args = parse_args()
    dataset = load_dataset(args.dataset, args.dataset_config, split=args.split)
    texts = [str(text) for text in dataset[args.text_column]]
    batches = [texts[i : i + args.batch_size] for i in range(0, len(texts), args.batch_size)]

    for batch in batches[: args.warmup_batches]:
        post_batch(args.url, batch, args.timeout)

    started = time.perf_counter()
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        responses = list(
            executor.map(
                lambda batch: post_batch(args.url, batch, args.timeout),
                batches,
            )
        )
    elapsed = time.perf_counter() - started

    embeddings = np.asarray([row for response in responses for row in response], dtype=np.float32)
    result: dict[str, object] = {
        "label": args.label,
        "url": args.url,
        "dataset": args.dataset,
        "dataset_config": args.dataset_config,
        "split": args.split,
        "text_column": args.text_column,
        "num_texts": len(texts),
        "batch_size": args.batch_size,
        "workers": args.workers,
        "warmup_batches": args.warmup_batches,
        "embedding_shape": list(embeddings.shape),
        "elapsed_seconds": elapsed,
        "texts_per_second": len(texts) / elapsed,
    }

    if args.save_embeddings:
        args.save_embeddings.parent.mkdir(parents=True, exist_ok=True)
        np.save(args.save_embeddings, embeddings)
        result["embeddings_path"] = str(args.save_embeddings)

    if args.compare_to:
        result["embedding_comparison"] = compare_embeddings(args.compare_to, embeddings)

    serialized = json.dumps(result, indent=2)
    print(serialized)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(serialized + "\n")


if __name__ == "__main__":
    main()