当前位置：首页 > news >正文

【RAG】【embeddings42】Amazon SageMaker 嵌入端点集成案例

news 2026/8/12 3:41:00

案例目标

本案例展示如何通过 LlamaIndex 与 Amazon SageMaker 嵌入端点进行交互，实现文本向量化功能。Amazon SageMaker 端点是一种完全托管的资源，支持部署机器学习模型，对新数据进行预测。

通过本案例，您将学习如何：

配置和连接到 Amazon SageMaker 嵌入端点
使用 SageMakerEmbedding 类生成文本嵌入向量
处理单个和批量文本的向量化请求
验证嵌入向量的维度和数值

技术栈与核心依赖

llama-index-embeddings-sagemaker-endpoint Amazon SageMaker AWS SDK (boto3) Python 3.7+

核心组件说明

SageMakerEmbedding- LlamaIndex 中用于与 Amazon SageMaker 嵌入端点交互的类
SageMaker Endpoint- Amazon 提供的完全托管的机器学习模型部署服务
AWS Credentials- 用于访问 AWS 服务的身份验证信息

环境配置

1. 安装依赖

pip install llama-index-embeddings-sagemaker-endpoint
pip install llama-index

2. 准备 SageMaker 嵌入端点

在使用本案例之前，您需要：

在 AWS 控制台中创建一个 SageMaker 嵌入端点
获取端点名称和区域信息
配置适当的 AWS 访问凭证

3. 配置 AWS 凭证

您可以通过以下两种方式配置 AWS 凭证：

方式一：使用 AWS 配置文件

from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding

ENDPOINT_NAME = "<-YOUR-ENDPOINT-NAME->"
PROFILE_NAME = "<-YOUR-PROFILE-NAME->"
embed_model = SageMakerEmbedding(
endpoint_name=ENDPOINT_NAME, profile_name=PROFILE_NAME
)

方式二：直接传递凭证参数

from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding

ENDPOINT_NAME = "<-YOUR-ENDPOINT-NAME->"
AWS_ACCESS_KEY_ID = "<-YOUR-AWS-ACCESS-KEY-ID->"
AWS_SECRET_ACCESS_KEY = "<-YOUR-AWS-SECRET-ACCESS-KEY->"
AWS_SESSION_TOKEN = "<-YOUR-AWS-SESSION-TOKEN->"
REGION_NAME = "<-YOUR-ENDPOINT-REGION-NAME->"

embed_model = SageMakerEmbedding(
endpoint_name=ENDPOINT_NAME,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_session_token=AWS_SESSION_TOKEN,
region_name=REGION_NAME,
)

案例实现

1. 导入必要模块

from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding

2. 初始化 SageMaker 嵌入模型

# 使用凭证参数初始化
embed_model = SageMakerEmbedding(
endpoint_name=ENDPOINT_NAME,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_session_token=AWS_SESSION_TOKEN,
region_name=REGION_NAME,
)

注意：请将 ENDPOINT_NAME、AWS_ACCESS_KEY_ID 等占位符替换为您的实际值。

3. 生成单个文本的嵌入

embeddings = embed_model.get_text_embedding(
"An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models, specifically LLM (Large Language Models), for making predictions on new data."
)

4. 检查嵌入结果

# 查看嵌入向量的前5个值
embeddings[:5]
# 输出示例: [0.021565623581409454, 0.019147753715515137, -0.056700449436903, -0.022435730323195457, -0.03930153697729111]

5. 生成批量文本的嵌入

embeddings = embed_model.get_text_embedding_batch(
[
"An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models",
"Sagemaker is integrated with llamaIndex",
]
)

# 检查结果数量
len(embeddings)
# 输出: 2