Require page rendering for page-level image embeddings

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()`).
 

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

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(

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

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()

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