MLflow Transformers Flavor 中的任务

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

概述
原生 Transformers 任务类型
用于 OpenAI 兼容推理的高级任务
- 输入和输出格式
- 使用 llm/v1 任务的代码示例
Databricks 模型服务上的预留吞吐量
常见问题
- 如何覆盖 OpenAI 兼容推理的默认查询参数？

概述

在 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/chat`	`text-generation (文本生成)`	Chat API 规范	以 JSON 格式返回 Chat Completion 对象。
`llm/v1/completions`	`text-generation (文本生成)`	Completions API 规范	以 JSON 格式返回 Completion 对象。
`llm/v1/embeddings`	`feature-extraction (特征提取)`	Embeddings API 规范	返回一个 Embedding 对象列表。此外，模型还会返回 `usage` 字段，其中包含用于生成嵌入的 token 数量。

注意

Completion API 被视为旧版，但 MLflow 为向后兼容性仍然支持它。我们建议使用 Chat API 以兼容 OpenAI 和其他模型提供商的最新 API。

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

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

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 模型服务上的预留吞吐量

Databricks 模型服务上的预留吞吐量是一项能力，可优化基础模型的推理性能，并提供性能保证。要在 Databricks 模型服务上服务 Transformers 模型并获得预留吞吐量，请在记录模型时指定 llm/v1/xxx 任务类型。MLflow 会记录启用 Databricks 模型服务上的预留吞吐量所需的元数据。

提示

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

常见问题

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

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

在推理时：您可以将参数作为输入字典的一部分传递，在调用 predict() 方法时，就像传递输入消息一样。
在记录模型时：您可以通过在记录模型时保存 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。但是，此行为可能不稳定，因为分词器在不同句子中生成的相同序列的 token ID 并不总是相同的，特别是对于基于 sentence-piece 的分词器。

概述​

原生 Transformers 任务类型​

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

输入和输出格式​

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

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

常见问题​

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

概述