Added support for models on AWS Bedrock and Sagemaker #772

drawal1 · 2024-04-04T19:37:32Z

This check-in replaces the old Bedrock module which was incomplete and obsolete. The new code refactors the module into aws_providers and aws_models and allows extending to new models.
Supported aws providers - Bedrock and Sagemaker
Supported models - Mistral, Anthropic, Llama (extensible design to add custom models)

"""Test harness: Testing the AWS modules"""
import dspy
from dspy import import AWSAnthropic, AWSMeta, AWSMistral, Bedrock, Sagemaker

class QASignature(dspy.Signature):
    """answer the question"""

    question = dspy.InputField()
    answer = dspy.OutputField()


def test_aws_models(lm):
    """Test the models on the given AWS provider"""

    predict_func = dspy.ChainOfThought(QASignature)
    with dspy.context(lm=lm):
        answer = predict_func(question="What is the capital of France")
        assert "paris" in str(answer).lower()
        print(answer)


if __name__ == "__main__":
    # NOTE: Configure your AWS credentials with the AWS CLI before running this test!

    bedrock = Bedrock(region_name="us-west-2")
    test_aws_models(AWSMistral(bedrock, model="mistral.mixtral-8x7b-instruct-v0:1"))
    test_aws_models(AWSMistral(bedrock, "mistral.mistral-7b-instruct-v0:2"))
    test_aws_models(AWSAnthropic(bedrock, "anthropic.claude-3-haiku-20240307-v1:0"))
    test_aws_models(AWSAnthropic(bedrock, "anthropic.claude-3-sonnet-20240229-v1:0"))
    # this is slower than molasses and generates irrelevant content after the answer!!
    # You may have to wait for 10-15min before it returns
    test_aws_models(AWSMeta(bedrock, "meta.llama2-70b-chat-v1"))

    # NOTE: Configure your Sagemaker endpoints before running this test!
    sagemaker = Sagemaker(region_name="us-west-2")
    # NOTE: Replace model value below with your own endpoint name
    test_aws_models(AWSMistral(sagemaker, model="g5-48xlarge-mixtral-8x7b"))

xiaochuan-du

do you think the lm.py need to be updated as well? otherwise, the llm_model.inspect_history(n=1) may throw errors

drawal1 · 2024-04-05T14:39:25Z

Good catch!

arnavsinghvi11 · 2024-04-08T17:53:23Z

dsp/modules/__init__.py

-from .cache_utils import *
-from .clarifai import *
-from .cohere import *
+from .cache_utils import *  # noqa: F403


@drawal1 can you remove all extra comments here?

arnavsinghvi11 · 2024-04-08T17:54:31Z

dsp/modules/aws_models.py

+from typing import Any
+
+from dsp.modules.lm import LM
+from shared.src.models.utils.aws_providers import AWSProvider, Bedrock, Sagemaker


where is this import coming from? can you wrap any external imports in a try/except block as done here?

arnavsinghvi11 · 2024-04-08T17:56:43Z

dsp/modules/aws_models.py

+        right now.
+        """
+        if not only_completed:
+            raise NotImplementedError("Error, only_completed not yet supported!")


minor styling but can you follow this here?

arnavsinghvi11 · 2024-04-08T17:57:34Z

dsp/modules/aws_models.py

+
+
+class AWSModel(LM):
+    """This class adds support for an AWS model."""


can you add more clear documentation on how this class works with the other classes?

arnavsinghvi11 · 2024-04-08T17:59:58Z

dsp/modules/aws_providers.py

+            profile_name (str, optional): boto3 credentials profile.
+            batch_n_enabled (bool): If False, call the LM N times rather than batching.
+        """
+        import boto3  # pylint: disable=import-outside-toplevel


same comment as above - wrap external imports in a try/except block as done here?

arnavsinghvi11 · 2024-04-08T18:01:03Z

dspy/__init__.py

-from .primitives import *
-from .retrieve import *
-from .signatures import *
+from .predict import *  # noqa: F403


same comment as above - remove any extra comments

arnavsinghvi11 · 2024-04-08T18:01:23Z

dspy/__init__.py

 Google = dsp.Google

 HFClientTGI = dsp.HFClientTGI
-HFClientVLLM = HFClientVLLM


why is this deleted?

arnavsinghvi11 · 2024-04-08T18:01:41Z

pyproject.toml

 pgvector = {version = "^0.2.5", optional = true}
 structlog = "^24.1.0"

-


remove change

arnavsinghvi11 · 2024-04-08T18:02:06Z

dsp/modules/__init__.py

+from .aws_models import AWSAnthropic, AWSLlama2, AWSMistral, AWSModel
+from .aws_providers import Bedrock, Sagemaker
 from .azure_openai import AzureOpenAI
-from .bedrock import *


why is importing bedrock removed from here?

arnavsinghvi11 · 2024-04-08T18:05:47Z

Hi @drawal1, left some comments on the PR. It's also failing some of the CI tests. Please run ruff check . --fix-only and address the other errors found in the build tests.

It would be great if you could add both documentation and integration tests for the AWS model support as well. Feel free to follow these LM docs for reference.

This was referenced Apr 4, 2024

Added support for models on AWS Bedrock and Sagemaker #751

Closed

Fixing claude v2 in AWS Bedrock #745

Closed

Add AWS Bedrock Claude v3 support #746

Closed

xiaochuan-du reviewed Apr 5, 2024

View reviewed changes

This was referenced Apr 5, 2024

Bedrock LM produces extraneous key error in unclear circumstances #365

Closed

Refactor Bedrock class to support Claude-3 model and improve error handling #720

Closed

arnavsinghvi11 requested changes Apr 8, 2024

View reviewed changes

drawal1 closed this Apr 9, 2024

drawal1 force-pushed the main branch from a56aea1 to 696f2d2 Compare April 9, 2024 16:23

drawal1 mentioned this pull request Apr 9, 2024

Added support for models on AWS Bedrock and Sagemaker #795

Merged



		class AWSModel(LM):
		"""This class adds support for an AWS model."""

		pgvector = {version = "^0.2.5", optional = true}
		structlog = "^24.1.0"

Added support for models on AWS Bedrock and Sagemaker #772

Added support for models on AWS Bedrock and Sagemaker #772

Uh oh!

Conversation

drawal1 commented Apr 4, 2024

Uh oh!

xiaochuan-du left a comment

Choose a reason for hiding this comment

Uh oh!

drawal1 commented Apr 5, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

arnavsinghvi11 commented Apr 8, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants