chore!: Rename model_name to model in the Gradient integration (#228)

* rename model_name into model * fix tests * rename model_name into model for text embedder * fix tests
deepset-ai · Jan 18, 2024 · 065a00f · 065a00f
1 parent c3912e7
commit 065a00f
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 12 deletions.
diff --git a/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py b/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py
@@ -28,7 +28,7 @@ class GradientDocumentEmbedder:
     embedder = GradientDocumentEmbedder(
         access_token=gradient_access_token,
         workspace_id=gradient_workspace_id,
-        model_name="bge_large"))
+        model="bge_large"))
     p = Pipeline()
     p.add_component(embedder, name="document_embedder")
     p.add_component(instance=GradientDocumentEmbedder(
@@ -41,7 +41,7 @@ class GradientDocumentEmbedder:
     def __init__(
         self,
         *,
-        model_name: str = "bge-large",
+        model: str = "bge-large",
         batch_size: int = 32_768,
         access_token: Optional[str] = None,
         workspace_id: Optional[str] = None,
@@ -51,7 +51,7 @@ def __init__(
         """
         Create a GradientDocumentEmbedder component.
 
-        :param model_name: The name of the model to use.
+        :param model: The name of the model to use.
         :param batch_size: Update cycle for tqdm progress bar, default is to update every 32_768 docs.
         :param access_token: The Gradient access token. If not provided it's read from the environment
                              variable GRADIENT_ACCESS_TOKEN.
@@ -62,7 +62,7 @@ def __init__(
         """
         self._batch_size = batch_size
         self._host = host
-        self._model_name = model_name
+        self._model_name = model
         self._progress_bar = progress_bar
 
         self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id)
@@ -77,7 +77,7 @@ def to_dict(self) -> dict:
         """
         Serialize the component to a Python dictionary.
         """
-        return default_to_dict(self, workspace_id=self._gradient.workspace_id, model_name=self._model_name)
+        return default_to_dict(self, workspace_id=self._gradient.workspace_id, model=self._model_name)
 
     def warm_up(self) -> None:
         """

diff --git a/integrations/gradient/src/gradient_haystack/embedders/gradient_text_embedder.py b/integrations/gradient/src/gradient_haystack/embedders/gradient_text_embedder.py
@@ -13,7 +13,7 @@ class GradientTextEmbedder:
     embedder = GradientTextEmbedder(
         access_token=gradient_access_token,
         workspace_id=gradient_workspace_id,
-        model_name="bge_large")
+        model="bge_large")
     p = Pipeline()
     p.add_component(instance=embedder, name="text_embedder")
     p.add_component(instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever")
@@ -25,23 +25,23 @@ class GradientTextEmbedder:
     def __init__(
         self,
         *,
-        model_name: str = "bge-large",
+        model: str = "bge-large",
         access_token: Optional[str] = None,
         workspace_id: Optional[str] = None,
         host: Optional[str] = None,
     ) -> None:
         """
         Create a GradientTextEmbedder component.
 
-        :param model_name: The name of the model to use.
+        :param model: The name of the model to use.
         :param access_token: The Gradient access token. If not provided it's read from the environment
                              variable GRADIENT_ACCESS_TOKEN.
         :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment
                              variable GRADIENT_WORKSPACE_ID.
         :param host: The Gradient host. By default it uses https://api.gradient.ai/.
         """
         self._host = host
-        self._model_name = model_name
+        self._model_name = model
 
         self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id)
 
@@ -55,7 +55,7 @@ def to_dict(self) -> dict:
         """
         Serialize the component to a Python dictionary.
         """
-        return default_to_dict(self, workspace_id=self._gradient.workspace_id, model_name=self._model_name)
+        return default_to_dict(self, workspace_id=self._gradient.workspace_id, model=self._model_name)
 
     def warm_up(self) -> None:
         """

diff --git a/integrations/gradient/tests/test_gradient_document_embedder.py b/integrations/gradient/tests/test_gradient_document_embedder.py
@@ -54,7 +54,7 @@ def test_to_dict(self):
         data = component.to_dict()
         assert data == {
             "type": "gradient_haystack.embedders.gradient_document_embedder.GradientDocumentEmbedder",
-            "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"},
+            "init_parameters": {"workspace_id": workspace_id, "model": "bge-large"},
         }
 
     def test_warmup(self):

diff --git a/integrations/gradient/tests/test_gradient_text_embedder.py b/integrations/gradient/tests/test_gradient_text_embedder.py
@@ -53,7 +53,7 @@ def test_to_dict(self):
         data = component.to_dict()
         assert data == {
             "type": "gradient_haystack.embedders.gradient_text_embedder.GradientTextEmbedder",
-            "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"},
+            "init_parameters": {"workspace_id": workspace_id, "model": "bge-large"},
         }
 
     def test_warmup(self):