求助：如何在Azure ML中注册SentenceTransformers预训练模型？

核心问题在于SentenceTransformers并非MLFlow原生支持的模型类型，直接用常规MLFlow注册流程会失败，需要通过自定义PyFunc包装器来适配，下面是完整的操作步骤：

1. 编写自定义MLFlow PyFunc包装器 #

创建一个Python脚本（比如sentence_transformer_wrapper.py），实现对SentenceTransformers模型的加载和预测逻辑：

import mlflow.pyfunc
from sentence_transformers import SentenceTransformer

class SentenceTransformerModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 从指定路径加载预训练模型
        self.model = SentenceTransformer(context.artifacts["model_path"])

    def predict(self, context, model_input):
        # 处理输入并返回句子嵌入结果
        sentences = model_input["sentences"].tolist()
        embeddings = self.model.encode(sentences)
        return embeddings

2. 保存模型并注册到Azure ML #

本地保存包装后的模型 #

先在本地环境安装依赖：pip install mlflow sentence-transformers torch azure-ai-ml，然后执行以下代码保存模型：

import mlflow
import os
from sentence_transformer_wrapper import SentenceTransformerModel

# 下载预训练模型到本地
model_name = "all-MiniLM-L6-v2"
local_model_path = f"./{model_name}"
model = SentenceTransformer(model_name)
model.save(local_model_path)

# 保存MLFlow包装后的模型
mlflow_pyfunc_path = "./sentence_transformer_mlflow"
mlflow.pyfunc.save_model(
    path=mlflow_pyfunc_path,
    python_model=SentenceTransformerModel(),
    artifacts={"model_path": local_model_path}
)

注册到Azure ML工作区 #

使用Azure ML的Python SDK v2注册模型：

from azure.ai.ml import MLClient
from azure.ai.ml.entities import Model
from azure.identity import DefaultAzureCredential

# 初始化MLClient
ml_client = MLClient(
    credential=DefaultAzureCredential(),
    subscription_id="<你的订阅ID>",
    resource_group_name="<资源组名称>",
    workspace_name="<工作区名称>"
)

# 注册模型
model = Model(
    path=mlflow_pyfunc_path,
    name="sentence-transformers-minilm",
    description="SentenceTransformers预训练模型all-MiniLM-L6-v2，用于生成句子嵌入",
    type="mlflow_model"
)
ml_client.models.create_or_update(model)

3. 部署到Azure ML在线端点 #

创建环境配置 #

创建一个conda.yaml文件，定义模型运行所需的依赖：

name: sentence-transformers-env
channels:
  - conda-forge
dependencies:
  - python=3.9
  - pip=23.0
  - pip:
      - sentence-transformers==2.2.2
      - torch==2.0.1
      - mlflow==2.4.1
      - azureml-mlflow==1.50.0

部署端点 #

使用Python SDK v2创建在线端点和部署：

from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment, Environment

# 创建在线端点
endpoint = ManagedOnlineEndpoint(
    name="sentence-transformers-endpoint",
    description="用于生成句子嵌入的在线端点",
    auth_mode="key"
)
ml_client.begin_create_or_update(endpoint).result()

# 创建环境
env = Environment(
    conda_file="conda.yaml",
    image="mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04:latest"
)

# 创建部署
deployment = ManagedOnlineDeployment(
    name="sentence-transformers-deployment",
    endpoint_name="sentence-transformers-endpoint",
    model="sentence-transformers-minilm:1",  # 替换为你的模型版本号
    environment=env,
    instance_type="Standard_DS3_v2",
    instance_count=1
)
ml_client.begin_create_or_update(deployment).result()

验证部署 #

部署完成后，发送测试请求验证：

import requests
import json

# 获取端点密钥和URL
endpoint = ml_client.online_endpoints.get("sentence-transformers-endpoint")
scoring_uri = endpoint.scoring_uri
key = ml_client.online_endpoints.get_keys("sentence-transformers-endpoint").primary_key

# 测试数据
data = {
    "input_data": {
        "columns": ["sentences"],
        "data": ["Hello world", "Azure ML is great"]
    }
}

# 发送请求
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
response = requests.post(scoring_uri, json=data, headers=headers)
print(response.json())

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