跳到主要内容

评估数据集 SDK 指南

通过实际工作流程和真实模式，掌握用于创建、演进和管理评估数据集的 API。

开始使用​

MLflow 提供了一个用于处理评估数据集的流畅 API，使常见工作流程变得简单直观。

from mlflow.genai.datasets import (
    create_dataset,
    get_dataset,
    search_datasets,
    set_dataset_tags,
    delete_dataset_tag,
)

您的数据集之旅​

遵循这个典型的工作流程来构建和演进您的评估数据集。

完整的开发工作流程

Continuous Improvement
创建/获取数据集
添加测试用例
运行评估
改进代码
测试与追踪
更新数据集
更新标签

步骤 1：创建您的数据集​

首先，使用 mlflow.genai.datasets.create_dataset() API 创建一个包含有意义元数据的新评估数据集。

from mlflow.genai.datasets import create_dataset

# Create a new dataset with tags for organization
dataset = create_dataset(
    name="customer_support_qa_v1",
    experiment_id=["0"],  # Link to experiments ("0" is default)
    tags={
        "version": "1.0",
        "purpose": "regression_testing",
        "model": "gpt-4",
        "team": "ml-platform",
        "status": "development",
    },
)

步骤 2：添加您的第一个测试用例​

通过从生产追踪数据和手动整理中添加测试用例来构建您的数据集。期望通常由主题领域专家 (SME) 定义，他们理解该领域并能确立构成正确行为的基准真相。

了解如何定义期望 → 期望是定义您的 AI 应生成什么的基准真相值。它们由审查输出并确立质量标准的主题领域专家添加。

import mlflow

# Search for production traces to build your dataset
# Request list format to work with individual Trace objects
production_traces = mlflow.search_traces(
    experiment_ids=["0"],  # Your production experiment
    filter_string="attributes.user_feedback = 'positive'",
    max_results=100,
    return_type="list",  # Returns list[Trace] for direct manipulation
)

# Subject matter experts add expectations to define correct behavior
for trace in production_traces:
    # Subject matter experts review traces and define what the output should satisfy
    mlflow.log_expectation(
        trace_id=trace.info.trace_id,
        name="quality_assessment",
        value={
            "should_match_production": True,
            "minimum_quality": 0.8,
            "response_time_ms": 2000,
            "contains_citation": True,
        },
    )

    # Can also add textual expectations
    mlflow.log_expectation(
        trace_id=trace.info.trace_id,
        name="expected_behavior",
        value="Response should provide step-by-step instructions with security considerations",
    )

# Add annotated traces to dataset (expectations are automatically included)
dataset.merge_records(production_traces)

步骤 3：演进您的数据集​

随着您发现边缘情况并加深理解，不断更新您的数据集。mlflow.entities.EvaluationDataset.merge_records() 方法智能地处理新记录和现有记录的更新。

# Capture a production failure
failure_case = {
    "inputs": {"question": "'; DROP TABLE users; --", "user_type": "malicious"},
    "expectations": {
        "handles_sql_injection": True,
        "returns_safe_response": True,
        "logs_security_event": True,
    },
    "source": {
        "source_type": "HUMAN",
        "source_data": {"discovered_by": "security_team"},
    },
    "tags": {"category": "security", "severity": "critical"},
}

# Add the new edge case
dataset.merge_records([failure_case])

# Update expectations for existing records
updated_records = []
for record in dataset.records:
    if "accuracy" in record.get("expectations", {}):
        # Raise the quality bar
        record["expectations"]["accuracy"] = max(
            0.9, record["expectations"]["accuracy"]
        )
        updated_records.append(record)

# Merge updates (intelligently handles duplicates)
dataset.merge_records(updated_records)

步骤 4：使用标签进行组织​

使用标签来追踪数据集的演进并实现强大的搜索。了解更多关于 mlflow.search_traces() 的信息，以便从生产数据构建您的数据集。

from mlflow.genai.datasets import set_dataset_tags

# Update dataset metadata
set_dataset_tags(
    dataset_id=dataset.dataset_id,
    tags={
        "status": "validated",
        "coverage": "comprehensive",
        "last_review": "2024-11-01",
    },
)

# Remove outdated tags
set_dataset_tags(
    dataset_id=dataset.dataset_id,
    tags={"development_only": None},  # Setting to None removes the tag
)

步骤 5：搜索和发现​

使用 mlflow.genai.datasets.search_datasets() 的强大搜索功能查找数据集。

from mlflow.genai.datasets import search_datasets

# Find datasets by experiment
datasets = search_datasets(experiment_ids=["0", "1"])  # Search in multiple experiments

# Search by name pattern
regression_datasets = search_datasets(filter_string="name LIKE '%regression%'")

# Complex search with tags
production_ready = search_datasets(
    filter_string="tags.status = 'validated' AND tags.coverage = 'comprehensive'",
    order_by=["last_update_time DESC"],
    max_results=10,
)

# The PagedList automatically handles pagination when iterating

通用筛选字符串示例​

以下是一些实用的筛选字符串示例，可帮助您找到合适的数据集。

筛选表达式描述用例
name = 'production_qa'精确名称匹配查找特定数据集
name LIKE '%test%'模式匹配查找所有测试数据集
tags.status = 'validated'标签相等性查找可用于生产的数据集
tags.version = '2.0' AND tags.team = 'ml'多个标签条件查找特定团队的版本
created_by = 'alice@company.com'创建者筛选按作者查找数据集
created_time > 1698800000000基于时间的筛选查找最近的数据集
tags.model = 'gpt-4' AND name LIKE '%eval%'组合条件特定模型的评估集
last_updated_by != 'bot@system'排除筛选排除自动更新

步骤 6：管理实验关联​

创建数据集后，可以使用 mlflow.genai.datasets.add_dataset_to_experiments() 和 mlflow.genai.datasets.remove_dataset_from_experiments() 将其动态关联到实验。

此功能支持几个重要的用例：

  • 跨团队协作：通过添加其实验 ID 来跨团队共享数据集。
  • 生命周期管理：随着项目的成熟，删除过时的实验关联。
  • 项目重组：随着项目结构的演变，动态重组数据集。
from mlflow.genai.datasets import (
    add_dataset_to_experiments,
    remove_dataset_from_experiments,
)

# Add dataset to additional experiments
dataset = add_dataset_to_experiments(
    dataset_id="d-1a2b3c4d5e6f7890abcdef1234567890", experiment_ids=["3", "4", "5"]
)
print(f"Dataset now linked to experiments: {dataset.experiment_ids}")

# Remove dataset from specific experiments
dataset = remove_dataset_from_experiments(
    dataset_id="d-1a2b3c4d5e6f7890abcdef1234567890", experiment_ids=["3"]
)
print(f"Updated experiment associations: {dataset.experiment_ids}")

活动记录模式​

EvaluationDataset 对象遵循活动记录模式——它既是数据容器，又提供了与后端交互的方法。

# Get a dataset
dataset = get_dataset(dataset_id="d-1a2b3c4d5e6f7890abcdef1234567890")

# The dataset object is "live" - it can fetch and update data
current_record_count = len(dataset.records)  # Lazy loads if needed

# Add new records directly on the object
new_records = [
    {
        "inputs": {"question": "What are your business hours?"},
        "expectations": {"mentions_hours": True, "includes_timezone": True},
    }
]
dataset.merge_records(new_records)  # Updates backend immediately

# Convert to DataFrame for analysis
df = dataset.to_df()
# Access auto-computed properties
schema = dataset.schema  # Field structure
profile = dataset.profile  # Dataset statistics

记录合并工作原理​

merge_records() 方法智能地处理新记录和现有记录的更新。记录根据其输入的哈希值进行匹配——如果已存在具有相同输入的记录，则会更新其期望和标签，而不是创建重复记录。

当您首次添加记录时，它们会与其输入、期望和元数据一起存储。

# Initial record
record_v1 = {
    "inputs": {"question": "What is MLflow?", "context": "ML platform overview"},
    "expectations": {"accuracy": 0.8, "mentions_tracking": True},
}

dataset.merge_records([record_v1])
# Creates a new record in the dataset

理解来源类型​

MLflow 通过来源类型跟踪评估数据集中每条记录的出处。这有助于您了解测试数据的来源，并按数据源分析性能。

来源类型行为

自动推断

当未提供显式来源时，MLflow 会根据记录的特征自动推断来源类型。

手动覆盖

您可以随时指定显式的来源信息来覆盖自动推断。

出处跟踪

来源类型允许按数据来源进行筛选和性能分析。

自动来源分配​

MLflow 会根据记录的特征自动分配来源类型。

从 MLflow 追踪创建的记录会自动分配 TRACE 来源类型。

# When adding traces directly (automatic TRACE source)
traces = mlflow.search_traces(experiment_ids=["0"], return_type="list")
dataset.merge_records(traces)  # All records get TRACE source type

# Or when using DataFrame from search_traces
traces_df = mlflow.search_traces(experiment_ids=["0"])  # Returns DataFrame
dataset.merge_records(
    traces_df
)  # Automatically detects traces and assigns TRACE source

手动来源指定​

您可以为任何记录显式指定来源类型和元数据。当未提供显式来源时，MLflow 会在将记录发送到后端之前，根据以下规则自动推断来源类型：

  • 包含期望的记录 → 推断为 HUMAN 来源（表示手动标注或基准真相）。
  • 仅包含输入的记录（无期望）→ 推断为 CODE 来源（表示程序化生成）。
  • 来自追踪的记录 → 始终标记为 TRACE 来源（无论是否有期望）。

此推断在 merge_records() 方法中在客户端执行，然后在记录发送到追踪后端。您可以通过提供显式的来源信息来覆盖此自动推断。

# Specify HUMAN source for manually curated test cases
human_curated = {
    "inputs": {"question": "What are your business hours?"},
    "expectations": {"accuracy": 1.0, "includes_timezone": True},
    "source": {
        "source_type": "HUMAN",
        "source_data": {"curator": "support_team", "date": "2024-11-01"},
    },
}

# Specify DOCUMENT source for data from documentation
from_docs = {
    "inputs": {"question": "How to install MLflow?"},
    "expectations": {"mentions_pip": True, "mentions_conda": True},
    "source": {
        "source_type": "DOCUMENT",
        "source_data": {"document_id": "install_guide", "page": 1},
    },
}

# Specify CODE source for programmatically generated data
generated = {
    "inputs": {"question": f"Test question {i}" for i in range(100)},
    "source": {
        "source_type": "CODE",
        "source_data": {"generator": "test_suite_v2", "seed": 42},
    },
}

dataset.merge_records([human_curated, from_docs, generated])

可用来源类型​

来源类型支持对您的评估结果进行强大的筛选和分析。您可以按数据来源分析性能，以了解您的模型在人类整理的测试用例与生成的测试用例，或生产追踪与文档示例上的表现是否不同。

TRACE

通过 MLflow 追踪捕获的生产数据 - 添加追踪时自动分配。

HUMAN

主题领域专家标注 - 为包含期望的记录推断。

CODE

程序化生成的测试 - 为不包含期望的记录推断。

DOCUMENT

来自文档或规范的测试用例 - 必须显式指定。

UNSPECIFIED

来源未知或未提供 - 用于遗留或导入数据。

按来源分析数据​

# Convert dataset to DataFrame for analysis
df = dataset.to_df()

# Check source type distribution
source_distribution = df["source_type"].value_counts()
print("Data sources in dataset:")
for source_type, count in source_distribution.items():
    print(f"  {source_type}: {count} records")

搜索筛选参考​

在您的筛选字符串中使用这些字段。注意：流畅 API 返回一个 PagedList，可以直接迭代 - 分页在您迭代结果时自动处理。

字段类型示例
name字符串name = 'production_tests'
tags.<key>字符串tags.status = 'validated'
created_by字符串created_by = 'alice@company.com'
last_updated_by字符串last_updated_by = 'bob@company.com'
created_timetimestampcreated_time > 1698800000000
last_update_timetimestamplast_update_time > 1698800000000

筛选运算符​

  • =, !=：精确匹配
  • LIKE, ILIKE：使用 % 通配符的模式匹配（ILIKE 区分大小写）。
  • >, <, >=, <=：数字/时间戳比较。
  • AND：组合条件（OR 目前不支持用于评估数据集）。
# Complex filter example
datasets = search_datasets(
    filter_string="""
        tags.status = 'production'
        AND name LIKE '%customer%'
        AND created_time > 1698800000000
    """,
    order_by=["last_update_time DESC"],
)

使用客户端 API​

对于应用程序和高级用例，您还可以使用 MlflowClient API，它提供与面向对象接口相同的 funkcionalita。

from mlflow import MlflowClient

client = MlflowClient()

# Create a dataset
dataset = client.create_dataset(
    name="customer_support_qa",
    experiment_id=["0"],
    tags={"version": "1.0", "team": "ml-platform"},
)

客户端 API 提供与流畅 API 相同的功能，但更适合：

  • 需要显式客户端管理的出产应用程序。
  • 需要自定义追踪 URI 或身份验证的场景。
  • 与现有的基于 MLflow 客户端的工作流程集成。

后续步骤​

端到端工作流程

了解从应用程序构建到生产的完整评估驱动开发工作流程。

查看完整工作流程 →

运行评估

使用您的数据集系统地评估和改进您的 GenAI 应用程序。

开始评估 →

定义期望

了解如何为您的测试数据添加基准真相期望以进行质量验证。

设置期望 →

捕获追踪

对您的应用程序进行仪器化，以捕获生产数据用于构建数据集。

启用追踪 →