| from __future__ import annotations |
|
|
| import random |
| from pathlib import Path |
| from typing import Optional |
|
|
| from .datasets.local import IMRNNDataset, load_beir_directory |
|
|
| DatasetSplit = IMRNNDataset |
|
|
|
|
| def download_beir_dataset(dataset_name: str, datasets_dir: Path) -> Path: |
| try: |
| from beir import util |
| except ImportError as exc: |
| raise ImportError('BEIR download support requires `pip install "imrnns[eval]"`.') from exc |
| url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset_name}.zip" |
| return Path(util.download_and_unzip(url, str(datasets_dir))) |
|
|
|
|
| def _dataset_path(dataset_name: str, datasets_dir: Path) -> Path: |
| candidate = datasets_dir / dataset_name |
| return candidate if candidate.is_dir() else download_beir_dataset(dataset_name, datasets_dir) |
|
|
|
|
| def _limit(dataset: DatasetSplit, max_queries: Optional[int]) -> DatasetSplit: |
| if max_queries is None: |
| return dataset |
| query_ids = sorted(dataset.queries)[:max_queries] |
| return DatasetSplit( |
| corpus=dataset.corpus, |
| queries={query_id: dataset.queries[query_id] for query_id in query_ids}, |
| qrels={query_id: dataset.qrels[query_id] for query_id in query_ids}, |
| ) |
|
|
|
|
| def load_beir_source( |
| dataset_name: str, |
| datasets_dir: Path, |
| max_queries: Optional[int] = None, |
| source_split: str = "test", |
| ) -> DatasetSplit: |
| data_path = _dataset_path(dataset_name, datasets_dir) |
| if dataset_name.lower() == "msmarco" and source_split == "test": |
| source_split = "train" |
| return _limit(load_beir_directory(data_path, split=source_split), max_queries) |
|
|
|
|
| def _subset(base: DatasetSplit, query_ids: list[str]) -> DatasetSplit: |
| return DatasetSplit( |
| corpus=base.corpus, |
| queries={query_id: base.queries[query_id] for query_id in query_ids}, |
| qrels={query_id: base.qrels[query_id] for query_id in query_ids}, |
| ) |
|
|
|
|
| def _partition(query_ids: list[str], fractions: tuple[float, ...], seed: int) -> list[list[str]]: |
| shuffled = sorted(query_ids) |
| random.Random(seed).shuffle(shuffled) |
| boundaries: list[int] = [] |
| running = 0.0 |
| for fraction in fractions[:-1]: |
| running += fraction |
| boundaries.append(round(len(shuffled) * running)) |
| starts = [0, *boundaries] |
| ends = [*boundaries, len(shuffled)] |
| return [shuffled[start:end] for start, end in zip(starts, ends)] |
|
|
|
|
| def load_beir_splits( |
| dataset_name: str, |
| datasets_dir: Path, |
| max_queries: Optional[int] = None, |
| source_split: str = "test", |
| seed: int = 42, |
| ) -> dict[str, DatasetSplit]: |
| """Load deterministic train/validation/test splits. |
| |
| When official train and test qrels both exist, train is split 85/15 for |
| training/validation and the official test set is kept intact. Datasets with |
| only one qrels split use a 70/15/15 partition based on sorted IDs and an |
| explicit seed. |
| """ |
|
|
| data_path = _dataset_path(dataset_name, datasets_dir) |
| train_qrels = data_path / "qrels" / "train.tsv" |
| test_qrels = data_path / "qrels" / "test.tsv" |
| if train_qrels.is_file() and test_qrels.is_file() and dataset_name.lower() != "msmarco": |
| official_train = _limit(load_beir_directory(data_path, split="train"), max_queries) |
| official_test = _limit(load_beir_directory(data_path, split="test"), max_queries) |
| train_ids, validation_ids = _partition(list(official_train.queries), (0.85, 0.15), seed) |
| return { |
| "train": _subset(official_train, train_ids), |
| "validation": _subset(official_train, validation_ids), |
| "test": official_test, |
| } |
|
|
| base = load_beir_source( |
| dataset_name, |
| datasets_dir=datasets_dir, |
| max_queries=max_queries, |
| source_split=source_split, |
| ) |
| train_ids, validation_ids, test_ids = _partition(list(base.queries), (0.70, 0.15, 0.15), seed) |
| return { |
| "train": _subset(base, train_ids), |
| "validation": _subset(base, validation_ids), |
| "test": _subset(base, test_ids), |
| } |
|
|
