From 0c0e6874579b656bbef55c1ec954daad1b3d1b36 Mon Sep 17 00:00:00 2001 From: Kurt Heiss Date: Tue, 23 Jun 2026 09:47:31 -0700 Subject: [PATCH] examples: scaffold Helm metadata-filter path in notebook Follow-up to PR #2195: document job-based sidecar ingest for Helm deployments in the metadata filtering notebook. Part A keeps local GraphIngestor flow; Part B outlines POST /v1/ingest/sidecar and meta_dataframe_id for assignees to complete. --- ...ever_retriever_query_metadata_filter.ipynb | 959 +++++++++--------- 1 file changed, 505 insertions(+), 454 deletions(-) diff --git a/examples/nemo_retriever_retriever_query_metadata_filter.ipynb b/examples/nemo_retriever_retriever_query_metadata_filter.ipynb index bc939c2915..89fece2719 100644 --- a/examples/nemo_retriever_retriever_query_metadata_filter.ipynb +++ b/examples/nemo_retriever_retriever_query_metadata_filter.ipynb @@ -1,471 +1,522 @@ { - "cells": [ - { - "cell_type": "markdown", - "id": "6d0ef3d1", - "metadata": {}, - "source": [ - "# Metadata and SQL filters with `Retriever.query`\n", - "\n", - "This notebook shows how **chunk metadata** ends up in LanceDB and how to narrow search results when calling **`Retriever.query`**.\n", - "\n", - "**Two complementary mechanisms**\n", - "\n", - "1. **Server-side filter (`where`)** — Pass a Lance / DataFusion SQL predicate in `vdb_kwargs` on each `query` (or set defaults on the `Retriever`). The predicate runs inside LanceDB on table columns **`vector`**, **`text`**, **`metadata`**, **`source`**. This is implemented in `LanceDB.retrieval` as a `.where(...)` clause on the vector search.\n", - "2. **Client-side filter** — After retrieval, use `filter_hits_by_content_metadata` to keep rows whose parsed `content_metadata` matches arbitrary Python logic (see also `nemo_retriever_metadata_and_filtered_search.ipynb`).\n", - "\n", - "**Metadata shape** — Ingestion stores each chunk’s `content_metadata` as a **JSON string** in the `metadata` column (compact JSON: no spaces after `:`). Sidecar columns (`meta_dataframe` / `meta_fields`) are merged into that object before upload.\n", - "\n", - "**Prerequisites** — Same as the sibling metadata notebook: sample PDFs under `NEMO_RETRIEVER_ROOT/data/`, Python env with `nemo-retriever`, LanceDB, and resources for extract + embed. Use `local_ingest_embed_backend=\"hf\"` if you are not running a local vLLM embed server." - ] - }, - { - "cell_type": "markdown", - "id": "6c621d9c", - "metadata": {}, - "source": [ - "## Imports and paths" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "0d0b2c1b", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import shutil\n", - "from pathlib import Path\n", - "\n", - "import pandas as pd\n", - "\n", - "from nemo_retriever.graph_ingestor import GraphIngestor\n", - "from nemo_retriever.params import EmbedParams, ExtractParams\n", - "from nemo_retriever.retriever import Retriever\n", - "from nemo_retriever.vdb import IngestVdbOperator, filter_hits_by_content_metadata, parse_hit_content_metadata\n", - "\n", - "NEMO_RETRIEVER_ROOT = Path(os.environ.get(\"NEMO_RETRIEVER_ROOT\", \"/raid/nv-ingest\")).resolve()" - ] - }, - { - "cell_type": "markdown", - "id": "e2ff5798", - "metadata": {}, - "source": [ - "## Configuration\n", - "\n", - "Use the **same embedding model** for ingestion and for `Retriever` queries. Use a fresh LanceDB directory when you change documents or sidecar data (`overwrite=True` on upload)." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "1115f9d3", - "metadata": {}, - "outputs": [ + "cells": [ { - "data": { - "text/plain": [ - "['/raid/nv-ingest/data/woods_frost.pdf',\n", - " '/raid/nv-ingest/data/multimodal_test.pdf']" - ] - }, - "execution_count": 2, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "model_name = \"nvidia/llama-nemotron-embed-1b-v2\"\n", - "LANCEDB_URI = str(Path(\"./nemo_retriever_query_filter_lancedb\").resolve())\n", - "TABLE_NAME = \"nv-ingest\"\n", - "\n", - "pdf_a = NEMO_RETRIEVER_ROOT / \"data\" / \"woods_frost.pdf\"\n", - "pdf_b = NEMO_RETRIEVER_ROOT / \"data\" / \"multimodal_test.pdf\"\n", - "files = [str(p) for p in (pdf_a, pdf_b) if p.is_file()]\n", - "if len(files) < 2:\n", - " raise FileNotFoundError(\n", - " f\"Expected sample PDFs at {pdf_a} and {pdf_b}. Set NEMO_RETRIEVER_ROOT or copy data files.\"\n", - " )\n", - "\n", - "files" - ] - }, - { - "cell_type": "markdown", - "id": "c0dd3ad7", - "metadata": {}, - "source": [ - "## Sidecar metadata\n", - "\n", - "These columns are joined onto each chunk’s `content_metadata` using the document path in `source` (with `meta_join_key=\"auto\"`)." - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "79914d90", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Metadata and SQL filters with `Retriever.query`\n", + "\n", + "This notebook shows how **chunk metadata** ends up in LanceDB and how to narrow search results when calling **`Retriever.query`**.\n", + "\n", + "**Two complementary mechanisms**\n", + "\n", + "1. **Server-side filter (`where`)** — Pass a Lance / DataFusion SQL predicate in `vdb_kwargs` on each `query` (or set defaults on the `Retriever`). The predicate runs inside LanceDB on table columns **`vector`**, **`text`**, **`metadata`**, **`source`**. This is implemented in `LanceDB.retrieval` as a `.where(...)` clause on the vector search.\n", + "2. **Client-side filter** — After retrieval, use `filter_hits_by_content_metadata` to keep rows whose parsed `content_metadata` matches arbitrary Python logic (see also `nemo_retriever_metadata_and_filtered_search.ipynb`).\n", + "\n", + "**Metadata shape** — Ingestion stores each chunk’s `content_metadata` as a **JSON string** in the `metadata` column (compact JSON: no spaces after `:`). Sidecar columns (`meta_dataframe` / `meta_fields`) are merged into that object before upload.\n", + "\n", + "**Deployment paths**\n", + "\n", + "1. **Part A — Local in-process ingest** (below) uses `GraphIngestor` with a sidecar CSV on disk (`meta_dataframe` path). Best for laptop or single-node dev without the retriever Helm chart.\n", + "2. **Part B — Helm / retriever service** (end of notebook) uses job-based ingest against a running deployment: upload the sidecar with `POST /v1/ingest/sidecar`, pass `meta_dataframe_id` in `vdb_upload_params`, then query the same LanceDB table.\n", + "\n", + "The published docs hub at `vdbs.md#metadata-and-filtering` points here only; Helm deployment detail lives in this notebook, not on the docs page.\n", + "\n", + "**Prerequisites** — Sample PDFs under `NEMO_RETRIEVER_ROOT/data/`, Python env with `nemo-retriever`, LanceDB, and resources for extract + embed. Use `local_ingest_embed_backend=\"hf\"` if you are not running a local vLLM embed server. For Part B, also need a retriever service URL (Helm release) and API key if auth is enabled." + ], + "id": "6d0ef3d1" + }, { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
sourcemeta_ameta_b
0/raid/nv-ingest/data/woods_frost.pdfalpha5
1/raid/nv-ingest/data/multimodal_test.pdfbravo10
\n", - "
" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part A — Local in-process ingest" ], - "text/plain": [ - " source meta_a meta_b\n", - "0 /raid/nv-ingest/data/woods_frost.pdf alpha 5\n", - "1 /raid/nv-ingest/data/multimodal_test.pdf bravo 10" - ] - }, - "execution_count": 3, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "meta_df = pd.DataFrame(\n", - " {\n", - " \"source\": [str(pdf_a.resolve()), str(pdf_b.resolve())],\n", - " \"meta_a\": [\"alpha\", \"bravo\"],\n", - " \"meta_b\": [5, 10],\n", - " }\n", - ")\n", - "meta_path = Path(\"./nemo_retriever_query_filter_sidecar.csv\").resolve()\n", - "meta_df.to_csv(meta_path, index=False)\n", - "meta_df" - ] - }, - { - "cell_type": "markdown", - "id": "82f2b7c8", - "metadata": {}, - "source": [ - "## 1) Extract and embed (graph pipeline)" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "e63b774c", - "metadata": {}, - "outputs": [ + "id": "b212bcdc" + }, { - "name": "stderr", - "output_type": "stream", - "text": [ - "/opt/retriever_runtime/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Imports and paths" + ], + "id": "6c621d9c" }, { - "name": "stdout", - "output_type": "stream", - "text": [ - " -> Downloading/loading weights from HuggingFace: nvidia/nemotron-page-elements-v3\n", - " -> Weights cached at: /root/.cache/huggingface/hub/models--nvidia--nemotron-page-elements-v3/snapshots/df62dbb631502575ac4d43b44d700b1674ab1d56/nemotron_page_elements_v3/weights.pth\n" - ] + "cell_type": "code", + "metadata": {}, + "source": [ + "import os\n", + "import shutil\n", + "from pathlib import Path\n", + "\n", + "import pandas as pd\n", + "\n", + "from nemo_retriever.graph_ingestor import GraphIngestor\n", + "from nemo_retriever.params import EmbedParams, ExtractParams\n", + "from nemo_retriever.retriever import Retriever\n", + "from nemo_retriever.vdb import IngestVdbOperator, filter_hits_by_content_metadata, parse_hit_content_metadata\n", + "\n", + "NEMO_RETRIEVER_ROOT = Path(os.environ.get(\"NEMO_RETRIEVER_ROOT\", \"/raid/nv-ingest\")).resolve()" + ], + "execution_count": 1, + "outputs": [], + "id": "0d0b2c1b" }, { - "name": "stderr", - "output_type": "stream", - "text": [ - "install_pinned_hf_hub_download: module 'nemotron_ocr.inference.pipeline_v2' has no 'hf_hub_download' attribute; revision pinning was NOT applied.\n", - "/opt/retriever_runtime/lib/python3.12/site-packages/nemotron_ocr/inference/models/recognizer.py:90: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.norm_first was True\n", - " self.tx = nn.TransformerEncoder(\n", - "/opt/retriever_runtime/lib/python3.12/site-packages/nemotron_ocr/inference/models/relational.py:105: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.norm_first was True\n", - " nn.TransformerEncoder(\n", - "Pipeline stages: 71%|████████████████████████████████████████████████████████████████▎ | 5/7 [00:34<00:16, 8.03s/stage, _BatchEmbedActor]/opt/retriever_runtime/lib/python3.12/site-packages/transformers/utils/hub.py:110: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.\n", - " warnings.warn(\n", - "`torch_dtype` is deprecated! Use `dtype` instead!\n", - "Pipeline stages: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:39<00:00, 5.63s/stage, _BatchEmbedActor]\n" - ] + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configuration\n", + "\n", + "Use the **same embedding model** for ingestion and for `Retriever` queries. Use a fresh LanceDB directory when you change documents or sidecar data (`overwrite=True` on upload)." + ], + "id": "e2ff5798" }, { - "data": { - "text/plain": [ - "5" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "extract_params = ExtractParams(\n", - " extract_text=True,\n", - " extract_tables=True,\n", - " extract_charts=False,\n", - " extract_images=False,\n", - ")\n", - "embed_params = EmbedParams(\n", - " model_name=model_name,\n", - " embed_granularity=\"page\",\n", - " local_ingest_embed_backend=\"hf\",\n", - ")\n", - "\n", - "ingestor = (\n", - " GraphIngestor(run_mode=\"inprocess\", documents=files)\n", - " .extract(extract_params)\n", - " .embed(embed_params)\n", - ")\n", - "result_df = ingestor.ingest()\n", - "records = result_df.to_dict(\"records\")\n", - "len(records)" - ] - }, - { - "cell_type": "markdown", - "id": "90e0511b", - "metadata": {}, - "source": [ - "## 2) Upload to LanceDB (sidecar merged into `metadata` JSON)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0ba141e2", - "metadata": {}, - "outputs": [], - "source": [ - "if Path(LANCEDB_URI).exists():\n", - " shutil.rmtree(LANCEDB_URI)\n", - "\n", - "upload_vdb_kwargs = {\n", - " \"uri\": LANCEDB_URI,\n", - " \"table_name\": TABLE_NAME,\n", - " \"overwrite\": True,\n", - " \"meta_dataframe\": str(meta_path),\n", - " \"meta_source_field\": \"source\",\n", - " \"meta_fields\": [\"meta_a\", \"meta_b\"],\n", - " \"meta_join_key\": \"auto\",\n", - "}\n", - "\n", - "uploader = IngestVdbOperator(vdb_op=\"lancedb\", vdb_kwargs=upload_vdb_kwargs)\n", - "uploader.process(records)" - ] - }, - { - "cell_type": "markdown", - "id": "dad108ab", - "metadata": {}, - "source": [ - "## 3) Inspect metadata on hits from `Retriever.query`\n", - "\n", - "Each hit includes a **`metadata`** string field (JSON). Parse it with `parse_hit_content_metadata` to work with **`meta_a`**, **`meta_b`**, page numbers, etc." - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "c15a68d1", - "metadata": {}, - "outputs": [ + "cell_type": "code", + "metadata": {}, + "source": [ + "model_name = \"nvidia/llama-nemotron-embed-1b-v2\"\n", + "LANCEDB_URI = str(Path(\"./nemo_retriever_query_filter_lancedb\").resolve())\n", + "TABLE_NAME = \"nv-ingest\"\n", + "\n", + "pdf_a = NEMO_RETRIEVER_ROOT / \"data\" / \"woods_frost.pdf\"\n", + "pdf_b = NEMO_RETRIEVER_ROOT / \"data\" / \"multimodal_test.pdf\"\n", + "files = [str(p) for p in (pdf_a, pdf_b) if p.is_file()]\n", + "if len(files) < 2:\n", + " raise FileNotFoundError(\n", + " f\"Expected sample PDFs at {pdf_a} and {pdf_b}. Set NEMO_RETRIEVER_ROOT or copy data files.\"\n", + " )\n", + "\n", + "files" + ], + "execution_count": 2, + "outputs": [ + { + "output_type": "execute_result", + "data": { + "text/plain": [ + "['/raid/nv-ingest/data/woods_frost.pdf',\n", + " '/raid/nv-ingest/data/multimodal_test.pdf']" + ] + } + } + ], + "id": "1115f9d3" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "num hits: 5\n", - "parsed content_metadata keys (sample): ['meta_a', 'meta_b', 'page_number']\n", - "meta_a: bravo meta_b: 10\n", - "text preview: TestingDocument\n", - "A sample document with headings and placeholder text\n", - "Introduction\n", - "This is a placeholder document that can be used for any purpose. It contains some \n", - "headings and some placeholder t\n" - ] - } - ], - "source": [ - "base_vdb = {\"uri\": LANCEDB_URI, \"table_name\": TABLE_NAME}\n", - "\n", - "retriever = Retriever(\n", - " vdb_kwargs=base_vdb,\n", - " embed_kwargs={\n", - " \"model_name\": model_name,\n", - " \"embed_model_name\": model_name,\n", - " \"local_ingest_embed_backend\": \"hf\",\n", - " },\n", - " top_k=8,\n", - ")\n", - "\n", - "q = \"introduction summary table\"\n", - "hits = retriever.query(q, top_k=8)\n", - "print(\"num hits:\", len(hits))\n", - "if hits:\n", - " first = hits[0]\n", - " cm = parse_hit_content_metadata(first)\n", - " print(\"parsed content_metadata keys (sample):\", sorted(cm.keys())[:20])\n", - " print(\"meta_a:\", cm.get(\"meta_a\"), \"meta_b:\", cm.get(\"meta_b\"))\n", - " print(\"text preview:\", (first.get(\"text\") or \"\")[:200])" - ] - }, - { - "cell_type": "markdown", - "id": "cab9110b", - "metadata": {}, - "source": [ - "## 4) Server-side filter: `query(..., vdb_kwargs={..., \"where\": \"...\"})`\n", - "\n", - "Per-call **`vdb_kwargs`** are merged with the retriever’s defaults. Pass a SQL predicate as **`where`** (alias **`_filter`**) to restrict rows **before** `top_k` is applied in the database.\n", - "\n", - "**Examples**\n", - "\n", - "- Filter on **`text`**: simple equality or `LIKE`.\n", - "- Filter on **`metadata`**: the column is a string; for JSON substrings use `LIKE` with a fragment that matches **compact** JSON (e.g. `\"meta_a\":\"alpha\"`).\n", - "\n", - "Escape single quotes in SQL strings by doubling them (`''`)." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "0327e2c3", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Sidecar metadata\n", + "\n", + "These columns are joined onto each chunk’s `content_metadata` using the document path in `source` (with `meta_join_key=\"auto\"`)." + ], + "id": "c0dd3ad7" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "hits with server-side meta_a filter: 2\n", - "alpha # Collection Year 1 A Boy's Will 1913 2 North of Boston 1914 3 Mountain Interval 1916 4 New Hampshir\n", - "alpha Stopping by Woods on a Snowy Evening, By Robert Frost\n", - "Figure 1: Snowy Woods\n", - "Whose woods these are \n" - ] - } - ], - "source": [ - "# Restrict to chunks whose serialized content_metadata contains meta_a == \"alpha\"\n", - "# (matches ingestion JSON like {\"meta_a\":\"alpha\",...})\n", - "pred_alpha = \"metadata LIKE '%\\\"meta_a\\\":\\\"alpha\\\"%'\"\n", - "\n", - "hits_alpha = retriever.query(q, top_k=8, vdb_kwargs={**base_vdb, \"where\": pred_alpha})\n", - "print(\"hits with server-side meta_a filter:\", len(hits_alpha))\n", - "for h in hits_alpha[:3]:\n", - " print(parse_hit_content_metadata(h).get(\"meta_a\"), (h.get(\"text\") or \"\")[:100])" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "c57ddafb", - "metadata": {}, - "outputs": [], - "source": [ - "# Same predicate using the _filter alias (useful for Milvus-style call sites)\n", - "hits_alias = retriever.query(q, top_k=8, vdb_kwargs={**base_vdb, \"_filter\": pred_alpha})\n", - "assert len(hits_alias) == len(hits_alpha)" - ] - }, - { - "cell_type": "markdown", - "id": "ca2a01be", - "metadata": {}, - "source": [ - "## 5) Compare: server-side `where` vs client-side `filter_hits_by_content_metadata`\n", - "\n", - "- **`where`** reduces work in LanceDB and can shrink the candidate set before vector ranking (depending on planner behavior).\n", - "- **`filter_hits_by_content_metadata`** is flexible (any Python predicate) but runs after the vector search returns." - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "608b1c8b", - "metadata": {}, - "outputs": [ + "cell_type": "code", + "metadata": {}, + "source": [ + "meta_df = pd.DataFrame(\n", + " {\n", + " \"source\": [str(pdf_a.resolve()), str(pdf_b.resolve())],\n", + " \"meta_a\": [\"alpha\", \"bravo\"],\n", + " \"meta_b\": [5, 10],\n", + " }\n", + ")\n", + "meta_path = Path(\"./nemo_retriever_query_filter_sidecar.csv\").resolve()\n", + "meta_df.to_csv(meta_path, index=False)\n", + "meta_df" + ], + "execution_count": 3, + "outputs": [ + { + "output_type": "execute_result", + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
sourcemeta_ameta_b
0/raid/nv-ingest/data/woods_frost.pdfalpha5
1/raid/nv-ingest/data/multimodal_test.pdfbravo10
\n", + "
" + ], + "text/plain": [ + " source meta_a meta_b\n", + "0 /raid/nv-ingest/data/woods_frost.pdf alpha 5\n", + "1 /raid/nv-ingest/data/multimodal_test.pdf bravo 10" + ] + } + } + ], + "id": "79914d90" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1) Extract and embed (graph pipeline)" + ], + "id": "82f2b7c8" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "extract_params = ExtractParams(\n", + " extract_text=True,\n", + " extract_tables=True,\n", + " extract_charts=False,\n", + " extract_images=False,\n", + ")\n", + "embed_params = EmbedParams(\n", + " model_name=model_name,\n", + " embed_granularity=\"page\",\n", + " local_ingest_embed_backend=\"hf\",\n", + ")\n", + "\n", + "ingestor = (\n", + " GraphIngestor(run_mode=\"inprocess\", documents=files)\n", + " .extract(extract_params)\n", + " .embed(embed_params)\n", + ")\n", + "result_df = ingestor.ingest()\n", + "records = result_df.to_dict(\"records\")\n", + "len(records)" + ], + "execution_count": 4, + "outputs": [ + { + "output_type": "stream", + "text": [ + "/opt/retriever_runtime/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "output_type": "stream", + "text": [ + " -> Downloading/loading weights from HuggingFace: nvidia/nemotron-page-elements-v3\n", + " -> Weights cached at: /root/.cache/huggingface/hub/models--nvidia--nemotron-page-elements-v3/snapshots/df62dbb631502575ac4d43b44d700b1674ab1d56/nemotron_page_elements_v3/weights.pth\n" + ] + }, + { + "output_type": "stream", + "text": [ + "install_pinned_hf_hub_download: module 'nemotron_ocr.inference.pipeline_v2' has no 'hf_hub_download' attribute; revision pinning was NOT applied.\n", + "/opt/retriever_runtime/lib/python3.12/site-packages/nemotron_ocr/inference/models/recognizer.py:90: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.norm_first was True\n", + " self.tx = nn.TransformerEncoder(\n", + "/opt/retriever_runtime/lib/python3.12/site-packages/nemotron_ocr/inference/models/relational.py:105: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.norm_first was True\n", + " nn.TransformerEncoder(\n", + "Pipeline stages: 71%|████████████████████████████████████████████████████████████████▎ | 5/7 [00:34<00:16, 8.03s/stage, _BatchEmbedActor]/opt/retriever_runtime/lib/python3.12/site-packages/transformers/utils/hub.py:110: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.\n", + " warnings.warn(\n", + "`torch_dtype` is deprecated! Use `dtype` instead!\n", + "Pipeline stages: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:39<00:00, 5.63s/stage, _BatchEmbedActor]\n" + ] + }, + { + "output_type": "execute_result", + "data": { + "text/plain": [ + "5" + ] + } + } + ], + "id": "e63b774c" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2) Upload to LanceDB (sidecar merged into `metadata` JSON)" + ], + "id": "90e0511b" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "if Path(LANCEDB_URI).exists():\n", + " shutil.rmtree(LANCEDB_URI)\n", + "\n", + "upload_vdb_kwargs = {\n", + " \"uri\": LANCEDB_URI,\n", + " \"table_name\": TABLE_NAME,\n", + " \"overwrite\": True,\n", + " \"meta_dataframe\": str(meta_path),\n", + " \"meta_source_field\": \"source\",\n", + " \"meta_fields\": [\"meta_a\", \"meta_b\"],\n", + " \"meta_join_key\": \"auto\",\n", + "}\n", + "\n", + "uploader = IngestVdbOperator(vdb_op=\"lancedb\", vdb_kwargs=upload_vdb_kwargs)\n", + "uploader.process(records)" + ], + "execution_count": null, + "outputs": [], + "id": "0ba141e2" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3) Inspect metadata on hits from `Retriever.query`\n", + "\n", + "Each hit includes a **`metadata`** string field (JSON). Parse it with `parse_hit_content_metadata` to work with **`meta_a`**, **`meta_b`**, page numbers, etc." + ], + "id": "dad108ab" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "wide search then client filter: 3\n", - "single query with where: 3\n" - ] + "cell_type": "code", + "metadata": {}, + "source": [ + "base_vdb = {\"uri\": LANCEDB_URI, \"table_name\": TABLE_NAME}\n", + "\n", + "retriever = Retriever(\n", + " vdb_kwargs=base_vdb,\n", + " embed_kwargs={\n", + " \"model_name\": model_name,\n", + " \"embed_model_name\": model_name,\n", + " \"local_ingest_embed_backend\": \"hf\",\n", + " },\n", + " top_k=8,\n", + ")\n", + "\n", + "q = \"introduction summary table\"\n", + "hits = retriever.query(q, top_k=8)\n", + "print(\"num hits:\", len(hits))\n", + "if hits:\n", + " first = hits[0]\n", + " cm = parse_hit_content_metadata(first)\n", + " print(\"parsed content_metadata keys (sample):\", sorted(cm.keys())[:20])\n", + " print(\"meta_a:\", cm.get(\"meta_a\"), \"meta_b:\", cm.get(\"meta_b\"))\n", + " print(\"text preview:\", (first.get(\"text\") or \"\")[:200])" + ], + "execution_count": 6, + "outputs": [ + { + "output_type": "stream", + "text": [ + "num hits: 5\n", + "parsed content_metadata keys (sample): ['meta_a', 'meta_b', 'page_number']\n", + "meta_a: bravo meta_b: 10\n", + "text preview: TestingDocument\n", + "A sample document with headings and placeholder text\n", + "Introduction\n", + "This is a placeholder document that can be used for any purpose. It contains some \n", + "headings and some placeholder t\n" + ] + } + ], + "id": "c15a68d1" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4) Server-side filter: `query(..., vdb_kwargs={..., \"where\": \"...\"})`\n", + "\n", + "Per-call **`vdb_kwargs`** are merged with the retriever’s defaults. Pass a SQL predicate as **`where`** (alias **`_filter`**) to restrict rows **before** `top_k` is applied in the database.\n", + "\n", + "**Examples**\n", + "\n", + "- Filter on **`text`**: simple equality or `LIKE`.\n", + "- Filter on **`metadata`**: the column is a string; for JSON substrings use `LIKE` with a fragment that matches **compact** JSON (e.g. `\"meta_a\":\"alpha\"`).\n", + "\n", + "Escape single quotes in SQL strings by doubling them (`''`)." + ], + "id": "cab9110b" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Restrict to chunks whose serialized content_metadata contains meta_a == \"alpha\"\n", + "# (matches ingestion JSON like {\"meta_a\":\"alpha\",...})\n", + "pred_alpha = \"metadata LIKE '%\\\"meta_a\\\":\\\"alpha\\\"%'\"\n", + "\n", + "hits_alpha = retriever.query(q, top_k=8, vdb_kwargs={**base_vdb, \"where\": pred_alpha})\n", + "print(\"hits with server-side meta_a filter:\", len(hits_alpha))\n", + "for h in hits_alpha[:3]:\n", + " print(parse_hit_content_metadata(h).get(\"meta_a\"), (h.get(\"text\") or \"\")[:100])" + ], + "execution_count": 7, + "outputs": [ + { + "output_type": "stream", + "text": [ + "hits with server-side meta_a filter: 2\n", + "alpha # Collection Year 1 A Boy's Will 1913 2 North of Boston 1914 3 Mountain Interval 1916 4 New Hampshir\n", + "alpha Stopping by Woods on a Snowy Evening, By Robert Frost\n", + "Figure 1: Snowy Woods\n", + "Whose woods these are \n" + ] + } + ], + "id": "0327e2c3" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Same predicate using the _filter alias (useful for Milvus-style call sites)\n", + "hits_alias = retriever.query(q, top_k=8, vdb_kwargs={**base_vdb, \"_filter\": pred_alpha})\n", + "assert len(hits_alias) == len(hits_alpha)" + ], + "execution_count": 8, + "outputs": [], + "id": "c57ddafb" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5) Compare: server-side `where` vs client-side `filter_hits_by_content_metadata`\n", + "\n", + "- **`where`** reduces work in LanceDB and can shrink the candidate set before vector ranking (depending on planner behavior).\n", + "- **`filter_hits_by_content_metadata`** is flexible (any Python predicate) but runs after the vector search returns." + ], + "id": "ca2a01be" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "hits_wide = retriever.query(q, top_k=16)\n", + "client_only = filter_hits_by_content_metadata(\n", + " hits_wide, lambda m: m.get(\"meta_a\") == \"bravo\" and m.get(\"meta_b\", 0) >= 10\n", + ")\n", + "print(\"wide search then client filter:\", len(client_only))\n", + "\n", + "# Approximate the same intent in SQL on the JSON string (meta_b is stored as a number in JSON)\n", + "pred_bravo = \"metadata LIKE '%\\\"meta_a\\\":\\\"bravo\\\"%' AND metadata LIKE '%\\\"meta_b\\\":10%'\"\n", + "hits_sql = retriever.query(q, top_k=16, vdb_kwargs={**base_vdb, \"where\": pred_bravo})\n", + "print(\"single query with where:\", len(hits_sql))" + ], + "execution_count": 9, + "outputs": [ + { + "output_type": "stream", + "text": [ + "wide search then client filter: 3\n", + "single query with where: 3\n" + ] + } + ], + "id": "608b1c8b" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [], + "execution_count": null, + "outputs": [], + "id": "290964a0-7d68-418a-8872-7f6ef6b220ad" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part B — Helm / retriever service ingest\n", + "\n", + "Use this path when the NeMo Retriever Helm chart is running and ingest goes through the job API instead of local `GraphIngestor`.\n", + "\n", + "**Flow**\n", + "\n", + "1. Build the same sidecar CSV as Part A (`source`, `meta_a`, `meta_b`, …).\n", + "2. `POST /v1/ingest/sidecar` with the CSV file → receive `sidecar_id`.\n", + "3. `POST /v1/ingest/job` to open a job (`expected_documents` = number of PDFs).\n", + "4. For each PDF, `POST /v1/ingest/job/{job_id}/document` with:\n", + " - `file` — document bytes\n", + " - `metadata` JSON containing `pipeline.vdb_upload_params` with `meta_dataframe_id`, `meta_source_field`, `meta_fields`, and LanceDB `vdb_kwargs` (`uri`, `table_name`, …)\n", + "5. Poll `GET /v1/ingest/job/{job_id}` or SSE events until the job completes.\n", + "6. Run the same `Retriever.query` + `where` / `filter_hits_by_content_metadata` examples from Part A against the service LanceDB URI.\n", + "\n", + "**Notes**\n", + "\n", + "- Remote ingest does **not** accept a local `meta_dataframe` path; upload via `/v1/ingest/sidecar` and reference `meta_dataframe_id` only.\n", + "- Sidecars are in-memory on the service pod, scoped by bearer token when auth is on, and evicted after TTL (default 3600s).\n", + "- See `nemo_retriever/tests/test_service_sidecar.py` for request shapes.\n", + "\n", + "**TODO (@jioffe502, @jdye64):** Add runnable cells using `RetrieverServiceClient` / `ServiceIngestor` against a Helm deployment, including job creation, sidecar upload, document upload with `vdb_upload_params`, and filtered query validation." + ], + "id": "26aacc42" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Helm / service implementation — to be added on this branch.\n", + "# Planned outline:\n", + "# from nemo_retriever.service.client import RetrieverServiceClient\n", + "# from nemo_retriever.service.service_ingestor import ServiceIngestor\n", + "#\n", + "# client = RetrieverServiceClient(base_url=SERVICE_URL, api_key=API_KEY)\n", + "# sidecar_id = ... # POST /v1/ingest/sidecar with meta_path\n", + "# job_id = client.create_job(expected_documents=len(files))\n", + "# pipeline_spec = {\"vdb_upload_params\": {\"meta_dataframe_id\": sidecar_id, ...}}\n", + "# client.ingest_files(files, job_id=job_id, pipeline_spec=pipeline_spec)" + ], + "execution_count": null, + "outputs": [], + "id": "0d6bdaa2" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" } - ], - "source": [ - "hits_wide = retriever.query(q, top_k=16)\n", - "client_only = filter_hits_by_content_metadata(\n", - " hits_wide, lambda m: m.get(\"meta_a\") == \"bravo\" and m.get(\"meta_b\", 0) >= 10\n", - ")\n", - "print(\"wide search then client filter:\", len(client_only))\n", - "\n", - "# Approximate the same intent in SQL on the JSON string (meta_b is stored as a number in JSON)\n", - "pred_bravo = \"metadata LIKE '%\\\"meta_a\\\":\\\"bravo\\\"%' AND metadata LIKE '%\\\"meta_b\\\":10%'\"\n", - "hits_sql = retriever.query(q, top_k=16, vdb_kwargs={**base_vdb, \"where\": pred_bravo})\n", - "print(\"single query with where:\", len(hits_sql))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "290964a0-7d68-418a-8872-7f6ef6b220ad", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.13" - } - }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat": 4, + "nbformat_minor": 5 }