forked from langchain-ai/langchain
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Bedrock llm and embeddings (langchain-ai#5464)
# Bedrock LLM and Embeddings This PR adds a new LLM and an Embeddings class for the [Bedrock](https://aws.amazon.com/bedrock) service. The PR also includes example notebooks for using the LLM class in a conversation chain and embeddings usage in creating an embedding for a query and document. **Note**: AWS is doing a private release of the Bedrock service on 05/31/2023; users need to request access and added to an allowlist in order to start using the Bedrock models and embeddings. Please use the [Bedrock Home Page](https://aws.amazon.com/bedrock) to request access and to learn more about the models available in Bedrock. <!-- For a quicker response, figure out the right person to tag with @ @hwchase17 - project lead Tracing / Callbacks - @agola11 Async - @agola11 DataLoaders - @eyurtsev Models - @hwchase17 - @agola11 Agents / Tools / Toolkits - @vowelparrot VectorStores / Retrievers / Memory - @dev2049 -->
- Loading branch information
1 parent
e188bac
commit 92dc497
Showing
7 changed files
with
533 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,24 @@ | ||
# Amazon Bedrock | ||
|
||
>[Amazon Bedrock](https://aws.amazon.com/bedrock/) is a fully managed service that makes FMs from leading AI startups and Amazon available via an API, so you can choose from a wide range of FMs to find the model that is best suited for your use case. | ||
## Installation and Setup | ||
|
||
```bash | ||
pip install boto3 | ||
``` | ||
|
||
## LLM | ||
|
||
See a [usage example](../modules/models/llms/integrations/bedrock.ipynb). | ||
|
||
```python | ||
from langchain import Bedrock | ||
``` | ||
|
||
## Text Embedding Models | ||
|
||
See a [usage example](../modules/models/text_embedding/examples/bedrock.ipynb). | ||
```python | ||
from langchain.embeddings import BedrockEmbeddings | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,86 @@ | ||
{ | ||
"cells": [ | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"# Amazon Bedrock" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"[Amazon Bedrock](https://aws.amazon.com/bedrock/) is a fully managed service that makes FMs from leading AI startups and Amazon available via an API, so you can choose from a wide range of FMs to find the model that is best suited for your use case" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"%pip install boto3" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": { | ||
"tags": [] | ||
}, | ||
"outputs": [], | ||
"source": [ | ||
"from langchain.llms.bedrock import Bedrock\n", | ||
"\n", | ||
"llm = Bedrock(credentials_profile_name=\"bedrock-admin\", model_id=\"amazon.titan-tg1-large\")" | ||
] | ||
}, | ||
{ | ||
"cell_type": "markdown", | ||
"metadata": {}, | ||
"source": [ | ||
"### Using in a conversation chain" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"from langchain.chains import ConversationChain\n", | ||
"from langchain.memory import ConversationBufferMemory\n", | ||
"\n", | ||
"conversation = ConversationChain(\n", | ||
" llm=llm,\n", | ||
" verbose=True,\n", | ||
" memory=ConversationBufferMemory()\n", | ||
")\n", | ||
"\n", | ||
"conversation.predict(input=\"Hi there!\")" | ||
] | ||
} | ||
], | ||
"metadata": { | ||
"kernelspec": { | ||
"display_name": "Python 3 (ipykernel)", | ||
"language": "python", | ||
"name": "python3" | ||
}, | ||
"language_info": { | ||
"codemirror_mode": { | ||
"name": "ipython", | ||
"version": 3 | ||
}, | ||
"file_extension": ".py", | ||
"mimetype": "text/x-python", | ||
"name": "python", | ||
"nbconvert_exporter": "python", | ||
"pygments_lexer": "ipython3", | ||
"version": "3.10.11" | ||
} | ||
}, | ||
"nbformat": 4, | ||
"nbformat_minor": 4 | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,75 @@ | ||
{ | ||
"cells": [ | ||
{ | ||
"cell_type": "markdown", | ||
"id": "75e378f5-55d7-44b6-8e2e-6d7b8b171ec4", | ||
"metadata": {}, | ||
"source": [ | ||
"# Bedrock Embeddings" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"id": "2dbe40fa-7c0b-4bcb-a712-230bf613a42f", | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"%pip install boto3" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"id": "282239c8-e03a-4abc-86c1-ca6120231a20", | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"from langchain.embeddings import BedrockEmbeddings\n", | ||
"\n", | ||
"embeddings = BedrockEmbeddings(credentials_profile_name=\"bedrock-admin\")" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"id": "19a46868-4bed-40cd-89ca-9813fbfda9cb", | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"embeddings.embed_query(\"This is a content of the document\")" | ||
] | ||
}, | ||
{ | ||
"cell_type": "code", | ||
"execution_count": null, | ||
"id": "cf0349c4-6408-4342-8691-69276a388784", | ||
"metadata": {}, | ||
"outputs": [], | ||
"source": [ | ||
"embeddings.embed_documents([\"This is a content of the document\"])" | ||
] | ||
} | ||
], | ||
"metadata": { | ||
"kernelspec": { | ||
"display_name": "Python 3 (ipykernel)", | ||
"language": "python", | ||
"name": "python3" | ||
}, | ||
"language_info": { | ||
"codemirror_mode": { | ||
"name": "ipython", | ||
"version": 3 | ||
}, | ||
"file_extension": ".py", | ||
"mimetype": "text/x-python", | ||
"name": "python", | ||
"nbconvert_exporter": "python", | ||
"pygments_lexer": "ipython3", | ||
"version": "3.10.11" | ||
} | ||
}, | ||
"nbformat": 4, | ||
"nbformat_minor": 5 | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,157 @@ | ||
import json | ||
import os | ||
from typing import Any, Dict, List, Optional | ||
|
||
from pydantic import BaseModel, Extra, root_validator | ||
|
||
from langchain.embeddings.base import Embeddings | ||
|
||
|
||
class BedrockEmbeddings(BaseModel, Embeddings): | ||
"""Embeddings provider to invoke Bedrock embedding models. | ||
To authenticate, the AWS client uses the following methods to | ||
automatically load credentials: | ||
https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html | ||
If a specific credential profile should be used, you must pass | ||
the name of the profile from the ~/.aws/credentials file that is to be used. | ||
Make sure the credentials / roles used have the required policies to | ||
access the Bedrock service. | ||
""" | ||
|
||
""" | ||
Example: | ||
.. code-block:: python | ||
from langchain.bedrock_embeddings import BedrockEmbeddings | ||
region_name ="us-east-1" | ||
credentials_profile_name = "default" | ||
model_id = "amazon.titan-e1t-medium" | ||
be = BedrockEmbeddings( | ||
credentials_profile_name=credentials_profile_name, | ||
region_name=region_name, | ||
model_id=model_id | ||
) | ||
""" | ||
|
||
client: Any #: :meta private: | ||
|
||
region_name: Optional[str] = None | ||
"""The aws region e.g., `us-west-2`. Fallsback to AWS_DEFAULT_REGION env variable | ||
or region specified in ~/.aws/config in case it is not provided here. | ||
""" | ||
|
||
credentials_profile_name: Optional[str] = None | ||
"""The name of the profile in the ~/.aws/credentials or ~/.aws/config files, which | ||
has either access keys or role information specified. | ||
If not specified, the default credential profile or, if on an EC2 instance, | ||
credentials from IMDS will be used. | ||
See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html | ||
""" | ||
|
||
model_id: str = "amazon.titan-e1t-medium" | ||
"""Id of the model to call, e.g., amazon.titan-e1t-medium, this is | ||
equivalent to the modelId property in the list-foundation-models api""" | ||
|
||
model_kwargs: Optional[Dict] = None | ||
"""Key word arguments to pass to the model.""" | ||
|
||
class Config: | ||
"""Configuration for this pydantic object.""" | ||
|
||
extra = Extra.forbid | ||
|
||
@root_validator() | ||
def validate_environment(cls, values: Dict) -> Dict: | ||
"""Validate that AWS credentials to and python package exists in environment.""" | ||
try: | ||
import boto3 | ||
|
||
if values["credentials_profile_name"] is not None: | ||
session = boto3.Session(profile_name=values["credentials_profile_name"]) | ||
else: | ||
# use default credentials | ||
session = boto3.Session() | ||
|
||
client_params = {} | ||
if values["region_name"]: | ||
client_params["region_name"] = values["region_name"] | ||
|
||
values["client"] = session.client("bedrock", **client_params) | ||
|
||
except ImportError: | ||
raise ModuleNotFoundError( | ||
"Could not import boto3 python package. " | ||
"Please install it with `pip install boto3`." | ||
) | ||
except Exception as e: | ||
raise ValueError( | ||
"Could not load credentials to authenticate with AWS client. " | ||
"Please check that credentials in the specified " | ||
"profile name are valid." | ||
) from e | ||
|
||
return values | ||
|
||
def _embedding_func(self, text: str) -> List[float]: | ||
"""Call out to Bedrock embedding endpoint.""" | ||
# replace newlines, which can negatively affect performance. | ||
text = text.replace(os.linesep, " ") | ||
_model_kwargs = self.model_kwargs or {} | ||
|
||
input_body = {**_model_kwargs} | ||
input_body["inputText"] = text | ||
body = json.dumps(input_body) | ||
content_type = "application/json" | ||
accepts = "application/json" | ||
|
||
embeddings = [] | ||
try: | ||
response = self.client.invoke_model( | ||
body=body, | ||
modelId=self.model_id, | ||
accept=accepts, | ||
contentType=content_type, | ||
) | ||
response_body = json.loads(response.get("body").read()) | ||
embeddings = response_body.get("embedding") | ||
except Exception as e: | ||
raise ValueError(f"Error raised by inference endpoint: {e}") | ||
|
||
return embeddings | ||
|
||
def embed_documents( | ||
self, texts: List[str], chunk_size: int = 1 | ||
) -> List[List[float]]: | ||
"""Compute doc embeddings using a Bedrock model. | ||
Args: | ||
texts: The list of texts to embed. | ||
chunk_size: Bedrock currently only allows single string | ||
inputs, so chunk size is always 1. This input is here | ||
only for compatibility with the embeddings interface. | ||
Returns: | ||
List of embeddings, one for each text. | ||
""" | ||
results = [] | ||
for text in texts: | ||
response = self._embedding_func(text) | ||
results.append(response) | ||
return results | ||
|
||
def embed_query(self, text: str) -> List[float]: | ||
"""Compute query embeddings using a Bedrock model. | ||
Args: | ||
text: The text to embed. | ||
Returns: | ||
Embeddings for the text. | ||
""" | ||
return self._embedding_func(text) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.