Upload folder using huggingface_hub

src/imrnns/__init__.py

@@ -1,15 +1,16 @@
 """IMRNNs package."""
 
 from .api import cache_embeddings, evaluate, run, train
+from .adapter import IMRNNAdapter, RetrievalResult
 from .hub import DEFAULT_REPO_ID, download_checkpoint, get_download_count, load_pretrained
-from .model import BiHyperNetIR, HyperNet, IMRNN, ModelConfig
+from .model import IMRNN, ModelConfig
 
 __all__ = [
-    "BiHyperNetIR",
     "DEFAULT_REPO_ID",
-    "HyperNet",
+    "IMRNNAdapter",
     "IMRNN",
     "ModelConfig",
+    "RetrievalResult",
     "cache_embeddings",
     "download_checkpoint",
     "evaluate",

src/imrnns/adapter.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+import torch
+from sentence_transformers import SentenceTransformer
+
+from .encoders import EncoderSpec
+from .hub import DEFAULT_REPO_ID, load_pretrained
+from .model import IMRNN
+
+
+@dataclass(frozen=True)
+class RetrievalResult:
+    rank: int
+    index: int
+    text: str
+    score: float
+
+
+def _format_query(text: str, encoder_spec: EncoderSpec) -> str:
+    return f"{encoder_spec.query_prefix}{text}" if encoder_spec.query_prefix else text
+
+
+def _format_document(text: str, encoder_spec: EncoderSpec) -> str:
+    return f"{encoder_spec.passage_prefix}{text}" if encoder_spec.passage_prefix else text
+
+
+class IMRNNAdapter:
+    """Inference wrapper for applying a pretrained IMRNN adapter to a base retriever."""
+
+    def __init__(
+        self,
+        *,
+        model: IMRNN,
+        encoder: SentenceTransformer,
+        encoder_spec: EncoderSpec,
+        metadata: dict[str, Any],
+        device: str,
+    ) -> None:
+        self.model = model
+        self.encoder = encoder
+        self.encoder_spec = encoder_spec
+        self.metadata = metadata
+        self.device = device
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        *,
+        encoder: str,
+        dataset: str,
+        repo_id: str = DEFAULT_REPO_ID,
+        device: str = "cpu",
+    ) -> "IMRNNAdapter":
+        model, metadata, encoder_spec = load_pretrained(
+            encoder=encoder,
+            dataset=dataset,
+            repo_id=repo_id,
+            device=device,
+        )
+        encoder_model = SentenceTransformer(encoder_spec.model_name, device=device)
+        return cls(
+            model=model,
+            encoder=encoder_model,
+            encoder_spec=encoder_spec,
+            metadata=metadata,
+            device=device,
+        )
+
+    def score(self, query: str, documents: Sequence[str], top_k: int | None = None) -> list[RetrievalResult]:
+        if not documents:
+            return []
+
+        formatted_query = _format_query(query, self.encoder_spec)
+        formatted_documents = [_format_document(document, self.encoder_spec) for document in documents]
+
+        with torch.no_grad():
+            query_embedding = self.encoder.encode(
+                [formatted_query],
+                convert_to_tensor=True,
+                show_progress_bar=False,
+                device=self.device,
+            )[0].to(self.device)
+            document_embeddings = self.encoder.encode(
+                formatted_documents,
+                convert_to_tensor=True,
+                show_progress_bar=False,
+                device=self.device,
+            ).to(self.device)
+            _, _, scores = self.model.score_candidates(query_embedding, document_embeddings)
+
+        ranked_indices = torch.argsort(scores, descending=True).tolist()
+        if top_k is not None:
+            ranked_indices = ranked_indices[:top_k]
+
+        return [
+            RetrievalResult(
+                rank=rank,
+                index=index,
+                text=documents[index],
+                score=float(scores[index].item()),
+            )
+            for rank, index in enumerate(ranked_indices, start=1)
+        ]

src/imrnns/api.py

@@ -9,7 +9,7 @@ from .checkpoints import default_checkpoint_name, load_model, save_checkpoint
 from .data import ContrastiveCachedDataset, load_cached_split
 from .encoders import get_encoder_spec
 from .evaluation import evaluate_model
-from .model import BiHyperNetIR, ModelConfig
+from .model import IMRNN, ModelConfig
 from .training import TrainingConfig, train_model
 
 
@@ -66,7 +66,7 @@ def train(
     val_split = load_cached_split(cache_dir, "val", beir_source, encoder_spec, device)
     test_split = load_cached_split(cache_dir, "test", beir_source, encoder_spec, device)
 
-    model = BiHyperNetIR(
+    model = IMRNN(
         ModelConfig(
             input_dim=encoder_spec.embedding_dim,
             output_dim=output_dim,

src/imrnns/checkpoints.py

@@ -7,7 +7,7 @@ from typing import Any
 import torch
 
 from .encoders import normalize_encoder_name
-from .model import BiHyperNetIR, ModelConfig
+from .model import IMRNN, ModelConfig
 
 
 def default_checkpoint_name(encoder: str, dataset: str) -> str:
@@ -29,7 +29,7 @@ def sanitize_legacy_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
 
 def save_checkpoint(
     path: Path,
-    model: BiHyperNetIR,
+    model: IMRNN,
     metadata: dict[str, Any],
 ) -> None:
     payload = {
@@ -56,9 +56,9 @@ def load_model(
     checkpoint_path: Path,
     model_config: ModelConfig,
     device: str,
-) -> tuple[BiHyperNetIR, dict[str, Any], list[str], list[str]]:
+) -> tuple[IMRNN, dict[str, Any], list[str], list[str]]:
     state_dict, metadata = load_checkpoint(checkpoint_path)
-    model = BiHyperNetIR(model_config)
+    model = IMRNN(model_config)
     missing, unexpected = model.load_state_dict(state_dict, strict=False)
     model.to(device)
     model.eval()

src/imrnns/cli.py

@@ -19,7 +19,7 @@ from .checkpoints import default_checkpoint_name, load_model, save_checkpoint
 from .data import ContrastiveCachedDataset, load_cached_split
 from .encoders import get_encoder_spec
 from .evaluation import evaluate_model
-from .model import BiHyperNetIR, ModelConfig
+from .model import IMRNN, ModelConfig
 from .training import TrainingConfig, train_model
 
 
@@ -85,7 +85,7 @@ def _command_cache(args: argparse.Namespace) -> int:
 
 def _command_train(args: argparse.Namespace) -> int:
     encoder_spec, cache_dir, train_split, val_split, test_split = _load_training_inputs(args)
-    model = BiHyperNetIR(
+    model = IMRNN(
         ModelConfig(
             input_dim=encoder_spec.embedding_dim,
             output_dim=args.output_dim,

src/imrnns/evaluation.py

@@ -10,7 +10,7 @@ import torch.nn.functional as F
 from tqdm import tqdm
 
 from .data import CachedSplit
-from .model import BiHyperNetIR
+from .model import IMRNN
 
 try:
     import faiss  # type: ignore
@@ -73,7 +73,7 @@ def _compute_metrics(ranked_doc_ids: list[str], qrel: dict[str, int], k_values:
 
 
 def evaluate_model(
-    model: BiHyperNetIR,
+    model: IMRNN,
     cached_split: CachedSplit,
     device: str,
     feedback_k: int = 100,
@@ -120,10 +120,10 @@ def evaluate_model(
                 dim=0,
             ).to(device)
 
-            _, _, rerank_scores = model.encode_candidates(query_embedding.float().to(device), candidate_embeddings)
-            rerank_scores = rerank_scores.cpu().tolist()
+            _, _, adapted_scores = model.score_candidates(query_embedding.float().to(device), candidate_embeddings)
+            adapted_scores = adapted_scores.cpu().tolist()
             reranked = [
-                doc_id for doc_id, _ in sorted(zip(candidate_ids, rerank_scores), key=lambda item: item[1], reverse=True)
+                doc_id for doc_id, _ in sorted(zip(candidate_ids, adapted_scores), key=lambda item: item[1], reverse=True)
             ][:ranking_k]
 
             metrics = _compute_metrics(reranked, cached_split.split.qrels[qid], k_values)

src/imrnns/hub.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi, hf_hub_download
 
 from .checkpoints import default_checkpoint_name, load_model
 from .encoders import EncoderSpec, get_encoder_spec, normalize_encoder_name
-from .model import BiHyperNetIR, ModelConfig
+from .model import IMRNN, ModelConfig
 
 DEFAULT_REPO_ID = "yashsaxena21/IMRNNs"
 CONFIG_FILENAME = "config.json"
@@ -92,7 +92,7 @@ def load_pretrained(
     revision: Optional[str] = None,
     cache_dir: Optional[Path] = None,
     local_files_only: bool = False,
-) -> tuple[BiHyperNetIR, dict[str, Any], EncoderSpec]:
+) -> tuple[IMRNN, dict[str, Any], EncoderSpec]:
     encoder_spec = get_encoder_spec(encoder)
     pretrained = download_checkpoint(
         encoder=encoder,

src/imrnns/model.py

@@ -99,7 +99,7 @@ class IMRNN(nn.Module):
         scores = torch.einsum("bd,bkd->bk", F.normalize(modulated_queries, p=2, dim=-1), F.normalize(modulated_documents, p=2, dim=-1))
         return modulated_queries, modulated_documents, scores
 
-    def encode_candidates(
+    def score_candidates(
         self,
         query_embedding: torch.Tensor,
         candidate_document_embeddings: torch.Tensor,
@@ -112,5 +112,5 @@ class IMRNN(nn.Module):
         return modulated_query.squeeze(0), modulated_docs.squeeze(0), scores.squeeze(0)
 
 
-class BiHyperNetIR(IMRNN):
-    """Backward-compatible alias for legacy checkpoints and code paths."""
+BiHyperNetIR = IMRNN
+"""Backward-compatible alias retained for legacy checkpoints and code paths."""

src/imrnns/training.py

@@ -8,7 +8,7 @@ from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from .data import ContrastiveCachedDataset, collate_contrastive_batch
-from .model import BiHyperNetIR
+from .model import IMRNN
 
 
 class MultipleNegativesRankingLoss(torch.nn.Module):
@@ -45,7 +45,7 @@ def build_dataloader(dataset: ContrastiveCachedDataset, batch_size: int, shuffle
 
 
 def evaluate_loss(
-    model: BiHyperNetIR,
+    model: IMRNN,
     dataloader: DataLoader,
     device: str,
     loss_fn: MultipleNegativesRankingLoss,
@@ -67,7 +67,7 @@ def evaluate_loss(
 
 
 def train_model(
-    model: BiHyperNetIR,
+    model: IMRNN,
     train_dataset: ContrastiveCachedDataset,
     val_dataset: ContrastiveCachedDataset,
     config: TrainingConfig,