nv-ingest-api 2025.7.14.dev20250714__tar.gz → 2025.7.16.dev20250716__tar.gz

{nv_ingest_api-2025.7.14.dev20250714/src/nv_ingest_api.egg-info → nv_ingest_api-2025.7.16.dev20250716}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.7.14.dev20250714
+Version: 2025.7.16.dev20250716
 Summary: Python module with core document ingestion functions.
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License
@@ -217,6 +217,7 @@ Requires-Dist: backoff==2.2.1
 Requires-Dist: pandas>=2.0
 Requires-Dist: pydantic>2.0.0
 Requires-Dist: pydantic-settings>2.0.0
+Requires-Dist: tritonclient
 Dynamic: license-file
 
 # nv-ingest-api

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "pandas>=2.0",
     "pydantic>2.0.0",
     "pydantic-settings>2.0.0",
+    "tritonclient",
 ]
 
 [project.urls]

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/enums/common.py

@@ -52,6 +52,8 @@ class ContentDescriptionEnum(str, Enum):
         Description for image extracted from PDF document.
     PDF_INFOGRAPHIC : str
         Description for structured infographic extracted from PDF document.
+    PDF_PAGE_IMAGE : str
+        Description for a full-page image rendered from a PDF document.
     PDF_TABLE : str
         Description for structured table extracted from PDF document.
     PDF_TEXT : str
@@ -70,6 +72,7 @@ class ContentDescriptionEnum(str, Enum):
     PDF_CHART: str = "Structured chart extracted from PDF document."
     PDF_IMAGE: str = "Image extracted from PDF document."
     PDF_INFOGRAPHIC: str = "Structured infographic extracted from PDF document."
+    PDF_PAGE_IMAGE: str = "Full-page image rendered from a PDF document."
     PDF_TABLE: str = "Structured table extracted from PDF document."
     PDF_TEXT: str = "Unstructured text from PDF document."
     PPTX_IMAGE: str = "Image extracted from PPTX presentation."
@@ -94,6 +97,8 @@ class ContentTypeEnum(str, Enum):
         Represents image content.
     INFO_MSG : str
         Represents an informational message.
+    PAGE_IMAGE : str
+        Represents a full-page image rendered from a document.
     STRUCTURED : str
         Represents structured content.
     TEXT : str
@@ -111,6 +116,7 @@ class ContentTypeEnum(str, Enum):
     INFOGRAPHIC: str = "infographic"
     INFO_MSG: str = "info_message"
     NONE: str = "none"
+    PAGE_IMAGE: str = "page_image"
     STRUCTURED: str = "structured"
     TABLE: str = "table"
     TEXT: str = "text"

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/extract/pdf/engines/nemoretriever.py

@@ -40,6 +40,7 @@ from nv_ingest_api.internal.schemas.meta.metadata_schema import validate_metadat
 from nv_ingest_api.internal.primitives.nim.model_interface.yolox import (
     YOLOX_PAGE_IMAGE_PREPROC_WIDTH,
     YOLOX_PAGE_IMAGE_PREPROC_HEIGHT,
+    YOLOX_PAGE_IMAGE_FORMAT,
 )
 from nv_ingest_api.internal.schemas.extract.extract_pdf_schema import NemoRetrieverParseConfigSchema
 from nv_ingest_api.util.metadata.aggregators import (
@@ -355,7 +356,7 @@ def nemoretriever_parse_extractor(
                 img_numpy = crop_image(page_image, transformed_bbox)
 
                 if img_numpy is not None:
-                    base64_img = numpy_to_base64(img_numpy)
+                    base64_img = numpy_to_base64(img_numpy, format=YOLOX_PAGE_IMAGE_FORMAT)
                     image = Base64Image(
                         image=base64_img,
                         bbox=transformed_bbox,

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/extract/pdf/engines/pdf_helpers/__init__.py

@@ -4,20 +4,21 @@
 # Copyright (c) 2024, NVIDIA CORPORATION.
 
 import base64
+import inspect
 import io
-
-import pandas as pd
-from typing import Any, Dict, List, Optional
 import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
 
-from nv_ingest_api.internal.extract.pdf.engines import (
-    adobe_extractor,
-    llama_parse_extractor,
-    nemoretriever_parse_extractor,
-    pdfium_extractor,
-    tika_extractor,
-    unstructured_io_extractor,
-)
+import pandas as pd
+from nv_ingest_api.internal.extract.pdf.engines import adobe_extractor
+from nv_ingest_api.internal.extract.pdf.engines import llama_parse_extractor
+from nv_ingest_api.internal.extract.pdf.engines import nemoretriever_parse_extractor
+from nv_ingest_api.internal.extract.pdf.engines import pdfium_extractor
+from nv_ingest_api.internal.extract.pdf.engines import tika_extractor
+from nv_ingest_api.internal.extract.pdf.engines import unstructured_io_extractor
 from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
 
 # Import extraction functions for different engines.
@@ -43,6 +44,7 @@ def _work_extract_pdf(
     extract_infographics: bool,
     extract_tables: bool,
     extract_charts: bool,
+    extract_page_as_image: bool,
     extractor_config: dict,
     execution_trace_log=None,
 ) -> Any:
@@ -52,17 +54,25 @@ def _work_extract_pdf(
 
     extract_method = extractor_config["extract_method"]
     extractor_fn = EXTRACTOR_LOOKUP.get(extract_method, pdfium_extractor)
-    return extractor_fn(
-        pdf_stream,
-        extract_text,
-        extract_images,
-        extract_infographics,
-        extract_tables,
-        extract_charts,
-        extractor_config,
-        execution_trace_log,
+
+    extractor_fn_args = dict(
+        pdf_stream=pdf_stream,
+        extract_text=extract_text,
+        extract_images=extract_images,
+        extract_infographics=extract_infographics,
+        extract_tables=extract_tables,
+        extract_charts=extract_charts,
+        extractor_config=extractor_config,
+        execution_trace_log=execution_trace_log,
     )
 
+    if "extract_page_as_image" in inspect.signature(extractor_fn).parameters:
+        extractor_fn_args["extract_page_as_image"] = extract_page_as_image
+    elif extract_page_as_image:
+        logger.warning(f"`extract_page_as_image` is set to True, but {extract_method} does not support it.")
+
+    return extractor_fn(**extractor_fn_args)
+
 
 @unified_exception_handler
 def _orchestrate_row_extraction(
@@ -97,6 +107,7 @@ def _orchestrate_row_extraction(
         extract_tables = params.pop("extract_tables", False)
         extract_charts = params.pop("extract_charts", False)
         extract_infographics = params.pop("extract_infographics", False)
+        extract_page_as_image = params.pop("extract_page_as_image", False)
         extract_method = params.get("extract_method", "pdfium")
     except KeyError as e:
         raise ValueError(f"Missing required extraction flag: {e}")
@@ -137,6 +148,7 @@ def _orchestrate_row_extraction(
         extract_text=extract_text,
         extract_images=extract_images,
         extract_infographics=extract_infographics,
+        extract_page_as_image=extract_page_as_image,
         extract_tables=extract_tables,
         extract_charts=extract_charts,
         extractor_config=extractor_config,

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/extract/pdf/engines/pdfium.py

@@ -24,16 +24,19 @@ import numpy as np
 import pandas as pd
 import pypdfium2 as libpdfium
 
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
 from nv_ingest_api.internal.primitives.nim.default_values import YOLOX_MAX_BATCH_SIZE
 from nv_ingest_api.internal.primitives.nim.model_interface.yolox import (
     YOLOX_PAGE_IMAGE_PREPROC_WIDTH,
     YOLOX_PAGE_IMAGE_PREPROC_HEIGHT,
+    YOLOX_PAGE_IMAGE_FORMAT,
     get_yolox_model_name,
     YoloxPageElementsModelInterface,
 )
 from nv_ingest_api.internal.schemas.extract.extract_pdf_schema import PDFiumConfigSchema
 from nv_ingest_api.internal.enums.common import TableFormatEnum, TextTypeEnum, AccessLevelEnum
 from nv_ingest_api.util.metadata.aggregators import (
+    construct_image_metadata_from_base64,
     construct_image_metadata_from_pdf_image,
     extract_pdf_metadata,
     construct_text_metadata,
@@ -46,6 +49,7 @@ from nv_ingest_api.util.pdf.pdfium import (
     extract_image_like_objects_from_pdfium_page,
 )
 from nv_ingest_api.util.pdf.pdfium import pdfium_pages_to_numpy
+from nv_ingest_api.util.image_processing import scale_image_to_encoding_size
 from nv_ingest_api.util.image_processing.transforms import numpy_to_base64, crop_image
 
 logger = logging.getLogger(__name__)
@@ -186,7 +190,7 @@ def _extract_page_element_images(
             if cropped is None:
                 continue
 
-            base64_img = numpy_to_base64(cropped)
+            base64_img = numpy_to_base64(cropped, format=YOLOX_PAGE_IMAGE_FORMAT)
 
             bbox_in_orig_coord = (
                 int(w1) - pad_width,
@@ -384,6 +388,7 @@ def pdfium_extractor(
     extract_infographics: bool,
     extract_tables: bool,
     extract_charts: bool,
+    extract_page_as_image: bool,
     extractor_config: dict,
     execution_trace_log: Optional[List[Any]] = None,
 ) -> pd.DataFrame:
@@ -524,6 +529,24 @@ def pdfium_extractor(
                 )
                 extracted_data.extend(image_data)
 
+            # Full page image extraction
+            if extract_page_as_image:
+                page_text = _extract_page_text(page)
+                image, _ = pdfium_pages_to_numpy([page], scale_tuple=(16384, 16384), trace_info=execution_trace_log)
+                base64_image = numpy_to_base64(image[0])
+                if len(base64_image) > 2**24 - 1:
+                    base64_image, _ = scale_image_to_encoding_size(base64_image, max_base64_size=2**24 - 1)
+                image_meta = construct_image_metadata_from_base64(
+                    base64_image,
+                    page_idx,
+                    page_count,
+                    source_metadata,
+                    base_unified_metadata,
+                    subtype=ContentTypeEnum.PAGE_IMAGE,
+                    text=page_text,
+                )
+                extracted_data.append(image_meta)
+
             # If we want tables or charts, rasterize the page and store it
             if extract_tables or extract_charts or extract_infographics:
                 image, padding_offsets = pdfium_pages_to_numpy(
@@ -574,6 +597,7 @@ def pdfium_extractor(
                 execution_trace_log=execution_trace_log,
             )
             futures.append(future)
+
             pages_for_tables.clear()
 
         # Wait for all asynchronous jobs to complete.

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/primitives/nim/model_interface/nemoretriever_parse.py

@@ -120,6 +120,7 @@ class NemoRetrieverParseModelInterface(ModelInterface):
             logger.debug("Formatting input for HTTP NemoRetrieverParse model")
             # Prepare payload for HTTP request
 
+            ## TODO: Ask @Edward Kim if we want to switch to JPEG/PNG here
             if "images" in data:
                 base64_list = [numpy_to_base64(img) for img in data["images"]]
             else:

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/primitives/nim/model_interface/yolox.py

@@ -2,9 +2,7 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-
-import base64
-import io
+import os
 import logging
 import warnings
 from math import log
@@ -20,11 +18,11 @@ import packaging
 import pandas as pd
 import torch
 import torchvision
-from PIL import Image
 
 from nv_ingest_api.internal.primitives.nim import ModelInterface
 from nv_ingest_api.internal.primitives.nim.model_interface.helpers import get_model_name
 from nv_ingest_api.util.image_processing import scale_image_to_encoding_size
+from nv_ingest_api.util.image_processing.transforms import numpy_to_base64
 
 logger = logging.getLogger(__name__)
 
@@ -35,6 +33,7 @@ YOLOX_PAGE_MIN_SCORE = 0.1
 YOLOX_PAGE_NIM_MAX_IMAGE_SIZE = 512_000
 YOLOX_PAGE_IMAGE_PREPROC_HEIGHT = 1024
 YOLOX_PAGE_IMAGE_PREPROC_WIDTH = 1024
+YOLOX_PAGE_IMAGE_FORMAT = os.getenv("YOLOX_PAGE_IMAGE_FORMAT", "PNG")
 
 # yolox-page-elements-v1 contants
 YOLOX_PAGE_V1_NUM_CLASSES = 4
@@ -239,15 +238,11 @@ class YoloxModelInterfaceBase(ModelInterface):
                 # Convert to uint8 if needed.
                 if image.dtype != np.uint8:
                     image = (image * 255).astype(np.uint8)
-                # Convert the numpy array to a PIL Image.
-                image_pil = Image.fromarray(image)
-                original_size = image_pil.size
-
-                # Save the image to a buffer and encode to base64.
-                buffered = io.BytesIO()
-                image_pil.save(buffered, format="PNG")
-                image_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
 
+                # Get original size directly from numpy array (width, height)
+                original_size = (image.shape[1], image.shape[0])
+                # Convert numpy array directly to base64 using OpenCV
+                image_b64 = numpy_to_base64(image, format=YOLOX_PAGE_IMAGE_FORMAT)
                 # Scale the image if necessary.
                 scaled_image_b64, new_size = scale_image_to_encoding_size(
                     image_b64, max_base64_size=self.nim_max_image_size

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/schemas/meta/ingest_job_schema.py

@@ -107,6 +107,10 @@ class IngestTaskEmbedSchema(BaseModelNoExt):
     model_name: Optional[str] = None
     api_key: Optional[str] = None
     filter_errors: bool = False
+    text_elements_modality: Optional[str] = None
+    image_elements_modality: Optional[str] = None
+    structured_elements_modality: Optional[str] = None
+    audio_elements_modality: Optional[str] = None
 
 
 class IngestTaskVdbUploadSchema(BaseModelNoExt):
@@ -195,6 +199,7 @@ class IngestTaskSchema(BaseModelNoExt):
         validated_task_properties = expected_schema_cls(**task_properties)
         values["type"] = task_type  # ensure type is now always the enum
         values["task_properties"] = validated_task_properties
+
         return values
 
     @field_validator("type", mode="before")

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/schemas/transform/transform_text_embedding_schema.py

@@ -22,5 +22,9 @@ class TextEmbeddingSchema(BaseModel):
     input_type: str = Field(default="passage")
     raise_on_failure: bool = Field(default=False)
     truncate: str = Field(default="END")
+    text_elements_modality: str = Field(default="text")
+    image_elements_modality: str = Field(default="text")
+    structured_elements_modality: str = Field(default="text")
+    audio_elements_modality: str = Field(default="text")
 
     model_config = ConfigDict(extra="forbid")

{nv_ingest_api-2025.7.14.dev20250714 → nv_ingest_api-2025.7.16.dev20250716}/src/nv_ingest_api/internal/transform/embed_text.py

@@ -4,6 +4,7 @@
 
 import logging
 from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 from typing import Any, Dict, Tuple, Optional, Iterable, List
 
 import pandas as pd
@@ -19,6 +20,9 @@ from nv_ingest_api.util.schema.schema_validator import validate_schema
 logger = logging.getLogger(__name__)
 
 
+MULTI_MODAL_MODELS = ["llama-3.2-nemoretriever-1b-vlm-embed-v1"]
+
+
 # ------------------------------------------------------------------------------
 # Asynchronous Embedding Requests
 # ------------------------------------------------------------------------------
@@ -33,6 +37,7 @@ def _make_async_request(
     input_type: str,
     truncate: str,
     filter_errors: bool,
+    modalities: Optional[List[str]] = None,
 ) -> list:
     """
     Interacts directly with the NIM embedding service to calculate embeddings for a batch of prompts.
@@ -74,11 +79,18 @@ def _make_async_request(
             base_url=embedding_nim_endpoint,
         )
 
+        extra_body = {
+            "input_type": input_type,
+            "truncate": truncate,
+        }
+        if modalities:
+            extra_body["modality"] = modalities
+
         resp = client.embeddings.create(
             input=prompts,
             model=embedding_model,
             encoding_format=encoding_format,
-            extra_body={"input_type": input_type, "truncate": truncate},
+            extra_body=extra_body,
         )
 
         response["embedding"] = resp.data
@@ -110,6 +122,7 @@ def _async_request_handler(
     input_type: str,
     truncate: str,
     filter_errors: bool,
+    modalities: Optional[List[str]] = None,
 ) -> List[dict]:
     """
     Gathers calculated embedding results from the NIM embedding service concurrently.
@@ -138,6 +151,9 @@ def _async_request_handler(
     List[dict]
         A list of response dictionaries from the embedding service.
     """
+    if modalities is None:
+        modalities = [None] * len(prompts)
+
     with ThreadPoolExecutor() as executor:
         futures = [
             executor.submit(
@@ -150,8 +166,9 @@ def _async_request_handler(
                 input_type=input_type,
                 truncate=truncate,
                 filter_errors=filter_errors,
+                modalities=modality_batch,
             )
-            for prompt_batch in prompts
+            for prompt_batch, modality_batch in zip(prompts, modalities)
         ]
         results = [future.result() for future in futures]
 
@@ -167,6 +184,7 @@ def _async_runner(
     input_type: str,
     truncate: str,
     filter_errors: bool,
+    modalities: Optional[List[str]] = None,
 ) -> dict:
     """
     Concurrently launches all NIM embedding requests and flattens the results.
@@ -204,6 +222,7 @@ def _async_runner(
         input_type,
         truncate,
         filter_errors,
+        modalities=modalities,
     )
 
     flat_results = {"embeddings": [], "info_msgs": []}
@@ -263,7 +282,19 @@ def _add_embeddings(row, embeddings, info_msgs):
     return row
 
 
-def _get_pandas_text_content(row):
+def _format_image_input_string(image_b64: Optional[str]) -> str:
+    if not image_b64:
+        return
+    return f"data:image/png;base64,{image_b64}"
+
+
+def _format_text_image_pair_input_string(text: Optional[str], image_b64: Optional[str]) -> str:
+    if (not text) or (not text.strip()) or (not image_b64):
+        return
+    return f"{text.strip()} {_format_image_input_string(image_b64)}"
+
+
+def _get_pandas_text_content(row, modality="text"):
     """
     Extracts text content from a DataFrame row.
 
@@ -280,7 +311,7 @@ def _get_pandas_text_content(row):
     return row["content"]
 
 
-def _get_pandas_table_content(row):
+def _get_pandas_table_content(row, modality="text"):
     """
     Extracts table/chart content from a DataFrame row.
 
@@ -294,10 +325,19 @@ def _get_pandas_table_content(row):
     str
         The table/chart content from the row.
     """
-    return row.get("table_metadata", {}).get("table_content")
+    if modality == "text":
+        content = row.get("table_metadata", {}).get("table_content")
+    elif modality == "image":
+        content = _format_image_input_string(row.get("content"))
+    elif modality == "text_image":
+        text = row.get("table_metadata", {}).get("table_content")
+        image = row.get("content")
+        content = _format_text_image_pair_input_string(text, image)
+
+    return content
 
 
-def _get_pandas_image_content(row):
+def _get_pandas_image_content(row, modality="text"):
     """
     Extracts image caption content from a DataFrame row.
 
@@ -311,10 +351,28 @@ def _get_pandas_image_content(row):
     str
         The image caption from the row.
     """
-    return row.get("image_metadata", {}).get("caption")
+    subtype = row.get("content_metadata", {}).get("subtype")
+    if modality == "text":
+        if subtype == "page_image":
+            content = row.get("image_metadata", {}).get("text")
+        else:
+            content = row.get("image_metadata", {}).get("caption")
+    elif modality == "image":
+        content = _format_image_input_string(row.get("content"))
+    elif modality == "text_image":
+        if subtype == "page_image":
+            text = row.get("image_metadata", {}).get("text")
+        else:
+            text = row.get("image_metadata", {}).get("caption")
+        image = row.get("content")
+        content = _format_text_image_pair_input_string(text, image)
 
+    # A workaround to save memory.
+    row["content"] = ""
+    return content
 
-def _get_pandas_audio_content(row):
+
+def _get_pandas_audio_content(row, modality="text"):
     """
     A pandas UDF used to select extracted audio transcription to be used to create embeddings.
     """
@@ -408,6 +466,23 @@ def _concatenate_extractions_pandas(
 # ------------------------------------------------------------------------------
 
 
+def does_model_support_multimodal_embeddings(model: str) -> bool:
+    """
+    Checks if a given model supports multi-modal embeddings.
+
+    Parameters
+    ----------
+    model : str
+        The name of the model.
+
+    Returns
+    -------
+    bool
+        True if the model supports multi-modal embeddings, False otherwise.
+    """
+    return model in MULTI_MODAL_MODELS
+
+
 def transform_create_text_embeddings_internal(
     df_transform_ledger: pd.DataFrame,
     task_config: Dict[str, Any],
@@ -460,6 +535,15 @@ def transform_create_text_embeddings_internal(
         ContentTypeEnum.AUDIO: _get_pandas_audio_content,
         ContentTypeEnum.VIDEO: lambda x: None,  # Not supported yet.
     }
+    task_type_to_modality = {
+        ContentTypeEnum.TEXT: task_config.get("text_elements_modality") or transform_config.text_elements_modality,
+        ContentTypeEnum.STRUCTURED: (
+            task_config.get("structured_elements_modality") or transform_config.structured_elements_modality
+        ),
+        ContentTypeEnum.IMAGE: task_config.get("image_elements_modality") or transform_config.image_elements_modality,
+        ContentTypeEnum.AUDIO: task_config.get("audio_elements_modality") or transform_config.audio_elements_modality,
+        ContentTypeEnum.VIDEO: lambda x: None,  # Not supported yet.
+    }
 
     def _content_type_getter(row):
         return row["content_metadata"]["type"]
@@ -480,7 +564,7 @@ def transform_create_text_embeddings_internal(
         # Extract content and normalize empty or non-str to None
         extracted_content = (
             df_content["metadata"]
-            .apply(content_getter)
+            .apply(partial(content_getter, modality=task_type_to_modality[content_type]))
             .apply(lambda x: x.strip() if isinstance(x, str) and x.strip() else None)
         )
         df_content["_content"] = extracted_content
@@ -488,9 +572,15 @@ def transform_create_text_embeddings_internal(
         # Prepare batches for only valid (non-None) content
         valid_content_mask = df_content["_content"].notna()
         if valid_content_mask.any():
-            filtered_content_batches = _generate_batches(
-                df_content.loc[valid_content_mask, "_content"].tolist(), batch_size=transform_config.batch_size
-            )
+            filtered_content_list = df_content.loc[valid_content_mask, "_content"].tolist()
+            filtered_content_batches = _generate_batches(filtered_content_list, batch_size=transform_config.batch_size)
+
+            if model_name in MULTI_MODAL_MODELS:
+                modality_list = [task_type_to_modality[content_type]] * len(filtered_content_list)
+                modality_batches = _generate_batches(modality_list, batch_size=transform_config.batch_size)
+            else:
+                modality_batches = None
+
             content_embeddings = _async_runner(
                 filtered_content_batches,
                 api_key,
@@ -500,6 +590,7 @@ def transform_create_text_embeddings_internal(
                 transform_config.input_type,
                 transform_config.truncate,
                 False,
+                modalities=modality_batches,
             )
             # Build a simple row index -> embedding map
             embeddings_dict = dict(