Add rerank (#1565)

Co-authored-by: Jeremy Dyer <jdye64@gmail.com>

Author: jperez999

nemo_retriever/pyproject.toml

@@ -63,6 +63,7 @@ dependencies = [
   "nemotron-ocr>=0.dev0",
   "markitdown",
   "timm==1.0.22",
+  "tqdm",
   "accelerate==1.12.0",
   "albumentations==2.0.8",
   "scikit-learn>=1.6.0",

nemo_retriever/src/nemo_retriever/examples/batch_pipeline.py

@@ -425,6 +425,14 @@ def main(
         "--runtime-metrics-prefix",
         help="Optional filename prefix for per-run metrics artifacts.",
     ),
+    reranker: Optional[bool] = typer.Option(
+        False, "--reranker/--no-reranker", help="Enable a re-ranking stage with a cross-encoder model."
+    ),
+    reranker_model_name: str = typer.Option(
+        "nvidia/llama-nemotron-rerank-1b-v2",
+        "--reranker-model-name",
+        help="Cross-encoder model name for re-ranking stage (passed to .embed()).",
+    ),
     structured_elements_modality: Optional[str] = typer.Option(
         None,
         "--structured-elements-modality",
@@ -782,6 +790,7 @@ def _extract_params(batch_tuning: dict, **overrides: Any) -> ExtractParams:
             ks=(1, 5, 10),
             hybrid=hybrid,
             match_mode=recall_match_mode,
+            reranker=reranker_model_name if reranker else None,
         )
 
         # Capture recall only times.

nemo_retriever/src/nemo_retriever/ingest_modes/lancedb_utils.py

@@ -197,7 +197,6 @@ def lancedb_schema(vector_dim: int = 2048) -> Any:
             pa.field("pdf_basename", pa.string()),
             pa.field("page_number", pa.int32()),
             pa.field("source", pa.string()),
-            pa.field("source_id", pa.string()),
             pa.field("path", pa.string()),
             pa.field("text", pa.string()),
             pa.field("metadata", pa.string()),

nemo_retriever/src/nemo_retriever/model/local/__init__.py

@@ -17,6 +17,7 @@
     "NemotronTableStructureV1",
     "NemotronGraphicElementsV1",
     "NemotronParseV12",
+    "NemotronRerankV2",
     "ParakeetCTC1B1ASR",
 ]
 
@@ -42,6 +43,10 @@ def __getattr__(name: str):
         from .nemotron_parse_v1_2 import NemotronParseV12
 
         return NemotronParseV12
+    if name == "NemotronRerankV2":
+        from .nemotron_rerank_v2 import NemotronRerankV2
+
+        return NemotronRerankV2
     if name == "ParakeetCTC1B1ASR":
         from .parakeet_ctc_1_1b_asr import ParakeetCTC1B1ASR
 

nemo_retriever/src/nemo_retriever/model/local/nemotron_rerank_v2.py

@@ -0,0 +1,210 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Local wrapper for nvidia/llama-nemotron-rerank-1b-v2 cross-encoder reranker."""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from nemo_retriever.utils.hf_cache import configure_global_hf_cache_base
+from ..model import BaseModel, RunMode
+
+
+_DEFAULT_MODEL = "nvidia/llama-nemotron-rerank-1b-v2"
+_DEFAULT_MAX_LENGTH = 512
+_DEFAULT_BATCH_SIZE = 32
+
+
+def _prompt_template(query: str, passage: str) -> str:
+    """Format a (query, passage) pair as the model expects."""
+    return f"question:{query} \n \n passage:{passage}"
+
+
+class NemotronRerankV2(BaseModel):
+    """
+    Local cross-encoder reranker wrapping nvidia/llama-nemotron-rerank-1b-v2.
+
+    The model scores (query, document) pairs and returns raw logits; higher
+    values indicate greater relevance.  It is fine-tuned from
+    meta-llama/Llama-3.2-1B with bi-directional attention and supports 26
+    languages with sequences up to 8 192 tokens.
+
+    Example::
+
+        reranker = NemotronRerankV2()
+        scores = reranker.score("What is ML?", ["Machine learning is…", "Paris is…"])
+        # scores -> [20.6, -23.1]  (higher = more relevant)
+    """
+
+    def __init__(
+        self,
+        model_name: str = _DEFAULT_MODEL,
+        device: Optional[str] = None,
+        hf_cache_dir: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        import torch
+        from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+        configure_global_hf_cache_base()
+
+        self._model_name = model_name
+        self._device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        kwargs: dict = {"trust_remote_code": True}
+        if hf_cache_dir:
+            kwargs["cache_dir"] = hf_cache_dir
+
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            padding_side="left",
+            **kwargs,
+        )
+        if self._tokenizer.pad_token is None:
+            self._tokenizer.pad_token = self._tokenizer.eos_token
+
+        self._model = (
+            AutoModelForSequenceClassification.from_pretrained(
+                model_name,
+                torch_dtype=torch.bfloat16,
+                **kwargs,
+            )
+            .eval()
+            .to(self._device)
+        )
+
+        if self._model.config.pad_token_id is None:
+            self._model.config.pad_token_id = self._tokenizer.eos_token_id
+
+    # ------------------------------------------------------------------
+    # BaseModel abstract properties
+    # ------------------------------------------------------------------
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    @property
+    def model_type(self) -> str:
+        return "reranker"
+
+    @property
+    def model_runmode(self) -> RunMode:
+        return "local"
+
+    @property
+    def input(self):
+        return "List[Tuple[str, str]]"
+
+    @property
+    def output(self):
+        return "List[float]"
+
+    @property
+    def input_batch_size(self) -> int:
+        return _DEFAULT_BATCH_SIZE
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def score(
+        self,
+        query: str,
+        documents: List[str],
+        *,
+        max_length: int = _DEFAULT_MAX_LENGTH,
+        batch_size: int = _DEFAULT_BATCH_SIZE,
+    ) -> List[float]:
+        """
+        Score relevance of *documents* to *query*.
+
+        Parameters
+        ----------
+        query:
+            The search query.
+        documents:
+            Candidate passages/documents to score.
+        max_length:
+            Tokenizer truncation length (default 512; max supported 8 192).
+        batch_size:
+            Number of (query, doc) pairs to process per GPU forward pass.
+
+        Returns
+        -------
+        List[float]
+            Raw logit scores aligned with *documents* (higher = more relevant).
+        """
+        import torch
+
+        if not documents:
+            return []
+
+        texts = [_prompt_template(query, d) for d in documents]
+        all_scores: List[float] = []
+
+        with torch.inference_mode():
+            for start in range(0, len(texts), batch_size):
+                chunk = texts[start : start + batch_size]
+                batch = self._tokenizer(
+                    chunk,
+                    padding=True,
+                    truncation=True,
+                    return_tensors="pt",
+                    max_length=max_length,
+                )
+                batch = {k: v.to(self._device) for k, v in batch.items()}
+                logits = self._model(**batch).logits
+                all_scores.extend(logits.view(-1).cpu().tolist())
+
+        return all_scores
+
+    def score_pairs(
+        self,
+        pairs: List[tuple],
+        *,
+        max_length: int = _DEFAULT_MAX_LENGTH,
+        batch_size: int = _DEFAULT_BATCH_SIZE,
+    ) -> List[float]:
+        """
+        Score a list of (query, document) pairs.
+
+        Parameters
+        ----------
+        pairs:
+            Sequence of ``(query, document)`` tuples.
+        max_length:
+            Tokenizer truncation length.
+        batch_size:
+            GPU forward-pass batch size.
+
+        Returns
+        -------
+        List[float]
+            Raw logit scores (higher = more relevant).
+        """
+        import torch
+
+        if not pairs:
+            return []
+
+        texts = [_prompt_template(q, d) for q, d in pairs]
+        all_scores: List[float] = []
+
+        with torch.inference_mode():
+            for start in range(0, len(texts), batch_size):
+                chunk = texts[start : start + batch_size]
+                batch = self._tokenizer(
+                    chunk,
+                    padding=True,
+                    truncation=True,
+                    return_tensors="pt",
+                    max_length=max_length,
+                )
+                batch = {k: v.to(self._device) for k, v in batch.items()}
+                logits = self._model(**batch).logits
+                all_scores.extend(logits.view(-1).cpu().tolist())
+
+        return all_scores