如何使用LlamaIndex在Amazon SageMaker部署预训练LLM？

1. 安装依赖 #

先安装必要的Python包：

pip install llama-index boto3

2. 定义内容处理器 #

你需要实现适配SageMaker端点的输入序列化和输出反序列化逻辑，示例如下（适配Hugging Face文本生成类模型的常见格式，需根据你的端点实际格式调整）：

from llama_index.llms.sagemaker import SageMakerEndpoint, SageMakerLLMPredictor
from typing import Optional, Dict
import json

class CustomContentHandler:
    def serialize(self, prompt: str, model_kwargs: Optional[Dict] = None) -> bytes:
        # 构造符合SageMaker端点要求的请求体
        payload = {
            "inputs": prompt,
            "parameters": model_kwargs or {"max_new_tokens": 256, "temperature": 0.7}
        }
        return json.dumps(payload).encode("utf-8")
    
    def deserialize(self, response: bytes) -> str:
        # 解析SageMaker返回的响应内容
        response_json = json.loads(response.decode("utf-8"))
        # 根据端点返回格式调整，示例为Hugging Face模型的返回结构
        return response_json[0]["generated_text"]

3. 初始化LlamaIndex的SageMaker LLM实例 #

import boto3

# 初始化SageMaker Runtime客户端，指定区域
sagemaker_client = boto3.client("sagemaker-runtime", region_name="me-central-1")

# 实例化内容处理器
content_handler = CustomContentHandler()

# 创建SageMaker LLM对象
llm = SageMakerEndpoint(
    endpoint_name="你的SageMaker端点名称",  # 替换为你的实际端点名
    predictor=SageMakerLLMPredictor(
        sagemaker_client=sagemaker_client,
        content_handler=content_handler
    ),
    model_kwargs={"max_new_tokens": 256, "temperature": 0.7}  # 默认模型参数
)

4. 基础推理与进阶使用 #

直接生成文本 #

response = llm.complete("请说明大语言模型在教育领域的应用")
print(response.text)

结合知识库构建问答系统（可选） #

如果需要基于自有文档构建问答，可按以下步骤操作：

from llama_index import VectorStoreIndex, SimpleDirectoryReader

# 加载本地文档（示例：加载data目录下的所有文件）
documents = SimpleDirectoryReader("data").load_data()
# 构建向量索引，关联我们的SageMaker LLM
index = VectorStoreIndex.from_documents(documents, llm=llm)
# 创建查询引擎
query_engine = index.as_query_engine()
# 发起查询
response = query_engine.query("文档中提到的大语言模型教育应用有哪些？")
print(response.response)

重要提示 #

• 格式适配：务必根据你的SageMaker端点实际的请求/响应格式修改CustomContentHandler中的序列化和反序列化逻辑，不同模型/容器的格式可能存在差异。
• 凭证配置：确保本地或运行环境已配置AWS凭证（可通过aws configure命令设置，或配置环境变量AWS_ACCESS_KEY_ID、AWS_SECRET_ACCESS_KEY）。
• 参数匹配：model_kwargs中的参数需要与你的模型支持的参数一致，比如部分模型可能使用max_length而非max_new_tokens。

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