Merge pull request #3 from waggle-sensor/feature-nrp-llm

Add NRP LLM caption function

Author: iperezx

.env.example

@@ -45,6 +45,11 @@ UNALLOWED_NODES="W042,N001,V012,W015,W01C,W01E,W024,W026,W02C,W02D,W02E,W02F,W03
 LOG_LEVEL='INFO'
 MONITOR_DATA_STREAM_INTERVAL=60
 MONITOR_DATA_STREAM_QUERY_DELAY_MINUTE=5
+LLM_RUN_MODE=TRITON
+# IF LLM_RUN_MODE is set to NRP
+# NRP_API_KEY=
+# NRP_API_ENDPOINT=
+# NRP_LLM_MODEL=
 
 #gardio-ui
 WEAVIATE_HOST=weaviate

benchmarking/benchmarks/INQUIRE/config.py

@@ -35,6 +35,9 @@ def __init__(self):
         self._weaviate_grpc_port = os.environ.get("WEAVIATE_GRPC_PORT", "50051")
         self._collection_name = os.environ.get("COLLECTION_NAME", "INQUIRE")
 
+        # model provider parameters
+        self._llm_model_provider = os.environ.get("LLM_MODEL_PROVIDER", "triton").lower()
+
         # Triton parameters
         self._triton_host = os.environ.get("TRITON_HOST", "triton")
         self._triton_port = os.environ.get("TRITON_PORT", "8001")
@@ -97,3 +100,9 @@ def __init__(self):
   keywords: <keyword1>, <keyword2>, ...
 """
         self.gemma3_prompt = os.environ.get("GEMMA3_PROMPT", default_prompt)
+
+    @staticmethod
+    def is_nrp_key_set():
+      """Check if NRP API key is set."""
+      if os.environ.get("NRP_API_KEY", "") == "":
+        raise ValueError("NRP_API_KEY is not set")

benchmarking/benchmarks/INQUIRE/model_provider.py

@@ -0,0 +1,62 @@
+from imsearch_eval.adapters import TritonModelProvider, NRPModelProvider
+from imsearch_eval.framework import Config
+import os
+from PIL import Image
+from typing import Optional
+
+class MixedModelProvider(NRPModelProvider):
+    """
+    Mixed model provider using NRPModelProvider and TritonModelProvider.
+
+    NRPModelProvider is used for caption generation and TritonModelProvider is used for embedding generation.
+    """
+
+    def __init__(
+        self,
+        api_key: str = os.environ.get("NRP_API_KEY"),
+        base_url: str = "https://ellm.nrp-nautilus.io/v1",
+        triton_model_provider: TritonModelProvider = None,
+        config: Config = None,
+        **client_kwargs,
+    ):
+        """
+        Initialize Mixed model provider.
+
+        Args:
+            api_key: NRP API token (defaults to environment variable "NRP_API_KEY").
+            base_url: Envoy gateway URL. Defaults to the NRP-managed LLM endpoint.
+            triton_model_provider: Triton model provider.
+            config: Config object.
+            **client_kwargs: Optional extra arguments passed to the NRPModelProvider.
+        """      
+        super().__init__(api_key=api_key, base_url=base_url, **client_kwargs)   
+        self.triton_model_provider = triton_model_provider
+        self.config = config
+        
+        # determine which model provider to use for caption generation
+        if self.config._llm_model_provider == "triton":
+            self.model_utils = self.triton_model_provider.model_utils
+        elif self.config._llm_model_provider == "nrp":
+            self.config.is_nrp_key_set()
+        else:
+            raise ValueError(f"Invalid model provider: {self.config._llm_model_provider} not supported")
+    
+
+    def get_embedding(
+        self, 
+        text: str, 
+        image: Optional[Image.Image] = None,
+        model_name: str = "clip"
+    ):
+        """
+        Get embedding for text and/or image using triton model provider.
+        
+        Args:
+            text: Text to embed
+            image: Optional PIL Image to embed
+            model_name: Name of the model to use ("clip", "colbert", "align")
+            
+        Returns:
+            Embedding vector (numpy array)
+        """
+        return self.triton_model_provider.get_embedding(text, image, model_name)

benchmarking/benchmarks/INQUIRE/requirements.txt

@@ -1,8 +1,9 @@
 # INQUIRE Benchmark Requirements
 # Core benchmarking framework
-imsearch_eval[weaviate] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.0
-imsearch_eval[triton] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.0
-imsearch_eval[huggingface] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.0
+imsearch_eval[weaviate] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.1
+imsearch_eval[triton] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.1
+imsearch_eval[huggingface] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.1
+imsearch_eval[nrp] @ git+https://github.com/waggle-sensor/imsearch_eval.git@0.1.1
 
 # Image processing
 Pillow>=10.0.0