如何用MLflow追踪与评估Amazon Bedrock上的LLM模型（如Claude）

MLflow目前没有针对Amazon Bedrock的官方直接集成，但可以通过自定义PyFunc模型包装实现对Bedrock上LLM的日志记录和性能评估，以下是完整的实现步骤和代码示例：

前置准备 #

1. 安装依赖：

pip install mlflow boto3 pandas

2. 配置AWS凭证：确保运行环境已通过环境变量、AWS CLI配置文件或IAM角色获得Bedrock模型调用权限（需包含bedrock:InvokeModel权限）。

完整代码实现 #

1. 自定义Bedrock Claude模型包装类 #

通过MLflow的PyFunc接口包装Bedrock模型调用逻辑，确保符合MLflow的模型规范：

import mlflow
import boto3
import json
import pandas as pd

class BedrockClaudeModel(mlflow.pyfunc.PythonModel):
    def __init__(self, model_id, system_prompt):
        self.model_id = model_id
        self.system_prompt = system_prompt
        self.bedrock_client = boto3.client("bedrock-runtime", region_name="us-east-1")

    def predict(self, context, model_input):
        # 处理输入：model_input应为包含"question"列的DataFrame
        results = []
        for question in model_input["question"]:
            # 构造Claude请求体
            request_body = json.dumps({
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 1024,
                "messages": [
                    {"role": "user", "content": f"{self.system_prompt}\n\nQuestion: {question}"}
                ]
            })

            # 调用Bedrock模型
            response = self.bedrock_client.invoke_model(
                modelId=self.model_id,
                contentType="application/json",
                accept="application/json",
                body=request_body
            )

            # 解析响应
            response_body = json.loads(response["body"].read())
            answer = response_body["content"][0]["text"].strip()
            results.append(answer)
        
        return pd.Series(results)

2. 模型日志记录与评估 #

在MLflow运行中完成模型注册和性能评估，逻辑与OpenAI集成类似：

# 准备评估数据集：需包含"question"（输入问题）和"ground_truth"（预期答案）列
eval_df = pd.DataFrame({
    "question": ["什么是MLflow？", "Amazon Bedrock主要用途是什么？"],
    "ground_truth": ["MLflow是一个开源的机器学习生命周期管理平台，用于追踪实验、打包模型和部署。", "Amazon Bedrock是AWS提供的托管式基础模型服务，允许开发者通过API调用各类大语言模型和生成式AI模型。"]
})

with mlflow.start_run() as run:
    # 定义系统提示词
    system_prompt = "请用两句话回答以下问题"

    # 初始化自定义模型
    claude_model = BedrockClaudeModel(
        model_id="anthropic.claude-3-5-sonnet-20240620-v1:0",
        system_prompt=system_prompt
    )

    # 日志记录模型到MLflow
    logged_model = mlflow.pyfunc.log_model(
        artifact_path="claude-bedrock-model",
        python_model=claude_model,
        input_example=pd.DataFrame({"question": ["示例问题"]})
    )

    # 评估模型性能
    evaluation_results = mlflow.evaluate(
        model=logged_model.model_uri,
        data=eval_df,
        targets="ground_truth",
        model_type="question-answering",
        evaluators="default"
    )

# 打印评估指标
print("评估指标：")
for metric_name, value in evaluation_results.metrics.items():
    print(f"{metric_name}: {value}")

关键说明 #

• 模型ID：Claude 3.5-sonnet的Bedrock模型ID需根据使用的AWS区域确认，示例中使用的是us-east-1区域的官方ID。
• 请求格式：Claude 3系列模型的请求体需遵循Anthropic的规范，包含anthropic_version字段，消息格式与OpenAI略有差异。
• 评估逻辑：MLflow的默认评估器会针对问答任务自动计算BLEU、ROUGE等文本相似度指标，无需额外配置。

内容的提问来源于stack exchange，提问作者Fabian Wörenkämper