nv-ingest-api 2025.4.21.dev20250421__py3-none-any.whl → 2025.4.23.dev20250423__py3-none-any.whl

nv_ingest_api/internal/primitives/nim/model_interface/helpers.py

@@ -0,0 +1,275 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import backoff
+import cv2
+import numpy as np
+import requests
+
+from nv_ingest_api.internal.primitives.nim.model_interface.decorators import multiprocessing_cache
+from nv_ingest_api.util.image_processing.transforms import pad_image, normalize_image
+from nv_ingest_api.util.string_processing import generate_url, remove_url_endpoints
+
+logger = logging.getLogger(__name__)
+
+
+def preprocess_image_for_paddle(array: np.ndarray, image_max_dimension: int = 960) -> np.ndarray:
+    """
+    Preprocesses an input image to be suitable for use with PaddleOCR by resizing, normalizing, padding,
+    and transposing it into the required format.
+
+    This function is intended for preprocessing images to be passed as input to PaddleOCR using GRPC.
+    It is not necessary when using the HTTP endpoint.
+
+    Steps:
+    -----
+    1. Resizes the image while maintaining aspect ratio such that its largest dimension is scaled to 960 pixels.
+    2. Normalizes the image using the `normalize_image` function.
+    3. Pads the image to ensure both its height and width are multiples of 32, as required by PaddleOCR.
+    4. Transposes the image from (height, width, channel) to (channel, height, width), the format expected by PaddleOCR.
+
+    Parameters:
+    ----------
+    array : np.ndarray
+        The input image array of shape (height, width, channels). It should have pixel values in the range [0, 255].
+
+    Returns:
+    -------
+    np.ndarray
+        A preprocessed image with the shape (channels, height, width) and normalized pixel values.
+        The image will be padded to have dimensions that are multiples of 32, with the padding color set to 0.
+
+    Notes:
+    -----
+    - The image is resized so that its largest dimension becomes 960 pixels, maintaining the aspect ratio.
+    - After normalization, the image is padded to the nearest multiple of 32 in both dimensions, which is
+      a requirement for PaddleOCR.
+    - The normalized pixel values are scaled between 0 and 1 before padding and transposing the image.
+    """
+    height, width = array.shape[:2]
+    scale_factor = image_max_dimension / max(height, width)
+    new_height = int(height * scale_factor)
+    new_width = int(width * scale_factor)
+    resized = cv2.resize(array, (new_width, new_height))
+
+    normalized = normalize_image(resized)
+
+    # PaddleOCR NIM (GRPC) requires input shapes to be multiples of 32.
+    new_height = (normalized.shape[0] + 31) // 32 * 32
+    new_width = (normalized.shape[1] + 31) // 32 * 32
+    padded, (pad_width, pad_height) = pad_image(
+        normalized, target_height=new_height, target_width=new_width, background_color=0, dtype=np.float32
+    )
+
+    # PaddleOCR NIM (GRPC) requires input to be (channel, height, width).
+    transposed = padded.transpose((2, 0, 1))
+
+    # Metadata can used for inverting transformations on the resulting bounding boxes.
+    metadata = {
+        "original_height": height,
+        "original_width": width,
+        "scale_factor": scale_factor,
+        "new_height": transposed.shape[1],
+        "new_width": transposed.shape[2],
+        "pad_height": pad_height,
+        "pad_width": pad_width,
+    }
+
+    return transposed, metadata
+
+
+def is_ready(http_endpoint: str, ready_endpoint: str) -> bool:
+    """
+    Check if the server at the given endpoint is ready.
+
+    Parameters
+    ----------
+    http_endpoint : str
+        The HTTP endpoint of the server.
+    ready_endpoint : str
+        The specific ready-check endpoint.
+
+    Returns
+    -------
+    bool
+        True if the server is ready, False otherwise.
+    """
+
+    # IF the url is empty or None that means the service was not configured
+    # and is therefore automatically marked as "ready"
+    if http_endpoint is None or http_endpoint == "":
+        return True
+
+    # If the url is for build.nvidia.com, it is automatically assumed "ready"
+    if "ai.api.nvidia.com" in http_endpoint:
+        return True
+
+    url = generate_url(http_endpoint)
+    url = remove_url_endpoints(url)
+
+    if not ready_endpoint.startswith("/") and not url.endswith("/"):
+        ready_endpoint = "/" + ready_endpoint
+
+    url = url + ready_endpoint
+
+    # Call the ready endpoint of the NIM
+    try:
+        # Use a short timeout to prevent long hanging calls. 5 seconds seems resonable
+        resp = requests.get(url, timeout=5)
+        if resp.status_code == 200:
+            # The NIM is saying it is ready to serve
+            return True
+        elif resp.status_code == 503:
+            # NIM is explicitly saying it is not ready.
+            return False
+        else:
+            # Any other code is confusing. We should log it with a warning
+            # as it could be something that might hold up ready state
+            logger.warning(f"'{url}' HTTP Status: {resp.status_code} - Response Payload: {resp.json()}")
+            return False
+    except requests.HTTPError as http_err:
+        logger.warning(f"'{url}' produced a HTTP error: {http_err}")
+        return False
+    except requests.Timeout:
+        logger.warning(f"'{url}' request timed out")
+        return False
+    except ConnectionError:
+        logger.warning(f"A connection error for '{url}' occurred")
+        return False
+    except requests.RequestException as err:
+        logger.warning(f"An error occurred: {err} for '{url}'")
+        return False
+    except Exception as ex:
+        # Don't let anything squeeze by
+        logger.warning(f"Exception: {ex}")
+        return False
+
+
+def _query_metadata(
+    http_endpoint: str,
+    field_name: str,
+    default_value: str,
+    retry_value: str = "",
+    metadata_endpoint: str = "/v1/metadata",
+) -> str:
+    if (http_endpoint is None) or (http_endpoint == ""):
+        return default_value
+
+    url = generate_url(http_endpoint)
+    url = remove_url_endpoints(url)
+
+    if not metadata_endpoint.startswith("/") and not url.endswith("/"):
+        metadata_endpoint = "/" + metadata_endpoint
+
+    url = url + metadata_endpoint
+
+    # Call the metadata endpoint of the NIM
+    try:
+        # Use a short timeout to prevent long hanging calls. 5 seconds seems reasonable
+        resp = requests.get(url, timeout=5)
+        if resp.status_code == 200:
+            field_value = resp.json().get(field_name, "")
+            if field_value:
+                return field_value
+            else:
+                # If the field is empty, retry
+                logger.warning(f"No {field_name} field in response from '{url}'. Retrying.")
+                return retry_value
+        else:
+            # Any other code is confusing. We should log it with a warning
+            logger.warning(f"'{url}' HTTP Status: {resp.status_code} - Response Payload: {resp.text}")
+            return retry_value
+    except requests.HTTPError as http_err:
+        logger.warning(f"'{url}' produced a HTTP error: {http_err}")
+        return retry_value
+    except requests.Timeout:
+        logger.warning(f"'{url}' request timed out")
+        return retry_value
+    except ConnectionError:
+        logger.warning(f"A connection error for '{url}' occurred")
+        return retry_value
+    except requests.RequestException as err:
+        logger.warning(f"An error occurred: {err} for '{url}'")
+        return retry_value
+    except Exception as ex:
+        # Don't let anything squeeze by
+        logger.warning(f"Exception: {ex}")
+        return retry_value
+
+
+@multiprocessing_cache(max_calls=100)  # Cache results first to avoid redundant retries from backoff
+@backoff.on_predicate(backoff.expo, max_time=30)
+def get_version(http_endpoint: str, metadata_endpoint: str = "/v1/metadata", version_field: str = "version") -> str:
+    """
+    Get the version of the server from its metadata endpoint.
+
+    Parameters
+    ----------
+    http_endpoint : str
+        The HTTP endpoint of the server.
+    metadata_endpoint : str, optional
+        The metadata endpoint to query (default: "/v1/metadata").
+    version_field : str, optional
+        The field containing the version in the response (default: "version").
+
+    Returns
+    -------
+    str
+        The version of the server, or an empty string if unavailable.
+    """
+    default_version = "1.0.0"
+
+    # TODO: Need a way to match NIM version to API versions.
+    if "ai.api.nvidia.com" in http_endpoint or "api.nvcf.nvidia.com" in http_endpoint:
+        return default_version
+
+    return _query_metadata(
+        http_endpoint,
+        field_name=version_field,
+        default_value=default_version,
+    )
+
+
+@multiprocessing_cache(max_calls=100)  # Cache results first to avoid redundant retries from backoff
+@backoff.on_predicate(backoff.expo, max_time=30)
+def get_model_name(
+    http_endpoint: str,
+    default_model_name,
+    metadata_endpoint: str = "/v1/metadata",
+    model_info_field: str = "modelInfo",
+) -> str:
+    """
+    Get the model name of the server from its metadata endpoint.
+
+    Parameters
+    ----------
+    http_endpoint : str
+        The HTTP endpoint of the server.
+    metadata_endpoint : str, optional
+        The metadata endpoint to query (default: "/v1/metadata").
+    model_info_field : str, optional
+        The field containing the model info in the response (default: "modelInfo").
+
+    Returns
+    -------
+    str
+        The model name of the server, or an empty string if unavailable.
+    """
+    if "ai.api.nvidia.com" in http_endpoint:
+        return http_endpoint.strip("/").strip("/chat/completions").split("/")[-1]
+
+    if "api.nvcf.nvidia.com" in http_endpoint:
+        return default_model_name
+
+    model_info = _query_metadata(
+        http_endpoint,
+        field_name=model_info_field,
+        default_value={"shortName": default_model_name},
+    )
+    short_name = model_info[0].get("shortName", default_model_name)
+    model_name = short_name.split(":")[0]
+
+    return model_name

nv_ingest_api/internal/primitives/nim/model_interface/nemoretriever_parse.py

@@ -0,0 +1,238 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+
+from nv_ingest_api.internal.primitives.nim import ModelInterface
+from nv_ingest_api.util.image_processing.transforms import numpy_to_base64
+
+ACCEPTED_TEXT_CLASSES = set(
+    [
+        "Text",
+        "Title",
+        "Section-header",
+        "List-item",
+        "TOC",
+        "Bibliography",
+        "Formula",
+        "Page-header",
+        "Page-footer",
+        "Caption",
+        "Footnote",
+        "Floating-text",
+    ]
+)
+ACCEPTED_TABLE_CLASSES = set(
+    [
+        "Table",
+    ]
+)
+ACCEPTED_IMAGE_CLASSES = set(
+    [
+        "Picture",
+    ]
+)
+ACCEPTED_CLASSES = ACCEPTED_TEXT_CLASSES | ACCEPTED_TABLE_CLASSES | ACCEPTED_IMAGE_CLASSES
+
+logger = logging.getLogger(__name__)
+
+
+class NemoRetrieverParseModelInterface(ModelInterface):
+    """
+    An interface for handling inference with a NemoRetrieverParse model.
+    """
+
+    def __init__(self, model_name: str = "nvidia/nemoretriever-parse"):
+        """
+        Initialize the instance with a specified model name.
+        Parameters
+        ----------
+        model_name : str, optional
+            The name of the model to be used, by default "nvidia/nemoretriever-parse".
+        """
+        self.model_name = model_name
+
+    def name(self) -> str:
+        """
+        Get the name of the model interface.
+
+        Returns
+        -------
+        str
+            The name of the model interface.
+        """
+        return "nemoretriever_parse"
+
+    def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Prepare input data for inference by resizing images and storing their original shapes.
+
+        Parameters
+        ----------
+        data : dict
+            The input data containing a list of images.
+
+        Returns
+        -------
+        dict
+            The updated data dictionary with resized images and original image shapes.
+        """
+
+        return data
+
+    def format_input(self, data: Dict[str, Any], protocol: str, max_batch_size: int, **kwargs) -> Any:
+        """
+        Format input data for the specified protocol.
+
+        Parameters
+        ----------
+        data : dict
+            The input data to format.
+        protocol : str
+            The protocol to use ("grpc" or "http").
+        **kwargs : dict
+            Additional parameters for HTTP payload formatting.
+
+        Returns
+        -------
+        Any
+            The formatted input data.
+
+        Raises
+        ------
+        ValueError
+            If an invalid protocol is specified.
+        """
+
+        # Helper function: chunk a list into sublists of length <= chunk_size.
+        def chunk_list(lst: list, chunk_size: int) -> List[list]:
+            return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]
+
+        if protocol == "grpc":
+            raise ValueError("gRPC protocol is not supported for NemoRetrieverParse.")
+        elif protocol == "http":
+            logger.debug("Formatting input for HTTP NemoRetrieverParse model")
+            # Prepare payload for HTTP request
+
+            if "images" in data:
+                base64_list = [numpy_to_base64(img) for img in data["images"]]
+            else:
+                base64_list = [numpy_to_base64(data["image"])]
+
+            formatted_batches = []
+            formatted_batch_data = []
+            b64_chunks = chunk_list(base64_list, max_batch_size)
+
+            for b64_chunk in b64_chunks:
+                payload = self._prepare_nemoretriever_parse_payload(b64_chunk)
+                formatted_batches.append(payload)
+                formatted_batch_data.append({})
+            return formatted_batches, formatted_batch_data
+
+        else:
+            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")
+
+    def parse_output(self, response: Any, protocol: str, data: Optional[Dict[str, Any]] = None, **kwargs) -> Any:
+        """
+        Parse the output from the model's inference response.
+
+        Parameters
+        ----------
+        response : Any
+            The response from the model inference.
+        protocol : str
+            The protocol used ("grpc" or "http").
+        data : dict, optional
+            Additional input data passed to the function.
+
+        Returns
+        -------
+        Any
+            The parsed output data.
+
+        Raises
+        ------
+        ValueError
+            If an invalid protocol is specified.
+        """
+
+        if protocol == "grpc":
+            raise ValueError("gRPC protocol is not supported for NemoRetrieverParse.")
+        elif protocol == "http":
+            logger.debug("Parsing output from HTTP NemoRetrieverParse model")
+            return self._extract_content_from_nemoretriever_parse_response(response)
+        else:
+            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")
+
+    def process_inference_results(self, output: Any, **kwargs) -> Any:
+        """
+        Process inference results for the NemoRetrieverParse model.
+
+        Parameters
+        ----------
+        output : Any
+            The raw output from the model.
+
+        Returns
+        -------
+        Any
+            The processed inference results.
+        """
+
+        return output
+
+    def _prepare_nemoretriever_parse_payload(self, base64_list: List[str]) -> Dict[str, Any]:
+        messages = []
+
+        for b64_img in base64_list:
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{b64_img}",
+                            },
+                        }
+                    ],
+                }
+            )
+        payload = {
+            "model": self.model_name,
+            "messages": messages,
+        }
+
+        return payload
+
+    def _extract_content_from_nemoretriever_parse_response(self, json_response: Dict[str, Any]) -> Any:
+        """
+        Extract content from the JSON response of a Deplot HTTP API request.
+
+        Parameters
+        ----------
+        json_response : dict
+            The JSON response from the Deplot API.
+
+        Returns
+        -------
+        Any
+            The extracted content from the response.
+
+        Raises
+        ------
+        RuntimeError
+            If the response does not contain the expected "choices" key or if it is empty.
+        """
+
+        if "choices" not in json_response or not json_response["choices"]:
+            raise RuntimeError("Unexpected response format: 'choices' key is missing or empty.")
+
+        tool_call = json_response["choices"][0]["message"]["tool_calls"][0]
+        return json.loads(tool_call["function"]["arguments"])