AWS SageMaker部署Hugging Face嵌入模型端点输出异常求助

在AWS SageMaker部署Hugging Face的sentence-transformers/all-MiniLM-L12-v2文本嵌入模型时，部署成功但输出不符合预期：期望每个输入字符串输出1×384的浮点列表，实际得到每个句子对应7×384的列表。当前代码如下：

import sagemaker
from sagemaker.huggingface.model import HuggingFaceModel

role = sagemaker.get_execution_role()

# Hub Model configuration
hub = {
  'HF_MODEL_ID':'sentence-transformers/all-MiniLM-L12-v2',
  'HF_TASK':'feature-extraction' # 当前使用的任务类型
}

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    env=hub,
    role=role,
    py_version='py36',
    transformers_version="4.6",
    pytorch_version="1.7",
)

predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.xlarge"
)

data = {
"inputs": ["This is an example sentence", "Each sentence is converted"]
}

result = predictor.predict(data)
print(len(result[0]))
print(result[0])

确实是任务类型选错了：

• feature-extraction任务会返回输入文本中每个token对应的嵌入向量，所以你看到的7×384是句子中包含的token数量（含[CLS]、[SEP]等特殊token）乘以模型维度。
• sentence-transformers系列模型专门用于生成句子级嵌入，需要对应使用sentence-similarity任务，该任务会自动对token嵌入做池化处理（如均值池化或取[CLS] token），最终输出单条1×384的句子嵌入。

修改hub配置中的HF_TASK为sentence-similarity即可，同时建议升级transformers和pytorch版本到更兼容的版本，避免版本兼容问题：

import sagemaker
from sagemaker.huggingface.model import HuggingFaceModel

role = sagemaker.get_execution_role()

# Hub Model configuration - 修改任务类型
hub = {
  'HF_MODEL_ID':'sentence-transformers/all-MiniLM-L12-v2',
  'HF_TASK':'sentence-similarity' # 改为句子相似度任务
}

# create Hugging Face Model Class - 升级版本
huggingface_model = HuggingFaceModel(
    env=hub,
    role=role,
    py_version='py38',
    transformers_version="4.26",
    pytorch_version="1.13",
)

predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.xlarge"
)

data = {
"inputs": ["This is an example sentence", "Each sentence is converted"]
}

result = predictor.predict(data)
print(len(result[0])) # 输出384
print(result[0]) # 1×384的浮点列表

如果坚持使用feature-extraction任务，也可以在预测后对token嵌入手动做均值池化（排除特殊token），但直接使用sentence-similarity任务是最简便且符合模型设计的方式。

内容的提问来源于stack exchange，提问作者MaxS.