nv-ingest-api 2025.5.11.dev20250511__tar.gz → 2025.5.12.dev20250512__tar.gz

{nv_ingest_api-2025.5.11.dev20250511/src/nv_ingest_api.egg-info → nv_ingest_api-2025.5.12.dev20250512}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.5.11.dev20250511
+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.11.dev20250511 → nv_ingest_api-2025.5.12.dev20250512}/src/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-2025.5.11.dev20250511 → nv_ingest_api-2025.5.12.dev20250512}/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)
 
+        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.11.dev20250511 → nv_ingest_api-2025.5.12.dev20250512/src/nv_ingest_api.egg-info}/PKG-INFO

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