Створення мультимодального RAG-конвеєра з NVIDIA NeMo Retriever, розміщеними NIM, LanceDB, повторним ранжуванням і генерацією з опорою на джерела
У цьому посібнику ми створимо вдосконалений мультимодальний конвеєр генерації з доповненням пошуку за допомогою NVIDIA NeMo Retriever. Спочатку ми налаштуємо середовище Python 3.12, встановимо необхідні пакети та виконаємо офлайн-вилучення тексту з PDF без використання графічного процесора або зовнішнього API-ключа. Потім розширимо робочий процес за допомогою розміщених кінцевих точок NVIDIA NIM для виявлення елементів сторінки, вилучення таблиць, діаграм та інфографіки, створення щільних векторних вкладень і збереження обробленого вмісту в LanceDB. Нарешті, ми реалізуємо щільний пошук, переранжування за допомогою vision-language-моделі, пошук із фільтрацією метаданих, генерацію обґрунтованих відповідей із вбудованими цитатами та полегшену оцінку 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-файлу та визначаємо його як вхідний документ для конвеєра. Потім виконуємо офлайн-вилучення тексту на базі ЦП за допомогою 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())
Ми створюємо конвеєр переранжування vision-language, який отримує ширший набір кандидатів і впорядковує результати відповідно до семантичної релевантності. Ми також застосовуємо текстовий фільтр, щоб звузити результати пошуку до фрагментів, які містять певний вміст документа. Ми безпосередньо перевіряємо таблицю 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, релізу продукту, вебінару тощо? Зв’яжіться з нами
Перекладено автоматично з англійської. Оригінал статті — за посиланням нижче.