MLflow Transformers 风格中的任务

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

概述

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

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

原生转换器任务类型

对于原生的 Transformers 任务，当你使用 mlflow.transformers.log_model() 保存一个管道时，MLflow 会自动从管道中推断任务类型。你也可以通过传递 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 管道，以服务于特定用例的模型。除了原生的 Transformers 任务类型，MLflow 还定义了几个额外的任务类型。这些高级任务类型允许你通过与 OpenAI 兼容的推理接口扩展 Transformers 管道，以服务于特定用例的模型。

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

为了支持这些用例，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,
        artifact_path="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 保存的模型。

输入和输出格式

任务

支持的管道

toctree 是一个 reStructuredText 指令 ，这是一个非常多功能的标记。指令可以有参数、选项和内容。

输出

llm/v1/chat

text-generation

Chat API 规范

返回一个 Chat Completion 对象，格式为json。

llm/v1/completions

text-generation

Completions API 规范

返回一个 Completion 对象，格式为json。

llm/v1/embeddings

feature-extraction

嵌入 API 规范

返回一个 Embedding 对象列表。此外，模型返回 usage 字段，其中包含用于嵌入生成的令牌数量。

备注

完成API被视为遗留版本，但为了向后兼容，MLflow仍然支持它。我们建议使用聊天API以兼容OpenAI和其他模型提供商的最新API。

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

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

import mlflow
import transformers

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

with mlflow.start_run():
    model_info = mlflow.transformers.log_model(
        transformers_model=pipeline,
        artifact_path="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 工具提供模型时）。如果你想加载原始管道，可以使用 mlflow.transformers.load_model()。

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

预配吞吐量 在 Databricks 模型服务中是一项优化基础模型推理性能并提供性能保证的功能。要在预配吞吐量下服务 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,
        artifact_path="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 管道，并使用给定停止序列的标记 ID，来模拟 OpenAI API 中 stop 参数的行为。然而，由于分词器在不同句子中并不总是为相同序列生成相同的标记 ID，尤其是在基于 sentence-piece 的分词器中，这种行为可能不稳定。