nv-ingest-api 2025.5.10.dev20250510__py3-none-any.whl → 2025.5.12.dev20250512__py3-none-any.whl

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

@@ -5,6 +5,7 @@
 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.
@@ -22,20 +23,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 +57,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 +105,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/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)
 
+        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}")
         # TODO(self.client.has_error(response)) => raise error
-
-        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.10.dev20250510.dist-info → nv_ingest_api-2025.5.12.dev20250512.dist-info}/METADATA

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

{nv_ingest_api-2025.5.10.dev20250510.dist-info → nv_ingest_api-2025.5.12.dev20250512.dist-info}/RECORD

@@ -46,7 +46,7 @@ nv_ingest_api/internal/primitives/control_message_task.py,sha256=nWVB3QsP6p8BKwH
 nv_ingest_api/internal/primitives/ingest_control_message.py,sha256=rvipBiiUaHuRhupFCFDCG8rv0PylSJibCiJ7rDeb98A,8514
 nv_ingest_api/internal/primitives/nim/__init__.py,sha256=i_i_fBR2EcRCh2Y19DF6GM3s_Q0VPgo_thPnhEIJUyg,266
 nv_ingest_api/internal/primitives/nim/default_values.py,sha256=W92XjfyeC6uuVxut6J7p00x1kpNsnXIDb97gSVytZJk,380
-nv_ingest_api/internal/primitives/nim/nim_client.py,sha256=7C_t3BnYz_hL2H8RmvOShLCKlfYmwIREC6vnOnzOHWA,14483
+nv_ingest_api/internal/primitives/nim/nim_client.py,sha256=lEP-PBp921--pxQzeVxxafR2BhONpli2Ad8oa0XLR4Y,14920
 nv_ingest_api/internal/primitives/nim/nim_model_interface.py,sha256=wMEgoi79YQn_4338MVemkeZgM1J-vnz0aZWpvqDhib4,2392
 nv_ingest_api/internal/primitives/nim/model_interface/__init__.py,sha256=wQSlVx3T14ZgQAt-EPzEczQusXVW0W8yynnUaFFGE3s,143
 nv_ingest_api/internal/primitives/nim/model_interface/cached.py,sha256=b1HX-PY1ExW5V6pXC1ZiHdobeG_BmbPr3rBbVJef13s,11003
@@ -56,7 +56,7 @@ nv_ingest_api/internal/primitives/nim/model_interface/helpers.py,sha256=x35a9AyT
 nv_ingest_api/internal/primitives/nim/model_interface/nemoretriever_parse.py,sha256=MFWPqMTXs_MZG3ripRR21o7f_mVeoE46Q10yvJ8KNr0,7023
 nv_ingest_api/internal/primitives/nim/model_interface/paddle.py,sha256=rSUPwl5XOrqneoS6aKhatVjrNBg_LhP3nwUWS_aTwz0,17950
 nv_ingest_api/internal/primitives/nim/model_interface/parakeet.py,sha256=OYg4AGki_wm--Np9VlSm0eZC-r54GbDOISbe9v0B9fw,12967
-nv_ingest_api/internal/primitives/nim/model_interface/text_embedding.py,sha256=oai0-moKDJOnOMfTaGQf-vo6qMRD6pbcf7_XRIt-oJ8,4934
+nv_ingest_api/internal/primitives/nim/model_interface/text_embedding.py,sha256=8ld_if6N3pe3W7NA8Xwm-ndCq53s_v3LmmoyQHnxxEo,5071
 nv_ingest_api/internal/primitives/nim/model_interface/vlm.py,sha256=qJ382PU1ZrIM-SR3cqIhtY_W2rmHec2HIa2aUB2SvaU,6031
 nv_ingest_api/internal/primitives/nim/model_interface/yolox.py,sha256=exN0pKTBXd3pb5kKP96jinTYisgz1Y7EyWmWUuDNnCY,49312
 nv_ingest_api/internal/primitives/tracing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -145,8 +145,8 @@ nv_ingest_api/util/service_clients/redis/redis_client.py,sha256=Xa9eeI3kfDBDlLsG
 nv_ingest_api/util/service_clients/rest/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nv_ingest_api/util/service_clients/rest/rest_client.py,sha256=K8hzIV4EcV-97G0SboY6LHMhWLx87l9wCI2CdWw9W_E,21734
 nv_ingest_api/util/string_processing/__init__.py,sha256=mkwHthyS-IILcLcL1tJYeF6mpqX3pxEw5aUzDGjTSeU,1411
-nv_ingest_api-2025.5.10.dev20250510.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-nv_ingest_api-2025.5.10.dev20250510.dist-info/METADATA,sha256=PP9KnoMh2nnyGqZ5KXrfj7bUO5D3R1NE091c-bKKbOo,13889
-nv_ingest_api-2025.5.10.dev20250510.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-nv_ingest_api-2025.5.10.dev20250510.dist-info/top_level.txt,sha256=abjYMlTJGoG5tOdfIB-IWvLyKclw6HLaRSc8MxX4X6I,14
-nv_ingest_api-2025.5.10.dev20250510.dist-info/RECORD,,
+nv_ingest_api-2025.5.12.dev20250512.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+nv_ingest_api-2025.5.12.dev20250512.dist-info/METADATA,sha256=rdw7szDF5gzTKuYwhI7F3ZdLRpw8RBXo7WuE7p1j6rc,13889
+nv_ingest_api-2025.5.12.dev20250512.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+nv_ingest_api-2025.5.12.dev20250512.dist-info/top_level.txt,sha256=abjYMlTJGoG5tOdfIB-IWvLyKclw6HLaRSc8MxX4X6I,14
+nv_ingest_api-2025.5.12.dev20250512.dist-info/RECORD,,