Изграждане на мултимодален RAG конвейер с NVIDIA NeMo Retriever, хоствани NIMs, LanceDB, повторно класиране и генериране с опора на източници
В този урок изграждаме усъвършенстван мултимодален конвейер за генериране с допълнено извличане с помощта на NVIDIA NeMo Retriever. Започваме с конфигуриране на среда с Python 3.12, инсталиране на необходимите пакети и офлайн извличане на текст от PDF файлове, без да разчитаме на GPU или външен API ключ. След това разширяваме работния процес с хоствани крайни точки NVIDIA NIM за откриване на елементи на страницата, извличане на таблици, диаграми и инфографики, генериране на плътни векторни ембединг представяния и съхраняване на обработеното съдържание в LanceDB. Накрая реализираме плътно извличане, визуално-езиково пренареждане, търсене с филтриране по метаданни, генериране на обосновани отговори с вградени цитирания и олекотена оценка recall-at-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 файл и го задаваме като входен документ за конвейера. След това извършваме базирано на CPU офлайн извличане на текст с PDFium и проверяваме извлечените редове, колони и съдържание на страниците.
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}")
Зареждаме сигурно API ключа на NVIDIA и дефинираме хостваните крайни точки NIM за откриване на оформление, OCR, извличане на таблици, графичен анализ, ембединг, пренареждане и генериране. Създаваме мултимодален конвейер за приемане, който извлича текст, таблици, диаграми и инфографики, като прилага разделяне на части с отчитане на токените и дедупликация на съдържанието. Генерираме ембединг представяния за извлеченото съдържание и качваме получените вектори и метаданни в таблица на 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, за да генерираме отговори, обосновани единствено от предоставения контекст. Включваме номерирани препратки към източниците и метаданни за страниците, така че генерираните отговори да могат да бъдат проследени до оригиналния документ. Завършваме с изчисляване на recall at k за малък набор от очаквани отговори и отчитаме финалните артефакти на векторната база данни и Markdown.
В заключение създадохме цялостна мултимодална RAG система, която преобразува структурирано и неструктурирано PDF съдържание в база от знания, готова за търсене и цитиране. Използвахме NeMo Retriever за координиране на извличането, дедупликацията, разделянето на части, създаването на ембединг представяния, индексирането във векторна база данни, извличането и пренареждането, като запазихме Colab средата олекотена чрез делегиране на инференцията на моделите към хоствани NVIDIA NIM услуги. Генерирахме и обосновани отговори с езиков модел Nemotron и измерихме ефективността на извличането с помощта на прост тест recall-at-k. Със завършването на този работен процес създадохме многократно използваема основа за изграждане на приложения за интелигентна обработка на документи, които обработват текст, таблици, диаграми и визуални елементи чрез единен конвейер за извличане.
Разгледайте ПЪЛНИТЕ КОДОВЕ тук. Също така можете да ни последвате в Twitter и не забравяйте да се присъедините към нашия ML SubReddit с над 150 хил. членове и да се абонирате за нашия информационен бюлетин. Чакайте! В Telegram ли сте? Вече можете да се присъедините към нас и в Telegram.
Искате да си партнирате с нас за популяризиране на вашето GitHub хранилище, Hugging Face страница, продуктово представяне, уебинар и др.? Свържете се с нас
Преведено автоматично от английски. Оригиналната статия е на връзката по-долу.