MLflow Transformers Flavor

MLflow Transformers Flavor 提供了与 Hugging Face Transformers 库的原生集成，支持 NLP、音频、视觉和多模态任务的模型日志记录、加载和推理。

主要特性

Pipeline and Component Logging: 保存完整的 pipeline 或单个模型组件
PyFunc Integration: 使用标准化的推理接口部署模型
PEFT Support: 对参数高效微调（LoRA、QLoRA 等）的原生支持
Prompt Templates: 与 pipeline 一起保存和管理 prompt 模板
Automatic Metadata Logging: 自动记录模型卡片和元数据
Flexible Inference Configuration: 通过 model_config 和 signature 参数自定义模型行为

安装

bash
pip install mlflow transformers

基本用法

Logging a Pipeline

python
import mlflow
from transformers import pipeline

# Create a text generation pipeline
text_gen = pipeline("text-generation", model="gpt2")

# Log the pipeline
with mlflow.start_run():
    mlflow.transformers.log_model(
        transformers_model=text_gen,
        name="model",
    )

Loading and Inference

python
# Load as native transformers
model = mlflow.transformers.load_model("runs:/<run_id>/model")
result = model("Hello, how are you?")

# Load as PyFunc
pyfunc_model = mlflow.pyfunc.load_model("runs:/<run_id>/model")
result = pyfunc_model.predict("Hello, how are you?")

Autologging with HuggingFace Trainer

当使用 HuggingFace Trainer 类进行微调时，您可以通过在 TrainingArguments 中设置 report_to="mlflow" 来启用自动日志记录到 MLflow。

python
from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",
    report_to="mlflow",  # Enable MLflow logging
    # ... other training arguments
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    # ... other trainer arguments
)

trainer.train()

这将自动将训练指标、超参数和模型检查点记录到您活动的 MLflow 运行中。

Tutorials

Important Considerations

PyFunc Limitations

并非所有 pipeline 类型都支持 PyFunc 推理
某些输出（例如，附加分数、引用）可能未被捕获
支持音频和文本 LLM；视觉和多模态模型需要原生加载
请参阅指南了解支持的 pipeline 类型

Input/Output Types

PyFunc 的输入和输出格式可能与原生 pipeline 不同。请确保与您的数据处理工作流兼容。

模型配置

在同时提供 ModelSignature 和 model_config 时，ModelSignature 中的参数将覆盖 model_config 中的参数。

Working with Large Models

对于包含数十亿参数的模型，MLflow 提供了优化技术来减少内存使用并加快日志记录速度。请参阅大型模型指南。

Tasks

task 参数决定了输入/输出格式。MLflow 支持原生的 Transformers 任务，以及 llm/v1/chat 和 llm/v1/completions 等高级任务，用于兼容 OpenAI 的推理。请参阅任务指南。

了解更多

Detailed Guide

Comprehensive documentation covering pipelines, PyFunc, signatures, PEFT, and more

Learn more →

Transformers Documentation

Official Hugging Face Transformers documentation

Learn more →

主要特性​

安装​

基本用法​

Logging a Pipeline​

Loading and Inference​

Autologging with HuggingFace Trainer​

Tutorials​

快速入门​

Fine-Tuning​

Advanced Use Cases​

Important Considerations​

PyFunc Limitations​

Input/Output Types​

模型配置​

Working with Large Models​

Tasks​

了解更多​