From 7471f3361b20a25f097027b896f0f5472763f73d Mon Sep 17 00:00:00 2001 From: edknv Date: Tue, 18 Aug 2026 09:46:09 -0700 Subject: [PATCH] Require page rendering for page-level image embeddings --- docs/docs/extraction/embedding.md | 8 ++- .../nemo_retriever/graph/ingestor_runtime.py | 17 +++++ .../ingestor/branch_extraction.py | 16 ++++- nemo_retriever/tests/test_ingest_manifest.py | 39 ++++++++++- .../tests/test_pdf_render_params_wiring.py | 61 +++++++++++++++- nemo_retriever/tests/test_pipeline_graph.py | 69 +++++++++++++++++-- 6 files changed, 201 insertions(+), 9 deletions(-) diff --git a/docs/docs/extraction/embedding.md b/docs/docs/extraction/embedding.md index 90a56909a..cf2e13abf 100644 --- a/docs/docs/extraction/embedding.md +++ b/docs/docs/extraction/embedding.md @@ -65,7 +65,13 @@ For documents where the entire page layout is important (such as infographics, c you can configure NeMo Retriever Library to treat every page as a single image. The following example extracts and embeds each page as an image. -- The `embed` method processes the page images. +- Set `embed_modality="image"` to use the rendered page image as the embedding input. +- Set `embed_granularity="page"` to create one result row for each PDF page. + +These arguments work together. When you set both arguments, the pipeline +enables page-image rendering during extraction, creates one row for each page, +and embeds the full rendered page image. Either argument alone does not enable +the complete page-as-image workflow. For parameter details, refer to the [Python API guide](nemo-retriever-api-reference.md) (`create_ingestor` and `.embed()`). diff --git a/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py b/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py index 9c99d8537..14fe52bdc 100644 --- a/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py +++ b/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py @@ -67,6 +67,14 @@ def _local_embed_requested(params: Any) -> bool: return gpu_embed is not None and float(gpu_embed) > 0 +def _image_embedding_requires_page_image(params: Any | None) -> bool: + """Return whether embedding consumes a rendered image for each page row.""" + return getattr(params, "embed_granularity", None) == "page" and getattr(params, "embed_modality", None) in { + "image", + "text_image", + } + + def default_concurrency_node_names( extract_params: Any | None, embed_params: Any | None, @@ -697,6 +705,15 @@ def build_graph( if split_config is None: split_config = resolve_split_params(None) + # Page-level image modalities require a full page raster regardless of + # whether PDF extraction runs through the dedicated or auto-dispatch graph. + if ( + extract_params is not None + and _image_embedding_requires_page_image(embed_params) + and not extract_params.extract_page_as_image + ): + extract_params = extract_params.model_copy(update={"extract_page_as_image": True}) + # Video ingestion uses a dedicated chain so each stage (fan-out, ASR, # frame OCR, scene fusion) shows up as its own Ray Data MapBatches op. # The audio-only shortcut below would otherwise short-circuit to a diff --git a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py index 370c708f3..a7ed9caf4 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py @@ -7,7 +7,7 @@ from __future__ import annotations import logging -from dataclasses import dataclass +from dataclasses import dataclass, replace from io import BytesIO from typing import Any, Callable @@ -18,6 +18,7 @@ build_graph, build_post_extract_graph, default_concurrency_node_names, + _image_embedding_requires_page_image, ) from nemo_retriever.ingestor.manifest import ( ExtractionBranchPlan, @@ -192,7 +193,7 @@ def _should_reshape_content_before_embed(self) -> bool: return any(branch.family in {"pdf", "image"} for branch in self.branches) def _resolve_branch(self, branch: ExtractionBranchPlan) -> ResolvedExtractionInputs: - return resolve_branch_extraction_inputs( + resolved = resolve_branch_extraction_inputs( branch, extract_params=self.extract_params, text_params=self.text_params, @@ -203,6 +204,17 @@ def _resolve_branch(self, branch: ExtractionBranchPlan) -> ResolvedExtractionInp video_text_dedup_params=self.video_text_dedup_params, av_fuse_params=self.av_fuse_params, ) + if ( + branch.family == "pdf" + and resolved.extract_params is not None + and _image_embedding_requires_page_image(self.embed_params) + and not resolved.extract_params.extract_page_as_image + ): + resolved = replace( + resolved, + extract_params=resolved.extract_params.model_copy(update={"extract_page_as_image": True}), + ) + return resolved def _build_extraction_only_graph(self, effective_extraction: ResolvedExtractionInputs) -> Any: return build_graph( diff --git a/nemo_retriever/tests/test_ingest_manifest.py b/nemo_retriever/tests/test_ingest_manifest.py index 7332b1756..3a13f7cd0 100644 --- a/nemo_retriever/tests/test_ingest_manifest.py +++ b/nemo_retriever/tests/test_ingest_manifest.py @@ -25,7 +25,7 @@ plan_extraction_branches, resolve_branch_extraction_inputs, ) -from nemo_retriever.common.params import ASRParams +from nemo_retriever.common.params import ASRParams, EmbedParams, ExtractParams def _resolve_plan( @@ -453,6 +453,43 @@ def fake_post_graph(**kwargs: Any) -> Graph: assert post_calls[0]["reshape_content_before_embed"] is False +@pytest.mark.parametrize("modality", ["image", "text_image"]) +def test_mixed_branch_image_embedding_enables_pdf_page_raster(monkeypatch, tmp_path, modality: str) -> None: + pdf = tmp_path / "manual.pdf" + text_file = tmp_path / "notes.txt" + pdf.write_bytes(b"pdf") + text_file.write_text("notes", encoding="utf-8") + extraction_calls: list[dict[str, Any]] = [] + + def fake_build_graph(**kwargs: Any) -> Graph: + extraction_calls.append(kwargs) + return _graph_with(_TagOperator(tag=kwargs["extraction_mode"])) + + monkeypatch.setattr("nemo_retriever.ingestor.branch_extraction.build_graph", fake_build_graph) + monkeypatch.setattr( + "nemo_retriever.ingestor.branch_extraction.build_post_extract_graph", + lambda **_kwargs: _graph_with(_PostOperator()), + ) + + ( + GraphIngestor(run_mode="inprocess", show_progress=False) + .files([str(pdf), str(text_file)]) + .extract( + ExtractParams( + extract_images=False, + extract_tables=False, + extract_charts=False, + extract_page_as_image=False, + ) + ) + .embed(EmbedParams(embed_modality=modality, embed_granularity="page")) + .ingest() + ) + + pdf_call = next(call for call in extraction_calls if call["extraction_mode"] == "pdf") + assert pdf_call["extract_params"].extract_page_as_image is True + + class _FakeDataset: def __init__(self, columns: list[str]) -> None: self.columns = columns diff --git a/nemo_retriever/tests/test_pdf_render_params_wiring.py b/nemo_retriever/tests/test_pdf_render_params_wiring.py index d26e35397..0efd27636 100644 --- a/nemo_retriever/tests/test_pdf_render_params_wiring.py +++ b/nemo_retriever/tests/test_pdf_render_params_wiring.py @@ -6,9 +6,14 @@ from __future__ import annotations +from pathlib import Path + +import numpy as np import pandas as pd +import pytest -from nemo_retriever.common.params import ExtractParams +from nemo_retriever import create_ingestor +from nemo_retriever.common.params import EmbedParams, ExtractParams from nemo_retriever.graph.graph_pipeline_registry import get_node_kwargs from nemo_retriever.graph.ingestor_runtime import build_graph from nemo_retriever.operators.graph_ops import multi_type_extract_operator as multi_type_module @@ -63,3 +68,57 @@ def __init__(self, **kwargs) -> None: assert captured["render_mode"] == "full_dpi" assert captured["image_format"] == "png" assert captured["jpeg_quality"] == 73 + + +@pytest.mark.parametrize("modality", ["image", "text_image"]) +def test_sdk_page_image_embedding_materializes_pdf_rasters_and_vectors(monkeypatch, modality: str) -> None: + class _FakeVLEmbedder: + def embed_images(self, images: list[str], batch_size: int = 8) -> np.ndarray: + assert all(isinstance(image, str) and image for image in images) + return np.ones((len(images), 2048), dtype=np.float32) + + def embed_text_image(self, texts: list[str], images: list[str], batch_size: int = 8) -> np.ndarray: + assert all(isinstance(text, str) and text for text in texts) + assert all(isinstance(image, str) and image for image in images) + return np.ones((len(images), 2048), dtype=np.float32) + + monkeypatch.setattr( + "nemo_retriever.models.create_local_embedder", + lambda *_args, **_kwargs: _FakeVLEmbedder(), + ) + document = Path(__file__).resolve().parents[2] / "data" / "multimodal_test.pdf" + + result = ( + create_ingestor(run_mode="inprocess", allow_no_gpu=True) + .files([str(document)]) + .extract( + ExtractParams( + extract_images=False, + extract_tables=False, + extract_charts=False, + extract_page_as_image=False, + use_page_elements=False, + ) + ) + .embed( + EmbedParams( + model_name="nvidia/llama-nemotron-embed-vl-1b-v2", + embed_model_name="nvidia/llama-nemotron-embed-vl-1b-v2", + local_ingest_embed_backend="hf", + embed_modality=modality, + embed_granularity="page", + ) + ) + .ingest() + ) + + assert len(result) == 3 + assert ( + result["page_image"] + .map(lambda page_image: isinstance(page_image, dict) and bool(page_image.get("image_b64"))) + .all() + ) + assert result["_image_b64"].map(lambda image_b64: isinstance(image_b64, str) and bool(image_b64)).all() + assert result["_contains_embeddings"].all() + assert result["text_embeddings_1b_v2_has_embedding"].all() + assert (result["text_embeddings_1b_v2_dim"] == 2048).all() diff --git a/nemo_retriever/tests/test_pipeline_graph.py b/nemo_retriever/tests/test_pipeline_graph.py index 9517339cc..7b2ce95eb 100644 --- a/nemo_retriever/tests/test_pipeline_graph.py +++ b/nemo_retriever/tests/test_pipeline_graph.py @@ -37,17 +37,21 @@ from nemo_retriever.common.ray_resource_hueristics import Resources -def _graph_node_names(graph: Graph) -> list[str]: - names: list[str] = [] +def _graph_nodes(graph: Graph) -> list[Node]: + nodes: list[Node] = [] def visit(node: Node) -> None: - names.append(getattr(node.operator, "name", node.name)) + nodes.append(node) for child in node.children: visit(child) for root in graph.roots: visit(root) - return names + return nodes + + +def _graph_node_names(graph: Graph) -> list[str]: + return [getattr(node.operator, "name", node.name) for node in _graph_nodes(graph)] def test_post_extract_graph_uses_explicit_content_reshape_flag() -> None: @@ -72,6 +76,63 @@ def test_text_build_graph_does_not_use_modal_content_reshape() -> None: assert "ExplodeContentToRows" not in _graph_node_names(graph) +@pytest.mark.parametrize("modality", ["image", "text_image"]) +def test_pdf_image_embedding_enables_page_raster(modality: str) -> None: + graph = build_graph( + extraction_mode="pdf", + extract_params=ExtractParams( + extract_images=False, + extract_tables=False, + extract_charts=False, + extract_page_as_image=False, + ), + embed_params=EmbedParams( + embed_modality=modality, + embed_granularity="page", + local_ingest_embed_backend="hf", + ), + ) + + pdf_extract_node = next( + node for node in _graph_nodes(graph) if node.operator.__class__.__name__ == "PDFExtractionActor" + ) + + assert pdf_extract_node.operator_kwargs["extract_page_as_image"] is True + + +def test_pdf_text_embedding_preserves_disabled_page_raster() -> None: + graph = build_graph( + extraction_mode="pdf", + extract_params=ExtractParams( + extract_images=False, + extract_tables=False, + extract_charts=False, + extract_page_as_image=False, + ), + embed_params=EmbedParams(embed_modality="text", embed_granularity="page"), + ) + + pdf_extract_node = next( + node for node in _graph_nodes(graph) if node.operator.__class__.__name__ == "PDFExtractionActor" + ) + + assert pdf_extract_node.operator_kwargs["extract_page_as_image"] is False + + +@pytest.mark.parametrize("modality", ["image", "text_image"]) +def test_auto_image_page_embedding_enables_page_raster(modality: str) -> None: + graph = build_graph( + extraction_mode="auto", + extract_params=ExtractParams(extract_page_as_image=False), + embed_params=EmbedParams(embed_modality=modality, embed_granularity="page"), + ) + + operator = graph.roots[0].operator + + assert isinstance(operator, MultiTypeExtractOperator) + assert operator.extract_params.extract_page_as_image is True + + def test_batch_graph_forwards_resolvable_hosted_parse_contract() -> None: from nemo_retriever.operators.extract.parse.nemotron_parse import _resolve_nemotron_parse_contract