This repository contains the MTEB scores and usage examples of Bedrock Titan Text Embeddings v2. You can use the embedding model either via the Bedrock InvokeModel API or via Bedrock's batch jobs. For RAG use cases we recommend the former to embed queries during search (latency optimized) and the latter to index corpus (throughput optimized).
1import json
2import boto3
3class TitanEmbeddings(object):
4 accept = "application/json"
5 content_type = "application/json"
6
7 def __init__(self, model_id="amazon.titan-embed-text-v2:0"):
8 self.bedrock = boto3.client(service_name='bedrock-runtime')
9 self.model_id = model_id
10 def __call__(self, text, dimensions, normalize=True):
11 """
12 Returns Titan Embeddings
13 Args:
14 text (str): text to embed
15 dimensions (int): Number of output dimensions.
16 normalize (bool): Whether to return the normalized embedding or not.
17 Return:
18 List[float]: Embedding
19
20 """
21 body = json.dumps({
22 "inputText": text,
23 "dimensions": dimensions,
24 "normalize": normalize
25 })
26 response = self.bedrock.invoke_model(
27 body=body, modelId=self.model_id, accept=self.accept, contentType=self.content_type
28 )
29 response_body = json.loads(response.get('body').read())
30 return response_body['embedding']
31
32if __name__ == '__main__':
33 """
34 Entrypoint for Amazon Titan Embeddings V2 - Text example.
35 """
36 dimensions = 1024
37 normalize = True
38
39 titan_embeddings_v2 = TitanEmbeddings(model_id="amazon.titan-embed-text-v2:0")
40
41 input_text = "What are the different services that you offer?"
42 embedding = titan_embeddings_v2(input_text, dimensions, normalize)
43
44 print(f"{input_text=}")
45 print(f"{embedding[:10]=}")
46
1import requests
2from aws_requests_auth.boto_utils import BotoAWSRequestsAuth
3
4region = "us-east-1"
5base_uri = f"bedrock.{region}.amazonaws.com"
6batch_job_uri = f"https://{base_uri}/model-invocation-job/"
7
8# For details on how to set up an IAM role for batch inference, see
9# https://docs.aws.amazon.com/bedrock/latest/userguide/batch-inference-permissions.html
10role_arn = "arn:aws:iam::111122223333:role/my-batch-inference-role"
11
12payload = {
13 "inputDataConfig": {
14 "s3InputDataConfig": {
15 "s3Uri": "s3://my-input-bucket/batch-input/",
16 "s3InputFormat": "JSONL"
17 }
18 },
19 "jobName": "embeddings-v2-batch-job",
20 "modelId": "amazon.titan-embed-text-v2:0",
21 "outputDataConfig": {
22 "s3OutputDataConfig": {
23 "s3Uri": "s3://my-output-bucket/batch-output/"
24 }
25 },
26 "roleArn": role_arn
27}
28
29request_auth = BotoAWSRequestsAuth(
30 aws_host=base_uri,
31 aws_region=region,
32 aws_service="bedrock"
33)
34
35
36response= requests.request("POST", batch_job_uri, json=payload, auth=request_auth)
37print(response.json())
