nv-ingest-api 2025.5.11.dev20250511__tar.gz → 2025.5.13.dev20250513__tar.gz

{nv_ingest_api-2025.5.11.dev20250511/src/nv_ingest_api.egg-info → nv_ingest_api-2025.5.13.dev20250513}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.5.11.dev20250511
+Version: 2025.5.13.dev20250513
 Summary: Python module with core document ingestion functions.
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/interface/transform.py

@@ -207,7 +207,7 @@ def transform_image_create_vlm_caption(
         "api_key": api_key,
         "prompt": prompt,
         "endpoint_url": endpoint_url,
-        "model_name": model_name,
+        "image_caption_model_name": model_name,
     }
     filtered_task_config: Dict[str, str] = {k: v for k, v in task_config.items() if v is not None}
 

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/extract/docx/docx_extractor.py

@@ -7,7 +7,7 @@ import base64
 import functools
 import io
 import logging
-from typing import Optional, Dict, Any, Union
+from typing import Optional, Dict, Any, Union, Tuple
 
 import pandas as pd
 from pydantic import BaseModel
@@ -146,7 +146,7 @@ def extract_primitives_from_docx_internal(
     task_config: Union[Dict[str, Any], BaseModel],
     extraction_config: DocxExtractorSchema,
     execution_trace_log: Optional[Dict[str, Any]] = None,
-) -> pd.DataFrame:
+) -> Tuple[pd.DataFrame, Union[Dict, None]]:
     """
     Processes a pandas DataFrame containing DOCX files encoded in base64, extracting text from
     each document and replacing the original content with the extracted text.
@@ -202,4 +202,4 @@ def extract_primitives_from_docx_internal(
     else:
         extracted_df = pd.DataFrame({"document_type": [], "metadata": [], "uuid": []})
 
-    return extracted_df
+    return extracted_df, {}

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/extract/image/image_extractor.py

@@ -16,7 +16,7 @@ import pandas as pd
 from pydantic import BaseModel
 
 from nv_ingest_api.internal.extract.image.image_helpers.common import unstructured_image_extractor
-from nv_ingest_api.internal.schemas.extract.extract_image_schema import ImageExtractorSchema
+from nv_ingest_api.internal.schemas.extract.extract_image_schema import ImageConfigSchema
 from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
 
 logger = logging.getLogger(__name__)
@@ -26,7 +26,7 @@ logger = logging.getLogger(__name__)
 def _decode_and_extract_from_image(
     base64_row: pd.Series,
     task_config: Dict[str, Any],
-    validated_extraction_config: ImageExtractorSchema,
+    validated_extraction_config: ImageConfigSchema,
     execution_trace_log: Optional[List[Any]] = None,
 ) -> Any:
     """
@@ -106,10 +106,10 @@ def _decode_and_extract_from_image(
 
         logger.debug(
             f"decode_and_extract: Extracting image content using image_extraction_config: "
-            f"{validated_extraction_config.image_extraction_config}"
+            f"{validated_extraction_config}"
         )
-        if validated_extraction_config.image_extraction_config is not None:
-            extract_params["image_extraction_config"] = validated_extraction_config.image_extraction_config
+        if validated_extraction_config is not None:
+            extract_params["image_extraction_config"] = validated_extraction_config
 
         if execution_trace_log is not None:
             extract_params["trace_info"] = execution_trace_log

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/extract/pdf/engines/nemoretriever.py

@@ -476,7 +476,7 @@ def _extract_text_and_bounding_boxes(
 
 def _create_clients(nemoretriever_parse_config):
     model_interface = nemoretriever_parse_utils.NemoRetrieverParseModelInterface(
-        model_name=nemoretriever_parse_config.model_name,
+        model_name=nemoretriever_parse_config.nemoretriever_parse_model_name,
     )
     nemoretriever_parse_client = create_inference_client(
         nemoretriever_parse_config.nemoretriever_parse_endpoints,

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/extract/pptx/engines/pptx_helper.py

@@ -17,7 +17,6 @@
 
 import logging
 import io
-import operator
 import re
 import uuid
 from collections import defaultdict
@@ -155,6 +154,12 @@ def _finalize_images(
             extracted_data.append(image_entry)
 
 
+def _safe_position(shape):
+    top = shape.top if shape.top is not None else float("inf")
+    left = shape.left if shape.left is not None else float("inf")
+    return (top, left)
+
+
 # -----------------------------------------------------------------------------
 # Helper Function: Recursive Image Extraction
 # -----------------------------------------------------------------------------
@@ -283,7 +288,7 @@ def python_pptx(
 
     for slide_idx, slide in enumerate(presentation.slides):
         # Obtain a flat list of shapes (ungrouped) sorted by top then left.
-        shapes = sorted(ungroup_shapes(slide.shapes), key=operator.attrgetter("top", "left"))
+        shapes = sorted(ungroup_shapes(slide.shapes), key=_safe_position)
 
         page_nearby_blocks = {
             "text": {"content": [], "bbox": []},
@@ -656,21 +661,43 @@ def get_bbox(
     shape_object: Optional[Slide] = None,
     text_depth: Optional[TextTypeEnum] = None,
 ):
-    bbox = (-1, -1, -1, -1)
-    if text_depth == TextTypeEnum.DOCUMENT:
-        bbox = (-1, -1, -1, -1)
-    elif text_depth == TextTypeEnum.PAGE:
-        top = left = 0
-        width = presentation_object.slide_width
-        height = presentation_object.slide_height
-        bbox = (top, left, top + height, left + width)
-    elif shape_object:
-        top = shape_object.top
-        left = shape_object.left
-        width = shape_object.width
-        height = shape_object.height
-        bbox = (top, left, top + height, left + width)
-    return bbox
+    """
+    Safely computes bounding box for a slide, shape, or document.
+    Ensures that missing or None values are gracefully handled.
+
+    Returns
+    -------
+    Tuple[int, int, int, int]
+        Bounding box as (top, left, bottom, right).
+        Defaults to (-1, -1, -1, -1) if invalid or unsupported.
+    """
+    try:
+        if text_depth == TextTypeEnum.DOCUMENT:
+            return (-1, -1, -1, -1)
+
+        elif text_depth == TextTypeEnum.PAGE and presentation_object:
+            top = left = 0
+            width = presentation_object.slide_width
+            height = presentation_object.slide_height
+            return (top, left, top + height, left + width)
+
+        elif shape_object:
+            top = shape_object.top if shape_object.top is not None else -1
+            left = shape_object.left if shape_object.left is not None else -1
+            width = shape_object.width if shape_object.width is not None else -1
+            height = shape_object.height if shape_object.height is not None else -1
+
+            # If all are valid, return normally, else return placeholder
+            if -1 in [top, left, width, height]:
+                return (-1, -1, -1, -1)
+
+            return (top, left, top + height, left + width)
+
+    except Exception as e:
+        logger.warning(f"get_bbox: Failed to compute bbox due to {e}")
+        return (-1, -1, -1, -1)
+
+    return (-1, -1, -1, -1)
 
 
 def ungroup_shapes(shapes):

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/extract/pptx/pptx_extractor.py

@@ -184,4 +184,4 @@ def extract_primitives_from_pptx_internal(
     else:
         extracted_df = pd.DataFrame({"document_type": [], "metadata": [], "uuid": []})
 
-    return extracted_df
+    return extracted_df, {}

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/primitives/nim/model_interface/text_embedding.py

@@ -5,9 +5,9 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from nv_ingest_api.internal.primitives.nim import ModelInterface
+import numpy as np
 
 
-# Assume ModelInterface is defined elsewhere in the project.
 class EmbeddingModelInterface(ModelInterface):
     """
     An interface for handling inference with an embedding model endpoint.
@@ -22,20 +22,13 @@ class EmbeddingModelInterface(ModelInterface):
 
     def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Prepare input data for embedding inference. Ensures that a 'prompts' key is provided
-        and that its value is a list.
-
-        Raises
-        ------
-        KeyError
-            If the 'prompts' key is missing.
+        Prepare input data for embedding inference. Returns a list of strings representing the text to be embedded.
         """
         if "prompts" not in data:
             raise KeyError("Input data must include 'prompts'.")
-        # Ensure the prompts are in list format.
         if not isinstance(data["prompts"], list):
             data["prompts"] = [data["prompts"]]
-        return data
+        return {"prompts": data["prompts"]}
 
     def format_input(
         self, data: Dict[str, Any], protocol: str, max_batch_size: int, **kwargs
@@ -63,29 +56,32 @@ class EmbeddingModelInterface(ModelInterface):
               - payloads is a list of JSON-serializable payload dictionaries.
               - batch_data_list is a list of dictionaries containing the key "prompts" corresponding to each batch.
         """
-        if protocol != "http":
-            raise ValueError("EmbeddingModelInterface only supports HTTP protocol.")
-
-        prompts = data.get("prompts", [])
 
         def chunk_list(lst, chunk_size):
+            lst = lst["prompts"]
             return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]
 
-        batches = chunk_list(prompts, max_batch_size)
-        payloads = []
-        batch_data_list = []
-        for batch in batches:
-            payload = {
-                "model": kwargs.get("model_name"),
-                "input": batch,
-                "encoding_format": kwargs.get("encoding_format", "float"),
-                "extra_body": {
-                    "input_type": kwargs.get("input_type", "query"),
+        batches = chunk_list(data, max_batch_size)
+        if protocol == "http":
+            payloads = []
+            batch_data_list = []
+            for batch in batches:
+                payload = {
+                    "model": kwargs.get("model_name"),
+                    "input": batch,
+                    "encoding_format": kwargs.get("encoding_format", "float"),
+                    "input_type": kwargs.get("input_type", "passage"),
                     "truncate": kwargs.get("truncate", "NONE"),
-                },
-            }
-            payloads.append(payload)
-            batch_data_list.append({"prompts": batch})
+                }
+                payloads.append(payload)
+                batch_data_list.append({"prompts": batch})
+        elif protocol == "grpc":
+            payloads = []
+            batch_data_list = []
+            for batch in batches:
+                text_np = np.array([[text.encode("utf-8")] for text in batch], dtype=np.object_)
+                payloads.append(text_np)
+                batch_data_list.append({"prompts": batch})
         return payloads, batch_data_list
 
     def parse_output(self, response: Any, protocol: str, data: Optional[Dict[str, Any]] = None, **kwargs) -> Any:
@@ -108,16 +104,17 @@ class EmbeddingModelInterface(ModelInterface):
         list
             A list of generated embeddings extracted from the response.
         """
-        if protocol != "http":
-            raise ValueError("EmbeddingModelInterface only supports HTTP protocol.")
-        if isinstance(response, dict):
-            embeddings = response.get("data")
-            if not embeddings:
-                raise RuntimeError("Unexpected response format: 'data' key is missing or empty.")
-            # Each item in embeddings is expected to have an 'embedding' field.
-            return [item.get("embedding", None) for item in embeddings]
-        else:
-            return [str(response)]
+        if protocol == "http":
+            if isinstance(response, dict):
+                embeddings = response.get("data")
+                if not embeddings:
+                    raise RuntimeError("Unexpected response format: 'data' key is missing or empty.")
+                # Each item in embeddings is expected to have an 'embedding' field.
+                return [item.get("embedding", None) for item in embeddings]
+            else:
+                return [str(response)]
+        elif protocol == "grpc":
+            return [res.flatten() for res in response]
 
     def process_inference_results(self, output: Any, protocol: str, **kwargs) -> Any:
         """

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/primitives/nim/model_interface/yolox.py

@@ -709,7 +709,13 @@ def postprocess_results(
             raise ValueError(f"Error in postprocessing {result.shape} and {original_image_shape}: {e}")
 
         for box, score, label in zip(bboxes, scores, labels):
-            class_name = class_labels[int(label)]
+            # TODO(Devin): Sometimes we get back unexpected class labels?
+            if (label < 0) or (label >= len(class_labels)):
+                logger.warning(f"Invalid class label {label} found in postprocessing")
+                continue
+            else:
+                class_name = class_labels[int(label)]
+
             annotation_dict[class_name].append([round(float(x), 4) for x in np.concatenate((box, [score]))])
 
         out.append(annotation_dict)

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/primitives/nim/nim_client.py

@@ -129,7 +129,7 @@ class NimClient:
         """
         if self.protocol == "grpc":
             logger.debug("Performing gRPC inference for a batch...")
-            response = self._grpc_infer(batch_input, model_name)
+            response = self._grpc_infer(batch_input, model_name, **kwargs)
             logger.debug("gRPC inference received response for a batch")
         elif self.protocol == "http":
             logger.debug("Performing HTTP inference for a batch...")
@@ -221,7 +221,7 @@ class NimClient:
 
         return all_results
 
-    def _grpc_infer(self, formatted_input: np.ndarray, model_name: str) -> np.ndarray:
+    def _grpc_infer(self, formatted_input: np.ndarray, model_name: str, **kwargs) -> np.ndarray:
         """
         Perform inference using the gRPC protocol.
 
@@ -238,16 +238,24 @@ class NimClient:
             The output of the model as a numpy array.
         """
 
-        input_tensors = [grpcclient.InferInput("input", formatted_input.shape, datatype="FP32")]
-        input_tensors[0].set_data_from_numpy(formatted_input)
+        parameters = kwargs.get("parameters", {})
+        output_names = kwargs.get("outputs", ["output"])
+        dtype = kwargs.get("dtype", "FP32")
+        input_name = kwargs.get("input_name", "input")
 
-        outputs = [grpcclient.InferRequestedOutput("output")]
-        response = self.client.infer(model_name=model_name, inputs=input_tensors, outputs=outputs)
-        logger.debug(f"gRPC inference response: {response}")
+        input_tensors = grpcclient.InferInput(input_name, formatted_input.shape, datatype=dtype)
+        input_tensors.set_data_from_numpy(formatted_input)
 
-        # TODO(self.client.has_error(response)) => raise error
+        outputs = [grpcclient.InferRequestedOutput(output_name) for output_name in output_names]
+        response = self.client.infer(
+            model_name=model_name, parameters=parameters, inputs=[input_tensors], outputs=outputs
+        )
+        logger.debug(f"gRPC inference response: {response}")
 
-        return response.as_numpy("output")
+        if len(outputs) == 1:
+            return response.as_numpy(outputs[0].name())
+        else:
+            return [response.as_numpy(output.name()) for output in outputs]
 
     def _http_infer(self, formatted_input: dict) -> dict:
         """

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/primitives/tracing/tagging.py

@@ -31,13 +31,15 @@ def traceable(trace_name=None):
 
     Notes
     -----
-    The decorated function must accept a IngestControlMessage object as its first argument. The
-    IngestControlMessage object must implement `has_metadata`, `get_metadata`, and `set_metadata`
-    methods used by the decorator to check for the trace tagging flag and to add trace metadata.
+    The decorated function must accept a IngestControlMessage object as one of its arguments.
+    For a regular function, this is expected to be the first argument; for a class method,
+    this is expected to be the second argument (after 'self'). The IngestControlMessage object
+    must implement `has_metadata`, `get_metadata`, and `set_metadata` methods used by the decorator
+    to check for the trace tagging flag and to add trace metadata.
 
     The trace metadata added by the decorator includes two entries:
-    - 'trace::entry::<trace_name>': The monotonic timestamp marking the function's entry.
-    - 'trace::exit::<trace_name>': The monotonic timestamp marking the function's exit.
+    - 'trace::entry::<trace_name>': The timestamp marking the function's entry.
+    - 'trace::exit::<trace_name>': The timestamp marking the function's exit.
 
     Example
     -------
@@ -47,23 +49,25 @@ def traceable(trace_name=None):
     ... def process_message(message):
     ...     pass
 
-    Applying the decorator with a custom trace name:
-
-    >>> @traceable(custom_trace_name="CustomTraceName")
-    ... def process_message(message):
-    ...     pass
-
-    In both examples, `process_message` will have entry and exit timestamps added to the
-    IngestControlMessage's metadata if 'config::add_trace_tagging' is True.
+    Applying the decorator with a custom trace name on a class method:
 
+    >>> class Processor:
+    ...     @traceable(trace_name="CustomTrace")
+    ...     def process(self, message):
+    ...         pass
     """
 
     def decorator_trace_tagging(func):
         @functools.wraps(func)
         def wrapper_trace_tagging(*args, **kwargs):
-            # Assuming the first argument is always the message
             ts_fetched = datetime.now()
-            message = args[0]
+            # Determine which argument is the message.
+            if hasattr(args[0], "has_metadata"):
+                message = args[0]
+            elif len(args) > 1 and hasattr(args[1], "has_metadata"):
+                message = args[1]
+            else:
+                raise ValueError("traceable decorator could not find a message argument with 'has_metadata()'")
 
             do_trace_tagging = (message.has_metadata("config::add_trace_tagging") is True) and (
                 message.get_metadata("config::add_trace_tagging") is True
@@ -79,7 +83,7 @@ def traceable(trace_name=None):
                     message.set_timestamp(f"trace::entry::{trace_prefix}_channel_in", ts_send)
                     message.set_timestamp(f"trace::exit::{trace_prefix}_channel_in", ts_fetched)
 
-            # Call the decorated function
+            # Call the decorated function.
             result = func(*args, **kwargs)
 
             if do_trace_tagging:

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/schemas/extract/extract_pdf_schema.py

@@ -131,7 +131,7 @@ class NemoRetrieverParseConfigSchema(BaseModel):
     nemoretriever_parse_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
     nemoretriever_parse_infer_protocol: str = ""
 
-    model_name: str = "nvidia/nemoretriever-parse"
+    nemoretriever_parse_model_name: str = "nvidia/nemoretriever-parse"
 
     timeout: float = 300.0
 

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/schemas/meta/ingest_job_schema.py

@@ -76,7 +76,7 @@ class IngestTaskCaptionSchema(BaseModelNoExt):
     api_key: Optional[str] = None
     endpoint_url: Optional[str] = None
     prompt: Optional[str] = None
-    model_name: Optional[str] = None
+    caption_model_name: Optional[str] = None
 
 
 class IngestTaskFilterParamsSchema(BaseModelNoExt):
@@ -104,7 +104,7 @@ class IngestTaskDedupSchema(BaseModelNoExt):
 
 class IngestTaskEmbedSchema(BaseModelNoExt):
     endpoint_url: Optional[str] = None
-    model_name: Optional[str] = None
+    embedding_model_name: Optional[str] = None
     api_key: Optional[str] = None
     filter_errors: bool = False
 

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/schemas/transform/transform_image_caption_schema.py

@@ -10,6 +10,6 @@ class ImageCaptionExtractionSchema(BaseModel):
     api_key: str = "api_key"
     endpoint_url: str = "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-11b-vision-instruct/chat/completions"
     prompt: str = "Caption the content of this image:"
-    model_name: str = "meta/llama-3.2-11b-vision-instruct"
+    image_caption_model_name: str = "meta/llama-3.2-11b-vision-instruct"
     raise_on_failure: bool = False
     model_config = ConfigDict(extra="forbid")

{nv_ingest_api-2025.5.11.dev20250511 → nv_ingest_api-2025.5.13.dev20250513}/src/nv_ingest_api/internal/transform/caption_image.py

@@ -173,7 +173,7 @@ def transform_image_create_vlm_caption_internal(
     api_key: str = task_config.get("api_key") or transform_config.api_key
     prompt: str = task_config.get("prompt") or transform_config.prompt
     endpoint_url: str = task_config.get("endpoint_url") or transform_config.endpoint_url
-    model_name: str = task_config.get("model_name") or transform_config.model_name
+    model_name: str = task_config.get("image_caption_model_name") or transform_config.image_caption_model_name
 
     # Create a mask for rows where the content type is "image".
     df_mask: pd.Series = df_transform_ledger["metadata"].apply(