跳到主要内容

MLflow Transformers Flavor 中的任务

本页面概述了如何在 MLflow Transformers Flavor 中使用 task 参数来控制模型的推理接口。

概述​

在 MLflow Transformers Flavor 中，task 在确定模型的输入和输出格式方面起着至关重要的作用。task 是 Transformers 库中的一个基本概念，它描述了每个模型的 API（输入和输出）的结构，并用于确定我们希望为任何给定模型显示哪个推理 API 和小部件。

MLflow 利用这个概念来确定模型的输入和输出格式，持久化正确的模型签名，并为服务不同类型的模型提供一致的Pyfunc 推理 API。此外，除了原生的 Transformers 任务类型之外，MLflow 还定义了一些额外的任务类型来支持更复杂的使用场景，例如聊天风格的应用程序。

原生 Transformers 任务类型​

对于原生 Transformers 任务，当您使用mlflow.transformers.log_model()保存 pipeline 时，MLflow 会自动从 pipeline 中推断出任务类型。您还可以通过传递 task 参数来显式指定任务类型。支持的任务类型的完整列表可在Transformers 文档中找到，但请注意，并非所有任务类型都受 MLflow 支持。

import mlflow
import transformers

pipeline = transformers.pipeline("text-generation", model="gpt2")

with mlflow.start_run():
    model_info = mlflow.transformers.save_model(
        transformers_model=pipeline,
        artifact_path="model",
        save_pretrained=False,
    )

print(f"Inferred task: {model_info.flavors['transformers']['task']}")
# >> Inferred task: text-generation

用于 OpenAI 兼容推理的高级任务​

除了原生的 Transformers 任务类型之外，MLflow 还定义了一些额外的任务类型。这些高级任务类型允许您使用 OpenAI 兼容的推理接口扩展 Transformers pipeline，从而为特定用例提供模型服务。

例如，Transformers text-generation pipeline 输入和输出单个字符串或字符串列表。但是，在服务模型时，通常需要具有更结构化的输入和输出格式。例如，在聊天风格的应用程序中，输入可能是一系列消息。

为了支持这些用例，MLflow 定义了一组以 llm/v1 为前缀的高级任务类型

  • "llm/v1/chat" 用于聊天风格的应用程序
  • "llm/v1/completions" 用于通用补全
  • "llm/v1/embeddings" 用于生成文本嵌入

使用这些高级任务类型的必要步骤是在记录模型时将 task 参数指定为 llm/v1 任务。

import mlflow

with mlflow.start_run():
    mlflow.transformers.log_model(
        transformers_model=pipeline,
        name="model",
        task="llm/v1/chat",  # <= Specify the llm/v1 task type
        # Optional, recommended for large models to avoid creating a local copy of the model weights
        save_pretrained=False,
    )
注意

此功能仅在 MLflow 2.11.0 及更高版本中可用。此外，llm/v1/chat 任务类型仅适用于使用 transformers >= 4.34.0 保存的模型。

输入和输出格式​

任务支持的 pipeline输入输出
llm/v1/chattext-generation (文本生成)Chat API 规范以 json 格式返回 Chat Completion 对象。
llm/v1/completionstext-generation (文本生成)Completions API 规范以 json 格式返回 Completion 对象。
llm/v1/embeddingsfeature-extraction (特征提取)Embeddings API 规范返回 Embedding 对象列表。此外，该模型还会返回 usage 字段，其中包含用于生成嵌入的 token 数量。
注意

Completion API 被认为是遗留的，但为了向后兼容，MLflow 仍然支持它。我们建议使用 Chat API 以便与 OpenAI 和其他模型提供商的最新 API 兼容。

使用 llm/v1 任务的代码示例​

以下代码片段演示了如何使用 llm/v1/chat 任务类型记录 Transformers pipeline，并使用该模型进行聊天风格的推理。查看notebook 教程，以查看更多实际示例！

import mlflow
import transformers

pipeline = transformers.pipeline("text-generation", "gpt2")

with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model=pipeline,
        name="model",
        task="llm/v1/chat",
        input_example={
            "messages": [
                {"role": "system", "content": "You are a bot."},
                {"role": "user", "content": "Hello, how are you?"},
            ]
        },
        save_pretrained=False,
    )

# Model metadata logs additional field "inference_task"
print(model_info.flavors["transformers"]["inference_task"])
# >> llm/v1/chat

# The original native task type is also saved
print(model_info.flavors["transformers"]["task"])
# >> text-generation

# Model signature is set to the chat API spec
print(model_info.signature)
# >> inputs:
# >>   ['messages': Array({content: string (required), name: string (optional), role: string (required)}) (required), 'temperature': double (optional), 'max_tokens': long (optional), 'stop': Array(string) (optional), 'n': long (optional), 'stream': boolean (optional)]
# >> outputs:
# >>   ['id': string (required), 'object': string (required), 'created': long (required), 'model': string (required), 'choices': Array({finish_reason: string (required), index: long (required), message: {content: string (required), name: string (optional), role: string (required)} (required)}) (required), 'usage': {completion_tokens: long (required), prompt_tokens: long (required), total_tokens: long (required)} (required)]
# >> params:
# >>     None

# The model can be served with the OpenAI-compatible inference API
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
prediction = pyfunc_model.predict(
    {
        "messages": [
            {"role": "system", "content": "You are a bot."},
            {"role": "user", "content": "Hello, how are you?"},
        ],
        "temperature": 0.5,
        "max_tokens": 200,
    }
)
print(prediction)
# >> [{'choices': [{'finish_reason': 'stop',
# >>               'index': 0,
# >>               'message': {'content': 'I'm doing well, thank you for asking.', 'role': 'assistant'}}],
# >>   'created': 1719875820,
# >>   'id': '355c4e9e-040b-46b0-bf22-00e93486100c',
# >>   'model': 'gpt2',
# >>   'object': 'chat.completion',
# >>   'usage': {'completion_tokens': 7, 'prompt_tokens': 13, 'total_tokens': 20}}]

请注意，输入和输出修改仅在模型使用 mlflow.pyfunc.load_model() 加载时（例如，当使用 mlflow models serve CLI 工具服务模型时）才适用。如果您只想加载原始 pipeline，则可以使用 mlflow.transformers.load_model()。

Databricks 模型服务上的预配置吞吐量​

预配置吞吐量 on Databricks Model Serving 是一种优化基础模型推理性能并提供性能保证的功能。要使用预配置吞吐量提供 Transformers 模型服务，请在记录模型时指定 llm/v1/xxx 任务类型。MLflow 会记录所需的元数据，以便在 Databricks 模型服务上启用预配置吞吐量。

提示

在记录大型模型时，您可以使用 save_pretrained=False 以避免创建模型权重的本地副本，从而节省时间和磁盘空间。有关更多详细信息，请参阅文档。

常见问题​

如何覆盖 OpenAI 兼容推理的默认查询参数？​

当服务使用 llm/v1 任务类型保存的模型时，MLflow 对 temperature 和 stop 等参数使用与 OpenAI API 相同的默认值。您可以通过在推理时传递值，或者在记录模型时设置不同的默认值来覆盖它们。

  1. 在推理时：您可以像传递输入消息一样，在调用 predict() 方法时将参数作为输入字典的一部分传递。
  2. 在记录模型时：您可以通过在记录模型时保存 model_config 参数来覆盖参数的默认值。
with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model=pipeline,
        name="model",
        task="llm/v1/chat",
        model_config={
            "temperature": 0.5,  # <= Set the default temperature
            "stop": ["foo", "bar"],  # <= Set the default stop sequence
        },
        save_pretrained=False,
    )
注意

stop 参数可用于为 llm/v1/chat 和 llm/v1/completions 任务指定停止序列。我们通过将 stopping_criteria 传递给 Transformers pipeline 来模拟 OpenAI API 中 stop 参数的行为，其中包含给定停止序列的 token ID。但是，该行为可能不稳定，因为 tokenizer 并不总是为不同句子中的相同序列生成相同的 token ID，特别是对于基于 sentence-piece 的 tokenizer。