面向 Amazon SageMaker 的 Hugging Face Embedding 容器介紹

我們很高興地宣佈，面向 Amazon SageMaker 的全新 Hugging Face Embedding 容器現已正式可用 (GA)。AWS 客戶現在可以在 SageMaker 上高效地部署 Embedding 模型，以構建包括檢索增強生成 (RAG) 在內的生成式 AI 應用。

在這篇部落格中，我們將向您展示如何使用新的 Hugging Face Embedding 容器，將開放的 Embedding 模型（如 Snowflake/snowflake-arctic-embed-l、BAAI/bge-large-en-v1.5 或 sentence-transformers/all-MiniLM-L6-v2）部署到 Amazon SageMaker 進行推理。我們將部署 Snowflake/snowflake-arctic-embed-m-v1.5，這是用於檢索的最佳開放 Embedding 模型之一 —— 您可以在 MTEB 排行榜上檢視其排名。

本示例涵蓋以下內容

• 1. 設定開發環境
• 2. 獲取新的 Hugging Face Embedding 容器
• 3. 將 Snowflake Arctic 部署到 Amazon SageMaker
• 4. 執行並評估推理效能
• 5. 刪除模型和終端節點

什麼是 Hugging Face Embedding 容器？

Hugging Face Embedding 容器是一個全新、專用的推理容器，用於在安全且受管理的環境中輕鬆部署 Embedding 模型。該深度學習容器 (DLC) 由 文字嵌入推理 (Text Embedding Inference, TEI) 提供支援，TEI 是一個用於部署和服務 Embedding 模型的極快且記憶體高效的解決方案。TEI 為最流行的模型（包括 FlagEmbedding、Ember、GTE 和 E5）實現了高效能提取。TEI 實現了許多特性，例如：

• 無模型圖編譯步驟
• 小巧的 Docker 映象和快速的啟動時間
• 基於 token 的動態批處理
• 使用 Flash Attention、Candle 和 cuBLASLt 優化了 Transformers 推理程式碼
• Safetensors 權重載入
• 生產就緒（透過 Open Telemetry 進行分散式追蹤，Prometheus 指標）

TEI 支援以下模型架構

• BERT/CamemBERT，例如 BAAI/bge-large-en-v1.5 或 Snowflake/snowflake-arctic-embed-m-v1.5
• RoBERTa，例如 sentence-transformers/all-roberta-large-v1
• XLM-RoBERTa，例如 sentence-transformers/paraphrase-xlm-r-multilingual-v1
• NomicBert，例如 jinaai/jina-embeddings-v2-base-en
• JinaBert，例如 nomic-ai/nomic-embed-text-v1.5

讓我們開始吧！

1. 設定開發環境

我們將使用 sagemaker Python SDK 將 Snowflake Arctic 部署到 Amazon SageMaker。我們需要確保已配置好 AWS 賬戶並安裝了 sagemaker Python SDK。

!pip install "sagemaker>=2.221.1" --upgrade --quiet

如果您要在本地環境中使用 Sagemaker，您需要一個具有 Sagemaker 所需許可權的 IAM 角色。您可以在此處瞭解更多資訊。

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it does not exist
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker session region: {sess.boto_region_name}")

2. 獲取新的 Hugging Face Embedding 容器

與部署常規的 Hugging Face 模型相比，我們首先需要獲取容器的 URI，並將其提供給我們的 HuggingFaceModel 模型類，其中 `image_uri` 指向該映象。為了在 Amazon SageMaker 中獲取新的 Hugging Face Embedding 容器，我們可以使用 sagemaker SDK 提供的 `get_huggingface_llm_image_uri` 方法。此方法允許我們獲取所需 Hugging Face Embedding 容器的 URI。需要注意的是，TEI 有 2 個不同的版本，分別用於 CPU 和 GPU，因此我們建立一個輔助函式，根據例項型別來獲取正確的映象 URI。

from sagemaker.huggingface import get_huggingface_llm_image_uri

# retrieve the image uri based on instance type
def get_image_uri(instance_type):
  key = "huggingface-tei" if instance_type.startswith("ml.g") or instance_type.startswith("ml.p") else "huggingface-tei-cpu"
  return get_huggingface_llm_image_uri(key, version="1.2.3")

3. 將 Snowflake Arctic 部署到 Amazon SageMaker

要將 Snowflake/snowflake-arctic-embed-m-v1.5 部署到 Amazon SageMaker，我們建立一個 HuggingFaceModel 模型類，並定義我們的終端節點配置，包括 HF_MODEL_ID、instance_type 等。我們將使用一個 c6i.2xlarge 例項型別，它有 4 個 Intel Ice-Lake vCPU，8GB 記憶體，每小時成本約為 0.204 美元。

import json
from sagemaker.huggingface import HuggingFaceModel

# sagemaker config
instance_type = "ml.g5.xlarge"

# Define Model and Endpoint configuration parameter
config = {
  'HF_MODEL_ID': "Snowflake/snowflake-arctic-embed-m-v1.5", # model_id from hf.co/models
}

# create HuggingFaceModel with the image uri
emb_model = HuggingFaceModel(
  role=role,
  image_uri=get_image_uri(instance_type),
  env=config
)

建立 HuggingFaceModel 後，我們可以使用 deploy 方法將其部署到 Amazon SageMaker。我們將使用 ml.c6i.2xlarge 例項型別部署模型。

# Deploy model to an endpoint
# https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy
emb = emb_model.deploy(
  initial_instance_count=1,
  instance_type=instance_type,
)

SageMaker 現在將建立我們的終端節點並將模型部署到該節點上。這可能需要約 5 分鐘。

4. 執行並評估推理效能

在我們的終端節點部署完成後，我們可以在其上執行推理。我們將使用 predictor 的 predict 方法在我們的終端節點上執行推理。

data = {
  "inputs": "the mesmerizing performances of the leads keep the film grounded and keep the audience riveted .",
}
 
res = emb.predict(data=data)
 
 
# print some results
print(f"length of embeddings: {len(res[0])}")
print(f"first 10 elements of embeddings: {res[0][:10]}")

太棒了！既然我們能生成 Embedding，讓我們來測試一下模型的效能。

我們將向我們的終端節點發送 3,900 個請求，並使用 10 個併發執行緒。我們將測量終端節點的平均延遲和吞吐量。我們將傳送一個包含 256 個 token 的輸入，總共約 100 萬個 token。我們決定使用 256 個 token 作為輸入長度，以在較短和較長的輸入之間找到一個平衡點。

注意：在執行負載測試時，請求從歐洲傳送，而終端節點部署在 us-east-1。這會給請求增加網路開銷延遲。

import threading
import time
number_of_threads = 10
number_of_requests = int(3900 // number_of_threads)
print(f"number of threads: {number_of_threads}")
print(f"number of requests per thread: {number_of_requests}")
 
def send_requests():
    for _ in range(number_of_requests):
        # input counted at https://huggingface.co/spaces/Xenova/the-tokenizer-playground for 100 tokens
        emb.predict(data={"inputs": "Hugging Face is a company and a popular platform in the field of natural language processing (NLP) and machine learning. They are known for their contributions to the development of state-of-the-art models for various NLP tasks and for providing a platform that facilitates the sharing and usage of pre-trained models. One of the key offerings from Hugging Face is the Transformers library, which is an open-source library for working with a variety of pre-trained transformer models, including those for text generation, translation, summarization, question answering, and more. The library is widely used in the research and development of NLP applications and is supported by a large and active community. Hugging Face also provides a model hub where users can discover, share, and download pre-trained models. Additionally, they offer tools and frameworks to make it easier for developers to integrate and use these models in their own projects. The company has played a significant role in advancing the field of NLP and making cutting-edge models more accessible to the broader community. Hugging Face also provides a model hub where users can discover, share, and download pre-trained models. Additionally, they offer tools and frameworks to make it easier for developers and ma"})
 
# Create multiple threads
threads = [threading.Thread(target=send_requests) for _ in range(number_of_threads) ]
# start all threads
start = time.time()
[t.start() for t in threads]
# wait for all threads to finish
[t.join() for t in threads]
print(f"total time: {round(time.time() - start)} seconds")

傳送 3,900 個請求或嵌入 100 萬個 token 大約耗時 841 秒。這意味著我們每秒可以執行約 5 個請求。但請記住，這包括了從歐洲到 us-east-1 的網路延遲。當我們透過 CloudWatch 檢視終端節點的延遲時，可以看到我們的 Embedding 模型在 10 個併發請求下的延遲為 2 秒。對於一個每月花費約 150 美元的小型且陳舊的 CPU 例項來說，這非常令人印象深刻。您可以將模型部署到 GPU 例項以獲得更快的推理時間。

注意：我們在一個帶有 1 塊 NVIDIA A10G GPU 的 `ml.g5.xlarge` 例項上運行了相同的測試。嵌入 100 萬個 token 大約耗時 30 秒。這意味著我們每秒可以執行約 130 個請求。在 10 個併發請求下，終端節點的延遲為 4 毫秒。在 Amazon SageMaker 上，`ml.g5.xlarge` 的成本約為每小時 1.408 美元。

GPU 例項比 CPU 例項快得多，但也更昂貴。如果您想批次處理 Embedding，可以使用 GPU 例項。如果您想以低成本執行一個小型的終端節點，可以使用 CPU 例項。我們計劃在未來為 Hugging Face Embedding 容器做一個專門的基準測試。

print(f"https://console.aws.amazon.com/cloudwatch/home?region={sess.boto_region_name}#metricsV2:graph=~(metrics~(~(~'AWS*2fSageMaker~'ModelLatency~'EndpointName~'{emb.endpoint_name}~'VariantName~'AllTraffic))~view~'timeSeries~stacked~false~region~'{sess.boto_region_name}~start~'-PT5M~end~'P0D~stat~'Average~period~30);query=~'*7bAWS*2fSageMaker*2cEndpointName*2cVariantName*7d*20{emb.endpoint_name}")

5. 刪除模型和終端節點

為了清理資源，我們可以刪除模型和終端節點

emb.delete_model()
emb.delete_endpoint()

總結

新的 Hugging Face Embedding 容器讓您可以輕鬆地將 Snowflake/snowflake-arctic-embed-l 等開放 Embedding 模型部署到 Amazon SageMaker 進行推理。我們逐步講解了如何設定開發環境、獲取容器、部署模型以及評估其推理效能。

有了這個新容器，客戶現在可以輕鬆部署高效能的 Embedding 模型，從而能夠以更高的效率建立複雜的生成式 AI 應用。我們很期待看到您使用面向 Amazon SageMaker 的新 Hugging Face Embedding 容器構建出的作品。如果您有任何問題或反饋，請告訴我們。