import argparse import json import sys import time from collections import defaultdict from datetime import date from pathlib import Path from typing import Any REPO_ROOT = Path(__file__).resolve().parents[2] PACKAGE_PARENT = REPO_ROOT.parent if str(PACKAGE_PARENT) not in sys.path: sys.path.insert(0, str(PACKAGE_PARENT)) from rusBeIR.beir.datasets.data_loader_hf import HFDataLoader from rusBeIR.beir.retrieval.evaluation import EvaluateRetrieval from rusBeIR.beir.retrieval.search.base import BaseSearch DATASETS_PATH = REPO_ROOT/"leaderboard"/ "data"/"datasets.json" DEFAULT_OUTPUT = REPO_ROOT/"leaderboard"/"data"/"results.jsonl" def load_datasets() -> list[dict[str, Any]]: with DATASETS_PATH.open("r", encoding="utf-8") as file: return json.load(file) def select_device(requested: str) -> str: if requested != "auto": return requested try: import torch return "cuda" if torch.cuda.is_available() else "cpu" except Exception: return "cpu" def safe_name(value: str) -> str: return value.replace("/", "__").replace(":", "_") def raw_results_path(results_dir: Path, model_name: str, dataset: dict[str, Any]) -> Path: return results_dir/f"results_{model_name}_{dataset['name']}_{dataset['split']}.json" def merge_metrics(*metric_groups: dict[str, float]) -> dict[str, float]: merged: dict[str, float] = {} for group in metric_groups: merged.update(group) return merged def average_metrics(per_dataset: dict[str, dict[str, float]]) -> dict[str, float]: values_by_metric: dict[str, list[float]] = defaultdict(list) for metrics in per_dataset.values(): for key, value in metrics.items(): values_by_metric[key].append(float(value)) return { key: round(sum(values) / len(values), 5) for key, values in sorted(values_by_metric.items()) if values } def retrieve(model, corpus: dict[str, dict[str, str]], queries: dict[str, str], top_k: int, query_batch_size: int, num_workers: int) -> dict[str, dict[str, float]]: if isinstance(model, BaseSearch): return EvaluateRetrieval(retriever=model, k_values=[top_k]).retrieve(corpus, queries) return model.retrieve(queries, corpus, top_n=top_k, data_batch_size=query_batch_size, num_workers=num_workers) def build_dense_model(args: argparse.Namespace, device: str): from rusBeIR.retrieval.models.dense.DenseHFModels import DenseHFModels class ConfigurableDenseModel(DenseHFModels): def __init__(self, model_name: str, maxlen: int, batch_size: int, device: str, pooling_method: str, query_prefix: str, passage_prefix: str, model_sep: str, padding_side: str | None, model_loader: str): self.model_loader = model_loader super().__init__(model_name=model_name, maxlen=maxlen, batch_size=batch_size, device=device, model_sep=model_sep, padding_side=padding_side) self.pooling_method = pooling_method self.query_prefix = query_prefix self.passage_prefix = passage_prefix def load_model(self, model_name: str, device: str = "cuda"): if self.model_loader == "t5-encoder": from transformers import AutoTokenizer, T5EncoderModel model = T5EncoderModel.from_pretrained(model_name).to(device) tokenizer = AutoTokenizer.from_pretrained(model_name) model.eval() return model, tokenizer return super().load_model(model_name, device) def encode_queries(self, queries, pooling_method=None, prefix=None): return super().encode_queries(queries, pooling_method=pooling_method or self.pooling_method, prefix=self.query_prefix if prefix is None else prefix) def encode_corpus(self, corpus, pooling_method=None, prefix=None): return super().encode_corpus(corpus, pooling_method=pooling_method or self.pooling_method, prefix=self.passage_prefix if prefix is None else prefix) return ConfigurableDenseModel( model_name=args.model_id, maxlen=args.maxlen, batch_size=args.batch_size, device=device, pooling_method=args.pooling, query_prefix=args.query_prefix, passage_prefix=args.passage_prefix, model_sep=args.model_sep, padding_side=args.padding_side, model_loader=args.model_loader, ) def build_model(args: argparse.Namespace): if args.model_type == "dense": device = select_device(args.device) return build_dense_model(args, device), device if args.model_type == "reranker": from rusBeIR.beir.reranking import Rerank from rusBeIR.beir.reranking.models import CrossEncoder device = select_device(args.device) cross_encoder = CrossEncoder(args.model_id, device=device) return Rerank(cross_encoder, batch_size=args.rerank_batch_size, max_length=args.rerank_max_length), device if args.sparse_model == "bm25s": from rusBeIR.retrieval.models.sparse.bm25s import BM25s return BM25s(method=args.bm25_method, k1=args.bm25_k1, b=args.bm25_b), "cpu" if args.sparse_model == "tfidf": from rusBeIR.retrieval.models.sparse.tfidf import TfidfSearch return TfidfSearch( lowercase=not args.no_lowercase, ngram_range=(args.tfidf_min_ngram, args.tfidf_max_ngram), max_features=args.tfidf_max_features, ), "cpu" raise SystemExit(f"Unknown sparse model: {args.sparse_model}") def evaluate_dataset(model, model_name: str, dataset: dict[str, Any], k_values: list[int], text_type: str, query_batch_size: int, num_workers: int, raw_results_dir: Path | None, model_type: str, first_stage_results_dir: Path | None, first_stage_model_name: str | None, rerank_top_k: int | None) -> dict[str, float]: corpus, queries, qrels = HFDataLoader(hf_repo=dataset["hf_repo"], hf_repo_qrels=dataset["qrels_repo"], streaming=False, keep_in_memory=False, text_type=text_type).load(split=dataset["split"]) if model_type == "reranker": if first_stage_results_dir is None or first_stage_model_name is None: raise ValueError("Reranker evaluation requires first-stage results.") first_stage_path = raw_results_path(first_stage_results_dir, first_stage_model_name, dataset) if not first_stage_path.exists(): raise FileNotFoundError(f"First-stage results not found: {first_stage_path}") with first_stage_path.open("r", encoding="utf-8") as file: first_stage_results = json.load(file) results = model.rerank(corpus, queries, first_stage_results, top_k=rerank_top_k or max(k_values)) else: results = retrieve(model=model, corpus=corpus, queries=queries, top_k=max(k_values), query_batch_size=query_batch_size, num_workers=num_workers) if raw_results_dir is not None: raw_results_dir.mkdir(parents=True, exist_ok=True) output_path = raw_results_path(raw_results_dir, model_name, dataset) with output_path.open("w", encoding="utf-8") as file: json.dump(results, file, ensure_ascii=False) retriever = EvaluateRetrieval(k_values=k_values) ndcg, map_scores, recall, precision = retriever.evaluate(qrels=qrels, results=results, k_values=k_values) mrr = retriever.evaluate_custom(qrels, results, k_values, "mrr") return merge_metrics(ndcg, map_scores, recall, precision, mrr) def read_jsonl(path: Path) -> list[dict[str, Any]]: if not path.exists(): return[] records = [] with path.open("r", encoding="utf-8") as file: for line_n, line in enumerate(file, start=1): line = line.strip() if not line or line.startswith("#"): continue try: records.append(json.loads(line)) except json.JSONDecodeError as error: raise ValueError(f"Invalid JSON in {path} at line {line_n}: {error}") return records def find_existing_record(path: Path, model_id: str) -> dict[str, Any] | None: for record in reversed(read_jsonl(path)): if record.get("model_id") == model_id: return record return None def insert_jsonl_record(path: Path, record: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) records = [item for item in read_jsonl(path) if item.get("model_id") != record["model_id"]] records.append(record) tmp_path = path.with_name(f"{path.name}.tmp") with tmp_path.open("w", encoding="utf-8") as file: for item in records: file.write(json.dumps(item, ensure_ascii=False, sort_keys=True)) file.write("\n") tmp_path.replace(path) def build_result_record(args, record_model_id: str, model_name: str, organization: str, hardware: str, started: float, per_dataset: dict[str, dict[str, float]]) -> dict[str, Any]: record = { "model_id": record_model_id, "model_name": model_name, "organization": organization, "type": args.model_type, "date": date.today().isoformat(), "verified": False, "hardware": hardware, "runtime_seconds": round(time.time() - started, 2), "source_url": args.source_url, "scores": { "average": average_metrics(per_dataset), "datasets": per_dataset, }, "notes": args.notes, } if args.model_type == "reranker": record["base_model_id"] = args.first_stage_model_id record["rerank_top_k"] = args.rerank_top_k or max(args.k_values) return record def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--model-id", default="", help="Hugging Face model id for dense and reranker models. Defaults to sparse model name for sparse runs.") parser.add_argument("--model-name", default="", help="Display name. Defaults to the last segment of --model-id.") parser.add_argument("--model-type", default="dense", choices=["dense", "sparse", "reranker"]) parser.add_argument("--sparse-model", default="bm25s", choices=["bm25s", "tfidf"], help="Sparse retriever to use when --model-type sparse.") parser.add_argument("--first-stage-model-id", "--base-model-id", dest="first_stage_model_id", default="", help="Model id/name whose raw retrieval results will be reranked.") parser.add_argument("--first-stage-results-dir", "--base-results-dir", dest="first_stage_results_dir", type=Path, default=None, help="Directory with first-stage raw result JSON files.") parser.add_argument("--rerank-top-k", type=int, default=None, help="Number of first-stage hits to rerank. Defaults to max(k-values).") parser.add_argument("--rerank-batch-size", type=int, default=32) parser.add_argument("--rerank-max-length", type=int, default=512) parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="JSONL file to update with the leaderboard row.") parser.add_argument("--raw-results-dir", type=Path, default=None, help="Optional directory for raw retrieval results.") parser.add_argument("--datasets", nargs="*", default=None, help="Dataset names to evaluate. Defaults to all official datasets.") parser.add_argument("--limit-datasets", type=int, default=None, help="Evaluate only the first N selected datasets.") parser.add_argument("--text-type", choices=["processed_text", "text"], default="text") parser.add_argument("--k-values", nargs="+", type=int, default=[1, 3, 5, 10, 100]) parser.add_argument("--device", default="auto", help="auto, cpu, cuda, cuda:0, etc.") parser.add_argument("--maxlen", type=int, default=512) parser.add_argument("--batch-size", type=int, default=128, help="Tokenizer/model batch size.") parser.add_argument("--query-batch-size", type=int, default=16, help="Query scoring batch size.") parser.add_argument("--num-workers", type=int, default=4) parser.add_argument("--pooling", choices=["average", "cls", "pooler", "last_token"], default="average") parser.add_argument("--model-loader", choices=["auto", "t5-encoder"], default="auto", help="Model loader for dense encoders. Use t5-encoder for encoder-only T5 embedding models such as FRIDA.") parser.add_argument("--query-prefix", default="") parser.add_argument("--passage-prefix", default="") parser.add_argument("--model-sep", default="[SEP]") parser.add_argument("--padding-side", choices=["left", "right"], default=None, help="Tokenizer padding side. Qwen3 Embedding models should use left padding with last_token pooling.") parser.add_argument("--bm25-method", default="lucene") parser.add_argument("--bm25-k1", type=float, default=None) parser.add_argument("--bm25-b", type=float, default=None) parser.add_argument("--tfidf-min-ngram", type=int, default=1) parser.add_argument("--tfidf-max-ngram", type=int, default=1) parser.add_argument("--tfidf-max-features", type=int, default=None) parser.add_argument("--no-lowercase", action="store_true") parser.add_argument("--source-url", default="") parser.add_argument("--notes", default="") parser.add_argument("--resume", action="store_true", help="Reuse existing per-dataset scores for this model_id from --output and skip already computed datasets.") return parser.parse_args() def main() -> None: args = parse_args() started = time.time() requested = set(args.datasets or []) datasets = [dataset for dataset in load_datasets() if dataset.get("official", True) and (not requested or dataset["name"] in requested)] if args.limit_datasets is not None: datasets = datasets[: args.limit_datasets] if not datasets: raise SystemExit("No datasets selected.") if args.model_type in {"dense", "reranker"} and not args.model_id: raise SystemExit("--model-id is required for dense and reranker models.") if args.model_type == "sparse" and not args.model_id: args.model_id = args.sparse_model first_stage_model_name = None if args.model_type == "reranker": if not args.first_stage_model_id: raise SystemExit("--first-stage-model-id is required for reranker evaluation.") if args.first_stage_results_dir is None: raise SystemExit("--first-stage-results-dir is required for reranker evaluation.") first_stage_model_name = safe_name(args.first_stage_model_id) missing_results = [ raw_results_path(args.first_stage_results_dir, first_stage_model_name, dataset) for dataset in datasets if not raw_results_path(args.first_stage_results_dir, first_stage_model_name, dataset).exists() ] if missing_results: missing = "\n".join(str(path) for path in missing_results) raise SystemExit(f"Missing first-stage result files:\n{missing}") record_model_id = args.model_id if args.model_type == "reranker": record_model_id = f"{args.first_stage_model_id}+{args.model_id}" model_name = args.model_name or record_model_id.split("/")[-1] organization = args.model_id.split("/", 1)[0] if "/" in args.model_id else "" model, hardware = build_model(args) per_dataset: dict[str, dict[str, float]] = {} if args.resume: existing_record = find_existing_record(args.output, record_model_id) existing_scores = (existing_record or {}).get("scores", {}).get("datasets", {}) if isinstance(existing_scores, dict): per_dataset.update(existing_scores) model_file_name = safe_name(args.model_id) if args.model_type == "reranker": model_file_name = f"{model_file_name}__rerank_{first_stage_model_name}" for dataset in datasets: dataset_name = dataset["name"] if args.resume and dataset_name in per_dataset: print(f"Skipping {dataset_name} ({dataset['split']}): already present in {args.output}", flush=True) continue print(f"Evaluating {dataset_name} ({dataset['split']})", flush=True) per_dataset[dataset_name] = evaluate_dataset( model=model, model_name=model_file_name, dataset=dataset, k_values=args.k_values, text_type=args.text_type, query_batch_size=args.query_batch_size, num_workers=args.num_workers, raw_results_dir=args.raw_results_dir, model_type=args.model_type, first_stage_results_dir=args.first_stage_results_dir, first_stage_model_name=first_stage_model_name, rerank_top_k=args.rerank_top_k ) checkpoint = build_result_record(args, record_model_id, model_name, organization, hardware, started, per_dataset) insert_jsonl_record(args.output, checkpoint) print(f"Saved checkpoint for {dataset_name} to {args.output}", flush=True) record = build_result_record(args, record_model_id, model_name, organization, hardware, started, per_dataset) insert_jsonl_record(args.output, record) print(json.dumps(record, ensure_ascii=False, indent=2), flush=True) if __name__ == "__main__": main()