nv-ingest-api 26.1.0rc4__py3-none-any.whl

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

@@ -0,0 +1,270 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Dict, Any, Optional, List
+
+import numpy as np
+import logging
+
+from nv_ingest_api.internal.primitives.nim import ModelInterface
+from nv_ingest_api.util.image_processing.transforms import base64_to_numpy
+
+logger = logging.getLogger(__name__)
+
+
+class DeplotModelInterface(ModelInterface):
+    """
+    An interface for handling inference with a Deplot model, supporting both gRPC and HTTP protocols,
+    now updated to handle multiple base64 images ('base64_images').
+    """
+
+    def name(self) -> str:
+        """
+        Get the name of the model interface.
+
+        Returns
+        -------
+        str
+            The name of the model interface ("Deplot").
+        """
+        return "Deplot"
+
+    def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Prepare input data by decoding one or more base64-encoded images into NumPy arrays.
+
+        Parameters
+        ----------
+        data : dict
+            The input data containing either 'base64_image' (single image)
+            or 'base64_images' (multiple images).
+
+        Returns
+        -------
+        dict
+            The updated data dictionary with 'image_arrays': a list of decoded NumPy arrays.
+        """
+
+        # Handle a single base64_image or multiple base64_images
+        if "base64_images" in data:
+            base64_list = data["base64_images"]
+            if not isinstance(base64_list, list):
+                raise ValueError("The 'base64_images' key must contain a list of base64-encoded strings.")
+            image_arrays = [base64_to_numpy(b64) for b64 in base64_list]
+
+        elif "base64_image" in data:
+            # Fallback for single image
+            image_arrays = [base64_to_numpy(data["base64_image"])]
+        else:
+            raise KeyError("Input data must include 'base64_image' or 'base64_images'.")
+
+        data["image_arrays"] = image_arrays
+
+        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 (gRPC or HTTP) for Deplot.
+        For HTTP, we now construct multiple messages—one per image batch—along with
+        corresponding batch data carrying the original image arrays and their dimensions.
+
+        Parameters
+        ----------
+        data : dict of str -> Any
+            The input data dictionary, expected to contain "image_arrays" (a list of np.ndarray).
+        protocol : str
+            The protocol to use, "grpc" or "http".
+        max_batch_size : int
+            The maximum number of images per batch.
+        kwargs : dict
+            Additional parameters to pass to the payload preparation (for HTTP).
+
+        Returns
+        -------
+        tuple
+            (formatted_batches, formatted_batch_data) where:
+              - For gRPC: formatted_batches is a list of NumPy arrays, each of shape (B, H, W, C)
+                with B <= max_batch_size.
+              - For HTTP: formatted_batches is a list of JSON-serializable payload dicts.
+              - In both cases, formatted_batch_data is a list of dicts containing:
+                    "image_arrays": the list of original np.ndarray images for that batch, and
+                    "image_dims":  a list of (height, width) tuples for each image in the batch.
+
+        Raises
+        ------
+        KeyError
+            If "image_arrays" is missing in the data dictionary.
+        ValueError
+            If the protocol is invalid, or if no valid images are found.
+        """
+        if "image_arrays" not in data:
+            raise KeyError("Expected 'image_arrays' in data. Call prepare_data_for_inference first.")
+
+        image_arrays = data["image_arrays"]
+        # Compute image dimensions from each image array.
+        image_dims = [(img.shape[0], img.shape[1]) for img in image_arrays]
+
+        # 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":
+            logger.debug("Formatting input for gRPC Deplot model (potentially batched).")
+            processed = []
+            for arr in image_arrays:
+                # Ensure each image has shape (1, H, W, C)
+                if arr.ndim == 3:
+                    arr = np.expand_dims(arr, axis=0)
+                arr = arr.astype(np.float32)
+                arr /= 255.0  # Normalize to [0,1]
+                processed.append(arr)
+
+            if not processed:
+                raise ValueError("No valid images found for gRPC formatting.")
+
+            formatted_batches = []
+            formatted_batch_data = []
+            proc_chunks = chunk_list(processed, max_batch_size)
+            orig_chunks = chunk_list(image_arrays, max_batch_size)
+            dims_chunks = chunk_list(image_dims, max_batch_size)
+
+            for proc_chunk, orig_chunk, dims_chunk in zip(proc_chunks, orig_chunks, dims_chunks):
+                # Concatenate along the batch dimension to form a single input.
+                batched_input = np.concatenate(proc_chunk, axis=0)
+                formatted_batches.append(batched_input)
+                formatted_batch_data.append({"image_arrays": orig_chunk, "image_dims": dims_chunk})
+            return formatted_batches, formatted_batch_data
+
+        elif protocol == "http":
+            logger.debug("Formatting input for HTTP Deplot model (multiple messages).")
+            if "base64_images" in data:
+                base64_list = data["base64_images"]
+            else:
+                base64_list = [data["base64_image"]]
+
+            formatted_batches = []
+            formatted_batch_data = []
+            b64_chunks = chunk_list(base64_list, max_batch_size)
+            orig_chunks = chunk_list(image_arrays, max_batch_size)
+            dims_chunks = chunk_list(image_dims, max_batch_size)
+
+            for b64_chunk, orig_chunk, dims_chunk in zip(b64_chunks, orig_chunks, dims_chunks):
+                payload = self._prepare_deplot_payload(
+                    base64_list=b64_chunk,
+                    max_tokens=kwargs.get("max_tokens", 500),
+                    temperature=kwargs.get("temperature", 0.5),
+                    top_p=kwargs.get("top_p", 0.9),
+                )
+                formatted_batches.append(payload)
+                formatted_batch_data.append({"image_arrays": orig_chunk, "image_dims": dims_chunk})
+            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 model's inference response.
+        """
+        if protocol == "grpc":
+            logger.debug("Parsing output from gRPC Deplot model (batched).")
+            # Each batch element might be returned as a list of bytes. Combine or keep separate as needed.
+            results = []
+            for item in response:
+                # If item is [b'...'], decode and join
+                if isinstance(item, list):
+                    joined_str = " ".join(o.decode("utf-8") for o in item)
+                    results.append(joined_str)
+                else:
+                    # single bytes or str
+                    val = item.decode("utf-8") if isinstance(item, bytes) else str(item)
+                    results.append(val)
+            return results  # Return a list of strings, one per image.
+
+        elif protocol == "http":
+            logger.debug("Parsing output from HTTP Deplot model.")
+            return self._extract_content_from_deplot_response(response)
+        else:
+            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")
+
+    def process_inference_results(self, output: Any, protocol: str, **kwargs) -> Any:
+        """
+        Process inference results for the Deplot model.
+
+        Parameters
+        ----------
+        output : Any
+            The raw output from the model.
+        protocol : str
+            The protocol used for inference (gRPC or HTTP).
+
+        Returns
+        -------
+        Any
+            The processed inference results.
+        """
+
+        # For Deplot, the output is the chart content as a string
+        return output
+
+    @staticmethod
+    def _prepare_deplot_payload(
+        base64_list: list,
+        max_tokens: int = 500,
+        temperature: float = 0.5,
+        top_p: float = 0.9,
+    ) -> Dict[str, Any]:
+        """
+        Prepare an HTTP payload for Deplot that includes one message per image,
+        matching the original single-image style:
+
+            messages = [
+              {
+                "role": "user",
+                "content": "Generate ... <img src=\"data:image/png;base64,...\" />"
+              },
+              {
+                "role": "user",
+                "content": "Generate ... <img src=\"data:image/png;base64,...\" />"
+              },
+              ...
+            ]
+
+        If your backend expects multiple messages in a single request, this keeps
+        the same structure as the single-image code repeated N times.
+        """
+        messages = []
+        # Note: deplot NIM currently only supports a single message per request
+        for b64_img in base64_list:
+            messages.append(
+                {
+                    "role": "user",
+                    "content": (
+                        "Generate the underlying data table of the figure below: "
+                        f'<img src="data:image/png;base64,{b64_img}" />'
+                    ),
+                }
+            )
+
+        payload = {
+            "model": "google/deplot",
+            "messages": messages,  # multiple user messages now
+            "max_tokens": max_tokens,
+            "stream": False,
+            "temperature": temperature,
+            "top_p": top_p,
+        }
+        return payload
+
+    @staticmethod
+    def _extract_content_from_deplot_response(json_response: Dict[str, Any]) -> Any:
+        """
+        Extract content from the JSON response of a Deplot HTTP API request.
+        The original code expected a single choice with a single textual content.
+        """
+        if "choices" not in json_response or not json_response["choices"]:
+            raise RuntimeError("Unexpected response format: 'choices' key is missing or empty.")
+
+        # If the service only returns one textual result, we return that one.
+        return json_response["choices"][0]["message"]["content"]

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

@@ -0,0 +1,338 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Optional
+
+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
+
+cv2.setNumThreads(1)
+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 preprocess_image_for_ocr(
+    array: np.ndarray,
+    target_height: Optional[int] = None,
+    target_width: Optional[int] = None,
+    pad_how: str = "bottom_right",
+    normalize: bool = False,
+    channel_first: bool = False,
+) -> np.ndarray:
+    """
+    Preprocesses an input image to be suitable for use with NemoRetriever-OCR.
+
+    This function is intended for preprocessing images to be passed as input to NemoRetriever-OCR using GRPC.
+    It is not necessary when using the HTTP endpoint.
+
+    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).
+    """
+    height, width = array.shape[:2]
+
+    if target_height is None:
+        target_height = height
+
+    if target_width is None:
+        target_width = width
+
+    padded, (pad_width, pad_height) = pad_image(
+        array,
+        target_height=target_height,
+        target_width=target_width,
+        background_color=255,
+        dtype=np.float32,
+        how=pad_how,
+    )
+
+    if normalize:
+        padded = padded / 255.0
+
+    if channel_first:
+        # NemoRetriever-OCR NIM (GRPC) requires input to be (channel, height, width).
+        padded = padded.transpose((2, 0, 1))
+
+    # Metadata can used for inverting transformations on the resulting bounding boxes.
+    metadata = {
+        "original_height": height,
+        "original_width": width,
+        "new_height": target_height,
+        "new_width": target_width,
+        "pad_height": pad_height,
+        "pad_width": pad_width,
+    }
+
+    return padded, 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