跳到主要内容

MLflow 中的 Sentence Transformers

注意

sentence_transformers Flavor 正在积极开发中，并标记为实验性。公共 API 可能会发生变化，并且随着 Flavor 新增功能，可能会添加新特性。

sentence_transformers 模型 Flavor 通过 mlflow.sentence_transformers.save_model() 和 mlflow.sentence_transformers.log_model() 函数，能够以 MLflow 格式记录 sentence-transformers 模型。使用这些函数还会为 MLflow 模型添加 python_function Flavor，使得模型可以通过 mlflow.pyfunc.load_model() 被解释为用于推理的通用 Python 函数。此外，可以使用 mlflow.sentence_transformers.load_model() 以原生 sentence-transformers 格式加载带有 sentence_transformers Flavor 的已保存或已记录的 MLflow 模型。

Sentence Transformers 教程​

想直接查看一些可用的示例和教程，了解如何利用 MLflow 使用此库？

查看教程

PyFunc 的输入和输出类型​

sentence_transformers python_function (pyfunc) 模型 Flavor 标准化了句子嵌入和计算语义相似性的过程。这种标准化通过将 sentence_transformers 所需的数据结构调整为兼容 JSON 序列化和转换为 Pandas DataFrame 的格式，从而支持服务和批量推理。

注意

sentence_transformers Flavor 支持用于多种任务的模型，例如嵌入生成、语义相似性计算和释义挖掘。具体的输入和输出类型将取决于所执行的模型和任务。

保存和记录 Sentence Transformers 模型​

您可以在 MLflow 中保存和记录 sentence-transformers 模型。以下是保存和记录模型的示例：

import mlflow
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("model_name")

# Saving the model
mlflow.sentence_transformers.save_model(model=model, path="path/to/save/directory")

# Logging the model
with mlflow.start_run():
    mlflow.sentence_transformers.log_model(
        sentence_transformers_model=model, artifact_path="model_artifact_path"
    )

保存具有 OpenAI 兼容推理接口的 Sentence Transformers 模型​

注意

此功能仅在 MLflow 2.11.0 及更高版本中可用。

MLflow 的 sentence_transformers Flavor 允许您在使用 mlflow.sentence_transformers.save_model() 和 mlflow.sentence_transformers.log_model() 保存模型时，通过字符串值 "llm/v1/embeddings" 传入 task 参数。

例如

import mlflow
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

mlflow.sentence_transformers.save_model(
    model=model, path="path/to/save/directory", task="llm/v1/embeddings"
)

当 task 设置为 "llm/v1/embeddings" 时，MLflow 会为您处理以下事项：

  • 为模型设置兼容嵌入的签名
  • 执行数据预处理和后处理，以确保输入和输出符合与 OpenAI API 规范兼容的嵌入 API 规范。

请注意，这些修改仅在通过 mlflow.pyfunc.load_model() 加载模型时（例如，使用 mlflow models serve CLI 工具提供模型服务时）适用。如果您只想加载基本管道，随时可以通过 mlflow.sentence_transformers.load_model() 进行加载。

除了 sentence-transformers Flavor，transformers Flavor 也支持 OpenAI 兼容的推理接口 ("llm/v1/chat" 和 "llm/v1/completions")。有关更多信息，请参阅Transformers Flavor 指南。

自定义 Python 函数实现​

除了使用预构建的模型外，您还可以使用 sentence_transformers Flavor 创建自定义 Python 函数。以下是用于比较文本文档之间相似度的自定义实现的示例：

import mlflow
from mlflow.pyfunc import PythonModel
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer, util


class DocumentSimilarityModel(PythonModel):
    def load_context(self, context):
        """Load the model context for inference."""
        self.model = SentenceTransformer.load(context.artifacts["model_path"])

    def predict(self, context, model_input):
        """Predict method for comparing similarity between documents."""
        if isinstance(model_input, pd.DataFrame) and model_input.shape[1] == 2:
            documents = model_input.values
        else:
            raise ValueError("Input must be a DataFrame with exactly two columns.")

        # Compute embeddings for each document separately
        embeddings1 = self.model.encode(documents[:, 0], convert_to_tensor=True)
        embeddings2 = self.model.encode(documents[:, 1], convert_to_tensor=True)

        # Calculate cosine similarity
        similarity_scores = util.cos_sim(embeddings1, embeddings2)

        return pd.DataFrame(similarity_scores.numpy(), columns=["similarity_score"])


# Example model saving and loading
model = SentenceTransformer("all-MiniLM-L6-v2")
model_path = "/tmp/sentence_transformers_model"
model.save(model_path)

# Example usage
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="document_similarity_model",
        python_model=DocumentSimilarityModel(),
        artifacts={"model_path": model_path},
    )

loaded = mlflow.pyfunc.load_model(model_info.model_uri)

# Test prediction
df = pd.DataFrame(
    {
        "doc1": ["Sentence Transformers is a wonderful package!"],
        "doc2": ["MLflow is pretty great too!"],
    }
)

result = loaded.predict(df)
print(result)

这将为传入的文档生成相似度分数，如下所示：

    similarity_score
0          0.275423