mlflow.transformers

MLflow 模块,提供对 HuggingFace/transformer 的支持。

mlflow.transformers.autolog(log_input_examples=False, log_model_signatures=False, log_models=False, log_datasets=False, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False, extra_tags=None)[source]

注意

自动记录已知与以下软件包版本兼容: 4.38.2 <= transformers <= 4.55.0。在该范围之外的包版本中使用时,自动记录可能无法成功。

此 autologging 集成仅用于禁用在基于 transformers 的模型训练和评估期间创建的无关子模型的错误 autologging。autologging 功能尚未在 transformers flavor 中完全实现。

mlflow.transformers.generate_signature_output(pipeline, data, model_config=None, params=None, flavor_config=None)[source]

用于生成响应输出以便提取用于模型保存和记录的输出签名。此函数模拟将已保存的模型或管道加载为 pyfunc 模型,而无需写入磁盘。

Parameters
  • pipeline – 一个 transformers pipeline 对象。请注意,不允许使用组件级或模型级的输入来提取输出示例。

  • data – 一个与给定的管道兼容的示例输入

  • model_config – 任何额外的模型配置,通过 kwargs 提供,用于指示来自 pipeline 推理调用的输出类型的格式。

  • params – 一个字典,包含要传递给 pipeline 用于推理的额外参数。

  • flavor_config – 模型的 flavor 配置。

Returns

来自 pyfunc 管道包装器的 predict 方法的输出

mlflow.transformers.get_default_conda_env(model)[source]
Returns

用于使用 transformers flavor 生成的 MLflow Models 的默认 Conda 环境,基于要记录的模型的模型实例框架类型。

mlflow.transformers.get_default_pip_requirements(model) list[str][source]
Parameters

model – 要保存的模型实例,以提供所需的底层深度学习执行框架依赖。请注意,这必须是实际的模型实例,而不是 Pipeline。

Returns

列出由 transformers flavor 生成的 MLflow 模型的默认 pip 依赖项。对 save_model()log_model() 的调用会至少在 pip 环境中包含这些依赖项。

mlflow.transformers.is_gpu_available()[source]
mlflow.transformers.load_model(model_uri: str, dst_path: Optional[str] = None, return_type='pipeline', device=None, **kwargs)[source]

注意

‘transformers’ MLflow Models 集成已知与 4.38.2 <= transformers <= 4.55.0 兼容。MLflow Models 与 transformers 的集成在与不在此版本范围内的包版本一起使用时可能无法成功。

从本地文件或运行中加载一个transformers对象。

Parameters
  • model_uri

    MLflow 模型的 URI 格式位置。例如:

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs:/<mlflow_run_id>/run-relative/path/to/model

    • mlflow-artifacts:/path/to/model

    有关受支持的 URI 方案的更多信息,请参见 Referencing Artifacts.

  • dst_path – 用于下载模型工件的本地文件系统路径。此目录在提供时必须已存在。若未指定,将创建一个本地输出路径。

  • return_type

    为已存储的 transformers 对象指定返回类型修饰符。 如果设置为 “components”,返回类型将是保存的 Pipeline 或预训练模型的各个独立组件的字典。 针对以 NLP 为中心的模型,组件通常会包含如下所示的返回表示(下面以文本分类为例):

    {"model": BertForSequenceClassification, "tokenizer": BertTokenizerFast}
    

    视觉模型将返回适当类型的 ImageProcessor 实例,而多模态模型则会连同模型一起返回 FeatureExtractorTokenizer。 将返回类型设置为 “components” 对于某些模型类型很有用,这些模型在某些用例中可能没有所需的 pipeline 返回类型。 如果设置为 “pipeline”,模型以及任何和所有所需的 TokenizerFeatureExtractorProcessorImageProcessor 对象将作为适当类型的 Pipeline 对象返回,该类型由模型实例类型所设置的 task 决定。要覆盖此行为,请在模型记录或保存时提供有效的 task 参数。默认值是 “pipeline”。

  • device – 要将模型加载到的设备。默认是 None。使用 0 将模型加载到默认 GPU。

  • kwargs – 可选的配置选项,用于加载一个 transformers 对象。有关参数及其用法的信息,请参见 transformers documentation

Returns

一个 transformers 模型实例或组件的字典

mlflow.transformers.log_model(transformers_model, artifact_path: str | None = None, processor=None, task: str | None = None, torch_dtype: torch.dtype | None = None, model_card=None, code_paths: list[str] | None = None, registered_model_name: str | None = None, signature: ModelSignature | None = None, input_example: ModelInputExample | None = None, await_registration_for=300, pip_requirements: list[str] | str | None = None, extra_pip_requirements: list[str] | str | None = None, conda_env=None, metadata: dict[str, Any] | None = None, model_config: dict[str, Any] | None = None, prompt_template: str | None = None, save_pretrained: bool = True, prompts: list[str | Prompt] | None = None, name: str | None = None, params: dict[str, Any] | None = None, tags: dict[str, Any] | None = None, model_type: str | None = None, step: int = 0, model_id: str | None = None, **kwargs)[source]

注意

‘transformers’ 的 MLflow Models 集成已知与 4.38.2 <= transformers <= 4.55.0 兼容。MLflow Models 与 transformers 的集成在使用超出此范围的包版本时可能无法成功。

将一个 transformers 对象记录为当前运行的 MLflow 工件。请注意,使用自定义代码记录 transformers 模型(即需要 trust_remote_code=True 的模型)需要 transformers >= 4.26.0

Parameters
  • transformers_model

    The transformers model to save. This can be one of the following format:

    1. A transformers Pipeline instance.

    2. A dictionary that maps required components of a pipeline to the named keys

      of [“model”, “image_processor”, “tokenizer”, “feature_extractor”]. The model key in the dictionary must map to a value that inherits from PreTrainedModel, TFPreTrainedModel, or FlaxPreTrainedModel. All other component entries in the dictionary must support the defined task type that is associated with the base model type configuration.

    3. A string that represents a path to a local/DBFS directory containing a model

      checkpoint. The directory must contain a config.json file that is required for loading the transformers model. This is particularly useful when logging a model that cannot be loaded into memory for serialization.

    An example of specifying a Pipeline from a default pipeline instantiation:

    from transformers import pipeline
    
    qa_pipe = pipeline("question-answering", "csarron/mobilebert-uncased-squad-v2")
    
    with mlflow.start_run():
        mlflow.transformers.log_model(
            transformers_model=qa_pipe,
            name="model",
        )
    

    An example of specifying component-level parts of a transformers model is shown below:

    from transformers import MobileBertForQuestionAnswering, AutoTokenizer
    
    architecture = "csarron/mobilebert-uncased-squad-v2"
    tokenizer = AutoTokenizer.from_pretrained(architecture)
    model = MobileBertForQuestionAnswering.from_pretrained(architecture)
    
    with mlflow.start_run():
        components = {
            "model": model,
            "tokenizer": tokenizer,
        }
        mlflow.transformers.log_model(
            transformers_model=components,
            name="model",
        )
    

    An example of specifying a local checkpoint path is shown below:

    with mlflow.start_run():
        mlflow.transformers.log_model(
            transformers_model="path/to/local/checkpoint",
            name="model",
        )
    

  • artifact_path – 已弃用。请改用 name

  • processor

    一个可选的 Processor 子类对象。一些模型架构,尤其是多模态类型,会使用 Processors 在单个入口点中将文本编码与图像或音频编码结合起来。

    注意

    如果在记录模型时提供了 processor,模型将无法作为 Pipeline 加载或用于 pyfunc 推理。

  • task – 模型在 transformers 中的特定任务类型。这些字符串用于使能够创建具有适当内部调用架构的 pipeline,以满足给定模型的需求。如果未指定此参数,将使用 transformers 库中的 pipeline 工具来推断正确的任务类型。如果指定的值在当前安装的 transformers 版本中不是受支持的类型,则会抛出 Exception。

  • torch_dtype – 应用于模型加载回时的 Pytorch dtype。这在你想以与模型训练时的 dtype 不同的特定 dtype 保存模型时很有用。如果未指定,将使用模型实例的当前 dtype。

  • model_card

    来自 huggingface-hub 的可选 ModelCard 实例。 如果提供,模型卡的内容将与提供的 transformers_model 一并保存。 如果未提供,则会尝试从所提供的基础预训练模型(或包含在提供的 Pipeline 中的模型)中获取该卡片。

    注意

    若要获取 ModelCard(如果未提供),必须安装 huggingface_hub 包,且版本必须 >=0.10.0

  • code_paths

    本地文件系统中指向 Python 文件依赖(或包含文件依赖的目录)路径的列表。这些文件在模型加载时会被预先添加到系统路径中。如果为某个模型声明了依赖文件且多个文件之间存在导入依赖关系,那么这些文件应从一个共同的根路径声明相对导入,以避免在加载模型时发生导入错误。

    有关 code_paths 功能、推荐的使用模式和限制的详细说明,请参阅 code_paths usage guide

  • registered_model_name – 如果提供,则在 registered_model_name 下创建一个模型版本,如果不存在具有该名称的注册模型,则同时创建一个注册模型。

  • signature

    一个描述模型输入和输出 Schema 的 Model Signature 对象。模型签名可以使用 infer_signature 函数从 mlflow.models.signature 推断得出。

    示例
    from mlflow.models import infer_signature
    from mlflow.transformers import generate_signature_output
    from transformers import pipeline
    
    en_to_de = pipeline("translation_en_to_de")
    
    data = "MLflow is great!"
    output = generate_signature_output(en_to_de, data)
    signature = infer_signature(data, output)
    
    with mlflow.start_run() as run:
        mlflow.transformers.log_model(
            transformers_model=en_to_de,
            name="english_to_german_translator",
            signature=signature,
            input_example=data,
        )
    
    model_uri = f"runs:/{run.info.run_id}/english_to_german_translator"
    loaded = mlflow.pyfunc.load_model(model_uri)
    
    print(loaded.predict(data))
    # MLflow 很棒!
    

    如果提供了 input_example 而未提供 signature,当且仅当 pipeline 类型为基于文本的模型(NLP)时,系统会自动推断出一个 signature 并将其应用到 MLmodel 文件中。如果 pipeline 类型不是受支持的类型,则该推断功能将无法正确工作,并且会发出警告。为确保记录精确的 signature,建议显式提供一个。

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame,然后使用 Pandas 的 split-oriented 格式序列化为 json,或者转换为一个 numpy array,其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时,输入示例用于推断模型签名。

  • await_registration_for – 等待模型版本完成创建并进入 READY 状态的秒数。默认情况下,该函数等待五分钟。指定 0 或 None 以跳过等待。

  • pip_requirements – 可以是 pip 需求字符串的可迭代对象 (e.g. ["transformers", "-r requirements.txt", "-c constraints.txt"]) 或本地文件系统上 pip requirements 文件的字符串路径 (e.g. "requirements.txt")。如果提供,则描述了运行该模型所需的环境。如果 None,则会由 mlflow.models.infer_pip_requirements() 从当前软件环境推断出默认的依赖列表。如果依赖推断失败,则回退为使用 get_default_pip_requirements。requirements 和 constraints 会被自动解析并分别写入 requirements.txtconstraints.txt 文件,并作为模型的一部分存储。这些依赖还会写入模型的 conda 环境 (conda.yaml) 文件的 pip 部分。

  • extra_pip_requirements

    要么是一个 pip 需求字符串的可迭代对象(例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]),要么是本地文件系统上 pip requirements 文件的字符串路径(例如 "requirements.txt")。如果提供,该参数描述了附加的 pip 依赖,这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txtconstraints.txt 文件,并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境(conda.yaml)文件的 pip 部分。

    警告

    以下参数不能同时指定:

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirementsextra_pip_requirements 指定 pip 依赖。

  • conda_env

    可以是 Conda 环境的字典表示,或者是指向 conda 环境 yaml 文件的路径。如果提供,它描述了该模型应在其中运行的环境。至少,它应当指定包含在 get_default_conda_env() 中的依赖项。如果 None,则会向模型添加一个 conda 环境,该环境的 pip 依赖由 mlflow.models.infer_pip_requirements() 推断而来。如果依赖推断失败,则回退为使用 get_default_pip_requirements。从 conda_env 中的 pip 依赖会被写入 pip 的 requirements.txt 文件,完整的 conda 环境会写入 conda.yaml。下面是 conda 环境的一个 示例 字典表示:

    {
        "name": "mlflow-env",
        "channels": ["conda-forge"],
        "dependencies": [
            "python=3.8.15",
            {
                "pip": [
                    "transformers==x.y.z"
                ],
            },
        ],
    }
    

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • model_config

    一个可应用于推理期间 pipeline 实例的有效覆盖项字典。 这些参数仅用于将模型作为 pyfunc 模型加载或在 Spark 中使用的情况。这些值不会应用于从 调用 mlflow.transformers.load_model() 返回的 Pipeline。

    警告

    如果提供的键与所提供任务的 Pipeline 实例不兼容,或不是对模型中任何可用参数的有效覆盖,则在运行时会引发异常。 在保存或记录之前验证此字典中的条目以确保它们有效非常重要。

    为问题生成模型提供覆盖项的示例:

    from transformers import pipeline, AutoTokenizer
    
    task = "text-generation"
    architecture = "gpt2"
    
    sentence_pipeline = pipeline(
        task=task,
        tokenizer=AutoTokenizer.from_pretrained(architecture),
        model=architecture,
    )
    
    # 验证覆盖项是否生效
    prompts = ["Generative models are", "I'd like a coconut so that I can"]
    
    # 在保存或记录之前对配置进行验证
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    
    # 验证不会抛出异常
    sentence_pipeline(prompts, **model_config)
    
    with mlflow.start_run():
        mlflow.transformers.log_model(
            transformers_model=sentence_pipeline,
            name="my_sentence_generator",
            task=task,
            model_config=model_config,
        )
    

  • prompt_template

    一个字符串,如果提供,将用于在推理之前格式化用户的输入。该字符串应包含一个占位符,{prompt},该占位符将被用户的输入替换。例如:"Answer the following question. Q: {prompt} A:".

    目前仅支持以下管道类型:

  • save_pretrained

    如果设置为 False,MLflow 将不会保存 Transformer 模型的权重文件,而是仅保存对 HuggingFace Hub 模型仓库及其提交哈希的引用。当你从 HuggingFace Hub 加载预训练模型并希望在不修改模型权重的情况下将其记录或保存到 MLflow 时,这很有用。在这种情况下,将此标志指定为 False 可以节省存储空间并减少保存模型的时间。有关更详细的用法,请参阅 Storage-Efficient Model Logging

    警告

    如果模型以 save_pretrained 设置为 False 的方式保存,则该模型无法注册到 MLflow Model Registry。为了将模型转换为可以注册的版本,你可以使用 mlflow.transformers.persist_pretrained_model() 从 HuggingFace Hub 下载模型权重并将其保存到现有的模型 artifacts。有关更详细的用法,请参阅 Transformers flavor documentation

    import mlflow.transformers
    
    model_uri = "YOUR_MODEL_URI_LOGGED_WITH_SAVE_PRETRAINED_FALSE"
    model = mlflow.transformers.persist_pretrained_model(model_uri)
    mlflow.register_model(model_uri, "model_name")
    

    重要

    当你保存 PEFT 模型时,MLflow 会将 save_pretrained 标志覆盖为 False,并且只存储 PEFT 适配器的权重。基础模型权重不会被保存,而是记录对 HuggingFace 仓库及其提交哈希的引用。

  • prompts

    在 MLflow Prompt Registry 中注册并与模型关联的 prompt URI 列表。 Each prompt URI should be in the form prompt://. 这些 prompt 应在与模型关联之前在 MLflow Prompt Registry 中注册。

    这将在模型与 prompt 之间创建相互链接。相关的 prompts 可以在存储于 MLmodel 文件中的模型元数据中看到。您也可以从 Prompt Registry 的 UI 导航到该模型。

    import mlflow
    
    prompt_template = "Hi, {name}! How are you doing today?"
    
    # 在 MLflow Prompt Registry 中注册一个 prompt
    mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")
    
    # 使用已注册的 prompt 记录模型
    with mlflow.start_run():
        model_info = mlflow.pyfunc.log_model(
            name=MyModel(),
            name="model",
            prompts=["prompt:/my_prompt/1"]
        )
    
    print(model_info.prompts)
    # 输出: ['prompt:/my_prompt/1']
    
    # 加载 prompt
    prompt = mlflow.genai.load_prompt(model_info.prompts[0])
    

  • name – 模型名称。

  • params – 一个用于与模型一同记录的参数字典。

  • tags – 一个要与模型一起记录的标签字典。

  • model_type – 模型的类型。

  • step – 在该步记录模型输出和指标

  • model_id – 模型的 ID。

  • kwargs – 用于 mlflow.models.model.Model 的附加参数

mlflow.transformers.persist_pretrained_model(model_uri: str) None[source]

将 Transformers 预训练模型权重持久化到指定 model_uri 的 artifacts 目录。该 API 主要用于更新以 save_pretrained=False 设置记录或保存的 MLflow Model。由于这些模型的完整预训练模型权重在 artifacts 中缺失,因此无法注册到 Databricks Workspace Model Registry。以此模式保存的 Transformers 模型只存储对 HuggingFace Hub 仓库的引用。该 API 会从 HuggingFace Hub 仓库下载模型权重并将其保存到给定 model_uri 的 artifacts 中,以便该模型可以注册到 Databricks Workspace Model Registry。

Parameters

model_uri – 现有 MLflow Model(Transformers flavor)的 URI。它必须在记录/保存时使用 save_pretrained=False。

示例:

import mlflow

# Saving a model with save_pretrained=False
with mlflow.start_run() as run:
    model = pipeline("question-answering", "csarron/mobilebert-uncased-squad-v2")
    mlflow.transformers.log_model(
        transformers_model=model, name="pipeline", save_pretrained=False
    )

# The model cannot be registered to the Model Registry as it is
try:
    mlflow.register_model(f"runs:/{run.info.run_id}/pipeline", "qa_pipeline")
except MlflowException as e:
    print(e.message)

# Use this API to persist the pretrained model weights
mlflow.transformers.persist_pretrained_model(f"runs:/{run.info.run_id}/pipeline")

# Now the model can be registered to the Model Registry
mlflow.register_model(f"runs:/{run.info.run_id}/pipeline", "qa_pipeline")
mlflow.transformers.save_model(transformers_model, path: str, processor=None, task: str | None = None, torch_dtype: torch.dtype | None = None, model_card=None, code_paths: list[str] | None = None, mlflow_model: Model | None = None, signature: ModelSignature | None = None, input_example: ModelInputExample | None = None, pip_requirements: list[str] | str | None = None, extra_pip_requirements: list[str] | str | None = None, conda_env=None, metadata: dict[str, Any] | None = None, model_config: dict[str, Any] | None = None, prompt_template: str | None = None, save_pretrained: bool = True, **kwargs) None[source]

注意

‘transformers’ MLflow Models 集成已知与 4.38.2 <= transformers <= 4.55.0 兼容。MLflow Models 与 transformers 的集成在与不在此版本范围内的包版本一起使用时可能无法成功。

将训练好的 transformers 模型保存到本地文件系统的某个路径上。注意,使用自定义代码保存 transformers 模型(即需要 trust_remote_code=True 的模型)需要 transformers >= 4.26.0

Parameters
  • transformers_model

    The transformers model to save. This can be one of the following format:

    1. A transformers Pipeline instance.

    2. A dictionary that maps required components of a pipeline to the named keys

      of [“model”, “image_processor”, “tokenizer”, “feature_extractor”]. The model key in the dictionary must map to a value that inherits from PreTrainedModel, TFPreTrainedModel, or FlaxPreTrainedModel. All other component entries in the dictionary must support the defined task type that is associated with the base model type configuration.

    3. A string that represents a path to a local/DBFS directory containing a model

      checkpoint. The directory must contain a config.json file that is required for loading the transformers model. This is particularly useful when logging a model that cannot be loaded into memory for serialization.

    An example of specifying a Pipeline from a default pipeline instantiation:

    from transformers import pipeline
    
    qa_pipe = pipeline("question-answering", "csarron/mobilebert-uncased-squad-v2")
    
    with mlflow.start_run():
        mlflow.transformers.save_model(
            transformers_model=qa_pipe,
            path="path/to/save/model",
        )
    

    An example of specifying component-level parts of a transformers model is shown below:

    from transformers import MobileBertForQuestionAnswering, AutoTokenizer
    
    architecture = "csarron/mobilebert-uncased-squad-v2"
    tokenizer = AutoTokenizer.from_pretrained(architecture)
    model = MobileBertForQuestionAnswering.from_pretrained(architecture)
    
    with mlflow.start_run():
        components = {
            "model": model,
            "tokenizer": tokenizer,
        }
        mlflow.transformers.save_model(
            transformers_model=components,
            path="path/to/save/model",
        )
    

    An example of specifying a local checkpoint path is shown below:

    with mlflow.start_run():
        mlflow.transformers.save_model(
            transformers_model="path/to/local/checkpoint",
            path="path/to/save/model",
        )
    

  • path – 序列化模型要保存到的本地目标路径。

  • processor

    一个可选的 Processor 子类对象。某些模型架构,尤其是多模态类型,使用 Processors 在单一入口点结合文本编码与图像或音频编码。

    注意

    如果在保存模型时提供了 processor,模型将无法作为 Pipeline 加载,也无法用于 pyfunc 推理。

  • task – 模型的 transformers 特定任务类型,或 MLflow 推理任务类型。 如果提供了 transformers 特定的任务类型,将使用这些字符串,以便可以创建具有适当内部调用架构的 pipeline 来满足给定模型的需求。 如果该参数以推理任务类型提供或未指定,transformers 库中的 pipeline 实用程序将用于推断正确的任务类型。如果指定的值不是受支持的类型,将抛出一个 Exception。

  • torch_dtype – 在将模型加载回时应用于模型的 Pytorch dtype。这在你希望以与模型训练时不同的特定 dtype 保存模型时很有用。如果未指定,将使用模型实例的当前 dtype。

  • model_card

    可选的来自 huggingface-hubModelCard 实例。若提供,model card 的内容将与所提供的 transformers_model 一并保存。如未提供,将尝试从所提供的基础预训练模型(或位于所提供的 Pipeline 中的模型)获取该卡。

    注意

    若要获取 ModelCard(如果未提供),必须安装 huggingface_hub 包,且版本必须 >=0.10.0

  • code_paths

    本地文件系统中指向 Python 文件依赖(或包含文件依赖的目录)路径的列表。这些文件在模型加载时会被预先添加到系统路径中。如果为某个模型声明了依赖文件且多个文件之间存在导入依赖关系,那么这些文件应从一个共同的根路径声明相对导入,以避免在加载模型时发生导入错误。

    有关 code_paths 功能、推荐的使用模式和限制的详细说明,请参阅 code_paths usage guide

  • mlflow_model – 一个 MLflow 模型对象,用于指定要将此模型添加到的 flavor。

  • signature

    一个描述模型输入和输出 Schema 的 Model Signature 对象。模型签名可以使用 infer_signature 函数(位于 mlflow.models.signature)推断得出。

    示例
    from mlflow.models import infer_signature
    from mlflow.transformers import generate_signature_output
    from transformers import pipeline
    
    en_to_de = pipeline("translation_en_to_de")
    
    data = "MLflow is great!"
    output = generate_signature_output(en_to_de, data)
    signature = infer_signature(data, output)
    
    mlflow.transformers.save_model(
        transformers_model=en_to_de,
        path="/path/to/save/model",
        signature=signature,
        input_example=data,
    )
    
    loaded = mlflow.pyfunc.load_model("/path/to/save/model")
    print(loaded.predict(data))
    # MLflow 很棒!
    

    如果提供了 input_example 而没有提供 signature,则当且仅当管道类型为基于文本的模型(NLP)时,签名将被自动推断并应用到 MLmodel 文件。如果管道类型不是受支持的类型,该推断功能将无法正确工作,并且会发出警告。为了确保记录精确的签名,建议显式提供一个签名。

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame,然后使用 Pandas 的 split-oriented 格式序列化为 json,或者转换为一个 numpy array,其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时,输入示例用于推断模型签名。

  • pip_requirements – 可以是 pip 依赖字符串的可迭代对象(例如 ["transformers", "-r requirements.txt", "-c constraints.txt"])或本地文件系统中 pip requirements 文件的字符串路径(例如 "requirements.txt")。如果提供,它描述了应在其中运行此模型的环境。如果 None,则由 mlflow.models.infer_pip_requirements() 从当前软件环境推断出默认的依赖列表。如果依赖推断失败,则回退到使用 get_default_pip_requirements。依赖和约束会被自动解析并分别写入 requirements.txtconstraints.txt 文件,并作为模型的一部分存储。依赖也会被写入模型的 conda 环境(conda.yaml)文件的 pip 部分。

  • extra_pip_requirements

    要么是一个 pip 需求字符串的可迭代对象(例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]),要么是本地文件系统上 pip requirements 文件的字符串路径(例如 "requirements.txt")。如果提供,该参数描述了附加的 pip 依赖,这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txtconstraints.txt 文件,并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境(conda.yaml)文件的 pip 部分。

    警告

    以下参数不能同时指定:

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirementsextra_pip_requirements 指定 pip 依赖。

  • conda_env

    可以是 Conda 环境的字典表示,也可以是指向 conda 环境 yaml 文件的路径。如果提供,该字段描述了运行此模型所需的环境。至少,它应当指定 get_default_conda_env() 中包含的依赖项。如果 None,则会向模型添加一个 conda 环境,该环境的 pip 依赖由 mlflow.models.infer_pip_requirements() 推断得到。如果依赖推断失败,则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖会被写入 pip requirements.txt 文件,完整的 conda 环境会被写入 conda.yaml。下面是一个 示例 字典表示的 conda 环境:

    {
        "name": "mlflow-env",
        "channels": ["conda-forge"],
        "dependencies": [
            "python=3.8.15",
            {
                "pip": [
                    "transformers==x.y.z"
                ],
            },
        ],
    }
    

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • model_config

    一个在推理期间可以应用到 pipeline 实例的有效覆盖项字典。 这些参数仅用于将模型作为 pyfunc Model 加载的情况,或用于 Spark。 这些值不会应用到通过调用 mlflow.transformers.load_model() 返回的 Pipeline 上。

    警告

    如果提供的键与所提供任务的 Pipeline 实例不兼容,或不是 Model 中任何可用参数的有效覆盖项, 将在运行时引发 Exception。务必在保存或记录之前验证此字典中的条目以确保它们有效。

    为问题生成模型提供覆盖项的示例:

    from transformers import pipeline, AutoTokenizer
    
    task = "text-generation"
    architecture = "gpt2"
    
    sentence_pipeline = pipeline(
        task=task,
        tokenizer=AutoTokenizer.from_pretrained(architecture),
        model=architecture,
    )
    
    # 验证覆盖项是否生效
    prompts = ["Generative models are", "I'd like a coconut so that I can"]
    
    # 在保存或记录之前对配置进行验证
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    
    # 验证不会抛出异常
    sentence_pipeline(prompts, **model_config)
    
    mlflow.transformers.save_model(
        transformers_model=sentence_pipeline,
        path="/path/for/model",
        task=task,
        model_config=model_config,
    )
    

  • prompt_template

    如果提供,该字符串将在推理之前用于格式化用户的输入。该字符串应包含一个占位符,{prompt},该占位符将被用户的输入替换。例如: "Answer the following question. Q: {prompt} A:".

    当前仅支持以下 pipeline 类型:

  • save_pretrained

    如果设置为 False,MLflow 将不会保存 Transformer 模型的权重文件,而是只保存对 HuggingFace Hub 模型仓库及其提交哈希的引用。当你从 HuggingFace Hub 加载预训练模型并希望将其记录或保存到 MLflow 而不修改模型权重时,这会很有用。在这种情况下,将此标志指定为 False 可以节省存储空间并减少保存模型的时间。有关更详细的用法,请参阅 Storage-Efficient Model Logging

    警告

    如果模型在保存时将 save_pretrained 设置为 False,则该模型不能注册到 MLflow Model Registry。为了将模型转换为可以注册的形式,你可以使用 mlflow.transformers.persist_pretrained_model() 从 HuggingFace Hub 下载模型权重并将其保存到现有的模型工件中。有关更详细的用法,请参阅 Transformers flavor documentation

    import mlflow.transformers
    
    model_uri = "YOUR_MODEL_URI_LOGGED_WITH_SAVE_PRETRAINED_FALSE"
    model = mlflow.transformers.persist_pretrained_model(model_uri)
    mlflow.register_model(model_uri, "model_name")
    

    重要

    当你保存 PEFT 模型时,MLflow 会将 save_pretrained 标志覆盖为 False 并仅存储 PEFT 适配器权重。基础模型权重不会被保存,而是记录 HuggingFace 仓库的引用及其提交哈希。

  • kwargs – 可选的用于 transformers 序列化的额外配置。