Merge pull request #1558 from weaviate/202502/nvidia-patches

Add `nvidia-reranker` factory function

Author: dirkkul

test/collection/test_config.py

@@ -1063,6 +1063,48 @@ def test_config_with_generative(
             "reranker-cohere": {},
         },
     ),
+    (
+        Configure.Reranker.voyageai(),
+        {
+            "reranker-voyageai": {},
+        },
+    ),
+    (
+        Configure.Reranker.voyageai(model="rerank-lite-1"),
+        {
+            "reranker-voyageai": {"model": "rerank-lite-1"},
+        },
+    ),
+    (
+        Configure.Reranker.jinaai(),
+        {
+            "reranker-jinaai": {},
+        },
+    ),
+    (
+        Configure.Reranker.jinaai(model="jina-reranker-v2-base-multilingual"),
+        {
+            "reranker-jinaai": {"model": "jina-reranker-v2-base-multilingual"},
+        },
+    ),
+    (
+        Configure.Reranker.nvidia(),
+        {
+            "reranker-nvidia": {},
+        },
+    ),
+    (
+        Configure.Reranker.nvidia(
+            model="nvidia/llama-3.2-nv-rerankqa-1b-v2",
+            base_url="https://integrate.api.nvidia.com/v1",
+        ),
+        {
+            "reranker-nvidia": {
+                "model": "nvidia/llama-3.2-nv-rerankqa-1b-v2",
+                "baseURL": "https://integrate.api.nvidia.com/v1",
+            },
+        },
+    ),
     (
         Configure.Reranker.transformers(),
         {

weaviate/collections/classes/config.py

@@ -215,13 +215,16 @@ class Rerankers(str, BaseEnum):
             Weaviate module backed by VoyageAI reranking models.
         `JINAAI`
             Weaviate module backed by JinaAI reranking models.
+        `NVIDIA`
+            Weaviate module backed by NVIDIA reranking models.
     """
 
     NONE = "none"
     COHERE = "reranker-cohere"
     TRANSFORMERS = "reranker-transformers"
     VOYAGEAI = "reranker-voyageai"
     JINAAI = "reranker-jinaai"
+    NVIDIA = "reranker-nvidia"
 
 
 class StopwordsPreset(str, BaseEnum):
@@ -650,6 +653,20 @@ class _RerankerVoyageAIConfig(_RerankerProvider):
     model: Optional[Union[RerankerVoyageAIModel, str]] = Field(default=None)
 
 
+class _RerankerNvidiaConfig(_RerankerProvider):
+    reranker: Union[Rerankers, _EnumLikeStr] = Field(
+        default=Rerankers.NVIDIA, frozen=True, exclude=True
+    )
+    model: Optional[str] = Field(default=None)
+    baseURL: Optional[AnyHttpUrl]
+
+    def _to_dict(self) -> Dict[str, Any]:
+        ret_dict = super()._to_dict()
+        if self.baseURL is not None:
+            ret_dict["baseURL"] = self.baseURL.unicode_string()
+        return ret_dict
+
+
 class _Generative:
     """Use this factory class to create the correct object for the `generative_config` argument in the `collections.create()` method.
 
@@ -1120,7 +1137,7 @@ def cohere(
     ) -> _RerankerProvider:
         """Create a `_RerankerCohereConfig` object for use when reranking using the `reranker-cohere` module.
 
-        See the [documentation](https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules/reranker-cohere)
+        See the [documentation](https://weaviate.io/developers/weaviate/model-providers/cohere/reranker)
         for detailed usage.
 
         Arguments:
@@ -1135,7 +1152,7 @@ def jinaai(
     ) -> _RerankerProvider:
         """Create a `_RerankerJinaAIConfig` object for use when reranking using the `reranker-jinaai` module.
 
-        See the [documentation](https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules/reranker-jinaai)
+        See the [documentation](https://weaviate.io/developers/weaviate/model-providers/jinaai/reranker)
         for detailed usage.
 
         Arguments:
@@ -1150,7 +1167,7 @@ def voyageai(
     ) -> _RerankerProvider:
         """Create a `_RerankerVoyageAIConfig` object for use when reranking using the `reranker-voyageai` module.
 
-        See the [documentation](https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules/reranker-voyageai)
+        See the [documentation](https://weaviate.io/developers/weaviate/model-providers/voyageai/reranker)
         for detailed usage.
 
         Arguments:
@@ -1159,6 +1176,24 @@ def voyageai(
         """
         return _RerankerVoyageAIConfig(model=model)
 
+    @staticmethod
+    def nvidia(
+        model: Optional[str] = None,
+        base_url: Optional[AnyHttpUrl] = None,
+    ) -> _RerankerProvider:
+        """Create a `_RerankerNvidiaConfig` object for use when reranking using the `reranker-nvidia` module.
+
+        See the [documentation](https://weaviate.io/developers/weaviate/model-providers/nvidia/reranker)
+        for detailed usage.
+
+        Arguments:
+            `model`
+                The model to use. Defaults to `None`, which uses the server-defined default
+            `baseurl`
+                The base URL to send the reranker requests to. Defaults to `None`, which uses the server-defined default.
+        """
+        return _RerankerNvidiaConfig(model=model, baseURL=base_url)
+
 
 class _CollectionConfigCreateBase(_ConfigCreateModel):
     description: Optional[str] = Field(default=None)