跳到主内容

MLflow 模型签名和输入示例指南

模型签名和输入示例介绍​

在 MLflow 中，*模型签名* 和 *模型输入示例* 的概念对于有效使用机器学习模型至关重要。这些组件不仅仅提供元数据；它们还为模型交互建立了关键准则，增强了在 MLflow 生态系统中的集成性和可用性。

模型签名​

MLflow 中的 模型签名 对于模型的清晰准确运行至关重要。它定义了模型输入和输出的预期格式，包括推理所需的任何附加参数。此规范充当明确的指南，确保模型与 MLflow 工具和外部服务的无缝集成。

模型输入示例​

作为模型签名的补充，模型输入示例 提供了有效模型输入的一个具体实例。这个实际示例对开发者而言非常有价值，它清晰地展示了进行有效模型测试和实际应用所需的数据格式和结构。

为何它们如此重要？​

模型签名和输入示例是稳健 ML 工作流的基础，它们为模型交互提供了蓝图，确保了一致性、准确性和易用性。它们充当了模型与其用户之间的契约，提供了预期数据格式的明确指南，从而防止了因不正确或意外输入可能导致的误解和错误。

模型签名​

在 MLflow 中，模型签名精确定义了模型输入、输出以及有效运行模型所需的任何附加参数的模式。此定义作为统一接口，指导用户正确准确地使用其模型。模型签名是 MLflow 生态系统不可或缺的一部分，它使 MLflow Tracking UI 和 Model Registry UI 都能够清晰地显示模型所需的输入、输出和参数。此外，MLflow 模型部署工具 利用这些签名来确保推理时使用的数据与模型已建立的规范一致，从而维护模型的完整性和性能。要深入了解这些签名如何强制执行数据准确性，请参阅签名强制执行部分。

在 MLflow 中，在模型中嵌入签名是一个简单的过程。使用 sklearn.log_model() 等函数记录或保存模型时，只需包含一个模型输入示例。此操作使 MLflow 能够自动推断模型的签名。有关此过程的详细说明可在如何记录带签名的模型部分找到。推断出的签名连同其他必要的模型元数据，以 JSON 格式存储在模型工件的MLmodel 文件中。如果需要向已记录或保存的模型添加签名，可以使用 set_signature() API。请参阅如何设置模型签名部分以获取实现此功能的详细指南。

模型签名组件​

MLflow 中模型签名的结构由三种不同的模式类型组成：(1) 输入、(2) 输出 和 (3) 参数 (params)。输入和输出模式分别指定了模型预期接收和产生的数据结构。这些模式可以根据不同模型的多种需求，定制以处理各种数据格式，包括列式数据和张量，以及数组 (Array) 和对象 (Object) 类型输入。

另一方面，参数 (params) 指的是在模型推理阶段至关重要的附加参数，这些参数通常是可选的。这些参数提供了额外的灵活性，允许对推理过程进行微调和自定义。

注意

在 MLflow 2.10.0 及更高版本中引入了模型签名处理对象 (Objects) 和数组 (Arrays) 的功能。在 2.10.0 之前的版本中，基于列的签名仅限于标量输入类型以及特定于列表和字典输入的一些条件类型，主要支持 transformers 风味。后续版本中的这一增强显著扩大了 MLflow 可以无缝适应的数据类型和结构的范围。

签名演练场​

为了帮助理解数据结构如何自动推断为有效签名，并提供大量有效签名的示例，我们创建了一个笔记本，您可以查看该笔记本，其中展示了不同的示例及其推断出的签名。您可以在此处查看笔记本。

或者，如果您想将笔记本下载到本地并使用您自己的数据结构进行测试，您可以在下方下载。

下载签名演练场笔记本

必需与可选输入字段​

在定义模型签名时，需要考虑某些决定签名强制执行的条件。其中最值得注意的一点是输入数据中的必需字段与可选字段的概念。

必需字段是指必须存在于输入数据中，模型才能进行预测的字段。如果缺少必需字段，签名强制执行验证将引发异常，指示缺少必需的输入字段。

要将字段配置为可选，在使用 mlflow.models.infer_signature() 函数时，必须为该字段传递值 None 或 np.nan。或者，您可以手动定义签名，并将该字段的 required 设置为 false。

模型签名类型​

MLflow 支持两种主要的签名类型：用于表格数据的基于列的签名，以及用于张量数据的基于张量的签名。

基于列的签名​

基于列的签名常用于需要表格数据输入的传统机器学习模型，例如 Pandas DataFrame。这些签名由一系列列组成，每列可能具有名称和指定的数据类型。输入和输出中每列的类型对应于支持的数据类型之一，并且可以为列可选地命名。此外，输入模式中的列可以指定为 optional，表明它们在模型输入中不是必需的，如有必要可以省略（详情请参阅可选列）。

支持的数据类型​

基于列的签名支持 MLflow DataType 规范中定义的数据原语

  • 字符串
  • 整数 1
  • 长整数
  • 浮点数
  • 双精度浮点数
  • 布尔值
  • 日期时间
输入 (Python)推断签名
from mlflow.models import infer_signature

infer_signature(model_input={"long_col": 1, "str_col": "a", "bool_col": True})
signature:
    input: '[
        {"name": "long_col", "type": "long",    "required": "true"},
        {"name": "str_col",  "type": "string",  "required": "true"},
        {"name": "bool_col", "type": "boolean", "required": "true"}
    ]'
    output: null
    params: null
注意

1 Python 通常将整数数据中的缺失值表示为浮点数，这会导致整数列中的类型变异以及 MLflow 中潜在的模式强制执行错误。为了避免此类问题，尤其是在 MLflow 中使用 Python 进行模型服务和 Spark 部署时，请将包含缺失值的整数列定义为双精度浮点数 (float64)。

基于列的签名还支持这些原语的复合数据类型。

  • 数组 (list, numpy 数组)
  • Spark ML 向量 (继承 Array[double] 类型)
  • 对象 (字典)
  • 任意类型
注意

AnyType 是一种特殊类型，可用于表示任何数据类型，包括 None 值。如果使用此类型，在 pyfunc predict 中的模式强制执行过程中不会对输入数据进行任何验证。函数 mlflow.models.infer_signature 仅在字段始终为 None 时才将字段推断为 AnyType（例如 {"a": None} --> ['a': Any (optional)]）；如果该字段具有其他有效类型，则会将其推断为可选字段（例如 [{"a": None}, {"a": "abc"}] --> ['a': string (optional)]）。如果某个字段明确接受多种有效类型（例如 [{"a": "string"}, {"a": 123}]），infer_signature 函数将会失败，您需要使用 ModelSignature 对象手动构造签名，并为该字段指定 AnyType 类型（例如 ModelSignature(Schema([ColSpec(AnyType(), "a", required=False)]))）。

警告
  • MLflow 2.10.0 版本引入了对数组 (Array) 和对象 (Object) 类型的支持。在早期版本的 MLflow 中将无法识别这些类型。如果您保存的模型使用了这些签名类型，应确保尝试加载这些模型的任何其他环境安装的 MLflow 版本至少为 2.10.0。
  • MLflow 2.15.0 版本引入了对 Spark ML 向量类型的支持，在早期版本的 MLflow 中将无法识别此类型。
  • MLflow 2.19.0 版本引入了对 AnyType 的支持。在早期版本的 MLflow 中将无法识别此类型。

通过查看签名示例笔记本，可以看到复合数据类型的更多示例。

输入 (Python)推断签名
from mlflow.models import infer_signature

infer_signature(
    model_input={
        # Python list
        "list_col": ["a", "b", "c"],
        # Numpy array
        "numpy_col": np.array([[1, 2], [3, 4]]),
        # Dictionary
        "obj_col": {"long_prop": 1, "array_prop": ["a", "b", "c"]},
    }
)
signature:
    input: '[
        {"list_col": Array(string) (required)},
        {"numpy_col": Array(Array(long)) (required)},
        {"obj_col": {array_prop: Array(string) (required), long_prop: long (required)} (required)}}
    ]'
    output: null
    params: null

可选列​

输入数据中包含 None 或 np.nan 值的列将被推断为 optional（即 required=False）

输入 (Python)推断签名
from mlflow.models import infer_signature

infer_signature(model_input=pd.DataFrame({"col": [1.0, 2.0, None]}))
signature:
    input: '[
        {"name": "col", "type": "double", "required": false}
    ]'
    output: null
    params: null
注意

嵌套数组可以包含空列表，但这不会使该列变为 optional，因为它表示一个空集 (∅)。在这种情况下，模式将从列表的其他元素推断，假设它们具有同质类型。如果您想使列可选，请传递 None 代替。

输入 (Python)推断签名
infer_signature(
    model_input={
        "list_with_empty": [["a", "b"], []],
        "list_with_none": [["a", "b"], None],
    }
)
signature:
    input: '[
        {"name": "list_with_empty", "type": "Array(str)", "required": "true" },
        {"name": "list_with_none" , "type": "Array(str)", "required": "false"},
    ]'
    output: null
    params: null

基于张量的签名​

基于张量的签名主要用于处理张量输入的模型，常见于涉及图像、音频数据及类似格式的深度学习应用。这些模式由一系列张量组成，每个张量可能具有名称并由特定的numpy 数据类型定义。

在基于张量的签名中，每个输入和输出张量由三个属性表征：一个 dtype（数据类型，与numpy 数据类型对齐）、一个 shape 以及一个可选的 name。重要的是要注意，基于张量的签名不接受可选输入。 shape 属性通常使用 -1 表示大小可能变化的任何维度，这在批量处理中很常见。

考虑一个在 MNIST 数据集上训练的分类模型。其模型签名的示例将包含一个输入张量，表示一张图像为 28 × 28 × 1 的 float32 数字数组。模型的输出可能是一个表示 10 个目标类别各自概率的张量。在这种情况下，表示批量大小的第一维度通常设置为 -1，允许模型处理不同大小的批量。

signature:
  inputs: '[{"name": "images", "type": "tensor", "tensor-spec": {"dtype": "uint8", "shape": [-1, 28, 28, 1]}}]'
  outputs: '[{"type": "tensor", "tensor-spec": {"shape": [-1, 10], "dtype": "float32"}}]'
  params: null

支持的数据类型​

基于张量的模式支持numpy 数据类型。

输入 (Python)推断签名
from mlflow.models import infer_signature

infer_signature(
    model_input=np.array(
        [
            [[1, 2, 3], [4, 5, 6]],
            [[7, 8, 9], [1, 2, 3]],
        ]
    )
)
signature:
    input: '[{"type": "tensor", "tensor-spec": {"dtype": "int64", "shape": [-1, 2, 3]}}]'
    output: None
    params: None
注意

基于张量的模式不支持可选输入。您可以传入包含 None 或 np.nan 值 的数组，但模式不会将其标记为可选。

带有推理参数的模型签名​

推理参数（或 'params'）是在推理阶段传递给模型的附加设置。常见示例包括语言学习模型 (LLMs) 中的 temperature 和 max_length 等参数。这些参数通常在训练期间不需要，但在推理时对调整模型行为起着至关重要的作用。随着基础模型的出现，这种配置变得越来越重要，因为同一模型可能需要针对不同的推理场景设置不同的参数。

MLflow 的 2.6.0 版本引入了在模型推理期间指定推理参数字典的功能。此功能增强了对推理结果的灵活性和控制力，使得对模型行为的调整更加细致入微。

要在推理时利用参数，必须将其纳入模型签名中。参数的模式被定义为一系列ParamSpec 元素，每个元素包含

  • 名称：参数的标识符，例如 temperature。
  • 类型：参数的数据类型，必须与支持的数据类型之一对齐。
  • 默认值：参数的默认值，如果在未提供特定值时确保有备选选项。
  • 形状：参数的形状，对于标量值通常为 None，对于列表为 (-1,)。

此功能标志着 MLflow 处理模型推理方式的重大进步，提供了一种更动态和适应性强的模型参数化方法。

signature:
    inputs: '[{"name": "input", "type": "string"}]'
    outputs: '[{"name": "output", "type": "string"}]'
    params: '[
        {
            "name": "temperature",
            "type": "float",
            "default": 0.5,
            "shape": null
        },
        {
            "name": "suppress_tokens",
            "type": "integer",
            "default": [101, 102],
              "shape": [-1]
        }
    ]'

推理参数在推理阶段以字典的形式提供给模型。每个参数值都会经过验证，以确保其与模型签名中指定的类型匹配。以下示例展示了在模型签名中定义参数的过程，并演示了它们在模型推理中的应用。

import mlflow
from mlflow.models import infer_signature


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params):
        return list(params.values())


params = {"temperature": 0.5, "suppress_tokens": [101, 102]}
# params' default values are saved with ModelSignature
signature = infer_signature(["input"], params=params)

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model=MyModel(), artifact_path="my_model", signature=signature
    )

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

# Not passing params -- predict with default values
loaded_predict = loaded_model.predict(["input"])
assert loaded_predict == [0.5, [101, 102]]

# Passing some params -- override passed-in params
loaded_predict = loaded_model.predict(["input"], params={"temperature": 0.1})
assert loaded_predict == [0.1, [101, 102]]

# Passing all params -- override all params
loaded_predict = loaded_model.predict(
    ["input"], params={"temperature": 0.5, "suppress_tokens": [103]}
)
assert loaded_predict == [0.5, [103]]

参数支持的数据类型​

MLflow 中定义的参数接受MLflow DataType 类型的值，包括这些数据类型的一维列表。目前，MLflow 只支持参数的 1D 列表。

注意

验证参数值时，值将被转换为 Python 原生类型。例如，np.float32(0.1) 将被转换为 float(0.1)。

签名强制执行​

MLflow 的模式强制执行会根据模型签名严格验证提供的输入和参数。如果输入不兼容，则会引发异常；如果参数不兼容，则会发出警告或引发异常。此强制执行在调用底层模型实现之前以及在整个模型推理过程中都会应用。但请注意，此强制执行仅适用于使用MLflow 模型部署工具或模型作为 python_function 加载的场景。对于以其原生格式加载的模型，例如通过 mlflow.sklearn.load_model() 加载的模型，不适用此强制执行。

名称顺序强制执行​

在 MLflow 中，会根据模型签名验证输入名称。缺少必需输入会触发异常，而缺少可选输入则不会。签名中未声明的输入将被忽略。当签名中的输入模式指定了输入名称时，匹配会按名称进行，并相应地重新排列输入。如果模式缺少输入名称，则匹配基于输入的顺序，MLflow 只检查提供的输入数量。

输入类型强制执行​

MLflow 强制执行模型签名中定义的输入类型。对于基于列的签名模型（例如使用 DataFrame 输入的模型），MLflow 会在必要时执行安全的类型转换，只允许无损转换。例如，允许将 int 转换为 long 或 int 转换为 double，但不允许将 long 转换为 double。如果类型无法兼容，MLflow 将会引发错误。

对于 PySpark DataFrame 输入，MLflow 会将 PySpark DataFrame 的一个样本转换为 Pandas DataFrame。MLflow 只会对数据行的一个子集强制执行模式。

对于基于张量的签名模型，类型检查更加严格。如果输入类型与模式中指定的类型不一致，则会抛出异常。

参数类型和形状强制执行​

在 MLflow 中，参数 (params) 的类型和形状会根据模型的签名进行仔细检查。在推理过程中，会验证每个参数的类型和形状，以确保其与签名中的规范一致。标量值预期形状为 None，而列表值应为 (-1,)。如果发现参数的类型或形状不兼容，MLflow 会引发异常。此外，还会根据签名中指定的类型对参数值进行验证检查。如果转换为指定类型失败，则会触发一个 MlflowException。有关有效参数的完整列表，请参阅模型推理参数部分。

重要

在推理期间接收到签名中未声明的参数的模型会为每个请求触发一个警告，并且任何无效参数都将被忽略。

处理带有缺失值的整数​

在 Python 中，带有缺失值的整数数据通常表示为浮点数。这会导致整数列数据类型的变异，由于整数和浮点数本质上不兼容，可能会在运行时导致模式强制执行错误。例如，如果训练数据中的列 'c' 全部为整数，它将被识别为整数。但是，如果 'c' 中引入了缺失值，它将表示为浮点数。如果模型的签名预期 'c' 是整数，MLflow 将由于无法将 float 转换为 int 而引发错误。为了缓解此问题，特别是因为 MLflow 使用 Python 进行模型服务和 Spark 部署，建议将带有缺失值的整数列定义为双精度浮点数 (float64)。

处理日期和时间戳​

Python 的日期时间类型带有内置精度，例如用于日精度的是 datetime64[D]，用于纳秒精度的是 datetime64[ns]。虽然这些精度细节在基于列的模型签名中被忽略，但在基于张量的签名中会被强制执行。

处理不规则数组 (Ragged Arrays)​

使用 infer_signature 时，numpy 中以 (-1,) 形状和 object dtype 为特征的不规则数组 (Ragged arrays) 会被自动管理。这会生成一个类似 Tensor('object', (-1,)) 的签名。为了获得更详细的表示，可以手动创建一个签名来反映不规则数组的特定性质，例如 Tensor('float64', (-1, -1, -1, 3))。然后，根据签名中详细说明的应用强制执行，以适应不规则输入数组。

如何记录带签名的模型​

在 MLflow 中为模型包含签名非常简单。在使用 log_model 或 save_model 函数（例如 sklearn.log_model()）时，只需提供一个模型输入示例。MLflow 将根据此输入示例和模型对给定示例的预测输出自动推断模型的签名。

提示

从 MLflow 2.20.0 开始，使用类型提示记录 PythonModel 时，模型签名会自动推断。有关详细信息，请查看带有类型提示的 PythonModel 指南。

或者，您可以显式地将签名对象附加到模型。这可以通过将 signature object 传递给您的 log_model 或 save_model 函数来实现。您可以手动创建模型签名对象，或使用 infer_signature 函数从包含有效模型输入（例如，减去目标列的训练数据集）、有效模型输出（例如，在训练数据集上进行的预测）和有效模型参数（例如，用于模型推理的参数字典，常见于transformers 的生成配置 (Generation Configs)）的数据集中生成签名。

注意

模型签名在MLflow 模型部署工具中起着至关重要的作用，特别是对于以 Python Function (PyFunc) 风味提供模型的服务。因此，在将签名附加到 log_model 或 save_model 调用时，务必确保签名准确反映模型 PyFunc 表示形式预期的输入和输出。如果模型作为 PyFunc 加载时的输入模式与用于测试的数据集不同，则此考虑尤其重要（例如pmdarima 模型风味的情况）。

基于列的签名示例​

以下示例演示了如何为在 Iris 数据集上训练的简单分类器存储模型签名

import pandas as pd
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier
import mlflow

iris = datasets.load_iris()
iris_train = pd.DataFrame(iris.data, columns=iris.feature_names)
clf = RandomForestClassifier(max_depth=7, random_state=0)

with mlflow.start_run():
    clf.fit(iris_train, iris.target)
    # Take the first row of the training dataset as the model input example.
    input_example = iris_train.iloc[[0]]
    # The signature is automatically inferred from the input example and its predicted output.
    mlflow.sklearn.log_model(clf, "iris_rf", input_example=input_example)

可以使用以下方式显式创建和记录相同的签名

from mlflow.models import ModelSignature, infer_signature
from mlflow.types.schema import Schema, ColSpec

# Option 1: Manually construct the signature object
input_schema = Schema(
    [
        ColSpec("double", "sepal length (cm)"),
        ColSpec("double", "sepal width (cm)"),
        ColSpec("double", "petal length (cm)"),
        ColSpec("double", "petal width (cm)"),
    ]
)
output_schema = Schema([ColSpec("long")])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)

# Option 2: Infer the signature
signature = infer_signature(iris_train, clf.predict(iris_train))

with mlflow.start_run():
    mlflow.sklearn.log_model(clf, "iris_rf", signature=signature)

基于张量的签名示例​

以下示例演示了如何为在 MNIST 数据集上训练的简单分类器存储模型签名

import tensorflow as tf
import mlflow

mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10),
    ]
)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])

with mlflow.start_run():
    model.fit(x_train, y_train, epochs=5)
    # Take the first three training examples as the model input example.
    input_example = x_train[:3, :]
    mlflow.tensorflow.log_model(model, "mnist_cnn", input_example=input_example)

可以使用以下方式显式创建和记录相同的签名

import numpy as np
from mlflow.models import ModelSignature, infer_signature
from mlflow.types.schema import Schema, TensorSpec

# Option 1: Manually construct the signature object
input_schema = Schema(
    [
        TensorSpec(np.dtype(np.float64), (-1, 28, 28, 1)),
    ]
)
output_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 10))])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)

# Option 2: Infer the signature
signature = infer_signature(testX, model.predict(testX))

with mlflow.start_run():
    mlflow.tensorflow.log_model(model, "mnist_cnn", signature=signature)

带参数的签名示例​

以下示例演示了如何为简单的 transformers 模型存储带有参数的模型签名

import mlflow
from mlflow.models import infer_signature
import transformers

architecture = "mrm8488/t5-base-finetuned-common_gen"
model = transformers.pipeline(
    task="text2text-generation",
    tokenizer=transformers.T5TokenizerFast.from_pretrained(architecture),
    model=transformers.T5ForConditionalGeneration.from_pretrained(architecture),
)
data = "pencil draw paper"

params = {
    "top_k": 2,
    "num_beams": 5,
    "max_length": 30,
    "temperature": 0.62,
    "top_p": 0.85,
    "repetition_penalty": 1.15,
    "begin_suppress_tokens": [1, 2, 3],
}

# infer signature with params
signature = infer_signature(
    data,
    mlflow.transformers.generate_signature_output(model, data),
    params,
)

# save model with signature
mlflow.transformers.save_model(
    model,
    "text2text",
    signature=signature,
)
pyfunc_loaded = mlflow.pyfunc.load_model("text2text")

# predict with params
result = pyfunc_loaded.predict(data, params=params)

可以使用以下方式显式创建相同的签名

from mlflow.models import ModelSignature
from mlflow.types.schema import ColSpec, ParamSchema, ParamSpec, Schema

input_schema = Schema([ColSpec(type="string")])
output_schema = Schema([ColSpec(type="string")])
params_schema = ParamSchema(
    [
        ParamSpec("top_k", "long", 2),
        ParamSpec("num_beams", "long", 5),
        ParamSpec("max_length", "long", 30),
        ParamSpec("temperature", "double", 0.62),
        ParamSpec("top_p", "double", 0.85),
        ParamSpec("repetition_penalty", "double", 1.15),
        ParamSpec("begin_suppress_tokens", "long", [1, 2, 3], (-1,)),
    ]
)
signature = ModelSignature(
    inputs=input_schema, outputs=output_schema, params=params_schema
)

GenAI 风味的模型签名示例​

像 langchain、OpenAI 和 transformers 这样的 GenAI 风味通常需要基于对象（字典）的模型签名。手动定义结构既复杂又容易出错。相反，您应该依赖 infer_signature 方法，根据输入示例自动生成模型签名。

记录 GenAI 风味模型时，您可以将输入示例传递给 log_model 方法的 input_example 参数。MLflow 将从提供的输入中推断模型签名，并将其应用于记录的模型。此外，此输入示例有助于验证模型能够成功运行预测，因此强烈建议在记录 GenAI 模型时始终包含输入示例。

注意

默认情况下，输入示例中的字段被推断为必需字段。要将字段标记为可选，您可以提供一个输入示例列表，其中某些示例包含可选字段的 None 值。（例如 [{"str": "a"}, {"str": None}]）

以 langchain 模型为例，您可以使用 models from code 功能，通过 2 个简单步骤记录带签名的模型

  1. 在一个名为 langchain_model.py 的独立 Python 文件中定义 langchain 模型

要成功记录 langchain 模型，强烈建议使用 models from code 功能。

import os
from operator import itemgetter

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import OpenAI

import mlflow

mlflow.set_experiment("Homework Helper")

mlflow.langchain.autolog()

prompt = PromptTemplate(
    template="You are a helpful tutor that evaluates my homework assignments and provides suggestions on areas for me to study further."
    " Here is the question: {question} and my answer which I got wrong: {answer}",
    input_variables=["question", "answer"],
)


def get_question(input):
    default = "What is your name?"
    if isinstance(input_data[0], dict):
        return input_data[0].get("content").get("question", default)
    return default


def get_answer(input):
    default = "My name is Bobo"
    if isinstance(input_data[0], dict):
        return input_data[0].get("content").get("answer", default)
    return default


model = OpenAI(temperature=0.95)

chain = (
    {
        "question": itemgetter("messages") | RunnableLambda(get_question),
        "answer": itemgetter("messages") | RunnableLambda(get_answer),
    }
    | prompt
    | model
    # Add this parser to convert the model output to a string
    # so that MLflow can correctly infer the output schema
    | StrOutputParser()
)

mlflow.models.set_model(chain)
注意

MLflow 只能推断某些数据类型的模型签名，因此对于 langchain 模型，必须添加一个解析器，例如 StrOutputParser，将输出解析为 MLflow 接受的格式。

  1. 使用有效的输入示例记录 langchain 模型
input_example = {
    "messages": [
        {
            "role": "user",
            "content": {
                "question": "What is the primary function of control rods in a nuclear reactor?",
                "answer": "To stir the primary coolant so that the neutrons are mixed well.",
            },
        },
    ]
}

chain_path = "langchain_model.py"
with mlflow.start_run():
    model_info = mlflow.langchain.log_model(
        lc_model=chain_path, artifact_path="model", input_example=input_example
    )

在此阶段，input_example 具有多种用途：它用于推断模型签名，验证模型在使用 PyFunc 风味重新加载时是否正常工作，并生成相应的 serving_input_example。生成的签名用于在部署之前验证模型的功能，并在模型部署后为调用者提供指导，以防他们使用无效数据结构发出请求。有关输入示例的优点和用法的更多信息，请参阅模型输入示例部分。要了解如何在部署前在本地验证模型，请参阅模型服务载荷示例部分。

如果您的模型有一个可选输入字段，您可以使用下面的 input_example 作为参考

from mlflow.models import infer_signature

input_example = {
    "messages": [
        {
            # specify name field in the first message
            "name": "userA",
            "role": "user",
            "content": {
                "question": "What is the primary function of control rods in a nuclear reactor?",
                "answer": "To stir the primary coolant so that the neutrons are mixed well.",
            },
        },
        {
            # no name field in the second message, so `name` will be inferred as optional
            "role": "user",
            "content": {
                "question": "What is MLflow?",
                "answer": "MLflow is an open-source platform",
            },
        },
    ]
}

print(infer_signature(input_example))

如果模型的输出包含 None 值，这些字段将被推断为 AnyType（自 MLflow 2.19.0 起），例如

from mlflow.models import infer_signature

data = [
    {
        "id": None,
        "object": "chat.completion",
        "created": 1731491873,
        "model": None,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": "MLflow",
                },
                "finish_reason": None,
            }
        ],
        "usage": {
            "prompt_tokens": None,
            "completion_tokens": None,
            "total_tokens": None,
        },
    }
]

print(infer_signature(data))
# inputs:
# [
#     'id': Any (optional),
#     'object': string (required),
#     'created': long (required),
#     'model': Any (optional),
#     'choices': Array({
#         'finish_reason': Any (optional),
#         'index': long (required),
#         'message': {
#             'content': string (required),
#             'role': string (required)
#         } (required)
#     }) (required),
#     'usage': {
#         'completion_tokens': Any (optional),
#         'prompt_tokens': Any (optional),
#         'total_tokens': Any (optional)
#     } (required)
# ]
# outputs:
# None
# params:
# None
  1. 重新加载模型并进行预测
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
result = loaded_model.predict(input_example)

print(result)

也支持使用 Dataclass 对象定义签名。将 Dataclass 实例传递给 infer_signature 将生成相应的模型签名等效项

from dataclasses import dataclass
from typing import List
from mlflow.models import infer_signature


@dataclass
class Message:
    role: str
    content: str


@dataclass
class ChatCompletionRequest:
    messages: List[Message]


chat_request = ChatCompletionRequest(
    messages=[
        Message(
            role="user",
            content="What is the primary function of control rods in a nuclear reactor?",
        ),
        Message(role="user", content="What is MLflow?"),
    ]
)

model_signature = infer_signature(
    chat_request,
    "Sample output as a string",
)
print(model_signature)
# inputs:
# ['messages': Array({content: string (required), role: string (required)}) (required)]
# outputs:
# [string (required)]
# params:
# None

如何设置模型签名​

模型可以保存时没有模型签名或带有错误的签名。要为现有已记录的模型添加或更新签名，请使用 mlflow.models.set_signature() API。以下是一些演示其用法的示例。

在已记录的模型上设置签名​

以下示例演示了如何在已记录的 sklearn 模型上设置模型签名。假设您已经记录了一个没有签名的 sklearn 模型，如下所示

import pandas as pd
from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier
import mlflow

X, y = datasets.load_iris(return_X_y=True, as_frame=True)
clf = RandomForestClassifier(max_depth=7, random_state=0)
with mlflow.start_run() as run:
    clf.fit(X, y)
    mlflow.sklearn.log_model(clf, "iris_rf")

您可以按如下方式在已记录的模型上设置签名

import pandas as pd
from sklearn import datasets
import mlflow
from mlflow.models.model import get_model_info
from mlflow.models import infer_signature, set_signature

# load the logged model
model_uri = f"runs:/{run.info.run_id}/iris_rf"
model = mlflow.pyfunc.load_model(model_uri)

# construct the model signature from test dataset
X_test, _ = datasets.load_iris(return_X_y=True, as_frame=True)
signature = infer_signature(X_test, model.predict(X_test))

# set the signature for the logged model
set_signature(model_uri, signature)

# now when you load the model again, it will have the desired signature
assert get_model_info(model_uri).signature == signature

请注意，也可以在 MLflow Tracking 之外保存的模型工件上设置模型签名。例如，您可以通过更改上一个代码片段中的 model_uri 变量指向模型的本地目录，轻松地在本地保存的 iris 模型上设置签名。

在注册的模型上设置签名​

由于 MLflow 模型注册表工件旨在只读，您无法直接在模型版本或由 models:/ URI 方案表示的模型工件上设置签名。相反，您应该首先在源模型工件上设置签名，然后使用更新的模型工件生成新的模型版本。以下示例说明了如何完成此操作。

假设您已经创建了以下没有签名的模型版本，如下所示

from sklearn.ensemble import RandomForestClassifier
import mlflow
from mlflow.client import MlflowClient

model_name = "add_signature_model"

with mlflow.start_run() as run:
    mlflow.sklearn.log_model(RandomForestClassifier(), "sklearn-model")

model_uri = f"runs:/{run.info.run_id}/sklearn-model"
mlflow.register_model(model_uri=model_uri, name=model_name)

要在模型版本上设置签名，请按如下方式创建带有新签名的副本模型版本

from sklearn.ensemble import RandomForestClassifier
import mlflow
from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository

client = mlflow.client.MlflowClient()
model_name = "add_signature_model"
model_version = 1
mv = client.get_model_version(name=model_name, version=model_version)

# set a dummy signature on the model version source
signature = infer_signature(np.array([1]))
set_signature(mv.source, signature)

# create a new model version with the updated source
client.create_model_version(name=model_name, source=mv.source, run_id=mv.run_id)

请注意，此过程会用新的模型签名覆盖模型版本 1 的源运行中的模型工件。

模型输入示例​

模型输入示例提供了有效模型输入的一个实例。输入示例与模型一起作为单独的工件存储，并在MLmodel 文件中引用。要在模型中包含输入示例，请将其添加到相应的 log_model 调用中，例如sklearn.log_model()。在未指定签名时，输入示例也用于在 log_model 调用中推断模型签名。

提示

在记录模型时包含输入示例提供了双重好处。首先，它有助于推断模型的签名。其次，同样重要的是，它验证了模型的要求。此输入示例用于使用即将记录的模型执行预测，从而提高了识别模型要求依赖项的准确性。强烈建议在记录模型时始终包含输入示例。

自 MLflow 2.16.0 起，记录带有输入示例的模型时，模型的工件目录中会保存两个文件

  • input_example.json：JSON 格式的输入示例。
  • serving_input_example.json：JSON 格式的输入示例，并经过额外转换，使其模式与查询部署的模型 REST 端点兼容。

以下示例演示了两个文件之间的区别

import mlflow


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        return model_input


with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model=MyModel(),
        artifact_path="model",
        input_example={"question": "What is MLflow?"},
    )

MLflow 记录的示例文件

文件名内容解释
input_example.json
{
  "question": "What is MLflow?"
}
原始格式的输入示例。
serving_input_example.json
{
  "inputs": {
    "question": "What is MLflow?"
  }
}
输入示例的 JSON 序列化版本，带有 MLflow 评分服务器在查询已部署的模型端点时所需的预定义键之一（dataframe_split、instances、inputs 或 dataframe_records）。
注意

在 MLflow 2.16.0 之前，保存时字典输入示例会被转换为 Pandas DataFrame 格式。在后续版本中，输入示例直接以其 JSON 序列化格式保存。对于 Pandas DataFrame，它会使用 to_dict(orient='split') 转换为字典格式并保存为 JSON 格式。langchain、openai、pyfunc 和 transformers 风味的 example_no_conversion 参数不再使用，可以安全移除，它将在未来版本中删除。

与模型签名类似，模型输入可以是基于列的（即 DataFrames）、基于张量的（即 numpy.ndarrays）或 JSON 对象（即 Python 字典）。我们支持使用元组将模型输入和参数组合，以提供带有参数的 input_example。请参见下面的示例

如何记录带基于列示例的模型​

对于接受基于列输入的模型，示例可以是一条记录或一批记录。示例输入可以是以下格式

  • Pandas DataFrame

给定的示例将使用 Pandas 的 split 格式序列化为 JSON。字节会进行 base64 编码。以下示例演示了如何使用模型的基于列的输入示例进行记录

import pandas as pd

input_example = pd.DataFrame(
    [
        {
            "sepal length (cm)": 5.1,
            "sepal width (cm)": 3.5,
            "petal length (cm)": 1.4,
            "petal width (cm)": 0.2,
        }
    ]
)
mlflow.sklearn.log_model(..., input_example=input_example)

如何记录带基于张量示例的模型​

对于接受基于张量输入的模型，示例必须是一批输入。默认情况下，轴 0 是批量轴，除非模型签名中另有指定。示例输入可以以下任何格式传入

  • numpy ndarray
  • Python dict 映射字符串到 numpy 数组
  • Scipy csr_matrix (稀疏矩阵)
  • Scipy csc_matrix (稀疏矩阵)

以下示例演示了如何使用模型的基于张量的输入示例进行记录

# each input has shape (4, 4)
input_example = np.array(
    [
        [[0, 0, 0, 0], [0, 134, 25, 56], [253, 242, 195, 6], [0, 93, 82, 82]],
        [[0, 23, 46, 0], [33, 13, 36, 166], [76, 75, 0, 255], [33, 44, 11, 82]],
    ],
    dtype=np.uint8,
)
mlflow.tensorflow.log_model(..., input_example=input_example)

如何使用 JSON 对象示例记录模型​

如果输入示例是 JSON 可序列化格式，我们支持按原样保存。输入示例可以是以下格式

  • dict（标量、字符串或标量值列表）
  • 列表
  • 标量

以下示例演示了如何使用模型的 JSON 对象输入示例进行记录

input_example = {
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "assistant", "content": "What would you like to ask?"},
        {"role": "user", "content": "Who owns MLflow?"},
    ]
}
mlflow.langchain.log_model(..., input_example=input_example)

如何记录包含参数的示例模型​

对于在推理期间需要额外参数的模型，您可以在保存或记录模型时包含包含参数的 input_example。为此，应将示例输入作为 tuple 提供。元组的第一个元素是输入数据示例，第二个元素是参数的 dict。有效参数的完整列表记录在模型推理参数部分。

  • Python tuple: (input_data, params)

以下示例演示了如何记录包含参数示例的模型

# input_example could be column-based or tensor-based example as shown above
# params must be a valid dictionary of params
input_data = "Hello, Dolly!"
params = {"temperature": 0.5, "top_k": 1}
input_example = (input_data, params)
mlflow.transformers.log_model(..., input_example=input_example)

模型服务载荷示例​

一旦 MLflow 模型部署到 REST 端点进行推理，请求载荷将以 JSON 序列化，并且可能与内存中的表示存在细微差异。为了验证您的模型适用于推理，您可以使用 serving_input_example.json 文件。提供 input_example 时，该文件会自动随模型一起记录，并包含给定输入示例的 JSON 格式，用于查询部署的模型端点。

以下示例演示了如何从已记录的模型中加载服务载荷

import mlflow
from mlflow.models.utils import load_serving_example

input_example = {
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "assistant", "content": "What would you like to ask?"},
        {"role": "user", "content": "Who owns MLflow?"},
    ]
}
model_info = mlflow.langchain.log_model(..., input_example=input_example)
print(f"model_uri: {model_info.model_uri}")
serving_example = load_serving_example(model_info.model_uri)
print(f"serving_example: {serving_example}")

您可以在服务之前验证输入示例是否有效

from mlflow.models import validate_serving_input

result = validate_serving_input(model_info.model_uri, serving_example)
print(f"prediction result: {result}")

在本地提供模型服务

mlflow models serve --model-uri "<YOUR_MODEL_URI>"

使用服务载荷示例验证模型推理

curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d 'YOUR_SERVING_EXAMPLE'