From d02d6bf8c6ef609cb42ec35ff55463148f711ac3 Mon Sep 17 00:00:00 2001 From: lilwilsond2 Date: Fri, 13 Oct 2023 03:25:26 -0500 Subject: [PATCH 01/12] init generator --- .../generators/gradient/__init__.py | 0 .../components/generators/gradient/base.py | 103 ++++++++++++++++++ 2 files changed, 103 insertions(+) create mode 100644 haystack/preview/components/generators/gradient/__init__.py create mode 100644 haystack/preview/components/generators/gradient/base.py diff --git a/haystack/preview/components/generators/gradient/__init__.py b/haystack/preview/components/generators/gradient/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py new file mode 100644 index 00000000000..41084d04c04 --- /dev/null +++ b/haystack/preview/components/generators/gradient/base.py @@ -0,0 +1,103 @@ +from typing import Optional, Dict, Any, overload + +import logging +from haystack.lazy_imports import LazyImport + +from haystack.preview import component, default_to_dict + +with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: + from gradientai import Gradient + +logger = logging.getLogger(__name__) + + +@component +class GradientGenerator: + @overload + def __init__( + self, + *, + access_token: Optional[str] = None, + base_model_slug: str, + host: Optional[str] = None, + max_generated_token_count: Optional[int] = None, + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + workspace_id: Optional[str] = None, + ) -> None: + ... + + @overload + def __init__( + self, + *, + access_token: Optional[str] = None, + host: Optional[str] = None, + max_generated_token_count: Optional[int] = None, + model_adapter_id: str, + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + workspace_id: Optional[str] = None, + ) -> None: + ... + + def __init__( + self, + *, + access_token: Optional[str] = None, + base_model_slug: Optional[str] = None, + host: Optional[str] = None, + max_generated_token_count: Optional[int] = None, + model_adapter_id: Optional[str] = None, + temperature: Optional[float] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + workspace_id: Optional[str] = None, + ) -> None: + self._access_token = access_token + self._base_model_slug = (base_model_slug,) + self._host = host + self._max_generated_token_count = max_generated_token_count + self._model_adapter_id = (model_adapter_id,) + self._temperature = temperature + self._top_k = top_k + self._top_p = top_p + self._workspace_id = workspace_id + + if (base_model_slug is None and model_adapter_id is None) or ( + isinstance(base_model_slug, str) and isinstance(model_adapter_id, str) + ): + raise ValueError("expected be provided exactly one of base_model_slug or model_adapter_id") + + self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) + if isinstance(base_model_slug, str): + self._model = self._gradient.get_base_model(base_model_slug=base_model_slug) + if isinstance(model_adapter_id, str): + self._model = self._gradient.get_model_adapter(model_adapter_id=model_adapter_id) + + def to_dict(self) -> Dict[str, Any]: + return default_to_dict( + self, + access_token=self._access_token, + base_model_slug=self._base_model_slug, + host=self._host, + max_generated_token_count=self._max_generated_token_count, + model_adapter_id=self._model_adapter_id, + temperature=self._temperature, + top_k=self._top_k, + top_p=self._top_p, + workspace_id=self._workspace_id, + ) + + @component.output_types(response=str) + def run(self, prompt: str): + resp = self._model.complete( + query=prompt, + max_generated_token_count=self._max_generated_token_count, + temperature=self._temperature, + top_k=self._top_k, + top_p=self._top_p, + ) + return {"response": resp.generated_output} From 336c28d8f47632e6050ce335904dc19cda2b9e44 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Thu, 26 Oct 2023 14:05:29 -0700 Subject: [PATCH 02/12] Gradient embeddings (#2) --- .../pipelines/test_gradient_rag_pipelines.py | 92 ++++++++++++ .../preview/components/embedders/__init__.py | 5 + .../embedders/gradient_document_embedder.py | 96 ++++++++++++ .../embedders/gradient_text_embedder.py | 71 +++++++++ .../preview/components/generators/__init__.py | 3 +- .../components/generators/gradient/base.py | 10 +- pyproject.toml | 3 + .../test_gradient_document_embedder.py | 139 ++++++++++++++++++ .../embedders/test_gradient_text_embedder.py | 111 ++++++++++++++ 9 files changed, 524 insertions(+), 6 deletions(-) create mode 100644 e2e/preview/pipelines/test_gradient_rag_pipelines.py create mode 100644 haystack/preview/components/embedders/gradient_document_embedder.py create mode 100644 haystack/preview/components/embedders/gradient_text_embedder.py create mode 100644 test/preview/components/embedders/test_gradient_document_embedder.py create mode 100644 test/preview/components/embedders/test_gradient_text_embedder.py diff --git a/e2e/preview/pipelines/test_gradient_rag_pipelines.py b/e2e/preview/pipelines/test_gradient_rag_pipelines.py new file mode 100644 index 00000000000..ca2113bbb6d --- /dev/null +++ b/e2e/preview/pipelines/test_gradient_rag_pipelines.py @@ -0,0 +1,92 @@ +import os +import json +import pytest + +from haystack.preview import Pipeline, Document +from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder +from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from haystack.preview.document_stores import InMemoryDocumentStore +from haystack.preview.components.writers import DocumentWriter +from haystack.preview.components.retrievers import InMemoryEmbeddingRetriever +from haystack.preview.components.generators.gradient.base import GradientGenerator +from haystack.preview.components.builders.answer_builder import AnswerBuilder +from haystack.preview.components.builders.prompt_builder import PromptBuilder + + +@pytest.mark.skipif( + not os.environ.get("GRADIENT_ACCESS_TOKEN", None), + reason="Export an env var called GRADIENT_ACCESS_TOKEN containing the Gradient access token to run this test.", +) +def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): + # Create the RAG pipeline + prompt_template = """ + Given these documents, answer the question.\nDocuments: + {% for doc in documents %} + {{ doc.text }} + {% endfor %} + + \nQuestion: {{question}} + \nAnswer: + """ + + gradient_access_token = os.environ.get("GRADIENT_ACCESS_TOKEN") + rag_pipeline = Pipeline() + embedder = GradientTextEmbedder(access_token=gradient_access_token) + rag_pipeline.add_component(instance=embedder, name="text_embedder") + rag_pipeline.add_component( + instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever" + ) + rag_pipeline.add_component(instance=PromptBuilder(template=prompt_template), name="prompt_builder") + rag_pipeline.add_component( + instance=GradientGenerator(access_token=gradient_access_token, base_model_slug="llama2-7b-chat"), name="llm" + ) + rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder") + rag_pipeline.connect("text_embedder", "retriever") + rag_pipeline.connect("retriever", "prompt_builder.documents") + rag_pipeline.connect("prompt_builder", "llm") + rag_pipeline.connect("llm.replies", "answer_builder.replies") + rag_pipeline.connect("retriever", "answer_builder.documents") + + # Draw the pipeline + rag_pipeline.draw(tmp_path / "test_gradient_embedding_rag_pipeline.png") + + # Serialize the pipeline to JSON + with open(tmp_path / "test_bm25_rag_pipeline.json", "w") as f: + json.dump(rag_pipeline.to_dict(), f) + + # Load the pipeline back + with open(tmp_path / "test_bm25_rag_pipeline.json", "r") as f: + rag_pipeline = Pipeline.from_dict(json.load(f)) + + # Populate the document store + documents = [ + Document(text="My name is Jean and I live in Paris."), + Document(text="My name is Mark and I live in Berlin."), + Document(text="My name is Giorgio and I live in Rome."), + ] + document_store = rag_pipeline.get_component("retriever").document_store + indexing_pipeline = Pipeline() + indexing_pipeline.add_component(instance=GradientDocumentEmbedder(), name="document_embedder") + indexing_pipeline.add_component(instance=DocumentWriter(document_store=document_store), name="document_writer") + indexing_pipeline.connect("document_embedder", "document_writer") + indexing_pipeline.run({"document_embedder": {"documents": documents}}) + + # Query and assert + questions = ["Who lives in Paris?", "Who lives in Berlin?", "Who lives in Rome?"] + answers_spywords = ["Jean", "Mark", "Giorgio"] + + for question, spyword in zip(questions, answers_spywords): + result = rag_pipeline.run( + { + "text_embedder": {"text": question}, + "prompt_builder": {"question": question}, + "answer_builder": {"query": question}, + } + ) + + assert len(result["answer_builder"]["answers"]) == 1 + generated_answer = result["answer_builder"]["answers"][0] + assert spyword in generated_answer.data + assert generated_answer.query == question + assert hasattr(generated_answer, "documents") + assert hasattr(generated_answer, "metadata") diff --git a/haystack/preview/components/embedders/__init__.py b/haystack/preview/components/embedders/__init__.py index a0840d7e0a4..4af150d9ce1 100644 --- a/haystack/preview/components/embedders/__init__.py +++ b/haystack/preview/components/embedders/__init__.py @@ -1,3 +1,6 @@ +from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder + from haystack.preview.components.embedders.sentence_transformers_text_embedder import SentenceTransformersTextEmbedder from haystack.preview.components.embedders.sentence_transformers_document_embedder import ( SentenceTransformersDocumentEmbedder, @@ -6,6 +9,8 @@ from haystack.preview.components.embedders.openai_text_embedder import OpenAITextEmbedder __all__ = [ + "GradientTextEmbedder", + "GradientDocumentEmbedder", "SentenceTransformersTextEmbedder", "SentenceTransformersDocumentEmbedder", "OpenAITextEmbedder", diff --git a/haystack/preview/components/embedders/gradient_document_embedder.py b/haystack/preview/components/embedders/gradient_document_embedder.py new file mode 100644 index 00000000000..7fd6084ca51 --- /dev/null +++ b/haystack/preview/components/embedders/gradient_document_embedder.py @@ -0,0 +1,96 @@ +import logging +from typing import List, Optional, Dict, Any + +from haystack.preview import component, Document, default_to_dict +from haystack.preview.lazy_imports import LazyImport + +with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: + from gradientai import Gradient + +logger = logging.getLogger(__name__) + + +@component +class GradientDocumentEmbedder: + """ + A component for computing Document embeddings using Gradient AI API.. + The embedding of each Document is stored in the `embedding` field of the Document. + """ + + def __init__( + self, + *, + model_name: str = "bge-large", + access_token: Optional[str] = None, + workspace_id: Optional[str] = None, + host: Optional[str] = None, + ) -> None: + """ + Create a GradientDocumentEmbedder component. + + :param model_name: 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._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) + + def _get_telemetry_data(self) -> Dict[str, Any]: + """ + Data that is sent to Posthog for usage analytics. + """ + return {"model": self._model_name} + + 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) + + def warm_up(self) -> None: + """ + Load the embedding model. + """ + if not hasattr(self, "_embedding_model"): + self._embedding_model = self._gradient.get_embeddings_model(slug=self._model_name) + + def _generate_embeddings(self, documents: List[Document], batch_size=100) -> List[List[float]]: + """ + Batches the documents and generates the embeddings. + """ + batches = [documents[i : i + batch_size] for i in range(0, len(documents), batch_size)] + + embeddings = [] + for batch in batches: + response = self._embedding_model.generate_embeddings(inputs=[{"input": doc.text} for doc in batch]) + embeddings.extend([e.embedding for e in response.embeddings]) + + return embeddings + + @component.output_types(documents=List[Document]) + def run(self, documents: List[Document]): + """ + Embed a list of Documents. + The embedding of each Document is stored in the `embedding` field of the Document. + + :param documents: A list of Documents to embed. + """ + if not isinstance(documents, list) or documents and not isinstance(documents[0], Document): + raise TypeError( + "GradientDocumentEmbedder expects a list of Documents as input." + "In case you want to embed a list of strings, please use the GradientTextEmbedder." + ) + + if not hasattr(self, "_embedding_model"): + raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") + + embeddings = self._generate_embeddings(documents=documents) + for doc, embedding in zip(documents, embeddings): + doc.embedding = embedding + + return {"documents": documents} diff --git a/haystack/preview/components/embedders/gradient_text_embedder.py b/haystack/preview/components/embedders/gradient_text_embedder.py new file mode 100644 index 00000000000..39f28f51952 --- /dev/null +++ b/haystack/preview/components/embedders/gradient_text_embedder.py @@ -0,0 +1,71 @@ +from typing import Any, Dict, List, Optional + +from haystack.preview import component, default_to_dict +from haystack.preview.lazy_imports import LazyImport + +with LazyImport(message="Run 'pip install gradientai'") as gradientai_import: + from gradientai import Gradient + + +@component +class GradientTextEmbedder: + """ + A component for embedding strings using models hosted on Gradient AI (https://gradient.ai). + """ + + def __init__( + self, + *, + model_name: 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 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._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) + + def _get_telemetry_data(self) -> Dict[str, Any]: + """ + Data that is sent to Posthog for usage analytics. + """ + return {"model": self._model_name} + + 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) + + def warm_up(self) -> None: + """ + Load the embedding model. + """ + if not hasattr(self, "_embedding_model"): + self._embedding_model = self._gradient.get_embeddings_model(slug=self._model_name) + + @component.output_types(embedding=List[float]) + def run(self, text: str): + """Generates an embedding for a single text.""" + if not isinstance(text, str): + raise TypeError( + "GradientTextEmbedder expects a string as an input." + "In case you want to embed a list of Documents, please use the GradientDocumentEmbedder." + ) + + if not hasattr(self, "_embedding_model"): + raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") + + result = self._embedding_model.generate_embeddings(inputs=[{"input": text}]) + return {"embedding": result.embeddings[0].embedding} diff --git a/haystack/preview/components/generators/__init__.py b/haystack/preview/components/generators/__init__.py index bc81975a3fa..8230c46bf9e 100644 --- a/haystack/preview/components/generators/__init__.py +++ b/haystack/preview/components/generators/__init__.py @@ -1,5 +1,6 @@ +from haystack.preview.components.generators.gradient.base import GradientGenerator from haystack.preview.components.generators.hugging_face_local import HuggingFaceLocalGenerator from haystack.preview.components.generators.hugging_face_tgi import HuggingFaceTGIGenerator from haystack.preview.components.generators.openai import GPTGenerator -__all__ = ["HuggingFaceLocalGenerator", "HuggingFaceTGIGenerator", "GPTGenerator"] +__all__ = ["GPTGenerator", "GradientGenerator", "HuggingFaceTGIGenerator", "HuggingFaceLocalGenerator"] diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py index 41084d04c04..199742096a6 100644 --- a/haystack/preview/components/generators/gradient/base.py +++ b/haystack/preview/components/generators/gradient/base.py @@ -1,4 +1,4 @@ -from typing import Optional, Dict, Any, overload +from typing import List, Optional, Dict, Any, overload import logging from haystack.lazy_imports import LazyImport @@ -57,10 +57,10 @@ def __init__( workspace_id: Optional[str] = None, ) -> None: self._access_token = access_token - self._base_model_slug = (base_model_slug,) + self._base_model_slug = base_model_slug self._host = host self._max_generated_token_count = max_generated_token_count - self._model_adapter_id = (model_adapter_id,) + self._model_adapter_id = model_adapter_id self._temperature = temperature self._top_k = top_k self._top_p = top_p @@ -91,7 +91,7 @@ def to_dict(self) -> Dict[str, Any]: workspace_id=self._workspace_id, ) - @component.output_types(response=str) + @component.output_types(replies=List[str]) def run(self, prompt: str): resp = self._model.complete( query=prompt, @@ -100,4 +100,4 @@ def run(self, prompt: str): top_k=self._top_k, top_p=self._top_p, ) - return {"response": resp.generated_output} + return {"replies": [resp.generated_output]} diff --git a/pyproject.toml b/pyproject.toml index 3d0ba57a8eb..beb8fc02ae0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -227,6 +227,9 @@ dev = [ # https://opendev.org/openstack/reno/src/branch/master/requirements.txt#L7 "dulwich>=0.21.0,<1.0.0", ] +gradient = [ + "gradientai==1.1.0" +] formatting = [ # Version specified following Black stability policy: diff --git a/test/preview/components/embedders/test_gradient_document_embedder.py b/test/preview/components/embedders/test_gradient_document_embedder.py new file mode 100644 index 00000000000..ba9568270f6 --- /dev/null +++ b/test/preview/components/embedders/test_gradient_document_embedder.py @@ -0,0 +1,139 @@ +import pytest +from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder +from unittest.mock import MagicMock, NonCallableMagicMock +import numpy as np + +from haystack.preview import Document + + +access_token = "access_token" +workspace_id = "workspace_id" +model = "bge-large" + + +def has_gradient(): + try: + import gradientai + + return True + except ModuleNotFoundError: + return False + + +@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") +class TestGradientDocumentEmbedder: + @pytest.mark.unit + def test_init_from_env(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", access_token) + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", workspace_id) + + embedder = GradientDocumentEmbedder() + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_without_access_token(self, monkeypatch): + monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=True) + + with pytest.raises(ValueError): + GradientDocumentEmbedder(workspace_id=workspace_id) + + @pytest.mark.unit + def test_init_without_workspace(self, monkeypatch): + monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=True) + + with pytest.raises(ValueError): + GradientDocumentEmbedder(access_token=access_token) + + @pytest.mark.unit + def test_init_from_params(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_from_params_precedence(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", "env_access_token") + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", "env_workspace_id") + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_to_dict(self): + component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + data = component.to_dict() + assert data == { + "type": "GradientDocumentEmbedder", + "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, + } + + @pytest.mark.unit + def test_warmup(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_warmup_doesnt_reload(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock(default_return_value="fake model") + embedder.warm_up() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_run_fail_if_not_warmed_up(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + + with pytest.raises(RuntimeError, match="warm_up()"): + embedder.run(documents=[Document(text=f"document number {i}") for i in range(5)]) + + @pytest.mark.unit + def test_run(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(5)] + ) + + documents = [Document(text=f"document number {i}") for i in range(5)] + + result = embedder.run(documents=documents) + + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) + + @pytest.mark.unit + def test_run_batch(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(110)] + ) + + documents = [Document(text=f"document number {i}") for i in range(110)] + + result = embedder.run(documents=documents) + + assert embedder._embedding_model.generate_embeddings.call_count == 2 + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) diff --git a/test/preview/components/embedders/test_gradient_text_embedder.py b/test/preview/components/embedders/test_gradient_text_embedder.py new file mode 100644 index 00000000000..3e2a6762f27 --- /dev/null +++ b/test/preview/components/embedders/test_gradient_text_embedder.py @@ -0,0 +1,111 @@ +import pytest +from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder +from unittest.mock import MagicMock, NonCallableMagicMock +import numpy as np + + +access_token = "access_token" +workspace_id = "workspace_id" +model = "bge-large" + + +def has_gradient(): + try: + import gradientai + + return True + except ModuleNotFoundError: + return False + + +@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") +class TestGradientTextEmbedder: + @pytest.mark.unit + def test_init_from_env(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", access_token) + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", workspace_id) + + embedder = GradientTextEmbedder() + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_without_access_token(self, monkeypatch): + monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=True) + + with pytest.raises(ValueError): + GradientTextEmbedder(workspace_id=workspace_id) + + @pytest.mark.unit + def test_init_without_workspace(self, monkeypatch): + monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=True) + + with pytest.raises(ValueError): + GradientTextEmbedder(access_token=access_token) + + @pytest.mark.unit + def test_init_from_params(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_init_from_params_precedence(self, monkeypatch): + monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", "env_access_token") + monkeypatch.setenv("GRADIENT_WORKSPACE_ID", "env_workspace_id") + + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + assert embedder is not None + assert embedder._gradient.workspace_id == workspace_id + assert embedder._gradient._api_client.configuration.access_token == access_token + + @pytest.mark.unit + def test_to_dict(self): + component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + data = component.to_dict() + assert data == { + "type": "GradientTextEmbedder", + "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, + } + + @pytest.mark.unit + def test_warmup(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_warmup_doesnt_reload(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._gradient.get_embeddings_model = MagicMock(default_return_value="fake model") + embedder.warm_up() + embedder.warm_up() + embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") + + @pytest.mark.unit + def test_run_fail_if_not_warmed_up(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + + with pytest.raises(RuntimeError, match="warm_up()"): + embedder.run(text="The food was delicious") + + @pytest.mark.unit + def test_run(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": 0}] + ) + + result = embedder.run(text="The food was delicious") + embedder._embedding_model.generate_embeddings.assert_called_once_with( + inputs=[{"input": "The food was delicious"}] + ) + + assert len(result["embedding"]) == 1024 # 1024 is the bge-large embedding size + assert all(isinstance(x, float) for x in result["embedding"]) From 43f744f6458792896b235b0b0948cb1a7fa17045 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Fri, 27 Oct 2023 13:31:11 -0700 Subject: [PATCH 03/12] apply pr requests --- .../pipelines/test_gradient_rag_pipelines.py | 4 +-- .../embedders/gradient_document_embedder.py | 8 +++++ .../embedders/gradient_text_embedder.py | 9 ++++++ .../components/generators/gradient/base.py | 31 ++++++++++++++++++- .../test_gradient_document_embedder.py | 11 +++++++ 5 files changed, 60 insertions(+), 3 deletions(-) diff --git a/e2e/preview/pipelines/test_gradient_rag_pipelines.py b/e2e/preview/pipelines/test_gradient_rag_pipelines.py index ca2113bbb6d..84f7f438b5d 100644 --- a/e2e/preview/pipelines/test_gradient_rag_pipelines.py +++ b/e2e/preview/pipelines/test_gradient_rag_pipelines.py @@ -14,8 +14,8 @@ @pytest.mark.skipif( - not os.environ.get("GRADIENT_ACCESS_TOKEN", None), - reason="Export an env var called GRADIENT_ACCESS_TOKEN containing the Gradient access token to run this test.", + not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None), + reason="Export an env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient access token to run this test.", ) def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): # Create the RAG pipeline diff --git a/haystack/preview/components/embedders/gradient_document_embedder.py b/haystack/preview/components/embedders/gradient_document_embedder.py index 7fd6084ca51..772d16cd018 100644 --- a/haystack/preview/components/embedders/gradient_document_embedder.py +++ b/haystack/preview/components/embedders/gradient_document_embedder.py @@ -15,6 +15,14 @@ class GradientDocumentEmbedder: """ A component for computing Document embeddings using Gradient AI API.. The embedding of each Document is stored in the `embedding` field of the Document. + + ```python + p = Pipeline() + p.add_component(instance=GradientDocumentEmbedder(), name="document_embedder") + p.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="document_writer") + p.connect("document_embedder", "document_writer") + p.run({"document_embedder": {"documents": documents}}) + ``` """ def __init__( diff --git a/haystack/preview/components/embedders/gradient_text_embedder.py b/haystack/preview/components/embedders/gradient_text_embedder.py index 39f28f51952..3e88d554a85 100644 --- a/haystack/preview/components/embedders/gradient_text_embedder.py +++ b/haystack/preview/components/embedders/gradient_text_embedder.py @@ -24,6 +24,15 @@ def __init__( """ Create a GradientTextEmbedder component. + ```python + p = Pipeline() + embedder = GradientTextEmbedder(access_token=gradient_access_token) + p.add_component(instance=embedder, name="text_embedder") + p.add_component(instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever") + p.connect("text_embedder", "retriever") + p.run("embed me!!!") + ``` + :param model_name: 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. diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py index 199742096a6..da594493206 100644 --- a/haystack/preview/components/generators/gradient/base.py +++ b/haystack/preview/components/generators/gradient/base.py @@ -13,6 +13,13 @@ @component class GradientGenerator: + """ + LLM Generator interfacing [Gradient AI](https://gradient.ai/). + + Queries the LLM using Gradient AI's SDK ('gradientai' package). + See [Gradient AI API](https://docs.gradient.ai/docs/sdk-quickstart) for more details. + """ + @overload def __init__( self, @@ -56,6 +63,21 @@ def __init__( top_p: Optional[float] = None, workspace_id: Optional[str] = None, ) -> None: + """ + Create a GradientGenerator component. + + :param access_token: The Gradient access token. If not provided it's read from the environment + variable GRADIENT_ACCESS_TOKEN. + :param base_model_slug: The base model slug to use. + :param host: The Gradient host. By default it uses https://api.gradient.ai/. + :param max_generated_token_count: The maximum number of tokens to generate. + :param model_adapter_id: The model adapter ID to use. + :param temperature: The temperature to use. + :param top_k: The top k to use. + :param top_p: The top p to use. + :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment + variable GRADIENT_WORKSPACE_ID. + """ self._access_token = access_token self._base_model_slug = base_model_slug self._host = host @@ -78,9 +100,11 @@ def __init__( self._model = self._gradient.get_model_adapter(model_adapter_id=model_adapter_id) def to_dict(self) -> Dict[str, Any]: + """ + Serialize this component to a dictionary. + """ return default_to_dict( self, - access_token=self._access_token, base_model_slug=self._base_model_slug, host=self._host, max_generated_token_count=self._max_generated_token_count, @@ -93,6 +117,11 @@ def to_dict(self) -> Dict[str, Any]: @component.output_types(replies=List[str]) def run(self, prompt: str): + """ + Queries the LLM with the prompt to produce replies. + + :param prompt: The prompt to be sent to the generative model. + """ resp = self._model.complete( query=prompt, max_generated_token_count=self._max_generated_token_count, diff --git a/test/preview/components/embedders/test_gradient_document_embedder.py b/test/preview/components/embedders/test_gradient_document_embedder.py index ba9568270f6..e68a2e552e9 100644 --- a/test/preview/components/embedders/test_gradient_document_embedder.py +++ b/test/preview/components/embedders/test_gradient_document_embedder.py @@ -137,3 +137,14 @@ def test_run_batch(self): assert isinstance(doc, Document) assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) + + @pytest.mark.unit + def test_run_empty(self): + from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + + result = embedder.run(documents=[]) + + assert result["documents"] == [] From 2eca3e687a9bc9201780a82836061e137712b5dd Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Tue, 31 Oct 2023 11:57:53 -0700 Subject: [PATCH 04/12] add release notes --- .../notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml diff --git a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml new file mode 100644 index 00000000000..44dcb7f1d2e --- /dev/null +++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml @@ -0,0 +1,4 @@ +--- +features: + - | + Adds integration with [Gradient AI](https://gradient.ai). From a6514b462e6479505749c49f480ce18c3e9f4e0f Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Tue, 31 Oct 2023 12:20:57 -0700 Subject: [PATCH 05/12] update texts (#5) --- e2e/preview/pipelines/test_gradient_rag_pipelines.py | 2 +- .../notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/e2e/preview/pipelines/test_gradient_rag_pipelines.py b/e2e/preview/pipelines/test_gradient_rag_pipelines.py index 84f7f438b5d..c7b7cbeed98 100644 --- a/e2e/preview/pipelines/test_gradient_rag_pipelines.py +++ b/e2e/preview/pipelines/test_gradient_rag_pipelines.py @@ -15,7 +15,7 @@ @pytest.mark.skipif( not os.environ.get("GRADIENT_ACCESS_TOKEN", None) or not os.environ.get("GRADIENT_WORKSPACE_ID", None), - reason="Export an env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient access token to run this test.", + reason="Export env variables called GRADIENT_ACCESS_TOKEN and GRADIENT_WORKSPACE_ID containing the Gradient configuration settings to run this test.", ) def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): # Create the RAG pipeline diff --git a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml index 44dcb7f1d2e..bb3917b0d39 100644 --- a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml +++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml @@ -1,4 +1,5 @@ --- -features: +preview: - | - Adds integration with [Gradient AI](https://gradient.ai). + Add changes to Haystack version 2, or remove this section. + Haystack version 2 can be found under haystack/preview. From 96826a2a84a6a35dc9eb19576947717d97ba746f Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Mon, 6 Nov 2023 15:34:13 -0800 Subject: [PATCH 06/12] fix tests after rebase --- e2e/preview/pipelines/test_gradient_rag_pipelines.py | 8 ++++---- .../components/embedders/gradient_document_embedder.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/e2e/preview/pipelines/test_gradient_rag_pipelines.py b/e2e/preview/pipelines/test_gradient_rag_pipelines.py index c7b7cbeed98..147e7fb8e6a 100644 --- a/e2e/preview/pipelines/test_gradient_rag_pipelines.py +++ b/e2e/preview/pipelines/test_gradient_rag_pipelines.py @@ -22,7 +22,7 @@ def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): prompt_template = """ Given these documents, answer the question.\nDocuments: {% for doc in documents %} - {{ doc.text }} + {{ doc.content }} {% endfor %} \nQuestion: {{question}} @@ -60,9 +60,9 @@ def test_gradient_embedding_retrieval_rag_pipeline(tmp_path): # Populate the document store documents = [ - Document(text="My name is Jean and I live in Paris."), - Document(text="My name is Mark and I live in Berlin."), - Document(text="My name is Giorgio and I live in Rome."), + Document(content="My name is Jean and I live in Paris."), + Document(content="My name is Mark and I live in Berlin."), + Document(content="My name is Giorgio and I live in Rome."), ] document_store = rag_pipeline.get_component("retriever").document_store indexing_pipeline = Pipeline() diff --git a/haystack/preview/components/embedders/gradient_document_embedder.py b/haystack/preview/components/embedders/gradient_document_embedder.py index 772d16cd018..1ddf0642a8c 100644 --- a/haystack/preview/components/embedders/gradient_document_embedder.py +++ b/haystack/preview/components/embedders/gradient_document_embedder.py @@ -75,7 +75,7 @@ def _generate_embeddings(self, documents: List[Document], batch_size=100) -> Lis embeddings = [] for batch in batches: - response = self._embedding_model.generate_embeddings(inputs=[{"input": doc.text} for doc in batch]) + response = self._embedding_model.generate_embeddings(inputs=[{"input": doc.content} for doc in batch]) embeddings.extend([e.embedding for e in response.embeddings]) return embeddings From 914dd89c698cdcba206cc816fd16471e1bca81cc Mon Sep 17 00:00:00 2001 From: hayden Date: Tue, 7 Nov 2023 16:44:16 -0600 Subject: [PATCH 07/12] Update add-gradient-ai-integration-045fd476e7d3aa6a.yaml (#9) --- .../notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml index bb3917b0d39..355a3874a57 100644 --- a/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml +++ b/releasenotes/notes/add-gradient-ai-integration-045fd476e7d3aa6a.yaml @@ -1,5 +1,4 @@ --- preview: - | - Add changes to Haystack version 2, or remove this section. - Haystack version 2 can be found under haystack/preview. + Adds integration with [Gradient AI](https://gradient.ai). From 64c373dcc5935c94951eb5ccaae75c6563d69f22 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Tue, 7 Nov 2023 19:20:24 -0800 Subject: [PATCH 08/12] apply the changes (#10) --- .../embedders/gradient_document_embedder.py | 10 ++- .../embedders/gradient_text_embedder.py | 26 ++++--- .../components/generators/gradient/base.py | 72 +++++++++---------- pyproject.toml | 3 - .../test_gradient_document_embedder.py | 16 +---- .../embedders/test_gradient_text_embedder.py | 39 ++++++---- 6 files changed, 88 insertions(+), 78 deletions(-) diff --git a/haystack/preview/components/embedders/gradient_document_embedder.py b/haystack/preview/components/embedders/gradient_document_embedder.py index 1ddf0642a8c..bfd34f4bb37 100644 --- a/haystack/preview/components/embedders/gradient_document_embedder.py +++ b/haystack/preview/components/embedders/gradient_document_embedder.py @@ -17,8 +17,13 @@ class GradientDocumentEmbedder: The embedding of each Document is stored in the `embedding` field of the Document. ```python + embedder = GradientDocumentEmbedder( + access_token=gradient_access_token, + workspace_id=gradient_workspace_id, + model_name="bge_large")) p = Pipeline() - p.add_component(instance=GradientDocumentEmbedder(), name="document_embedder") + p.add_component(embedder, name="document_embedder") + p.add_component(instance=GradientDocumentEmbedder( p.add_component(instance=DocumentWriter(document_store=InMemoryDocumentStore()), name="document_writer") p.connect("document_embedder", "document_writer") p.run({"document_embedder": {"documents": documents}}) @@ -43,6 +48,7 @@ def __init__( variable GRADIENT_WORKSPACE_ID. :param host: The Gradient host. By default it uses https://api.gradient.ai/. """ + gradientai_import.check() self._host = host self._model_name = model_name @@ -88,7 +94,7 @@ def run(self, documents: List[Document]): :param documents: A list of Documents to embed. """ - if not isinstance(documents, list) or documents and not isinstance(documents[0], Document): + if not isinstance(documents, list) or documents and any(not isinstance(doc, Document) for doc in documents): raise TypeError( "GradientDocumentEmbedder expects a list of Documents as input." "In case you want to embed a list of strings, please use the GradientTextEmbedder." diff --git a/haystack/preview/components/embedders/gradient_text_embedder.py b/haystack/preview/components/embedders/gradient_text_embedder.py index 3e88d554a85..8c89da08b4d 100644 --- a/haystack/preview/components/embedders/gradient_text_embedder.py +++ b/haystack/preview/components/embedders/gradient_text_embedder.py @@ -11,6 +11,18 @@ class GradientTextEmbedder: """ A component for embedding strings using models hosted on Gradient AI (https://gradient.ai). + + ```python + embedder = GradientTextEmbedder( + access_token=gradient_access_token, + workspace_id=gradient_workspace_id, + model_name="bge_large") + p = Pipeline() + p.add_component(instance=embedder, name="text_embedder") + p.add_component(instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever") + p.connect("text_embedder", "retriever") + p.run("embed me!!!") + ``` """ def __init__( @@ -24,15 +36,6 @@ def __init__( """ Create a GradientTextEmbedder component. - ```python - p = Pipeline() - embedder = GradientTextEmbedder(access_token=gradient_access_token) - p.add_component(instance=embedder, name="text_embedder") - p.add_component(instance=InMemoryEmbeddingRetriever(document_store=InMemoryDocumentStore()), name="retriever") - p.connect("text_embedder", "retriever") - p.run("embed me!!!") - ``` - :param model_name: 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. @@ -40,6 +43,7 @@ def __init__( variable GRADIENT_WORKSPACE_ID. :param host: The Gradient host. By default it uses https://api.gradient.ai/. """ + gradientai_import.check() self._host = host self._model_name = model_name @@ -77,4 +81,8 @@ def run(self, text: str): raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") result = self._embedding_model.generate_embeddings(inputs=[{"input": text}]) + + if (not result) or (result.embeddings is None) or (len(result.embeddings) == 0): + raise RuntimeError("The embedding model did not return any embeddings.") + return {"embedding": result.embeddings[0].embedding} diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py index da594493206..3a9fc63566e 100644 --- a/haystack/preview/components/generators/gradient/base.py +++ b/haystack/preview/components/generators/gradient/base.py @@ -18,37 +18,17 @@ class GradientGenerator: Queries the LLM using Gradient AI's SDK ('gradientai' package). See [Gradient AI API](https://docs.gradient.ai/docs/sdk-quickstart) for more details. - """ - - @overload - def __init__( - self, - *, - access_token: Optional[str] = None, - base_model_slug: str, - host: Optional[str] = None, - max_generated_token_count: Optional[int] = None, - temperature: Optional[float] = None, - top_k: Optional[int] = None, - top_p: Optional[float] = None, - workspace_id: Optional[str] = None, - ) -> None: - ... - @overload - def __init__( - self, - *, - access_token: Optional[str] = None, - host: Optional[str] = None, - max_generated_token_count: Optional[int] = None, - model_adapter_id: str, - temperature: Optional[float] = None, - top_k: Optional[int] = None, - top_p: Optional[float] = None, - workspace_id: Optional[str] = None, - ) -> None: - ... + ```python + llm = GradientGenerator( + access_token=gradient_access_token, + workspace_id=gradient_workspace_id, + base_model_slug="llama2-7b-chat") + llm.warm_up() + print(llm.run(prompt="What is the meaning of life?")) + # Output: {'replies': ['42']} + ``` + """ def __init__( self, @@ -78,6 +58,8 @@ def __init__( :param workspace_id: The Gradient workspace ID. If not provided it's read from the environment variable GRADIENT_WORKSPACE_ID. """ + gradientai_import.check() + self._access_token = access_token self._base_model_slug = base_model_slug self._host = host @@ -88,16 +70,20 @@ def __init__( self._top_p = top_p self._workspace_id = workspace_id - if (base_model_slug is None and model_adapter_id is None) or ( - isinstance(base_model_slug, str) and isinstance(model_adapter_id, str) - ): - raise ValueError("expected be provided exactly one of base_model_slug or model_adapter_id") + has_base_model_slug = base_model_slug is not None and base_model_slug != "" + has_model_adapter_id = model_adapter_id is not None and model_adapter_id != "" + + if not has_base_model_slug and not has_model_adapter_id: + raise ValueError("Either base_model_slug or model_adapter_id must be provided.") + if has_base_model_slug and has_model_adapter_id: + raise ValueError("Only one of base_model_slug or model_adapter_id must be provided.") + + if has_base_model_slug: + self._base_model_slug = base_model_slug + if has_model_adapter_id: + self._model_adapter_id = model_adapter_id self._gradient = Gradient(access_token=access_token, host=host, workspace_id=workspace_id) - if isinstance(base_model_slug, str): - self._model = self._gradient.get_base_model(base_model_slug=base_model_slug) - if isinstance(model_adapter_id, str): - self._model = self._gradient.get_model_adapter(model_adapter_id=model_adapter_id) def to_dict(self) -> Dict[str, Any]: """ @@ -115,6 +101,16 @@ def to_dict(self) -> Dict[str, Any]: workspace_id=self._workspace_id, ) + def warm_up(self): + """ + Initializes the LLM model instance if it doesn't exist. + """ + if not hasattr(self, "_model"): + if isinstance(self._base_model_slug, str): + self._model = self._gradient.get_base_model(base_model_slug=self.base_model_slug) + if isinstance(self._model_adapter_id, str): + self._model = self._gradient.get_model_adapter(model_adapter_id=self.model_adapter_id) + @component.output_types(replies=List[str]) def run(self, prompt: str): """ diff --git a/pyproject.toml b/pyproject.toml index beb8fc02ae0..3d0ba57a8eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -227,9 +227,6 @@ dev = [ # https://opendev.org/openstack/reno/src/branch/master/requirements.txt#L7 "dulwich>=0.21.0,<1.0.0", ] -gradient = [ - "gradientai==1.1.0" -] formatting = [ # Version specified following Black stability policy: diff --git a/test/preview/components/embedders/test_gradient_document_embedder.py b/test/preview/components/embedders/test_gradient_document_embedder.py index e68a2e552e9..297596df9e7 100644 --- a/test/preview/components/embedders/test_gradient_document_embedder.py +++ b/test/preview/components/embedders/test_gradient_document_embedder.py @@ -1,4 +1,5 @@ import pytest +from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess from haystack.preview.components.embedders.gradient_document_embedder import GradientDocumentEmbedder from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np @@ -11,16 +12,6 @@ model = "bge-large" -def has_gradient(): - try: - import gradientai - - return True - except ModuleNotFoundError: - return False - - -@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") class TestGradientDocumentEmbedder: @pytest.mark.unit def test_init_from_env(self, monkeypatch): @@ -96,8 +87,6 @@ def test_run_fail_if_not_warmed_up(self): @pytest.mark.unit def test_run(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess - embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( @@ -108,6 +97,7 @@ def test_run(self): result = embedder.run(documents=documents) + assert embedder._embedding_model.generate_embeddings.call_count == 1 assert isinstance(result["documents"], list) assert len(result["documents"]) == len(documents) for doc in result["documents"]: @@ -140,8 +130,6 @@ def test_run_batch(self): @pytest.mark.unit def test_run_empty(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess - embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() diff --git a/test/preview/components/embedders/test_gradient_text_embedder.py b/test/preview/components/embedders/test_gradient_text_embedder.py index 3e2a6762f27..fa553e27f1d 100644 --- a/test/preview/components/embedders/test_gradient_text_embedder.py +++ b/test/preview/components/embedders/test_gradient_text_embedder.py @@ -1,4 +1,5 @@ import pytest +from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess from haystack.preview.components.embedders.gradient_text_embedder import GradientTextEmbedder from unittest.mock import MagicMock, NonCallableMagicMock import numpy as np @@ -9,16 +10,6 @@ model = "bge-large" -def has_gradient(): - try: - import gradientai - - return True - except ModuleNotFoundError: - return False - - -@pytest.mark.skipif(not has_gradient(), reason="Gradient is not installed") class TestGradientTextEmbedder: @pytest.mark.unit def test_init_from_env(self, monkeypatch): @@ -93,9 +84,33 @@ def test_run_fail_if_not_warmed_up(self): embedder.run(text="The food was delicious") @pytest.mark.unit - def test_run(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess + def test_run_fail_when_no_embeddings_returned(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess(embeddings=[]) + + with pytest.raises(RuntimeError): + _result = embedder.run(text="The food was delicious") + embedder._embedding_model.generate_embeddings.assert_called_once_with( + inputs=[{"input": "The food was delicious"}] + ) + + @pytest.mark.unit + def test_run_empty_string(self): + embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) + embedder._embedding_model = NonCallableMagicMock() + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": 0}] + ) + + result = embedder.run(text="") + embedder._embedding_model.generate_embeddings.assert_called_once_with(inputs=[{"input": ""}]) + assert len(result["embedding"]) == 1024 # 1024 is the bge-large embedding size + assert all(isinstance(x, float) for x in result["embedding"]) + + @pytest.mark.unit + def test_run(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( From 5eead3d2dc92748b8af5a958b787d90dae848239 Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Tue, 14 Nov 2023 18:19:14 -0800 Subject: [PATCH 09/12] update tests --- .../embedders/gradient_document_embedder.py | 6 ++-- .../components/generators/gradient/base.py | 4 +-- .../test_gradient_document_embedder.py | 30 +++++++++++++++---- 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/haystack/preview/components/embedders/gradient_document_embedder.py b/haystack/preview/components/embedders/gradient_document_embedder.py index bfd34f4bb37..c1e2fa71a60 100644 --- a/haystack/preview/components/embedders/gradient_document_embedder.py +++ b/haystack/preview/components/embedders/gradient_document_embedder.py @@ -34,6 +34,7 @@ def __init__( self, *, model_name: str = "bge-large", + batch_size: int = 100, access_token: Optional[str] = None, workspace_id: Optional[str] = None, host: Optional[str] = None, @@ -49,6 +50,7 @@ def __init__( :param host: The Gradient host. By default it uses https://api.gradient.ai/. """ gradientai_import.check() + self._batch_size = batch_size self._host = host self._model_name = model_name @@ -73,7 +75,7 @@ def warm_up(self) -> None: if not hasattr(self, "_embedding_model"): self._embedding_model = self._gradient.get_embeddings_model(slug=self._model_name) - def _generate_embeddings(self, documents: List[Document], batch_size=100) -> List[List[float]]: + def _generate_embeddings(self, documents: List[Document], batch_size: int) -> List[List[float]]: """ Batches the documents and generates the embeddings. """ @@ -103,7 +105,7 @@ def run(self, documents: List[Document]): if not hasattr(self, "_embedding_model"): raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.") - embeddings = self._generate_embeddings(documents=documents) + embeddings = self._generate_embeddings(documents=documents, batch_size=self._batch_size) for doc, embedding in zip(documents, embeddings): doc.embedding = embedding diff --git a/haystack/preview/components/generators/gradient/base.py b/haystack/preview/components/generators/gradient/base.py index 3a9fc63566e..dd420a43f4a 100644 --- a/haystack/preview/components/generators/gradient/base.py +++ b/haystack/preview/components/generators/gradient/base.py @@ -107,9 +107,9 @@ def warm_up(self): """ if not hasattr(self, "_model"): if isinstance(self._base_model_slug, str): - self._model = self._gradient.get_base_model(base_model_slug=self.base_model_slug) + self._model = self._gradient.get_base_model(base_model_slug=self._base_model_slug) if isinstance(self._model_adapter_id, str): - self._model = self._gradient.get_model_adapter(model_adapter_id=self.model_adapter_id) + self._model = self._gradient.get_model_adapter(model_adapter_id=self._model_adapter_id) @component.output_types(replies=List[str]) def run(self, prompt: str): diff --git a/test/preview/components/embedders/test_gradient_document_embedder.py b/test/preview/components/embedders/test_gradient_document_embedder.py index 297596df9e7..75c57c61ff2 100644 --- a/test/preview/components/embedders/test_gradient_document_embedder.py +++ b/test/preview/components/embedders/test_gradient_document_embedder.py @@ -83,7 +83,7 @@ def test_run_fail_if_not_warmed_up(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) with pytest.raises(RuntimeError, match="warm_up()"): - embedder.run(documents=[Document(text=f"document number {i}") for i in range(5)]) + embedder.run(documents=[Document(content=f"document number {i}") for i in range(5)]) @pytest.mark.unit def test_run(self): @@ -93,7 +93,7 @@ def test_run(self): embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(5)] ) - documents = [Document(text=f"document number {i}") for i in range(5)] + documents = [Document(content=f"document number {i}") for i in range(5)] result = embedder.run(documents=documents) @@ -107,8 +107,6 @@ def test_run(self): @pytest.mark.unit def test_run_batch(self): - from gradientai.openapi.client.models.generate_embedding_success import GenerateEmbeddingSuccess - embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() @@ -116,7 +114,7 @@ def test_run_batch(self): embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(110)] ) - documents = [Document(text=f"document number {i}") for i in range(110)] + documents = [Document(content=f"document number {i}") for i in range(110)] result = embedder.run(documents=documents) @@ -128,6 +126,28 @@ def test_run_batch(self): assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) + @pytest.mark.unit + def test_run_custom_batch(self): + embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id, batch_size=20) + embedder._embedding_model = NonCallableMagicMock() + + DOCUMENT_COUNT = 101 + embedder._embedding_model.generate_embeddings.return_value = GenerateEmbeddingSuccess( + embeddings=[{"embedding": np.random.rand(1024).tolist(), "index": i} for i in range(DOCUMENT_COUNT)] + ) + + documents = [Document(content=f"document number {i}") for i in range(DOCUMENT_COUNT)] + + result = embedder.run(documents=documents) + + assert embedder._embedding_model.generate_embeddings.call_count == 6 + assert isinstance(result["documents"], list) + assert len(result["documents"]) == len(documents) + for doc in result["documents"]: + assert isinstance(doc, Document) + assert isinstance(doc.embedding, list) + assert isinstance(doc.embedding[0], float) + @pytest.mark.unit def test_run_empty(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) From b63786537e520e2d9da915fba7f9bc3faaccd1bd Mon Sep 17 00:00:00 2001 From: Mateusz Haligowski Date: Wed, 15 Nov 2023 22:15:18 -0800 Subject: [PATCH 10/12] add gradientai to ci --- .github/workflows/tests_preview.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/tests_preview.yml b/.github/workflows/tests_preview.yml index a69bb369eb3..076fc137448 100644 --- a/.github/workflows/tests_preview.yml +++ b/.github/workflows/tests_preview.yml @@ -116,7 +116,7 @@ jobs: python-version: ${{ env.PYTHON_VERSION }} - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run run: pytest -m "not integration" test/preview @@ -174,7 +174,7 @@ jobs: sudo apt install ffmpeg # for local Whisper tests - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run run: pytest --maxfail=5 -m "integration" test/preview @@ -230,7 +230,7 @@ jobs: colima start - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run Tika run: docker run -d -p 9998:9998 apache/tika:2.9.0.0 @@ -281,7 +281,7 @@ jobs: python-version: ${{ env.PYTHON_VERSION }} - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run run: pytest --maxfail=5 -m "integration" test/preview -k 'not tika' From d10e272fee1f5e9dd66969b154e885fc3e38ca2e Mon Sep 17 00:00:00 2001 From: ZanSara Date: Thu, 23 Nov 2023 10:31:07 +0000 Subject: [PATCH 11/12] install gradientai in the CI --- .github/workflows/e2e_preview.yml | 2 +- .github/workflows/linting_preview.yml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/e2e_preview.yml b/.github/workflows/e2e_preview.yml index 2e9763d2801..a695c8c34b9 100644 --- a/.github/workflows/e2e_preview.yml +++ b/.github/workflows/e2e_preview.yml @@ -36,7 +36,7 @@ jobs: sudo apt install ffmpeg # for local Whisper tests - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Run tests run: pytest e2e/preview diff --git a/.github/workflows/linting_preview.yml b/.github/workflows/linting_preview.yml index 1c299844369..08ecc2645e2 100644 --- a/.github/workflows/linting_preview.yml +++ b/.github/workflows/linting_preview.yml @@ -38,7 +38,7 @@ jobs: python-version: ${{ env.PYTHON_VERSION }} - name: Install Haystack - run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' + run: pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai - name: Mypy if: steps.files.outputs.any_changed == 'true' @@ -69,7 +69,7 @@ jobs: - name: Install Haystack run: | - pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' + pip install .[dev,preview,audio] langdetect transformers[torch,sentencepiece]==4.35.2 'sentence-transformers>=2.2.0' pypdf markdown-it-py mdit_plain tika 'azure-ai-formrecognizer>=3.2.0b2' gradientai pip install ./haystack-linter - name: Pylint From a8e89652a933cd0d4571354291a5b667609560b3 Mon Sep 17 00:00:00 2001 From: ZanSara Date: Thu, 23 Nov 2023 10:34:07 +0000 Subject: [PATCH 12/12] fix to_dict tests --- .../components/embedders/test_gradient_document_embedder.py | 2 +- .../preview/components/embedders/test_gradient_text_embedder.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/preview/components/embedders/test_gradient_document_embedder.py b/test/preview/components/embedders/test_gradient_document_embedder.py index 75c57c61ff2..b114d382e27 100644 --- a/test/preview/components/embedders/test_gradient_document_embedder.py +++ b/test/preview/components/embedders/test_gradient_document_embedder.py @@ -59,7 +59,7 @@ def test_to_dict(self): component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "GradientDocumentEmbedder", + "type": "haystack.preview.components.embedders.gradient_document_embedder.GradientDocumentEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } diff --git a/test/preview/components/embedders/test_gradient_text_embedder.py b/test/preview/components/embedders/test_gradient_text_embedder.py index fa553e27f1d..9eddff7eef3 100644 --- a/test/preview/components/embedders/test_gradient_text_embedder.py +++ b/test/preview/components/embedders/test_gradient_text_embedder.py @@ -57,7 +57,7 @@ def test_to_dict(self): component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() assert data == { - "type": "GradientTextEmbedder", + "type": "haystack.preview.components.embedders.gradient_text_embedder.GradientTextEmbedder", "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, }