Bedrock embeddings async methods (#9024)

## Description
This PR adds the `aembed_query` and `aembed_documents` async methods for
improving the embeddings generation for large documents. The
implementation uses asyncio tasks and gather to achieve concurrency as
there is no bedrock async API in boto3.

### Maintainers
@agola11 
@aarora79  

### Open questions
To avoid throttling from the Bedrock API, should there be an option to
limit the concurrency of the calls?
This commit is contained in:
Piyush Jain
2023-08-10 14:21:03 -07:00
committed by GitHub
parent 67ca187560
commit 8eea46ed0e
3 changed files with 60 additions and 15 deletions

View File

@@ -31,12 +31,11 @@
},
"outputs": [],
"source": [
"from langchain.llms.bedrock import Bedrock\n",
"from langchain.llms import Bedrock\n",
"\n",
"llm = Bedrock(\n",
" credentials_profile_name=\"bedrock-admin\",\n",
" model_id=\"amazon.titan-tg1-large\",\n",
" endpoint_url=\"custom_endpoint_url\",\n",
" model_id=\"amazon.titan-tg1-large\"\n",
")"
]
},

View File

@@ -20,7 +20,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 5,
"id": "282239c8-e03a-4abc-86c1-ca6120231a20",
"metadata": {},
"outputs": [],
@@ -28,7 +28,7 @@
"from langchain.embeddings import BedrockEmbeddings\n",
"\n",
"embeddings = BedrockEmbeddings(\n",
" credentials_profile_name=\"bedrock-admin\", endpoint_url=\"custom_endpoint_url\"\n",
" credentials_profile_name=\"bedrock-admin\", region_name=\"us-east-1\"\n",
")"
]
},
@@ -49,7 +49,29 @@
"metadata": {},
"outputs": [],
"source": [
"embeddings.embed_documents([\"This is a content of the document\"])"
"embeddings.embed_documents([\"This is a content of the document\", \"This is another document\"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9f6b364d",
"metadata": {},
"outputs": [],
"source": [
"# async embed query\n",
"await embeddings.aembed_query(\"This is a content of the document\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c9240a5a",
"metadata": {},
"outputs": [],
"source": [
"# async embed documents\n",
"await embeddings.aembed_documents([\"This is a content of the document\", \"This is another document\"])"
]
}
],
@@ -69,7 +91,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.9.13"
}
},
"nbformat": 4,