diff --git a/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py b/integrations/gradient/src/gradient_haystack/embedders/gradient_document_embedder.py index 551aa9dd5..34fc1f87f 100644 --- 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 index 2ddc229ce..ba753297c 100644 --- 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,7 +25,7 @@ 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, @@ -33,7 +33,7 @@ def __init__( """ 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 @@ -41,7 +41,7 @@ def __init__( :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 index 6e75360fe..dc59a76fb 100644 --- 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 index 7ae846e93..bd4b396ca 100644 --- 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):