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]

注意

Autologging 已知与以下软件包版本兼容:4.38.2 <= transformers <= 4.57.3。使用此范围之外的软件包版本进行 autologging 可能会失败。

此 autologging 集成仅用于禁用在训练和评估基于 transformer 的模型时创建的无关子模型的虚假 autologging。Transformers flavor 的 autologging 功能未完全实现。

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

用于生成响应输出来提取用于模型保存和日志记录的输出签名的实用程序。此函数模拟加载保存的模型或 pipeline 作为 pyfunc 模型,而无需写入磁盘。

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

  • data – 与给定 pipeline 兼容的示例输入

  • model_config – 任何其他模型配置,以 kwargs 的形式提供,用于告知 pipeline 推理调用的输出类型格式。

  • params – 要传递给 pipeline 进行推理的其他参数字典。

  • flavor_config – 模型的 flavor 配置。

返回

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

mlflow.transformers.get_default_conda_env(model)[source]
返回

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

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

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

返回

MLflow Models 使用 transformers flavor 生成的模型列表。调用 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.57.3。使用此范围之外的软件包版本进行 transformers 的 MLflow Models 集成可能会失败。

从本地文件或 run 加载 transformers 对象。

参数
  • 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 方案的更多信息,请参阅 引用 Artifacts

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

  • return_type

    用于存储的 transformers 对象的返回类型修改器。如果设置为“components”,则返回类型将是 Pipeline 或预训练模型的已保存单个组件的字典。面向 NLP 的模型组件通常会返回如下所示的文本分类示例的返回表示。

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

    视觉模型将返回相应类型的 ImageProcessor 实例,而多模态模型将同时返回 FeatureExtractorTokenizer 以及模型。返回“components”对于某些不具有所需 pipeline 返回类型的模型类型在某些用例中可能很有用。如果设置为“pipeline”,则模型以及任何必需的 TokenizerFeatureExtractorProcessorImageProcessor 对象将包含在由模型实例类型设置的 task 定义的相应类型的 Pipeline 对象中。要覆盖此行为,请在模型日志记录或保存期间提供有效的 task 参数。默认为“pipeline”。

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

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

返回

一个 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.57.3。使用此范围之外的软件包版本进行 transformers 的 MLflow Models 集成可能会失败。

transformers 对象记录为当前 run 的 MLflow artifact。请注意,使用自定义代码记录 transformers 模型(即需要 trust_remote_code=True 的模型)需要 transformers >= 4.26.0

参数
  • transformers_model

    要保存的 transformers 模型。它可以是以下格式之一:

    1. 一个 transformers Pipeline 实例。

    2. 一个将 pipeline 所需组件映射到命名键的字典

      的 ["model", "image_processor", "tokenizer", "feature_extractor"]。字典中的 model 键必须映射到继承自 PreTrainedModelTFPreTrainedModelFlaxPreTrainedModel 的值。字典中的所有其他组件条目都必须支持与基础模型类型配置关联的任务类型。

    3. 一个表示包含模型检查点的本地/DBFS 目录路径的字符串

      该目录必须包含加载 transformers 模型所需的 config.json 文件。这在记录无法加载到内存中进行序列化的模型时特别有用。

    从默认 pipeline 实例化指定 Pipeline 的示例

    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",
        )
    

    指定 transformer 模型组件级部分的示例

    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",
        )
    

    指定本地检查点路径的示例

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

  • artifact_path – Deprecated. Use name instead.

  • processor

    一个可选的 Processor 子类对象。某些模型架构,特别是多模态类型,会使用 Processors 将文本编码和图像或音频编码合并到一个入口点。

    注意

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

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

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

  • model_card

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

    注意

    为了能够获取 ModelCard(如果未提供),必须安装 huggingface_hub 包,并且版本必须是 >=0.10.0

  • code_paths

    A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

    For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.

  • registered_model_name – 如果提供,则在 registered_model_name 下创建一个模型版本,如果给定名称的注册模型不存在,也会创建该注册模型。

  • signature

    一个描述模型输入和输出 Schema 的模型签名对象。可以使用 mlflow.models.signatureinfer_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 ist großartig!
    

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

  • input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame,然后使用 Pandas 的面向拆分(split-oriented)格式序列化为 json,或者转换为 numpy 数组,其中示例将通过转换为列表来序列化为 json。字节将进行 base64 编码。当 signature 参数为 None 时,输入示例用于推断模型签名。

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

  • pip_requirements – pip 需求字符串的可迭代对象(例如 ["transformers", "-r requirements.txt", "-c constraints.txt"])或本地文件系统上 pip 需求文件的字符串路径(例如 "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.txt")。如果提供,这将描述附加到根据用户当前软件环境自动生成的默认 pip 需求集中的其他 pip 需求。需求和约束都会被自动解析并分别写入 requirements.txtconstraints.txt 文件,并作为模型的一部分存储。需求也会被写入模型 Conda 环境(conda.yaml)文件的 pip 部分。

    警告

    以下参数不能同时指定

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    此示例演示了如何使用 pip_requirementsextra_pip_requirements 指定 pip requirements。

  • conda_env

    Conda 环境的字典表示形式或 Conda 环境 yaml 文件的路径。如果提供,这将描述此模型应运行的环境。至少,它应该指定 get_default_conda_env() 中包含的依赖项。如果为 None,则会将一个带有 mlflow.models.infer_pip_requirements() 推断出的 pip 需求的 Conda 环境添加到模型中。如果需求推断失败,则回退使用 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 中任何可用参数的有效覆盖,则会在运行时引发异常。在保存或记录之前,验证此字典中的条目以确保它们有效非常重要。

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

    from transformers import pipeline, AutoTokenizer
    
    task = "text-generation"
    architecture = "gpt2"
    
    sentence_pipeline = pipeline(
        task=task,
        tokenizer=AutoTokenizer.from_pretrained(architecture),
        model=architecture,
    )
    
    # Validate that the overrides function
    prompts = ["Generative models are", "I'd like a coconut so that I can"]
    
    # validation of config prior to save or log
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    
    # Verify that no exceptions are thrown
    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:"

    目前,仅支持以下 pipeline 类型:

    以下示例显示了如何使用 prompt 模板记录 text-generation pipeline,并通过 python_function (pyfunc) flavor 来使用它。

    import mlflow
    from transformers import pipeline
    
    # Initialize a text-generation pipeline
    generator = pipeline("text-generation", model="gpt2")
    
    # Define a prompt template. The ``{prompt}`` placeholder will be replaced
    # with the raw user input at inference time.
    prompt_template = "Answer the following question concisely.\n\nQ: {prompt}\nA:"
    
    example_prompt = "What is MLflow?"
    
    # Log the model with the prompt template and an input example
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            transformers_model=generator,
            name="qa_text_generator",
            prompt_template=prompt_template,
            input_example=example_prompt,
        )
    
    # Load the model back as a pyfunc model
    loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
    
    # The input to ``predict`` is the raw question string; the prompt template
    # is applied internally before calling the underlying transformers pipeline.
    loaded_model.predict("What is experiment tracking?")
    

  • save_pretrained

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

    警告

    如果模型以 save_pretrained 设置为 False 的方式保存,则模型无法注册到 MLflow Model Registry。为了将模型转换为可注册的模型,您可以使用 mlflow.transformers.persist_pretrained_model() 从 HuggingFace Hub 下载模型权重并将其保存在现有模型 artifact 中。有关更详细的使用方法,请参阅 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 存储库及其 commit hash 的引用。

  • prompts

    与模型关联的 MLflow 提示注册表中已注册的提示 URI 列表。每个提示 URI 格式应为 prompt:/<name>/<version>。应在将提示与模型关联之前在 MLflow 提示注册表中注册这些提示。

    这将创建模型和提示之间的双向链接。关联的提示可以在 MLmodel 文件中存储的模型元数据中看到。从提示注册表 UI,您也可以导航到模型。

    import mlflow
    
    prompt_template = "Hi, {name}! How are you doing today?"
    
    # Register a prompt in the MLflow Prompt Registry
    mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")
    
    # Log a model with the registered prompt
    with mlflow.start_run():
        model_info = mlflow.pyfunc.log_model(
            name=MyModel(),
            name="model",
            prompts=["prompt:/my_prompt/1"]
        )
    
    print(model_info.prompts)
    # Output: ['prompt:/my_prompt/1']
    
    # Load the prompt
    prompt = mlflow.genai.load_prompt(model_info.prompts[0])
    

  • name – 模型名称。

  • params – 要与模型一起记录的参数字典。

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

  • model_type – 模型的类型。

  • step – 记录模型输出和指标的步骤

  • model_id – 模型的 ID。

  • kwargsmlflow.models.model.Model 的其他参数

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

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

参数

model_uri – Transformers flavor 的现有 MLflow 模型的 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.57.3。使用此范围之外的软件包版本进行 transformers 的 MLflow Models 集成可能会失败。

将经过训练的 Transformers 模型保存到本地文件系统上的路径。请注意,保存具有自定义代码的 Transformers 模型(即需要 trust_remote_code=True 的模型)需要 transformers >= 4.26.0

参数
  • transformers_model

    要保存的 transformers 模型。它可以是以下格式之一:

    1. 一个 transformers Pipeline 实例。

    2. 一个将 pipeline 所需组件映射到命名键的字典

      的 ["model", "image_processor", "tokenizer", "feature_extractor"]。字典中的 model 键必须映射到继承自 PreTrainedModelTFPreTrainedModelFlaxPreTrainedModel 的值。字典中的所有其他组件条目都必须支持与基础模型类型配置关联的任务类型。

    3. 一个表示包含模型检查点的本地/DBFS 目录路径的字符串

      该目录必须包含加载 transformers 模型所需的 config.json 文件。这在记录无法加载到内存中进行序列化的模型时特别有用。

    从默认 pipeline 实例化指定 Pipeline 的示例

    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",
        )
    

    指定 transformer 模型组件级部分的示例

    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",
        )
    

    指定本地检查点路径的示例

    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 工具来推断正确的任务类型。如果指定的值不是受支持的类型,将抛出异常。

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

  • model_card

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

    注意

    为了能够获取 ModelCard(如果未提供),必须安装 huggingface_hub 包,并且版本必须是 >=0.10.0

  • code_paths

    A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

    For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.

  • mlflow_model – 一个 MLflow 模型对象,指定该模型正在添加到的 flavor。

  • signature

    一个描述模型输入和输出 Schema 的模型签名对象。可以使用 mlflow.models.signatureinfer_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 ist großartig!
    

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

  • input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame,然后使用 Pandas 的面向拆分(split-oriented)格式序列化为 json,或者转换为 numpy 数组,其中示例将通过转换为列表来序列化为 json。字节将进行 base64 编码。当 signature 参数为 None 时,输入示例用于推断模型签名。

  • pip_requirements – pip 需求字符串的可迭代对象(例如 ["transformers", "-r requirements.txt", "-c constraints.txt"])或本地文件系统上 pip 需求文件的字符串路径(例如 "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.txt")。如果提供,这将描述附加到根据用户当前软件环境自动生成的默认 pip 需求集中的其他 pip 需求。需求和约束都会被自动解析并分别写入 requirements.txtconstraints.txt 文件,并作为模型的一部分存储。需求也会被写入模型 Conda 环境(conda.yaml)文件的 pip 部分。

    警告

    以下参数不能同时指定

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    此示例演示了如何使用 pip_requirementsextra_pip_requirements 指定 pip requirements。

  • conda_env

    Conda 环境的字典表示形式或 Conda 环境 yaml 文件的路径。如果提供,这将描述此模型应运行的环境。至少,它应该指定 get_default_conda_env() 中包含的依赖项。如果为 None,则会将一个带有 mlflow.models.infer_pip_requirements() 推断出的 pip 需求的 Conda 环境添加到模型中。如果需求推断失败,则回退使用 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 中任何可用参数的有效覆盖,则会在运行时引发异常。在保存或记录之前,验证此字典中的条目以确保它们有效非常重要。

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

    from transformers import pipeline, AutoTokenizer
    
    task = "text-generation"
    architecture = "gpt2"
    
    sentence_pipeline = pipeline(
        task=task,
        tokenizer=AutoTokenizer.from_pretrained(architecture),
        model=architecture,
    )
    
    # Validate that the overrides function
    prompts = ["Generative models are", "I'd like a coconut so that I can"]
    
    # validation of config prior to save or log
    model_config = {
        "top_k": 2,
        "num_beams": 5,
        "max_length": 30,
        "temperature": 0.62,
        "top_p": 0.85,
        "repetition_penalty": 1.15,
    }
    
    # Verify that no exceptions are thrown
    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 类型:

    以下示例显示了如何使用 prompt 模板记录 text-generation pipeline,并通过 python_function (pyfunc) flavor 来使用它。

    import mlflow
    from transformers import pipeline
    
    # Initialize a text-generation pipeline
    generator = pipeline("text-generation", model="gpt2")
    
    # Define a prompt template. The ``{prompt}`` placeholder will be replaced
    # with the raw user input at inference time.
    prompt_template = "Answer the following question concisely.\n\nQ: {prompt}\nA:"
    
    example_prompt = "What is MLflow?"
    
    # Log the model with the prompt template and an input example
    with mlflow.start_run():
        model_info = mlflow.transformers.log_model(
            transformers_model=generator,
            name="qa_text_generator",
            prompt_template=prompt_template,
            input_example=example_prompt,
        )
    
    # Load the model back as a pyfunc model
    loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
    
    # The input to ``predict`` is the raw question string; the prompt template
    # is applied internally before calling the underlying transformers pipeline.
    loaded_model.predict("What is experiment tracking?")
    

  • save_pretrained

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

    警告

    如果模型以 save_pretrained 设置为 False 的方式保存,则模型无法注册到 MLflow Model Registry。为了将模型转换为可注册的模型,您可以使用 mlflow.transformers.persist_pretrained_model() 从 HuggingFace Hub 下载模型权重并将其保存在现有模型 artifact 中。有关更详细的使用方法,请参阅 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 存储库及其 commit hash 的引用。

  • kwargs – Transformers 序列化的可选其他配置。