如何在AWS Bedrock中批量处理文本生成嵌入？

问题说明 #

现有AWS Bedrock Runtime的单次调用生成嵌入示例，是否支持直接传入CSV或DataFrame这类批量文本，实现批量生成嵌入？

回答 #

AWS Bedrock的嵌入模型（如Amazon Titan Text Embeddings V2）不直接支持CSV或DataFrame作为输入格式，但可以通过以下方式实现批量生成嵌入：

• 先从CSV或DataFrame中提取待处理的文本列表
• 使用Bedrock Runtime的invoke_model循环处理单条文本，或使用invoke_model_batch接口批量提交请求（后者效率更高）
• 注意控制单条请求的文本token数不超过模型限制，同时做好批量请求的错误重试与结果聚合

单次调用示例（中文注释） #

import boto3
import json

# 创建指定AWS区域的Bedrock Runtime客户端
client = boto3.client("bedrock-runtime", region_name="us-east-1")

# 指定模型ID，例如亚马逊Titan文本嵌入V2模型
model_id = "amazon.titan-embed-text-v2:0"

# 待转换为嵌入的文本
input_text = "请推荐与电影《盗梦空间》主题类似的书籍。"

# 构造模型请求参数
native_request = {"inputText": input_text}

# 将请求参数转为JSON格式
request = json.dumps(native_request)

# 调用模型
response = client.invoke_model(modelId=model_id, body=request)

# 解码模型返回的响应内容
model_response = json.loads(response["body"].read())

# 提取生成的嵌入向量和输入文本的token数量
embedding = model_response["embedding"]
input_token_count = model_response["inputTextTokenCount"]

print("\n你的输入文本:")
print(input_text)
print(f"输入token数量: {input_token_count}")
print(f"生成的嵌入向量长度: {len(embedding)}")
print("嵌入向量:")
print(embedding)

批量处理示例（基于DataFrame） #

import boto3
import json
import pandas as pd

# 创建Bedrock Runtime客户端
client = boto3.client("bedrock-runtime", region_name="us-east-1")
model_id = "amazon.titan-embed-text-v2:0"

# 读取CSV到DataFrame（假设CSV有一列名为text）
df = pd.read_csv("your_texts.csv")
text_list = df["text"].tolist()

# 批量生成嵌入的函数
def batch_generate_embeddings(texts):
    embeddings = []
    token_counts = []
    for text in texts:
        try:
            # 构造请求
            request = json.dumps({"inputText": text})
            response = client.invoke_model(modelId=model_id, body=request)
            model_response = json.loads(response["body"].read())
            embeddings.append(model_response["embedding"])
            token_counts.append(model_response["inputTextTokenCount"])
        except Exception as e:
            print(f"处理文本失败: {text}, 错误: {str(e)}")
            embeddings.append(None)
            token_counts.append(None)
    return embeddings, token_counts

# 执行批量处理
df["embedding"], df["token_count"] = batch_generate_embeddings(text_list)

# 保存结果到新CSV
df.to_csv("texts_with_embeddings.csv", index=False)
print("批量处理完成，结果已保存")

注意事项 #

• 使用invoke_model_batch接口可以进一步提升批量处理效率，该接口允许单次提交最多25条请求
• 不同模型的最大token限制不同，例如Amazon Titan Text Embeddings V2的最大输入token数为8192，需确保单条文本不超过此限制
• 批量处理时建议添加重试机制，避免因网络或服务临时异常导致任务中断

内容的提问来源于stack exchange，提问作者kishi