From 26649f7fa69ad9c97dd62fc81ec480a76c1363e9 Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 11:05:30 +0100 Subject: [PATCH 1/4] refactor: remove reimplementing exceptions --- .../generators/amazon_sagemaker/errors.py | 32 +------------------ 1 file changed, 1 insertion(+), 31 deletions(-) diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py index 6c13d0fcb..09f22bf52 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py @@ -1,46 +1,16 @@ -from typing import Optional - class SagemakerError(Exception): """ - Error generated by the Amazon Sagemaker integration. + Parent class for all exceptions raised by the Sagemaker component """ - def __init__( - self, - message: Optional[str] = None, - ): - super().__init__() - if message: - self.message = message - - def __getattr__(self, attr): - # If self.__cause__ is None, it will raise the expected AttributeError - getattr(self.__cause__, attr) - - def __str__(self): - return self.message - - def __repr__(self): - return str(self) - - class AWSConfigurationError(SagemakerError): """Exception raised when AWS is not configured correctly""" - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) - class SagemakerNotReadyError(SagemakerError): """Exception for issues that occur during Sagemaker inference""" - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) - class SagemakerInferenceError(SagemakerError): """Exception for issues that occur during Sagemaker inference""" - - def __init__(self, message: Optional[str] = None): - super().__init__(message=message) From 6b8118bc9729fe97703cb58f2669aaf12e92588a Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 11:30:08 +0100 Subject: [PATCH 2/4] docs: review docs --- .../generators/amazon_sagemaker/sagemaker.py | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py index c171ccdf6..b64902f00 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py @@ -31,22 +31,19 @@ class SagemakerGenerator: [SageMaker JumpStart foundation models documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-foundation-models-use.html). Usage example: - - Make sure your AWS credentials are set up correctly. You can use environment variables or a shared credentials file. - Then you can use the generator as follows: ```python + # Make sure your AWS credentials are set up correctly. You can use environment variables or a shared credentials file. + # Then you can use the generator as follows: from haystack_integrations.components.generators.amazon_sagemaker import SagemakerGenerator + generator = SagemakerGenerator(model="jumpstart-dft-hf-llm-falcon-7b-bf16") response = generator.run("What's Natural Language Processing? Be brief.") print(response) - ``` - ``` - >> {'replies': ['Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on - >> the interaction between computers and human language. It involves enabling computers to understand, interpret, - >> and respond to natural human language in a way that is both meaningful and useful.'], 'meta': [{}]} + >>> {'replies': ['Natural Language Processing (NLP) is a branch of artificial intelligence that focuses on + >>> the interaction between computers and human language. It involves enabling computers to understand, interpret, + >>> and respond to natural human language in a way that is both meaningful and useful.'], 'meta': [{}]} ``` """ - model_generation_keys: ClassVar = ["generated_text", "generation"] def __init__( @@ -73,7 +70,6 @@ def __init__( :param model: The name for SageMaker Model Endpoint. :param aws_custom_attributes: Custom attributes to be passed to SageMaker, for example `{"accept_eula": True}` in case of Llama-2 models. - :param generation_kwargs: Additional keyword arguments for text generation. For a list of supported parameters see your model's documentation page, for example here for HuggingFace models: https://huggingface.co/blog/sagemaker-huggingface-llm#4-run-inference-and-chat-with-our-model @@ -121,15 +117,15 @@ def resolve_secret(secret: Optional[Secret]) -> Optional[str]: def _get_telemetry_data(self) -> Dict[str, Any]: """ Returns data that is sent to Posthog for usage analytics. - :returns: a dictionary with following keys: - - model: The name of the model. - + :returns: A dictionary with the following keys: + - `model`: The name of the model. """ return {"model": self.model} def to_dict(self) -> Dict[str, Any]: """ Serializes the component to a dictionary. + :returns: Dictionary with serialized data. """ @@ -149,10 +145,11 @@ def to_dict(self) -> Dict[str, Any]: def from_dict(cls, data) -> "SagemakerGenerator": """ Deserializes the component from a dictionary. + :param data: Dictionary to deserialize from. :returns: - Deserialized component. + Deserialized component. """ deserialize_secrets_inplace( data["init_parameters"], @@ -170,6 +167,7 @@ def _get_aws_session( ): """ Creates an AWS Session with the given parameters. + Checks if the provided AWS credentials are valid and can be used to connect to AWS. :param aws_access_key_id: AWS access key ID. @@ -200,8 +198,10 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): :param prompt: The string prompt to use for text generation. :param generation_kwargs: Additional keyword arguments for text generation. These parameters will - potentially override the parameters passed in the `__init__` method. - + potentially override the parameters passed in the `__init__` method. + :raises ValueError: If the model response type is not a list of dictionaries or a single dictionary. + :raises SagemakerNotReadyError: If the SageMaker model is not ready to accept requests. + :raises SagemakerInferenceError: If the SageMaker Inference returns an error. :returns: A dictionary with the following keys: - `replies`: A list of strings containing the generated responses - `meta`: A list of dictionaries containing the metadata for each response. @@ -249,5 +249,5 @@ def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None): msg = f"Sagemaker model not ready: {res.text}" raise SagemakerNotReadyError(msg) from err - msg = f"SageMaker Inference returned an error. Status code: {res.status_code} Response body: {res.text}" - raise SagemakerInferenceError(msg, status_code=res.status_code) from err + msg = f"SageMaker Inference returned an error. Status code: {res.status_code}. Response body: {res.text}" + raise SagemakerInferenceError(msg) from err From 8ff9e8b81bb86cd4f96f8d67bb5916208b2252d2 Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 11:31:45 +0100 Subject: [PATCH 3/4] style: reformat --- .../components/generators/amazon_sagemaker/errors.py | 2 +- .../components/generators/amazon_sagemaker/sagemaker.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py index 09f22bf52..e518c3c6d 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py @@ -1,9 +1,9 @@ - class SagemakerError(Exception): """ Parent class for all exceptions raised by the Sagemaker component """ + class AWSConfigurationError(SagemakerError): """Exception raised when AWS is not configured correctly""" diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py index b64902f00..3cf71abe4 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py @@ -44,6 +44,7 @@ class SagemakerGenerator: >>> and respond to natural human language in a way that is both meaningful and useful.'], 'meta': [{}]} ``` """ + model_generation_keys: ClassVar = ["generated_text", "generation"] def __init__( From 13ce3c7536f2e107567ec241bbe0ce95b439bed5 Mon Sep 17 00:00:00 2001 From: Tobias Wochinger Date: Wed, 6 Mar 2024 12:02:02 +0100 Subject: [PATCH 4/4] style: shorten line --- .../components/generators/amazon_sagemaker/sagemaker.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py index 3cf71abe4..2a04d6a2a 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py @@ -32,8 +32,8 @@ class SagemakerGenerator: Usage example: ```python - # Make sure your AWS credentials are set up correctly. You can use environment variables or a shared credentials file. - # Then you can use the generator as follows: + # Make sure your AWS credentials are set up correctly. You can use environment variables or a shared credentials + # file. Then you can use the generator as follows: from haystack_integrations.components.generators.amazon_sagemaker import SagemakerGenerator generator = SagemakerGenerator(model="jumpstart-dft-hf-llm-falcon-7b-bf16")