Add source_id column back to lancedb

Author: jdye64

nemo_retriever/src/nemo_retriever/ingest_modes/batch.py

@@ -295,6 +295,9 @@ def extract(self, params: ExtractParams | None = None, **kwargs: Any) -> "BatchI
         This does not run extraction yet; it records configuration so the batch
         executor can build a concrete pipeline later.
 
+        If all input files have a ``.txt`` extension, the pipeline automatically
+        delegates to :meth:`extract_txt` with default :class:`TextChunkParams`.
+
         Resource-tuning kwargs (auto-detected from available resources if omitted):
 
         - ``pdf_split_batch_size``: Batch size for PDF split stage (default 1).
@@ -308,6 +311,13 @@ def extract(self, params: ExtractParams | None = None, **kwargs: Any) -> "BatchI
         - ``ocr_cpus_per_actor``: CPUs reserved per OCR actor (default 1).
         """
 
+        if self._input_documents and all(f.lower().endswith(".txt") for f in self._input_documents):
+            txt_params = TextChunkParams(
+                max_tokens=kwargs.pop("max_tokens", 1024),
+                overlap_tokens=kwargs.pop("overlap_tokens", 0),
+            )
+            return self.extract_txt(params=txt_params)
+
         resolved = _coerce_params(params, ExtractParams, kwargs)
         if (
             any(

nemo_retriever/src/nemo_retriever/ingest_modes/lancedb_utils.py

@@ -197,7 +197,9 @@ 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(
+                "source_id", pa.string()
+            ),  # Different than the source. Field contains path+page_number for aggregation tasks
             pa.field("path", pa.string()),
             pa.field("text", pa.string()),
             pa.field("metadata", pa.string()),

nemo_retriever/src/nemo_retriever/utils/hf_model_registry.py

@@ -28,6 +28,7 @@
     "nvidia/llama-nemotron-embed-vl-1b-v2": "859e1f2dac29c56c37a5279cf55f53f3e74efc6b",
     "meta-llama/Llama-3.2-1B": "4e20de362430cd3b72f300e6b0f18e50e7166e08",
     "intfloat/e5-large-unsupervised": "15af9288f69a6291f37bfb89b47e71abc747b206",
+    "nvidia/llama-nemotron-rerank-1b-v2": "aee9a1be0bbd89489f8bd0ec5763614c8bb85878",
 }
 
 

nemo_retriever/tests/test_lancedb_utils.py

@@ -198,6 +198,7 @@ def test_returns_schema_with_correct_fields(self):
         assert "text" in names
         assert "metadata" in names
         assert "source" in names
+        assert "source_id" in names
         assert len(names) == 10