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

模型签名和输入示例简介

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

模型签名

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

模型输入示例

补充模型签名,模型输入示例 提供了一个具体的实例,展示了有效的模型输入是什么样子的。这个实践示例对开发者来说是无价的,提供了对所需数据格式和结构的清晰理解,以便进行有效的模型测试和实际应用。

为什么这些很重要?

../../_images/signature-vs-no-signature.png

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

模型签名

在MLflow中,模型签名精确地定义了模型输入、输出以及任何其他有效模型操作所需的附加参数的架构。这个定义作为一个统一的接口,指导用户正确和准确地使用他们的模型。模型签名是MLflow生态系统的核心部分,使得MLflow跟踪UI和模型注册UI能够清晰地显示模型所需的输入、输出和参数。此外,MLflow模型部署工具 利用这些签名来确保推理时使用的数据与模型建立的规范一致,从而保持模型的完整性和性能。有关这些签名如何强制执行数据准确性的更多见解,请参阅 签名强制执行 部分。

在 MLflow 中,将签名嵌入到您的模型中是一个直接的过程。当使用 sklearn.log_model() 等函数来记录或保存模型时,只需包含一个 模型输入示例。此操作使 MLflow 能够 自动推断模型的签名。有关此过程的详细说明,请参阅 如何记录带有签名的模型 部分。推断的签名以及其他重要的模型元数据,以 JSON 格式存储在您的模型工件的 MLmodel 文件 中。如果需要为已经记录或保存的模型添加签名,可以使用 set_signature() API。有关实现此功能的详细指南,请参阅 如何在模型上设置签名 部分。

模型签名组件

MLflow 中模型签名的结构由三种不同的模式类型组成:(1) 输入(2) 输出**和**(3) 参数 (params)**输入**和**输出**模式分别指定了模型期望接收和生成的数据结构。这些可以定制以处理各种数据格式,包括列数据和张量,以及数组和对象类型的输入,以满足不同模型的多样化需求。

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

备注

MLflow 版本 2.10.0 及以后,模型签名中处理对象和数组的能力被引入。在 2.10.0 之前的版本中,基于列的签名仅限于标量输入类型和某些特定于列表和字典输入的条件类型,主要支持 transformers 风格。后期版本的这一增强功能显著扩大了 MLflow 可以无缝容纳的数据类型和结构的范畴。

签名游乐场

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

或者,如果你想在本地下载笔记本并用自己的数据结构进行测试,你可以从下面下载。

Download the Signature Playground Notebook

必填与可选输入字段

在定义模型签名时,有一些条件决定了签名的强制执行,这些条件是需要考虑的重要因素。其中最值得注意的是输入数据中必需与可选的概念。

必需字段是指在输入数据中必须存在的字段,以便模型能够进行预测。如果缺少必需字段,签名强制验证将引发异常,指出缺少必需的输入字段。

为了将一个字段配置为可选,在使用 mlflow.models.infer_signature() 函数时,必须为该字段传入 Nonenp.nan 的值。或者,您可以手动定义签名并将该字段的 required 字段设置为 false

模型签名类型

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

基于列的签名

基于列的签名通常用于需要表格数据输入的传统机器学习模型中,例如 Pandas DataFrame。这些签名由一系列列组成,每列可能具有名称和指定的数据类型。输入和输出中每列的类型对应于 支持的数据类型 之一,并且列可以可选地命名。此外,输入模式中的列可以被指定为 optional,表示它们在模型输入中的包含不是强制性的,必要时可以省略(更多详情请参见 可选列)。

支持的数据类型

基于列的签名支持在 MLflow 数据类型 规范中定义的数据原语:

  • 字符串

  • 整数 1

  • 浮动

  • 布尔

  • datetime

输入 (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 中潜在的架构强制错误。为了避免此类问题,特别是在使用 Python 进行模型服务和 Spark 部署时,请将包含缺失值的整数列定义为双精度浮点数(float64)。

基于列的签名也支持这些基本类型的复合数据类型。

  • 数组(列表,numpy 数组)

  • Spark ML 向量(它继承了 Array[double] 类型)

  • 对象 (字典)

警告

  • 在 MLflow 版本 2.10.0 中引入了对数组和对象类型的支持。这些类型在 MLflow 的早期版本中将无法识别。如果你正在保存一个使用这些签名类型的模型,你应该确保任何尝试加载这些模型的环境都安装了至少为 2.10.0 版本的 MLflow。

  • Spark ML 向量类型的支持在 MLflow 版本 2.15.0 中引入,这些类型在 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

可选列

输入数据中包含 Nonenp.nan 值的列将被推断为 可选`(即 `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

备注

嵌套数组可以包含一个空列表,但这并不会使列变为 可选 ,因为它代表一个空集(∅)。在这种情况下,模式将从列表的其他元素推断,假设它们具有同质类型。如果你想使一列可选,请传递 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 数据类型 <https://numpy.org/devdocs/user/basics.types.html>`_ 对齐),一个 **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

备注

基于张量的模式不支持可选输入。你可以传递一个包含 Nonenp.nan 值的数组,但模式不会将它们标记为可选。

带有推理参数的模型签名

推理参数(或称为’params’)是在推理阶段传递给模型的额外设置。常见的例子包括语言学习模型(LLMs)中的``temperature``和``max_length``等参数。这些参数在训练过程中通常不是必需的,但在推理时对模型的行为调整起着至关重要的作用。随着基础模型的使用,这种配置变得越来越重要,因为同一个模型在不同的推理场景中可能需要不同的参数设置。

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

要在推理时利用参数,它们必须被纳入 模型签名 。参数的架构被定义为一个 ParamSpec 元素序列,每个元素包括:

  • 名称: 参数的标识符,例如 temperature

  • 类型: 参数的数据类型,必须与 支持的数据类型 之一相匹配。

  • 默认值: 参数的默认值,确保在没有提供特定值时有一个回退选项。

  • shape: 参数的形状,通常标量值为 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, ctx, 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 数据类型 的值,包括这些数据类型的一维列表。目前,MLflow 仅支持参数的一维列表。

备注

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

签名强制执行

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

../../_images/signature-enforcement.png

名称排序强制执行

在 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将因无法将浮点数转换为整数而引发错误。为了缓解这个问题,特别是由于MLflow使用Python进行模型服务和Spark部署,建议将包含缺失值的整数列定义为双精度浮点数(float64)。

处理日期和时间戳

Python 的 datetime 类型自带内置精度,例如 datetime64[D] 表示天精度,datetime64[ns] 表示纳秒精度。尽管这种精度细节在基于列的模型签名中被忽略,但在基于张量的签名中是强制执行的。

处理不规则数组

numpy 中的不规则数组,其特征是形状为 (-1,) 且 dtype 为对象,在使用 infer_signature 时会自动管理。这会导致签名类似于 Tensor('object', (-1,))。为了更详细的表示,可以手动创建签名以反映不规则数组的特定性质,例如 Tensor('float64', (-1, -1, -1, 3))。然后根据签名中的详细信息应用执行,以适应不规则输入数组。

如何使用签名记录模型

在 MLflow 中为您的模型添加签名非常简单。只需在调用 log_model 或 save_model 函数时提供一个 模型输入示例 ,例如使用 sklearn.log_model() 。MLflow 将根据此输入示例和模型对该示例的预测输出自动推断模型的签名。

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

备注

模型签名在 MLflow 模型部署工具 中起着至关重要的作用,特别是在以 Python 函数 (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 风格的模型签名示例

GenAI 风格,如 langchain、OpenAI 和 transformers,通常需要基于对象(字典)的模型签名。手动定义结构复杂且容易出错。相反,您应该依赖 infer_signature 方法,根据输入示例自动生成模型签名。

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

备注

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

以 langchain 模型为例,你可以通过代码功能中的模型,使用两个简单的步骤记录一个带有签名的模型:

  1. 在名为 langchain_model.py 的单独 Python 文件中定义 langchain 模型:

要成功记录一个 langchain 模型,强烈建议使用 代码中的模型 功能。

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 flavor 重新加载时是否正常工作,并生成相应的 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))
  1. 加载模型并进行预测:

loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
result = loaded_model.predict(input_example)

print(result)

A 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 变量以指向模型的本地目录,从而轻松地为本地保存的鸢尾花模型设置签名。

在注册模型上设置签名

由于 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评分服务器在`查询已部署模型端点 <../../deployment/deploy-model-locally.html#local-inference-server-spec>`_时所需的预定义键之一(dataframe_splitinstancesinputsdataframe_records)。

备注

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

与模型签名类似,模型输入可以是基于列的(即DataFrame)、基于张量的(即numpy.ndarrays)或json对象(即python字典)。我们通过使用元组来组合模型输入和参数,从而支持带有参数的input_example。请参见以下示例:

如何使用基于列的示例记录模型

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

  • Pandas DataFrame

给定的示例将使用 Pandas 的面向分割格式序列化为 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 (标量、字符串或标量值列表)

  • list

  • scalars

以下示例展示了如何使用您的模型记录一个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'