用 NVIDIA NeMo Retriever + LanceDB 构建多模态 RAG 流水线

本教程将构建一套先进的多模态检索增强生成（RAG）管道，依托 NVIDIA NeMo Retriever 实现。首先配置 Python 3.12 环境，安装所需依赖包，并在不依赖 GPU 或外部 API Key 的情况下执行离线 PDF 文本提取。随后，通过 NVIDIA NIM 托管端点扩展工作流，完成页面元素检测、表格与图表及信息图提取、密集向量嵌入生成，并将处理后的内容存储至 LanceDB。最后实现密集检索、视觉语言重排序、元数据过滤搜索、带内联引用的 grounded 响应生成，以及轻量级的 recall@k 评估，以验证多模态文档内容的检索质量。

import sys, os, subprocess, textwrap, json, time, warnings
warnings.filterwarnings("ignore")
assert sys.version_info[:2] == (3, 12), (
   f"nemo-retriever requires Python 3.12.x (found {sys.version.split()[0]}). "
   "Colab's default runtime is 3.12; if you changed it, switch back."
)
def sh(cmd):
   print(f"$ {cmd}")
   subprocess.run(cmd, shell=True, check=False)
try:
   import nemo_retriever
   print("nemo-retriever already installed")
except ImportError:
   sh("pip install -q --ignore-installed PyJWT nemo-retriever openai")
import nemo_retriever
print("nemo-retriever version:", nemo_retriever.__version__)
from nemo_retriever import create_ingestor
try:
   from nemo_retriever.io import to_markdown, to_markdown_by_page
except ImportError:
   from nemo_retriever.common.io import to_markdown, to_markdown_by_page
try:
   from nemo_retriever.retriever import Retriever
except ImportError:
   from nemo_retriever.graph.retriever import Retriever
import pandas as pd
pd.set_option("display.max_colwidth", 160)
DOC = "multimodal_test.pdf"
if not os.path.exists(DOC):
   sh(f"curl -sL -o {DOC} "
      "https://raw.githubusercontent.com/NVIDIA/NeMo-Retriever/main/data/multimodal_test.pdf")
print("document:", DOC, os.path.getsize(DOC), "bytes")
DOCS = [DOC]
print("\n=== STAGE 1: offline text extraction (no API key) ===")
offline = (
   create_ingestor(run_mode="inprocess", allow_no_gpu=True)
   .files(DOCS)
   .extract(
       extract_text=True,
       extract_tables=False, extract_charts=False,
       extract_images=False, extract_infographics=False,
       use_page_elements=False,
       extract_page_as_image=False,
       method="pdfium",
   )
)
df_offline = offline.ingest()
print("rows:", df_offline.shape, "\ncolumns:", list(df_offline.columns))
print("\npage 1 text preview:\n", df_offline.iloc[0]["text"][:400])

配置 Python 3.12 环境，安装 NVIDIA NeMo Retriever 并导入所需的摄取和检索组件。下载示例多模态 PDF 并将其定义为管道的输入文档。随后使用 PDFium 执行基于 CPU 的离线文本提取，并检查提取的行、列和页面内容。

from getpass import getpass
if not os.environ.get("NVIDIA_API_KEY"):
   try:
       from google.colab import userdata
       os.environ["NVIDIA_API_KEY"] = userdata.get("NVIDIA_API_KEY")
   except Exception:
       os.environ["NVIDIA_API_KEY"] = getpass("NVIDIA_API_KEY (nvapi-...): ").strip()
API_KEY = os.environ.get("NVIDIA_API_KEY", "").strip()
HAVE_KEY = API_KEY.startswith("nvapi-")
print("API key present:", HAVE_KEY)
PAGE_ELEMENTS_URL   = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-page-elements-v3"
OCR_URL             = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-ocr-v1"
TABLE_STRUCT_URL    = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-table-structure-v1"
GRAPHIC_ELEM_URL    = "https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-graphic-elements-v1"
EMBED_URL           = "https://integrate.api.nvidia.com/v1/embeddings"
RERANK_URL          = "https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-nemotron-rerank-vl-1b-v2/reranking"
CHAT_URL            = "https://integrate.api.nvidia.com/v1"
EMBED_MODEL  = "nvidia/llama-nemotron-embed-1b-v2"
RERANK_MODEL = "nvidia/llama-nemotron-rerank-vl-1b-v2"
LLM_MODEL    = "nvidia/llama-3.3-nemotron-super-49b-v1.5"
LANCEDB_URI, TABLE = "./lancedb", "colab_demo"
df = df_offline
if HAVE_KEY:
   print("\n=== STAGE 2: multimodal ingest via hosted NIMs ===")
   ing = (
       create_ingestor(
           run_mode="inprocess",
           allow_no_gpu=True,
           error_policy="collect",
       )
       .files(DOCS)
       .extract(
           extract_text=True,
           extract_tables=True,
           extract_charts=True,
           extract_infographics=True,
           extract_images=False,
           method="pdfium",
           dpi=200,
           table_output_format="markdown",
           page_elements_invoke_url=PAGE_ELEMENTS_URL,
           ocr_invoke_url=OCR_URL,
           table_structure_invoke_url=TABLE_STRUCT_URL,
           graphic_elements_invoke_url=GRAPHIC_ELEM_URL,
           api_key=API_KEY,
           request_timeout_s=120.0,
           split_config={"text": {"max_tokens": 512, "overlap_tokens": 64}},
       )
       .dedup(content_hash=True, bbox_iou=True, iou_threshold=0.45)
       .embed(
           embedding_endpoint=EMBED_URL,
           model_name=EMBED_MODEL,
           embed_model_name=EMBED_MODEL,
           api_key=API_KEY,
           input_type="passage",
           inference_batch_size=16,
           nim_http_max_concurrent=8,
       )
       .vdb_upload(
           vdb_op="lancedb",
           vdb_kwargs={
               "uri": LANCEDB_URI,
               "table_name": TABLE,
               "overwrite": True,
               "create_index": True,
               "index_type": "IVF_HNSW_SQ",
               "metric": "l2",
           },
       )
   )
   t0 = time.time()
   df = ing.ingest(show_progress=True)
   print(f"ingested in {time.time()-t0:.1f}s -> {df.shape}")

安全加载 NVIDIA API Key 并定义托管 NIM 端点，用于布局检测、OCR、表格提取、图形分析、嵌入、重排序和生成。创建多模态摄取管道，执行文本、表格、图表和信息图的提取，同时应用基于 token 的分块和内容去重。为提取的内容生成嵌入，并将结果向量及元数据上传至 LanceDB 表。

print("\n=== Extraction inspection ===")
for col in ["tables", "charts", "infographics", "images"]:
   if col in df.columns:
       n = int(df[col].apply(lambda v: len(v) if isinstance(v, (list, tuple)) else 0).sum())
       print(f"  {col:<14} {n}")
pages = to_markdown_by_page(df)
print("\npages rendered to markdown:", list(pages.keys()))
print("\n--- page 1 markdown (first 900 chars) ---\n", pages[min(pages)][:900])
full_md = to_markdown(df)
if full_md:
   with open("extracted.md", "w") as f:
       f.write(full_md)
   print("\nfull document markdown -> extracted.md")
if HAVE_KEY:
   print("\n=== STAGE 3: dense retrieval ===")
   retriever = Retriever(
       run_mode="service",
       top_k=5,
       rerank=False,
       vdb_kwargs={"uri": LANCEDB_URI, "table_name": TABLE},
       embed_kwargs={
           "embedding_endpoint": EMBED_URL,
           "model_name": EMBED_MODEL,
           "embed_model_name": EMBED_MODEL,
           "api_key": API_KEY,
           "input_type": "query",
       },
   )
   QUERIES = [
       "Given their activities, which animal is responsible for the typos in my documents?",
       "What is the most expensive gadget and how much does it cost?",
       "Which animal is at the beach?",
   ]
   def show(hits, label=""):
       print(f"\n--- {label} ---")
       for i, h in enumerate(hits, 1):
           meta = h.get("metadata")
           if isinstance(meta, str):
               try: meta = json.loads(meta)
               except Exception: meta = {}
           page = (meta or {}).get("page_number", "?")
           score = h.get("_distance", h.get("rerank_score", ""))
           body = " ".join(str(h.get("text", "")).split())[:180]
           print(f" {i}. p{page} score={score}  {body}")
   show(retriever.query(QUERIES[0]), "single query")
   for q, hits in zip(QUERIES, retriever.queries(QUERIES, top_k=3)):
       show(hits, q[:60])

检查提取的多模态元素并将处理后的文档转换为页面级和全文档 Markdown。配置密集检索器，对用户查询进行嵌入并在 LanceDB 向量索引中搜索最相关的文档块。测试单独查询和批量查询，同时展示页码、相似度分数和检索到的文本预览。

if HAVE_KEY:
   print("\n=== STAGE 4: retrieve + VL rerank ===")
   reranking = Retriever(
       run_mode="service",
       top_k=5,
       rerank=True,
       vdb_kwargs={"uri": LANCEDB_URI, "table_name": TABLE},
       embed_kwargs={
           "embedding_endpoint": EMBED_URL, "model_name": EMBED_MODEL,
           "embed_model_name": EMBED_MODEL, "api_key": API_KEY, "input_type": "query",
       },
       rerank_kwargs={
           "model_name": RERANK_MODEL,
           "invoke_url": RERANK_URL,
           "api_key": API_KEY,
           "refine_factor": 4,
           "batch_size": 16,
       },
   )
   try:
       show(reranking.query(QUERIES[0]), "reranked")
   except Exception as e:
       print("rerank unavailable, dense results stand:", type(e).__name__, str(e)[:160])
if HAVE_KEY:
   print("\n=== STAGE 5: filtered retrieval ===")
   try:
       hits = retriever.query(
           "gadget costs",
           top_k=5,
           vdb_kwargs={"where": "text LIKE '%Cost%'"},
       )
       show(hits, "where: text LIKE '%Cost%'")
   except Exception as e:
       print("filter push-down failed:", type(e).__name__, str(e)[:160])
   import lancedb
   tbl = lancedb.connect(LANCEDB_URI).open_table(TABLE)
   print("\nrows in LanceDB:", tbl.count_rows())
   print(tbl.to_pandas()[["text"]].head(3).to_string())

创建视觉语言重排序管道，检索更广泛的候选集并根据语义相关性对结果重新排序。同时应用基于文本的过滤器，将检索结果收窄至文档中包含特定内容的块。直接检查 LanceDB 表以验证存储的记录数量并查看索引的文本。

if HAVE_KEY:
   print("\n=== STAGE 6: RAG answer ===")
   from openai import OpenAI
   client = OpenAI(base_url=CHAT_URL, api_key=API_KEY)
   def rag(question, k=5):
       hits = retriever.query(question, top_k=k)
       ctx = []
       for i, h in enumerate(hits, 1):
           meta = h.get("metadata")
           if isinstance(meta, str):
               try: meta = json.loads(meta)
               except Exception: meta = {}
           ctx.append(f"[{i}] (page {(meta or {}).get('page_number','?')})\n{h.get('text','')}")
       prompt = textwrap.dedent(f"""\
           Answer the question using ONLY the numbered context below.
           Cite the sources you used as [1], [2], etc. If the context is
           insufficient, say so plainly.
           Context:
           {chr(10).join(ctx)}
           Question: {question}
           """)
       r = client.chat.completions.create(
           model=LLM_MODEL,
           messages=[{"role": "user", "content": prompt}],
           temperature=0.0, max_tokens=512,
       )
       return r.choices[0].message.content, hits
   for q in QUERIES[:2]:
       try:
           ans, _ = rag(q)
           print(f"\nQ: {q}\nA: {ans}\n" + "-" * 70)
       except Exception as e:
           print("generation failed:", type(e).__name__, str(e)[:200])
if HAVE_KEY:
   print("\n=== Recall@k check ===")
   GOLD = [
       ("which animal is jumping onto a laptop", "Cat"),
       ("what does the chart show", "Gadgets"),
       ("which animal is at the beach", "Giraffe"),
   ]
   K = 5
   hit_lists = retriever.queries([q for q, _ in GOLD], top_k=K)
   got = sum(
       any(exp.lower() in str(h.get("text", "")).lower() for h in hits)
       for (_, exp), hits in zip(GOLD, hit_lists)
   )
   print(f"recall@{K} = {got}/{len(GOLD)} = {got/len(GOLD):.2f}")
print("\nDone. Artifacts: ./lancedb (vector table), ./extracted.md (markdown).")

将检索到的文档块与托管的 Nemotron 语言模型结合，生成仅基于提供的上下文进行 grounded 的回答。包含编号的来源引用和页面元数据，使生成的响应可追溯至原始文档。最后为一小部分预期答案计算 recall@k，并报告最终的向量数据库和 Markdown 产物。

总结

我们构建了一套完整的多模态 RAG 系统，将结构化和非结构化 PDF 内容转化为可搜索、可引用的知识。使用 NeMo Retriever 协调提取、去重、分块、嵌入、向量数据库索引、检索和重排序，同时通过将模型推理委托给托管的 NVIDIA NIM 服务保持 Colab 运行时的轻量。还使用 Nemotron 语言模型生成 grounded 回答，并通过简单的 recall@k 测试衡量检索效果。通过完成此工作流，我们建立了一个可复用的基础架构，用于构建文档智能应用——通过统一检索管道处理文本、表格、图表和视觉元素。

查看完整代码此处。此外，欢迎关注我们的 Twitter，别忘了加入我们的 150k+ ML SubReddit 并订阅我们的 Newsletter。等一下！你用 telegram 吗？现在也可以加入我们了。

需要与我们合作推广您的 GitHub 仓库、Hugging Face 页面、产品发布或网络研讨会吗？请联系我们

Sana Hassan，Marktechpost 咨询实习生，同时是马德拉斯理工学院的双学位学生，热衷于将技术和 AI 应用于解决现实世界的挑战。凭借对解决实际问题的浓厚兴趣，他将新鲜视角带入 AI 与现实生活解决方案的交汇处。