| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Dict, List, Union |
|
|
| import torch |
| from sentence_transformers import SentenceTransformer |
| from torch.utils.data import Dataset |
|
|
| from .beir_data import DatasetSplit |
| from .encoders import EncoderSpec |
|
|
|
|
| @dataclass(frozen=True) |
| class CachedSplit: |
| split: DatasetSplit |
| document_embeddings: dict[str, torch.Tensor] |
| query_embeddings: dict[str, torch.Tensor] |
| negatives: dict[str, list[str]] |
|
|
|
|
| def _query_cache_path(cache_dir: Path, split_name: str) -> Path: |
| return _split_dir(cache_dir, split_name) / "query_embeddings.pt" |
|
|
|
|
| def _split_dir(cache_dir: Path, split_name: str) -> Path: |
| return cache_dir / split_name |
|
|
|
|
| def load_document_embeddings(cache_dir: Path) -> dict[str, torch.Tensor]: |
| path = cache_dir / "documents" / "embeddings.pt" |
| if not path.is_file(): |
| raise FileNotFoundError(f"Document embedding cache is missing: {path}") |
| return torch.load(path, map_location="cpu", weights_only=True) |
|
|
|
|
| def load_negatives(cache_dir: Path, split_name: str) -> dict[str, list[str]]: |
| with open(_split_dir(cache_dir, split_name) / "negatives.json", encoding="utf-8") as handle: |
| return json.load(handle) |
|
|
|
|
| def encode_queries( |
| queries: dict[str, str], |
| encoder_spec: EncoderSpec, |
| cache_dir: Path, |
| split_name: str, |
| device: str, |
| batch_size: int = 64, |
| ) -> dict[str, torch.Tensor]: |
| cache_path = _query_cache_path(cache_dir, split_name) |
| if cache_path.exists(): |
| return torch.load(cache_path, map_location="cpu", weights_only=True) |
|
|
| cache_path.parent.mkdir(parents=True, exist_ok=True) |
| model = SentenceTransformer( |
| encoder_spec.model_name, |
| device=device, |
| revision=encoder_spec.revision, |
| ) |
| query_ids = sorted(queries) |
| texts = [encoder_spec.query_prefix + queries[qid] for qid in query_ids] |
| encoded = model.encode( |
| texts, |
| batch_size=batch_size, |
| convert_to_tensor=True, |
| show_progress_bar=True, |
| device=device, |
| ) |
| query_embeddings = {qid: embedding.cpu() for qid, embedding in zip(query_ids, encoded)} |
| torch.save(query_embeddings, cache_path) |
| return query_embeddings |
|
|
|
|
| def load_cached_split( |
| cache_dir: Path, |
| split_name: str, |
| dataset_source: DatasetSplit, |
| encoder_spec: EncoderSpec, |
| device: str, |
| ) -> CachedSplit: |
| negatives = load_negatives(cache_dir, split_name) |
| cached_qids = list(negatives.keys()) |
| filtered_queries = { |
| qid: dataset_source.queries[qid] |
| for qid in cached_qids |
| if qid in dataset_source.queries and qid in dataset_source.qrels |
| } |
| filtered_qrels = {qid: dataset_source.qrels[qid] for qid in filtered_queries} |
| filtered_split = DatasetSplit( |
| corpus=dataset_source.corpus, |
| queries=filtered_queries, |
| qrels=filtered_qrels, |
| ) |
| return CachedSplit( |
| split=filtered_split, |
| document_embeddings=load_document_embeddings(cache_dir), |
| query_embeddings=encode_queries(filtered_split.queries, encoder_spec, cache_dir, split_name, device), |
| negatives=negatives, |
| ) |
|
|
|
|
| class ContrastiveCachedDataset(Dataset): |
| def __init__( |
| self, |
| cached_split: CachedSplit, |
| num_negatives: int, |
| ) -> None: |
| if num_negatives <= 0: |
| raise ValueError("num_negatives must be positive.") |
| self.cached_split = cached_split |
| self.num_negatives = num_negatives |
| self.examples: list[tuple[str, str, list[str]]] = [] |
|
|
| for qid, qrel in cached_split.split.qrels.items(): |
| if qid not in cached_split.query_embeddings: |
| continue |
| positives = sorted( |
| ( |
| (doc_id, relevance) |
| for doc_id, relevance in qrel.items() |
| if relevance > 0 and doc_id in cached_split.document_embeddings |
| ), |
| key=lambda item: (-item[1], item[0]), |
| ) |
| negatives = [ |
| doc_id for doc_id in cached_split.negatives.get(qid, []) if doc_id in cached_split.document_embeddings |
| ] |
| if not positives or not negatives: |
| continue |
| self.examples.append((qid, positives[0][0], negatives[:num_negatives])) |
|
|
| def __len__(self) -> int: |
| return len(self.examples) |
|
|
| def __getitem__(self, index: int) -> dict[str, torch.Tensor]: |
| qid, positive_id, negative_ids = self.examples[index] |
| query_embedding = self.cached_split.query_embeddings[qid].float() |
| positive_embedding = self.cached_split.document_embeddings[positive_id].float() |
| normalized_negative_ids = list(negative_ids[: self.num_negatives]) |
| while len(normalized_negative_ids) < self.num_negatives: |
| normalized_negative_ids.append(normalized_negative_ids[-1]) |
|
|
| negative_embeddings = [ |
| self.cached_split.document_embeddings[doc_id].float() for doc_id in normalized_negative_ids |
| ] |
| documents = torch.stack([positive_embedding, *negative_embeddings], dim=0) |
| return { |
| "qid": qid, |
| "query_embedding": query_embedding, |
| "documents": documents, |
| } |
|
|
|
|
| def collate_contrastive_batch(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, Union[torch.Tensor, List[str]]]: |
| return { |
| "qids": [item["qid"] for item in batch], |
| "query_embeddings": torch.stack([item["query_embedding"] for item in batch], dim=0), |
| "documents": torch.stack([item["documents"] for item in batch], dim=0), |
| } |
|
|
