跳到主要内容

MLflow 用于生成式 AI 的模型注册表数据模型

MLflow 模型注册表提供了一个集中的模型存储、一套 API 和一个 UI，用于管理 AI 模型的整个生命周期。对于生成式 AI 应用程序，模型注册表作为生产就绪模型、智能体和 AI 应用程序的记录系统，实现协作模型管理和部署工作流。

概述​

模型注册表充当模型生命周期管理的中心枢纽，为生成式 AI 应用程序提供版本控制、阶段转换和部署协调。

模型注册表核心概念​

注册模型​

注册模型代表注册表中一个唯一的模型，可以有多个版本

模型版本​

注册模型的每个版本都包含特定的工件和元数据

阶段转换​

模型版本在其生命周期中经历定义的阶段

模型注册表数据模型​

注册表实体关系​

模型注册表维护模型、版本和部署阶段之间的关系

生成式 AI 模型注册工作流​

从开发到生产​

注册生成式 AI 模型的典型工作流

代码示例：注册生成式 AI 模型​

这里有一个在模型注册表中注册基于 LangChain 的聊天机器人的实际示例

import mlflow
import os

# Set up MLflow tracking
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("chatbot-development")

# Create the chatbot model script
# %%writefile "./chatbot_model.py"  # Uncomment if running in Jupyter

from langchain_openai import ChatOpenAI
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from mlflow.models import set_model
import os


class ChatbotModel:
    def __init__(self):
        # Initialize LangChain components
        self.llm = ChatOpenAI(
            model="gpt-4",
            temperature=0.7,
            max_tokens=1000,
            api_key=os.getenv("OPENAI_API_KEY"),  # Use environment variable
        )

        self.memory = ConversationBufferMemory()

        self.chatbot = ConversationChain(llm=self.llm, memory=self.memory, verbose=True)

    def invoke(self, input_data):
        """Handle both direct string input and structured input"""
        if isinstance(input_data, dict):
            query = input_data.get("query", input_data.get("input", ""))
        else:
            query = str(input_data)

        return self.chatbot.predict(input=query)


# Set the model for MLflow
set_model(ChatbotModel())

# Log the model using Models from Code
with mlflow.start_run() as run:
    # Log model parameters
    mlflow.log_params(
        {
            "model_type": "conversational",
            "base_model": "gpt-4",
            "temperature": 0.7,
            "max_tokens": 1000,
            "memory_type": "buffer",
        }
    )

    # Define input example
    input_example = {"query": "Hello, how can you help me?"}

    # Log the model using Models from Code
    model_info = mlflow.langchain.log_model(
        lc_model="chatbot_model.py",  # Path to the script file
        name="chatbot",
        input_example=input_example,
        registered_model_name="customer-support-chatbot",
        extra_pip_requirements=["langchain-openai", "langchain"],
    )

    # Add model metadata
    client = mlflow.tracking.MlflowClient()

    # Tag the registered model
    client.set_registered_model_tag(
        name="customer-support-chatbot", key="team", value="ai-support"
    )

    client.set_registered_model_tag(
        name="customer-support-chatbot", key="use_case", value="customer_service"
    )

    # Update model description
    client.update_registered_model(
        name="customer-support-chatbot",
        description="GPT-4 based chatbot for customer support with conversation memory",
    )

    # Add version-specific tags
    client.set_model_version_tag(
        name="customer-support-chatbot",
        version=model_info.registered_model_version,
        key="validation_status",
        value="pending",
    )

    print(f"Model registered with version: {model_info.registered_model_version}")

# Create an alias for easy reference
client.set_registered_model_alias(
    name="customer-support-chatbot",
    alias="champion",
    version=model_info.registered_model_version,
)

# Load model from registry for serving using alias
model_uri_alias = "models:/customer-support-chatbot@champion"
loaded_model = mlflow.langchain.load_model(model_uri_alias)

# Test the loaded model
test_query = {"query": "What's the weather like?"}
response = loaded_model.invoke(test_query)
print(f"Response: {response}")

生成式 AI 的模型注册表 API​

关键 API 操作​

模型注册表提供全面的 API，用于管理生成式 AI 模型

部署集成​

从注册表到生产​

模型注册表与部署平台无缝集成

生成式 AI 模型的最佳实践​

模型组织​

命名约定：使用描述性的、分层名称，指示模型的目的和类型（例如，nlp/chat/customer-support-v2）。

版本文档：记录版本之间的重大更改，包括提示修改、模型升级和性能改进。

元数据标准：为生成式 AI 模型建立一致的元数据模式，包括基础模型信息、上下文窗口和工具集成。

阶段管理​

暂存验证：在将生成式 AI 模型提升到生产环境之前，务必在暂存环境中用代表性工作负载进行验证。

A/B 测试：使用模型别名在生产环境中实现模型版本之间的平滑 A/B 测试。

回滚策略：将以前的生产版本保持在可访问状态，以便在出现问题时快速回滚。

协作工作流​

访问控制：明确定义模型注册、阶段转换和生产部署的权限。

审查流程：为生产升级实施审批工作流，特别是针对面向客户的生成式 AI 应用程序。

更改跟踪：使用标签和描述来跟踪模型被升级或存档的原因。

高级功能​

模型沿袭跟踪​

注册表维护从实验到生产的完整沿袭

与 MLflow 生态系统集成​

全面的平台集成​

模型注册表与 MLflow 的其他组件无缝集成

总结​

MLflow 模型注册表为管理生成式 AI 模型的整个生命周期提供了坚实的基础。通过集中模型管理、版本控制和部署工作流，它使团队能够自信地开发、验证和部署 AI 应用程序。注册表与 MLflow 更广泛生态系统的集成确保了生产生成式 AI 系统的全面跟踪、评估和监控能力。