From 2c5e5a7ebc9227834e5b4ad0213fee309d3f27df Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Thu, 23 Nov 2023 17:45:33 -0500 Subject: [PATCH 01/25] feat: add bedrock embeddings to embedding encoder --- .../nodes/retriever/_embedding_encoder.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index d07e7ab1b8..ceb55e9a9f 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -42,6 +42,8 @@ from haystack.modeling.infer import Inferencer from haystack.nodes.retriever._losses import _TRAINING_LOSSES +with LazyImport(message="Run 'pip install boto3'") as boto3_import: + import boto3 COHERE_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30)) COHERE_BACKOFF = int(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10)) @@ -434,6 +436,57 @@ def save(self, save_dir: Union[Path, str]): raise NotImplementedError(f"Saving is not implemented for {self.__class__}") +class _BedrockEmbeddingEncoder(_BaseEmbeddingEncoder): + def __init__(self, retriever: "EmbeddingRetriever"): + boto3_import.check() + + # See https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html for more details + # The maximum input text is 8K tokens and the maximum output vector length is 1536 + # Bedrock embeddings do not support batch operations + self.model: str = "amazon.titan-embed-text-v1" + self.client = retriever.client + + def embed(self, text: str) -> np.ndarray: + input_body = {} + input_body["inputText"] = text + body = json.dumps(input_body) + response = self.client.invoke_model( + body=body, + modelId=self.model, + accept="application/json", + contentType="application/json", + ) + + response_body = json.loads(response.get("body").read()) + return np.array(response_body.get("embedding")) + + def embed_queries(self, queries: List[str]) -> np.ndarray: + all_embeddings = [] + for q in queries: + generated_embeddings = self.embed(q) + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) + + def embed_documents(self, docs: List[Document]) -> np.ndarray: + return self.embed_queries([d.content for d in docs]) + + def train( + self, + training_data: List[Dict[str, Any]], + learning_rate: float = 2e-5, + n_epochs: int = 1, + num_warmup_steps: Optional[int] = None, + batch_size: int = 16, + train_loss: Literal["mnrl", "margin_mse"] = "mnrl", + num_workers: int = 0, + use_amp: bool = False, + **kwargs, + ): + raise NotImplementedError(f"Training is not implemented for {self.__class__}") + + def save(self, save_dir: Union[Path, str]): + raise NotImplementedError(f"Saving is not implemented for {self.__class__}") + _EMBEDDING_ENCODERS: Dict[str, Callable] = { "farm": _DefaultEmbeddingEncoder, "transformers": _DefaultEmbeddingEncoder, From bc370b861ca74f92ce8117a9b5aef75324e6622f Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Thu, 23 Nov 2023 23:10:56 -0500 Subject: [PATCH 02/25] reno & black --- haystack/nodes/retriever/_embedding_encoder.py | 8 +++----- .../aws-bedrock-embedding-encoder-a978884c1a2c8237.yaml | 4 ++++ 2 files changed, 7 insertions(+), 5 deletions(-) create mode 100644 releasenotes/notes/aws-bedrock-embedding-encoder-a978884c1a2c8237.yaml diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index ceb55e9a9f..7f8bc9d159 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -451,11 +451,8 @@ def embed(self, text: str) -> np.ndarray: input_body["inputText"] = text body = json.dumps(input_body) response = self.client.invoke_model( - body=body, - modelId=self.model, - accept="application/json", - contentType="application/json", - ) + body=body, modelId=self.model, accept="application/json", contentType="application/json" + ) response_body = json.loads(response.get("body").read()) return np.array(response_body.get("embedding")) @@ -487,6 +484,7 @@ def train( def save(self, save_dir: Union[Path, str]): raise NotImplementedError(f"Saving is not implemented for {self.__class__}") + _EMBEDDING_ENCODERS: Dict[str, Callable] = { "farm": _DefaultEmbeddingEncoder, "transformers": _DefaultEmbeddingEncoder, diff --git a/releasenotes/notes/aws-bedrock-embedding-encoder-a978884c1a2c8237.yaml b/releasenotes/notes/aws-bedrock-embedding-encoder-a978884c1a2c8237.yaml new file mode 100644 index 0000000000..caa1f479b8 --- /dev/null +++ b/releasenotes/notes/aws-bedrock-embedding-encoder-a978884c1a2c8237.yaml @@ -0,0 +1,4 @@ +--- +features: + - | + Adding Bedrock Embeddings Encoder to use as a retriever. From b28fbc3d3510b3f4bf4d7cbebaa2640c1b9474b4 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Fri, 24 Nov 2023 00:16:05 -0500 Subject: [PATCH 03/25] feat: refactoring for bedrock embedding encoder --- haystack/nodes/retriever/_embedding_encoder.py | 18 +++++++++++++++++- haystack/nodes/retriever/dense.py | 3 +++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 7f8bc9d159..b6b29834e1 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -444,7 +444,23 @@ def __init__(self, retriever: "EmbeddingRetriever"): # The maximum input text is 8K tokens and the maximum output vector length is 1536 # Bedrock embeddings do not support batch operations self.model: str = "amazon.titan-embed-text-v1" - self.client = retriever.client + self.aws_config = retriever.aws_config + self.client = self.initialize_boto3_client() + + def initialize_boto3_client(self): + if self.aws_config: + access_key_id = self.aws_config.get("aws_access_key_id") + secret_access_key = self.aws_config.get("aws_secret_access_key") + region = self.aws_config.get("region") + try: + return boto3.client( + 'bedrock-runtime', + aws_access_key_id=access_key_id, + aws_secret_access_key=secret_access_key, + region_name=region + ) + except Exception as e: + raise ValueError(f"AWS client error {e}") def embed(self, text: str) -> np.ndarray: input_body = {} diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index d4663e6cfd..026f2cd672 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1468,6 +1468,7 @@ def __init__( azure_deployment_name: Optional[str] = None, api_base: str = "https://api.openai.com/v1", openai_organization: Optional[str] = None, + aws_config: Optional[Dict[str, Any]]= None, ): """ :param document_store: An instance of DocumentStore from which to retrieve documents. @@ -1532,6 +1533,7 @@ def __init__( will not be used. :param api_base: The OpenAI API base URL, defaults to `"https://api.openai.com/v1"`. :param openai_organization: The OpenAI-Organization ID, defaults to `None`. For more details, see OpenAI + :param aws_config: The aws_config contains {aws_access_key, aws_secret_key, aws_region } to use with the boto3 client for an AWS Bedrock retriever. Defaults to 'None'. [documentation](https://platform.openai.com/docs/api-reference/requesting-organization). """ torch_and_transformers_import.check() @@ -1565,6 +1567,7 @@ def __init__( self.azure_base_url = azure_base_url self.azure_deployment_name = azure_deployment_name self.openai_organization = openai_organization + self.aws_config = aws_config self.model_format = ( self._infer_model_format(model_name_or_path=embedding_model, use_auth_token=use_auth_token) if model_format is None From 7eb3855d3e3b9fa1ae7770598767d45212e18c3c Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Fri, 24 Nov 2023 00:19:42 -0500 Subject: [PATCH 04/25] feat: bedrock embedding encoder --- haystack/nodes/retriever/_embedding_encoder.py | 4 ++-- haystack/nodes/retriever/dense.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index b6b29834e1..f87238e486 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -454,10 +454,10 @@ def initialize_boto3_client(self): region = self.aws_config.get("region") try: return boto3.client( - 'bedrock-runtime', + "bedrock-runtime", aws_access_key_id=access_key_id, aws_secret_access_key=secret_access_key, - region_name=region + region_name=region, ) except Exception as e: raise ValueError(f"AWS client error {e}") diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 026f2cd672..cd24896bdf 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1468,7 +1468,7 @@ def __init__( azure_deployment_name: Optional[str] = None, api_base: str = "https://api.openai.com/v1", openai_organization: Optional[str] = None, - aws_config: Optional[Dict[str, Any]]= None, + aws_config: Optional[Dict[str, Any]] = None, ): """ :param document_store: An instance of DocumentStore from which to retrieve documents. From 928920c4b76bb867755ab3cfb4b1066b6c7cfd88 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Fri, 24 Nov 2023 01:13:32 -0500 Subject: [PATCH 05/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 1 + haystack/nodes/retriever/dense.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index f87238e486..7e64d9e2e3 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -508,4 +508,5 @@ def save(self, save_dir: Union[Path, str]): "retribert": _RetribertEmbeddingEncoder, "openai": _OpenAIEmbeddingEncoder, "cohere": _CohereEmbeddingEncoder, + "bedrock": _BedrockEmbeddingEncoder, } diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index cd24896bdf..489dc099fb 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1895,6 +1895,8 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[ return "openai" if model_name_or_path in COHERE_EMBEDDING_MODELS: return "cohere" + if model_name_or_path == "bedrock": + return "bedrock" # Check if model name is a local directory with sentence transformers config file in it if Path(model_name_or_path).exists(): if Path(f"{model_name_or_path}/config_sentence_transformers.json").exists(): From cfa06d05dc475e05fad3761d9c1eff984272f23a Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Fri, 24 Nov 2023 01:39:34 -0500 Subject: [PATCH 06/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 2 ++ haystack/nodes/retriever/dense.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 7e64d9e2e3..97d09c2e74 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -461,6 +461,8 @@ def initialize_boto3_client(self): ) except Exception as e: raise ValueError(f"AWS client error {e}") + else: + raise ValueError(f"Please pass boto3.client(bedrock-runtime) credentials configuration") def embed(self, text: str) -> np.ndarray: input_body = {} diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 489dc099fb..054e678be1 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1895,7 +1895,7 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[ return "openai" if model_name_or_path in COHERE_EMBEDDING_MODELS: return "cohere" - if model_name_or_path == "bedrock": + elif model_name_or_path == "bedrock": return "bedrock" # Check if model name is a local directory with sentence transformers config file in it if Path(model_name_or_path).exists(): From 98f997a779509c754506634ff5c40d90fcbe1200 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Fri, 24 Nov 2023 01:48:02 -0500 Subject: [PATCH 07/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 2 +- haystack/nodes/retriever/dense.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 97d09c2e74..b205b3743c 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -462,7 +462,7 @@ def initialize_boto3_client(self): except Exception as e: raise ValueError(f"AWS client error {e}") else: - raise ValueError(f"Please pass boto3.client(bedrock-runtime) credentials configuration") + raise ValueError("Please pass boto3.client(bedrock-runtime) credentials configuration") def embed(self, text: str) -> np.ndarray: input_body = {} diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 054e678be1..5bb043ac76 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1893,12 +1893,12 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[ ) if valid_openai_model_name: return "openai" - if model_name_or_path in COHERE_EMBEDDING_MODELS: + elif model_name_or_path in COHERE_EMBEDDING_MODELS: return "cohere" elif model_name_or_path == "bedrock": return "bedrock" # Check if model name is a local directory with sentence transformers config file in it - if Path(model_name_or_path).exists(): + elif Path(model_name_or_path).exists(): if Path(f"{model_name_or_path}/config_sentence_transformers.json").exists(): return "sentence_transformers" # Check if sentence transformers config file in model hub From 6d11af33b08236ccd823b3405d88a984b310025a Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Fri, 24 Nov 2023 01:53:47 -0500 Subject: [PATCH 08/25] feat: bedrock refactoring --- haystack/nodes/retriever/dense.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 5bb043ac76..489dc099fb 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1893,12 +1893,12 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[ ) if valid_openai_model_name: return "openai" - elif model_name_or_path in COHERE_EMBEDDING_MODELS: + if model_name_or_path in COHERE_EMBEDDING_MODELS: return "cohere" - elif model_name_or_path == "bedrock": + if model_name_or_path == "bedrock": return "bedrock" # Check if model name is a local directory with sentence transformers config file in it - elif Path(model_name_or_path).exists(): + if Path(model_name_or_path).exists(): if Path(f"{model_name_or_path}/config_sentence_transformers.json").exists(): return "sentence_transformers" # Check if sentence transformers config file in model hub From 5739b5109c9666fdec8d96f429c3f1a0df4d2e42 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Sat, 25 Nov 2023 22:19:22 -0500 Subject: [PATCH 09/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index b205b3743c..8b3cee3d6d 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -449,20 +449,20 @@ def __init__(self, retriever: "EmbeddingRetriever"): def initialize_boto3_client(self): if self.aws_config: - access_key_id = self.aws_config.get("aws_access_key_id") - secret_access_key = self.aws_config.get("aws_secret_access_key") - region = self.aws_config.get("region") try: return boto3.client( "bedrock-runtime", - aws_access_key_id=access_key_id, - aws_secret_access_key=secret_access_key, - region_name=region, + aws_access_key_id=self.aws_config.get("aws_access_key_id"), + aws_secret_access_key=self.aws_config.get("aws_secret_access_key"), + region_name=self.aws_config.get("region"), ) except Exception as e: - raise ValueError(f"AWS client error {e}") + raise ValueError("Please pass boto3.client(bedrock-runtime) credentials configuration") else: - raise ValueError("Please pass boto3.client(bedrock-runtime) credentials configuration") + try: + return boto3.client("bedrock-runtime") + except Exception as e: + raise ValueError(f"AWS client error {e}") def embed(self, text: str) -> np.ndarray: input_body = {} From b80d03942a174d5d11ff4af7529a4d348c69aaa2 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Sat, 25 Nov 2023 22:26:28 -0500 Subject: [PATCH 10/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 8b3cee3d6d..8ace23506f 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -456,7 +456,7 @@ def initialize_boto3_client(self): aws_secret_access_key=self.aws_config.get("aws_secret_access_key"), region_name=self.aws_config.get("region"), ) - except Exception as e: + except Exception: raise ValueError("Please pass boto3.client(bedrock-runtime) credentials configuration") else: try: From 78a506d216fd1c3f141089cc3078e2201a427111 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Sun, 26 Nov 2023 10:42:59 -0500 Subject: [PATCH 11/25] feat: bedrock refactoring, add cohere --- .../nodes/retriever/_embedding_encoder.py | 72 +++++++++++++------ haystack/nodes/retriever/dense.py | 10 +-- 2 files changed, 55 insertions(+), 27 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 8ace23506f..dcc4e5f342 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -57,6 +57,11 @@ "embed-multilingual-v2.0", ] +BEDROCK_EMBEDDING_MODELS = [ + "amazon.titan-embed-text-v1", + "cohere.embed-english-v3", + "cohere.embed-multilingual-v3"] + class _DefaultEmbeddingEncoder(_BaseEmbeddingEncoder): def __init__(self, retriever: "EmbeddingRetriever"): @@ -442,38 +447,59 @@ def __init__(self, retriever: "EmbeddingRetriever"): # See https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html for more details # The maximum input text is 8K tokens and the maximum output vector length is 1536 - # Bedrock embeddings do not support batch operations - self.model: str = "amazon.titan-embed-text-v1" + # Bedrock Titan embeddings do not support batch operations self.aws_config = retriever.aws_config - self.client = self.initialize_boto3_client() + self.client = self.initialize_boto3_session().client("bedrock-runtime") + self.model: str = next( + (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "amazon.titan-embed-text-v1" + ) - def initialize_boto3_client(self): + def initialize_boto3_session(self) -> boto3.Session: if self.aws_config: - try: - return boto3.client( - "bedrock-runtime", - aws_access_key_id=self.aws_config.get("aws_access_key_id"), - aws_secret_access_key=self.aws_config.get("aws_secret_access_key"), - region_name=self.aws_config.get("region"), - ) - except Exception: - raise ValueError("Please pass boto3.client(bedrock-runtime) credentials configuration") + profile_name = self.aws_config.get('profile_name', None) + access_key = self.aws_config.get('aws_access_key_id', None) + secret_key = self.aws_config.get('aws_secret_access_key', None) + region = self.aws_config.get('region_name', None) + if profile_name: + try: + return boto3.Session(profile_name=profile_name, region_name=region) + except Exception as e: + raise ValueError(f"AWS client error {e}") + elif access_key and secret_key: + return boto3.Session(aws_access_key_id=access_key, aws_secret_access_key=secret_key, region_name=region) else: try: - return boto3.client("bedrock-runtime") + return boto3.Session() except Exception as e: raise ValueError(f"AWS client error {e}") def embed(self, text: str) -> np.ndarray: - input_body = {} - input_body["inputText"] = text - body = json.dumps(input_body) - response = self.client.invoke_model( - body=body, modelId=self.model, accept="application/json", contentType="application/json" - ) + if self.model == "amazon.titan-embed-text-v1": + input_body = {} + input_body["inputText"] = text + body = json.dumps(input_body) + response = self.client.invoke_model( + body=body, modelId=self.model, accept="application/json", contentType="application/json" + ) + + response_body = json.loads(response.get("body").read()) + return np.array(response_body.get("embedding")) + else: + coherePayload = json.dumps({ + 'texts': [text], + 'input_type': 'search_document', + 'truncate': 'END' + }) + response = self.client.invoke_model( + body=coherePayload, + modelId=self.model, + accept='application/json', + contentType='application/json' + ) - response_body = json.loads(response.get("body").read()) - return np.array(response_body.get("embedding")) + body = response.get('body').read().decode('utf-8') + response_body = json.loads(body) + return np.array(response_body['embeddings']) def embed_queries(self, queries: List[str]) -> np.ndarray: all_embeddings = [] @@ -510,5 +536,5 @@ def save(self, save_dir: Union[Path, str]): "retribert": _RetribertEmbeddingEncoder, "openai": _OpenAIEmbeddingEncoder, "cohere": _CohereEmbeddingEncoder, - "bedrock": _BedrockEmbeddingEncoder, + "bedrock": _BedrockEmbeddingEncoder } diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 489dc099fb..d1a0c88917 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -17,7 +17,7 @@ from haystack.schema import Document, FilterType from haystack.document_stores import BaseDocumentStore from haystack.nodes.retriever.base import BaseRetriever -from haystack.nodes.retriever._embedding_encoder import _EMBEDDING_ENCODERS, COHERE_EMBEDDING_MODELS +from haystack.nodes.retriever._embedding_encoder import _EMBEDDING_ENCODERS, COHERE_EMBEDDING_MODELS, BEDROCK_EMBEDDING_MODELS from haystack.utils.early_stopping import EarlyStopping from haystack.telemetry import send_event from haystack.lazy_imports import LazyImport @@ -1475,7 +1475,8 @@ def __init__( :param embedding_model: Local path or name of model in Hugging Face's model hub such as ``'sentence-transformers/all-MiniLM-L6-v2'``. The embedding model could also potentially be an OpenAI model ["ada", "babbage", "davinci", "curie"] or - a Cohere model ["embed-english-v2.0", "embed-english-light-v2.0", "embed-multilingual-v2.0"]. + a Cohere model ["embed-english-v2.0", "embed-english-light-v2.0", "embed-multilingual-v2.0"] or + an AWS Bedrock model ["amazon.titan-embed-text-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3"]. :param model_version: The version of model to use from the HuggingFace model hub. Can be tag name, branch name, or commit hash. :param use_gpu: Whether to use all available GPUs or the CPU. Falls back on CPU if no GPU is available. :param batch_size: Number of documents to encode at once. @@ -1490,6 +1491,7 @@ def __init__( 4. `retribert` : (will use `_RetribertEmbeddingEncoder` as embedding encoder) 5. `openai` : (will use `_OpenAIEmbeddingEncoder` as embedding encoder) 6. `cohere` : (will use `_CohereEmbeddingEncoder` as embedding encoder) + 7. `bedrock` : (will use `_BedrockEmbeddingEncoder` as embedding encoder) :param pooling_strategy: Strategy for combining the embeddings from the model (for farm / transformers models only). Options: @@ -1533,8 +1535,8 @@ def __init__( will not be used. :param api_base: The OpenAI API base URL, defaults to `"https://api.openai.com/v1"`. :param openai_organization: The OpenAI-Organization ID, defaults to `None`. For more details, see OpenAI - :param aws_config: The aws_config contains {aws_access_key, aws_secret_key, aws_region } to use with the boto3 client for an AWS Bedrock retriever. Defaults to 'None'. [documentation](https://platform.openai.com/docs/api-reference/requesting-organization). + :param aws_config: The aws_config contains {aws_access_key, aws_secret_key, aws_region, profile_name} to use with the boto3 Session for an AWS Bedrock retriever. Defaults to 'None'. """ torch_and_transformers_import.check() @@ -1895,7 +1897,7 @@ def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[ return "openai" if model_name_or_path in COHERE_EMBEDDING_MODELS: return "cohere" - if model_name_or_path == "bedrock": + if model_name_or_path in BEDROCK_EMBEDDING_MODELS: return "bedrock" # Check if model name is a local directory with sentence transformers config file in it if Path(model_name_or_path).exists(): From 998cfe369657dd7f3ab013ff299545f59b1b746a Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Sun, 26 Nov 2023 10:49:46 -0500 Subject: [PATCH 12/25] feat: bedrock refactoring, add cohere --- .../nodes/retriever/_embedding_encoder.py | 32 +++++++------------ haystack/nodes/retriever/dense.py | 6 +++- 2 files changed, 16 insertions(+), 22 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index dcc4e5f342..af4eeb751a 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -57,10 +57,7 @@ "embed-multilingual-v2.0", ] -BEDROCK_EMBEDDING_MODELS = [ - "amazon.titan-embed-text-v1", - "cohere.embed-english-v3", - "cohere.embed-multilingual-v3"] +BEDROCK_EMBEDDING_MODELS = ["amazon.titan-embed-text-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3"] class _DefaultEmbeddingEncoder(_BaseEmbeddingEncoder): @@ -454,12 +451,12 @@ def __init__(self, retriever: "EmbeddingRetriever"): (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "amazon.titan-embed-text-v1" ) - def initialize_boto3_session(self) -> boto3.Session: + def initialize_boto3_session(self): if self.aws_config: - profile_name = self.aws_config.get('profile_name', None) - access_key = self.aws_config.get('aws_access_key_id', None) - secret_key = self.aws_config.get('aws_secret_access_key', None) - region = self.aws_config.get('region_name', None) + profile_name = self.aws_config.get("profile_name", None) + access_key = self.aws_config.get("aws_access_key_id", None) + secret_key = self.aws_config.get("aws_secret_access_key", None) + region = self.aws_config.get("region_name", None) if profile_name: try: return boto3.Session(profile_name=profile_name, region_name=region) @@ -485,21 +482,14 @@ def embed(self, text: str) -> np.ndarray: response_body = json.loads(response.get("body").read()) return np.array(response_body.get("embedding")) else: - coherePayload = json.dumps({ - 'texts': [text], - 'input_type': 'search_document', - 'truncate': 'END' - }) + coherePayload = json.dumps({"texts": [text], "input_type": "search_document", "truncate": "END"}) response = self.client.invoke_model( - body=coherePayload, - modelId=self.model, - accept='application/json', - contentType='application/json' + body=coherePayload, modelId=self.model, accept="application/json", contentType="application/json" ) - body = response.get('body').read().decode('utf-8') + body = response.get("body").read().decode("utf-8") response_body = json.loads(body) - return np.array(response_body['embeddings']) + return np.array(response_body["embeddings"]) def embed_queries(self, queries: List[str]) -> np.ndarray: all_embeddings = [] @@ -536,5 +526,5 @@ def save(self, save_dir: Union[Path, str]): "retribert": _RetribertEmbeddingEncoder, "openai": _OpenAIEmbeddingEncoder, "cohere": _CohereEmbeddingEncoder, - "bedrock": _BedrockEmbeddingEncoder + "bedrock": _BedrockEmbeddingEncoder, } diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index d1a0c88917..5576c9cf03 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -17,7 +17,11 @@ from haystack.schema import Document, FilterType from haystack.document_stores import BaseDocumentStore from haystack.nodes.retriever.base import BaseRetriever -from haystack.nodes.retriever._embedding_encoder import _EMBEDDING_ENCODERS, COHERE_EMBEDDING_MODELS, BEDROCK_EMBEDDING_MODELS +from haystack.nodes.retriever._embedding_encoder import ( + _EMBEDDING_ENCODERS, + COHERE_EMBEDDING_MODELS, + BEDROCK_EMBEDDING_MODELS, +) from haystack.utils.early_stopping import EarlyStopping from haystack.telemetry import send_event from haystack.lazy_imports import LazyImport From a87395c88ffee73eb386de231b46cfb847a0f3db Mon Sep 17 00:00:00 2001 From: Stefano Fiorucci <44616784+anakin87@users.noreply.github.com> Date: Tue, 5 Dec 2023 13:08:21 +0100 Subject: [PATCH 13/25] pylint: disable too-many-return-statements in method --- haystack/nodes/retriever/dense.py | 1 + 1 file changed, 1 insertion(+) diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py index 5576c9cf03..3139ec1561 100644 --- a/haystack/nodes/retriever/dense.py +++ b/haystack/nodes/retriever/dense.py @@ -1894,6 +1894,7 @@ def _preprocess_documents(self, docs: List[Document]) -> List[Document]: @staticmethod def _infer_model_format(model_name_or_path: str, use_auth_token: Optional[Union[str, bool]]) -> str: + # pylint: disable=too-many-return-statements valid_openai_model_name = model_name_or_path in ["ada", "babbage", "davinci", "curie"] or any( m in model_name_or_path for m in ["-ada-", "-babbage-", "-davinci-", "-curie-"] ) From 5b6e2c32ad745d9ba0d55a054739922916f8dce4 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Wed, 6 Dec 2023 19:56:17 -0500 Subject: [PATCH 14/25] feat: bedrock refactoring --- .../nodes/retriever/_embedding_encoder.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index af4eeb751a..49794ee6c3 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -447,26 +447,26 @@ def __init__(self, retriever: "EmbeddingRetriever"): # Bedrock Titan embeddings do not support batch operations self.aws_config = retriever.aws_config self.client = self.initialize_boto3_session().client("bedrock-runtime") - self.model: str = next( - (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "amazon.titan-embed-text-v1" - ) + self.model: str = next((m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), None) + + if self.model is None: + raise ValueError("Model not found in the retriever's embedding models.") def initialize_boto3_session(self): if self.aws_config: - profile_name = self.aws_config.get("profile_name", None) - access_key = self.aws_config.get("aws_access_key_id", None) - secret_key = self.aws_config.get("aws_secret_access_key", None) - region = self.aws_config.get("region_name", None) - if profile_name: - try: - return boto3.Session(profile_name=profile_name, region_name=region) - except Exception as e: - raise ValueError(f"AWS client error {e}") - elif access_key and secret_key: - return boto3.Session(aws_access_key_id=access_key, aws_secret_access_key=secret_key, region_name=region) - else: + aws_profile_name = self.aws_config.get("profile_name", None) + aws_access_key_id = self.aws_config.get("aws_access_key_id", None) + aws_secret_access_key = self.aws_config.get("aws_secret_access_key", None) + aws_region_name = self.aws_config.get("region_name", None) + aws_session_token = self.aws_config.get("session_token", None) try: - return boto3.Session() + return boto3.Session( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + region_name=aws_region_name, + profile_name=aws_profile_name, + ) except Exception as e: raise ValueError(f"AWS client error {e}") @@ -482,7 +482,7 @@ def embed(self, text: str) -> np.ndarray: response_body = json.loads(response.get("body").read()) return np.array(response_body.get("embedding")) else: - coherePayload = json.dumps({"texts": [text], "input_type": "search_document", "truncate": "END"}) + coherePayload = json.dumps({"texts": [text], "input_type": "search_query", "truncate": "END"}) response = self.client.invoke_model( body=coherePayload, modelId=self.model, accept="application/json", contentType="application/json" ) From d50a03e9c8a805360b8f3e86fbb178e866aa2099 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Wed, 6 Dec 2023 20:09:58 -0500 Subject: [PATCH 15/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 49794ee6c3..76d69607f3 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -447,9 +447,9 @@ def __init__(self, retriever: "EmbeddingRetriever"): # Bedrock Titan embeddings do not support batch operations self.aws_config = retriever.aws_config self.client = self.initialize_boto3_session().client("bedrock-runtime") - self.model: str = next((m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), None) + self.model: str = next((m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "Not Found") - if self.model is None: + if self.model == "Not Found": raise ValueError("Model not found in the retriever's embedding models.") def initialize_boto3_session(self): From 8760b237773a0ea023ba5762f9704629828da826 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Wed, 6 Dec 2023 20:25:41 -0500 Subject: [PATCH 16/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 76d69607f3..b74fe4cce2 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -447,7 +447,9 @@ def __init__(self, retriever: "EmbeddingRetriever"): # Bedrock Titan embeddings do not support batch operations self.aws_config = retriever.aws_config self.client = self.initialize_boto3_session().client("bedrock-runtime") - self.model: str = next((m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "Not Found") + self.model: str = next( + (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "amazon.titan-embed-text-v1" + ) if self.model == "Not Found": raise ValueError("Model not found in the retriever's embedding models.") From 1a1bbcef5660788487f30ff8495f05cf1a97b7b2 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Wed, 6 Dec 2023 20:33:24 -0500 Subject: [PATCH 17/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index b74fe4cce2..d6c93df49a 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -451,9 +451,6 @@ def __init__(self, retriever: "EmbeddingRetriever"): (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "amazon.titan-embed-text-v1" ) - if self.model == "Not Found": - raise ValueError("Model not found in the retriever's embedding models.") - def initialize_boto3_session(self): if self.aws_config: aws_profile_name = self.aws_config.get("profile_name", None) From 92b410d8d090fd4db3001a8e763c0f22f5f143cb Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Wed, 6 Dec 2023 20:41:37 -0500 Subject: [PATCH 18/25] feat: bedrock refactoring --- haystack/nodes/retriever/_embedding_encoder.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index d6c93df49a..c871e9e7da 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -448,9 +448,12 @@ def __init__(self, retriever: "EmbeddingRetriever"): self.aws_config = retriever.aws_config self.client = self.initialize_boto3_session().client("bedrock-runtime") self.model: str = next( - (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "amazon.titan-embed-text-v1" + (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "Model Not Found" ) + if self.model == "Model Not Found": + raise ValueError("Model not found in Bedrock Embedding Models") + def initialize_boto3_session(self): if self.aws_config: aws_profile_name = self.aws_config.get("profile_name", None) From b8136437ba62f554b497315992fb2a58b3a5c490 Mon Sep 17 00:00:00 2001 From: jlonge4 Date: Thu, 7 Dec 2023 12:35:48 -0500 Subject: [PATCH 19/25] feat: bedrock refactoring --- .../nodes/retriever/_embedding_encoder.py | 26 ++++++++++++++----- 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index c871e9e7da..6110425543 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -472,7 +472,7 @@ def initialize_boto3_session(self): except Exception as e: raise ValueError(f"AWS client error {e}") - def embed(self, text: str) -> np.ndarray: + def embed(self, text: str, embed_type: Optional[str] = None) -> np.ndarray: if self.model == "amazon.titan-embed-text-v1": input_body = {} input_body["inputText"] = text @@ -484,7 +484,7 @@ def embed(self, text: str) -> np.ndarray: response_body = json.loads(response.get("body").read()) return np.array(response_body.get("embedding")) else: - coherePayload = json.dumps({"texts": [text], "input_type": "search_query", "truncate": "END"}) + coherePayload = json.dumps({"texts": [text], "input_type": embed_type, "truncate": "END"}) response = self.client.invoke_model( body=coherePayload, modelId=self.model, accept="application/json", contentType="application/json" ) @@ -496,12 +496,26 @@ def embed(self, text: str) -> np.ndarray: def embed_queries(self, queries: List[str]) -> np.ndarray: all_embeddings = [] for q in queries: - generated_embeddings = self.embed(q) - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) + if self.model == "amazon.titan-embed-text-v1": + generated_embeddings = self.embed(q) + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) + else: + generated_embeddings = self.embed(q, "search_query") + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) def embed_documents(self, docs: List[Document]) -> np.ndarray: - return self.embed_queries([d.content for d in docs]) + all_embeddings = [] + for d in docs: + if self.model == "amazon.titan-embed-text-v1": + generated_embeddings = self.embed(d) + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) + else: + generated_embeddings = self.embed(d, "search_document") + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) def train( self, From 7125d9b820574614bd04defdc54730846449a9c2 Mon Sep 17 00:00:00 2001 From: anakin87 Date: Fri, 15 Dec 2023 17:20:42 +0100 Subject: [PATCH 20/25] fix mypy and pylint errors --- .../nodes/retriever/_embedding_encoder.py | 52 +++++++++---------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 6110425543..193bc6550a 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -455,22 +455,24 @@ def __init__(self, retriever: "EmbeddingRetriever"): raise ValueError("Model not found in Bedrock Embedding Models") def initialize_boto3_session(self): - if self.aws_config: - aws_profile_name = self.aws_config.get("profile_name", None) - aws_access_key_id = self.aws_config.get("aws_access_key_id", None) - aws_secret_access_key = self.aws_config.get("aws_secret_access_key", None) - aws_region_name = self.aws_config.get("region_name", None) - aws_session_token = self.aws_config.get("session_token", None) - try: - return boto3.Session( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - region_name=aws_region_name, - profile_name=aws_profile_name, - ) - except Exception as e: - raise ValueError(f"AWS client error {e}") + if self.aws_config is None: + raise ValueError("`aws_config` is not set. To use Bedrock models, you should set `aws_config` when initializing the retriever.") + + aws_profile_name = self.aws_config.get("profile_name", None) + aws_access_key_id = self.aws_config.get("aws_access_key_id", None) + aws_secret_access_key = self.aws_config.get("aws_secret_access_key", None) + aws_region_name = self.aws_config.get("region_name", None) + aws_session_token = self.aws_config.get("session_token", None) + try: + return boto3.Session( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + region_name=aws_region_name, + profile_name=aws_profile_name, + ) + except Exception as e: + raise ValueError(f"AWS client error {e}") def embed(self, text: str, embed_type: Optional[str] = None) -> np.ndarray: if self.model == "amazon.titan-embed-text-v1": @@ -498,25 +500,21 @@ def embed_queries(self, queries: List[str]) -> np.ndarray: for q in queries: if self.model == "amazon.titan-embed-text-v1": generated_embeddings = self.embed(q) - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) else: generated_embeddings = self.embed(q, "search_query") - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) def embed_documents(self, docs: List[Document]) -> np.ndarray: all_embeddings = [] for d in docs: if self.model == "amazon.titan-embed-text-v1": - generated_embeddings = self.embed(d) - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) + generated_embeddings = self.embed(d.content) else: - generated_embeddings = self.embed(d, "search_document") - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) - + generated_embeddings = self.embed(d.content, "search_document") + all_embeddings.append(generated_embeddings) + return np.concatenate(all_embeddings) + def train( self, training_data: List[Dict[str, Any]], From 82496854564b00160c6036548e3e20f6e90e245d Mon Sep 17 00:00:00 2001 From: anakin87 Date: Fri, 15 Dec 2023 17:30:13 +0100 Subject: [PATCH 21/25] manually run precommit --- haystack/nodes/retriever/_embedding_encoder.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 193bc6550a..08823a95cd 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -456,8 +456,10 @@ def __init__(self, retriever: "EmbeddingRetriever"): def initialize_boto3_session(self): if self.aws_config is None: - raise ValueError("`aws_config` is not set. To use Bedrock models, you should set `aws_config` when initializing the retriever.") - + raise ValueError( + "`aws_config` is not set. To use Bedrock models, you should set `aws_config` when initializing the retriever." + ) + aws_profile_name = self.aws_config.get("profile_name", None) aws_access_key_id = self.aws_config.get("aws_access_key_id", None) aws_secret_access_key = self.aws_config.get("aws_secret_access_key", None) @@ -513,8 +515,8 @@ def embed_documents(self, docs: List[Document]) -> np.ndarray: else: generated_embeddings = self.embed(d.content, "search_document") all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) - + return np.concatenate(all_embeddings) + def train( self, training_data: List[Dict[str, Any]], From badc4f93f773962777eeb63e8aa5205d50ec0f5f Mon Sep 17 00:00:00 2001 From: tstadel Date: Sun, 17 Dec 2023 20:16:27 +0100 Subject: [PATCH 22/25] refactor init --- .../nodes/retriever/_embedding_encoder.py | 41 +++++++++---------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 08823a95cd..205129c0a2 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -440,38 +440,37 @@ def save(self, save_dir: Union[Path, str]): class _BedrockEmbeddingEncoder(_BaseEmbeddingEncoder): def __init__(self, retriever: "EmbeddingRetriever"): - boto3_import.check() - - # See https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html for more details - # The maximum input text is 8K tokens and the maximum output vector length is 1536 - # Bedrock Titan embeddings do not support batch operations - self.aws_config = retriever.aws_config - self.client = self.initialize_boto3_session().client("bedrock-runtime") - self.model: str = next( - (m for m in BEDROCK_EMBEDDING_MODELS if m in retriever.embedding_model), "Model Not Found" - ) + """Embedding Encoder for Bedrock models + See https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings.html for more details. + The maximum input text is 8K tokens and the maximum output vector length is 1536. + Titan embeddings do not support batch operations. - if self.model == "Model Not Found": - raise ValueError("Model not found in Bedrock Embedding Models") + :param retriever: EmbeddingRetriever object + """ + boto3_import.check() + if retriever.embedding_model not in BEDROCK_EMBEDDING_MODELS: + raise ValueError("Model not supported by Bedrock Embedding Encoder") + self.model = retriever.embedding_model + self.client = self.initialize_boto3_session(retriever.aws_config).client("bedrock-runtime") - def initialize_boto3_session(self): - if self.aws_config is None: + def initialize_boto3_session(self, aws_config: Dict[str, Any]): + if aws_config is None: raise ValueError( "`aws_config` is not set. To use Bedrock models, you should set `aws_config` when initializing the retriever." ) - aws_profile_name = self.aws_config.get("profile_name", None) - aws_access_key_id = self.aws_config.get("aws_access_key_id", None) - aws_secret_access_key = self.aws_config.get("aws_secret_access_key", None) - aws_region_name = self.aws_config.get("region_name", None) - aws_session_token = self.aws_config.get("session_token", None) + aws_access_key_id = aws_config.get("aws_access_key_id", None) + aws_secret_access_key = aws_config.get("aws_secret_access_key", None) + aws_session_token = aws_config.get("aws_session_token", None) + region_name = aws_config.get("region_name", None) + profile_name = aws_config.get("profile_name", None) try: return boto3.Session( aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, - region_name=aws_region_name, - profile_name=aws_profile_name, + region_name=region_name, + profile_name=profile_name, ) except Exception as e: raise ValueError(f"AWS client error {e}") From 1315ff7f5addaf5dfbad1f42d6fe9454c6199cce Mon Sep 17 00:00:00 2001 From: tstadel Date: Sun, 17 Dec 2023 20:18:31 +0100 Subject: [PATCH 23/25] fix cohere truncate and refactor embed --- .../nodes/retriever/_embedding_encoder.py | 75 ++++++++++--------- 1 file changed, 39 insertions(+), 36 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 205129c0a2..97e2a6a656 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -475,46 +475,49 @@ def initialize_boto3_session(self, aws_config: Dict[str, Any]): except Exception as e: raise ValueError(f"AWS client error {e}") - def embed(self, text: str, embed_type: Optional[str] = None) -> np.ndarray: - if self.model == "amazon.titan-embed-text-v1": - input_body = {} - input_body["inputText"] = text - body = json.dumps(input_body) - response = self.client.invoke_model( - body=body, modelId=self.model, accept="application/json", contentType="application/json" - ) - - response_body = json.loads(response.get("body").read()) - return np.array(response_body.get("embedding")) - else: - coherePayload = json.dumps({"texts": [text], "input_type": embed_type, "truncate": "END"}) - response = self.client.invoke_model( - body=coherePayload, modelId=self.model, accept="application/json", contentType="application/json" - ) - - body = response.get("body").read().decode("utf-8") - response_body = json.loads(body) - return np.array(response_body["embeddings"]) + def _embed_batch_cohere( + self, texts: List[str], input_type: Literal["search_query", "search_document"] + ) -> np.ndarray: + cohere_payload = {"texts": texts, "input_type": input_type, "truncate": "RIGHT"} + response = self._invoke_model(cohere_payload) + embeddings = np.array(response["embeddings"]) + return embeddings + + def _embed_titan(self, text: str) -> np.ndarray: + titan_payload = {"inputText": text} + response = self._invoke_model(titan_payload) + embeddings = np.array(response["embedding"]) + return embeddings + + def _invoke_model(self, payload: Dict[str, Any]) -> Dict[str, Any]: + body = json.dumps(payload) + response = self.client.invoke_model( + body=body, modelId=self.model, accept="application/json", contentType="application/json" + ) + body = response.get("body").read().decode("utf-8") + response_body = json.loads(body) + return response_body def embed_queries(self, queries: List[str]) -> np.ndarray: - all_embeddings = [] - for q in queries: - if self.model == "amazon.titan-embed-text-v1": - generated_embeddings = self.embed(q) - else: - generated_embeddings = self.embed(q, "search_query") - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) + if self.model == "amazon.titan-embed-text-v1": + all_embeddings = [] + for query in queries: + generated_embeddings = self._embed_titan(query) + all_embeddings.append(generated_embeddings) + return np.stack(all_embeddings) + else: + return self._embed_batch_cohere(queries, input_type="search_query") def embed_documents(self, docs: List[Document]) -> np.ndarray: - all_embeddings = [] - for d in docs: - if self.model == "amazon.titan-embed-text-v1": - generated_embeddings = self.embed(d.content) - else: - generated_embeddings = self.embed(d.content, "search_document") - all_embeddings.append(generated_embeddings) - return np.concatenate(all_embeddings) + if self.model == "amazon.titan-embed-text-v1": + all_embeddings = [] + for doc in docs: + generated_embeddings = self._embed_titan(doc.content) + all_embeddings.append(generated_embeddings) + return np.stack(all_embeddings) + else: + contents = [d.content for d in docs] + return self._embed_batch_cohere(contents, input_type="search_document") def train( self, From 6be5163b913f57bd5d52d5dc4e736408293f683b Mon Sep 17 00:00:00 2001 From: tstadel Date: Sun, 17 Dec 2023 20:25:22 +0100 Subject: [PATCH 24/25] fix mypy --- haystack/nodes/retriever/_embedding_encoder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index 97e2a6a656..c2cb9aabda 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -453,7 +453,7 @@ def __init__(self, retriever: "EmbeddingRetriever"): self.model = retriever.embedding_model self.client = self.initialize_boto3_session(retriever.aws_config).client("bedrock-runtime") - def initialize_boto3_session(self, aws_config: Dict[str, Any]): + def initialize_boto3_session(self, aws_config: Optional[Dict[str, Any]]): if aws_config is None: raise ValueError( "`aws_config` is not set. To use Bedrock models, you should set `aws_config` when initializing the retriever." From 9fdcb939b93f73dcccd3845e837aba0a6637b2ef Mon Sep 17 00:00:00 2001 From: tstadel Date: Sun, 17 Dec 2023 20:39:17 +0100 Subject: [PATCH 25/25] proper exception handing --- haystack/nodes/retriever/_embedding_encoder.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py index c2cb9aabda..f75b997123 100644 --- a/haystack/nodes/retriever/_embedding_encoder.py +++ b/haystack/nodes/retriever/_embedding_encoder.py @@ -16,7 +16,7 @@ HAYSTACK_REMOTE_API_MAX_RETRIES, HAYSTACK_REMOTE_API_TIMEOUT_SEC, ) -from haystack.errors import CohereError, CohereUnauthorizedError +from haystack.errors import AWSConfigurationError, CohereError, CohereUnauthorizedError from haystack.nodes.retriever._openai_encoder import _OpenAIEmbeddingEncoder from haystack.schema import Document from haystack.telemetry import send_event @@ -44,6 +44,7 @@ with LazyImport(message="Run 'pip install boto3'") as boto3_import: import boto3 + from botocore.exceptions import BotoCoreError COHERE_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30)) COHERE_BACKOFF = int(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10)) @@ -451,9 +452,9 @@ def __init__(self, retriever: "EmbeddingRetriever"): if retriever.embedding_model not in BEDROCK_EMBEDDING_MODELS: raise ValueError("Model not supported by Bedrock Embedding Encoder") self.model = retriever.embedding_model - self.client = self.initialize_boto3_session(retriever.aws_config).client("bedrock-runtime") + self.client = self._initialize_boto3_session(retriever.aws_config).client("bedrock-runtime") - def initialize_boto3_session(self, aws_config: Optional[Dict[str, Any]]): + def _initialize_boto3_session(self, aws_config: Optional[Dict[str, Any]]): if aws_config is None: raise ValueError( "`aws_config` is not set. To use Bedrock models, you should set `aws_config` when initializing the retriever." @@ -472,8 +473,10 @@ def initialize_boto3_session(self, aws_config: Optional[Dict[str, Any]]): region_name=region_name, profile_name=profile_name, ) - except Exception as e: - raise ValueError(f"AWS client error {e}") + except BotoCoreError as e: + raise AWSConfigurationError( + f"Failed to initialize the session with provided AWS credentials {aws_config}" + ) from e def _embed_batch_cohere( self, texts: List[str], input_type: Literal["search_query", "search_document"]