mlflow.data

mlflow.data 模块帮助你将模型训练和评估数据集记录到 MLflow Tracking 的运行中，以及从运行中检索数据集信息。它提供了以下重要接口：

以下示例展示了如何使用 mlflow.data 将训练数据集记录到运行中，从运行中检索有关数据集的信息，并加载数据集的源。

import mlflow.data
import pandas as pd
from mlflow.data.pandas_dataset import PandasDataset

# Construct a Pandas DataFrame using iris flower data from a web URL
dataset_source_url = "http://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
df = pd.read_csv(dataset_source_url)
# Construct an MLflow PandasDataset from the Pandas DataFrame, and specify the web URL
# as the source
dataset: PandasDataset = mlflow.data.from_pandas(df, source=dataset_source_url)

with mlflow.start_run():
    # Log the dataset to the MLflow Run. Specify the "training" context to indicate that the
    # dataset is used for model training
    mlflow.log_input(dataset, context="training")

# Retrieve the run, including dataset information
run = mlflow.get_run(mlflow.last_active_run().info.run_id)
dataset_info = run.inputs.dataset_inputs[0].dataset
print(f"Dataset name: {dataset_info.name}")
print(f"Dataset digest: {dataset_info.digest}")
print(f"Dataset profile: {dataset_info.profile}")
print(f"Dataset schema: {dataset_info.schema}")

# Load the dataset's source, which downloads the content from the source URL to the local
# filesystem
dataset_source = mlflow.data.get_source(dataset_info)
dataset_source.load()
class mlflow.data.dataset.Dataset(source: DatasetSource, name: str | None = None, digest: str | None = None)[源代码]

基类：object

表示一个用于 MLflow 跟踪的数据集，包括数据集的名称、摘要（哈希）、模式和概况，以及源信息（例如，数据集所源自的 S3 存储桶或托管的 Delta 表）。大多数数据集还公开了用于训练和评估的特征和目标。

property digest: str

数据集的唯一哈希或指纹，例如 "498c7496"。

property name: str

数据集的名称，例如 "iris_data"，"myschema.mycatalog.mytable@v1" 等。

abstract property profile: Any | None

数据集的可选汇总统计信息，例如表中的行数、每个表列的均值/中位数/标准差等。

abstract property schema: Any | None

可选的数据集模式，例如一个表示数据集特征和目标的 mlflow.types.Schema 实例。

property source: DatasetSource

关于数据集来源的信息，表示为一个 DatasetSource 的实例。例如，这可能是数据集所源自的 S3 位置或托管的 Delta 表的名称。

abstract to_dict() → Dict[str, str][源代码]

为数据集创建配置字典。

子类应重写此方法以在配置字典中提供额外的字段，例如，模式、配置文件等。

返回一个包含以下字段的字符串字典：名称、摘要、源、源类型。

to_json() → str[源代码]

获取 数据集 的 JSON 字符串表示。

返回:

表示 Dataset 的 JSON 字符串。

class mlflow.data.dataset_source.DatasetSource[源代码]

基类：object

表示在 MLflow Tracking 中使用的数据集的来源，提供诸如云存储位置、delta 表名称/版本等信息。

from_json(cls, source_json: str) → DatasetSource[源代码]
abstract classmethod from_dict(source_dict: Dict[Any, Any]) → DatasetSource[源代码]

从字典表示构建 DatasetSource 的一个实例。

参数:

source_dict – 数据集源的字典表示。

返回:

一个 DatasetSource 实例。

abstract load() → Any[源代码]

加载由 DatasetSource 引用的文件/对象。例如，根据 DatasetSource 的类型，这可能从 S3 下载源 CSV 文件到本地文件系统，加载源 Delta 表作为 Spark DataFrame 等。

返回:

下载的源文件，例如本地文件系统路径、Spark DataFrame等。

abstract to_dict() → Dict[str, Any][源代码]

获取 DatasetSource 的 JSON 兼容字典表示。

返回:

DatasetSource 的 JSON 兼容字典表示形式。

to_json() → str[源代码]

获取 DatasetSource 的 JSON 字符串表示。

返回:

表示 DatasetSource 的 JSON 字符串。

mlflow.data.get_source(dataset: Dataset | DatasetInput | Dataset) → DatasetSource[源代码]

获取指定数据集或数据集输入的源代码。

参数:

dataset – 一个 mlflow.data.dataset.Dataset 的实例，或 mlflow.entities.Dataset，或 mlflow.entities.DatasetInput。

返回:

一个 DatasetSource 的实例。

pandas

mlflow.data.from_pandas(df: DataFrame, source: str | DatasetSource | None = None, targets: str | None = None, name: str | None = None, digest: str | None = None, predictions: str | None = None) → PandasDataset[源代码]

从 Pandas DataFrame、可选的目标、可选的预测和来源构建一个 PandasDataset 实例。

参数:

  • df – 一个 Pandas DataFrame。

  • source – DataFrame 的来源，例如文件系统路径、S3 URI、HTTPS URL、带有版本的 delta 表名或 spark 表等。source 可以指定为 URI、类路径字符串或 DatasetSource 的实例。如果未指定，则假定来源是调用 from_pandas 的代码位置（例如笔记本单元格、脚本等）。

  • targets – 监督训练的可选目标列名。此列必须存在于数据框 (df) 中。

  • name – 数据集的名称。如果未指定，则会生成一个名称。

  • digest – 数据集摘要（哈希）。如果未指定，则自动计算摘要。

  • predictions – 模型评估的可选预测列名称。此列必须存在于数据框 (df) 中。

示例 
import mlflow
import pandas as pd

x = pd.DataFrame(
    [["tom", 10, 1, 1], ["nick", 15, 0, 1], ["juli", 14, 1, 1]],
    columns=["Name", "Age", "Label", "ModelOutput"],
)
dataset = mlflow.data.from_pandas(x, targets="Label", predictions="ModelOutput")
class mlflow.data.pandas_dataset.PandasDataset[源代码]

表示一个用于 MLflow 跟踪的 Pandas DataFrame。

property df: DataFrame

底层 pandas DataFrame。

property predictions: str | None

预测列的名称。如果没有预测列，则可能为 None。

property profile: Any | None

数据集的概况。如果无法计算概况，则可能为 None。

property schema: Schema | None

表示表格数据集的 mlflow.types.Schema 实例。如果无法从数据集中推断出模式，则可能为 None。

property source: DatasetSource

数据集的来源。

property targets: str | None

目标列的名称。如果没有目标列，则可能为 None。

to_dict() → Dict[str, str][源代码]

为数据集创建配置字典。

返回一个包含以下字段的字符串字典：名称、摘要、源、源类型、模式和配置文件。

NumPy

mlflow.data.from_numpy(features: ndarray | Dict[str, ndarray], source: str | DatasetSource | None = None, targets: ndarray | Dict[str, ndarray] | None = None, name: str | None = None, digest: str | None = None) → NumpyDataset[源代码]

从NumPy特征、可选目标和源构建一个 NumpyDataset 对象。如果源是类路径的，那么这将根据源路径构建一个DatasetSource对象。否则，假定源是一个DatasetSource对象。

参数:

  • features – NumPy 特性，表示为 np.ndarray 或命名 np.ndarray 的字典。

  • source – numpy 数据的来源，例如文件系统路径、S3 URI、HTTPS URL、带有版本的 delta 表名或 spark 表等。source 可以指定为 URI、类路径字符串或 DatasetSource 的实例。如果未指定，则假定来源是调用 from_numpy 的代码位置（例如笔记本单元格、脚本等）。

  • targets – 可选的 NumPy 目标，表示为一个 np.ndarray 或命名 np.ndarrays 的字典。

  • name – 数据集的名称。如果未指定，则会生成一个名称。

  • digest – 数据集摘要（哈希）。如果未指定，则自动计算摘要。

基本示例 
import mlflow
import numpy as np

x = np.random.uniform(size=[2, 5, 4])
y = np.random.randint(2, size=[2])
dataset = mlflow.data.from_numpy(x, targets=y)
字典示例 
import mlflow
import numpy as np

x = {
    "feature_1": np.random.uniform(size=[2, 5, 4]),
    "feature_2": np.random.uniform(size=[2, 5, 4]),
}
y = np.random.randint(2, size=[2])
dataset = mlflow.data.from_numpy(x, targets=y)
class mlflow.data.numpy_dataset.NumpyDataset[源代码]

表示一个用于 MLflow 跟踪的 NumPy 数据集。

property features: ndarray | Dict[str, ndarray]

数据集的特性。

property profile: Any | None

数据集的概况。如果无法计算概况，则可能为 None。

property schema: TensorDatasetSchema | None

MLflow TensorSpec 模式，表示数据集的特征和目标（可选）。

property source: DatasetSource

数据集的来源。

property targets: ndarray | Dict[str, ndarray] | None

数据集的目标。如果没有目标可用，则可能为 None。

to_dict() → Dict[str, str][源代码]

为数据集创建配置字典。

返回一个包含以下字段的字符串字典：名称、摘要、源、源类型、模式和配置文件。

Spark

mlflow.data.load_delta(path: str | None = None, table_name: str | None = None, version: str | None = None, targets: str | None = None, name: str | None = None, digest: str | None = None) → SparkDataset[源代码]

从 Delta 表加载一个 SparkDataset 以用于 MLflow 跟踪。

参数:

  • path – Delta 表的路径。必须指定 path 或 table_name 之一。

  • table_name – Delta 表的名称。必须指定 path 或 table_name 中的一个。

  • version – Delta 表的版本。如果未指定，将推断版本。

  • targets – 可选。包含监督学习目标（标签）的 Delta 表列的名称。

  • name – 数据集的名称。例如：“wiki_train”。如果未指定，将自动生成一个名称。

  • digest – 数据集的摘要（哈希，指纹）。如果未指定，则会自动计算摘要。

返回:

一个 SparkDataset 的实例。

mlflow.data.from_spark(df: pyspark.sql.DataFrame, path: str | None = None, table_name: str | None = None, version: str | None = None, sql: str | None = None, targets: str | None = None, name: str | None = None, digest: str | None = None, predictions: str | None = None) → SparkDataset[源代码]

给定一个 Spark DataFrame，构建一个 SparkDataset 对象，用于 MLflow 跟踪。

参数:

  • df – 用于构建 SparkDataset 的 Spark DataFrame。

  • path – DataFrame 最初来自的 Spark 或 Delta 源的路径。请注意，路径不必与 DataFrame 完全匹配，因为 DataFrame 可能已被 Spark 操作修改。这用于通过 SparkDataset.source.load() 请求时重新加载数据集。如果没有指定 path、table_name 或 sql，则使用 CodeDatasetSource，它将从运行上下文中获取信息。

  • table_name – DataFrame 最初来自的 Spark 或 Delta 表的名称。请注意，该表不必与 DataFrame 完全匹配，因为 DataFrame 可能已经通过 Spark 操作进行了修改。这用于通过 SparkDataset.source.load() 在请求时重新加载数据集。如果没有指定 path、table_name 或 sql，则使用 CodeDatasetSource，它将从运行上下文中获取信息。

  • version – 如果 DataFrame 最初来自 Delta 表，则指定 Delta 表的版本。这用于通过 SparkDataset.source.load() 在请求时重新加载数据集。如果指定了 sql，则不能指定 version。

  • sql – 用于构造DataFrame的原始Spark SQL语句。请注意，Spark SQL语句不必与DataFrame完全匹配，因为DataFrame可能已被Spark操作修改。这用于通过 SparkDataset.source.load() 在请求时重新加载数据集。如果未指定 path、table_name 或 sql 中的任何一个，则使用CodeDatasetSource，它将从运行上下文中获取信息。

  • targets – 可选。包含监督学习目标（标签）的数据框列的名称。

  • name – 数据集的名称。例如：“wiki_train”。如果未指定，将自动生成一个名称。

  • digest – 数据集的摘要（哈希，指纹）。如果未指定，则会自动计算摘要。

  • predictions – 可选。包含模型预测的列的名称，如果数据集包含模型预测。如果指定，此列必须存在于数据框（df）中。

返回:

一个 SparkDataset 的实例。

class mlflow.data.spark_dataset.SparkDataset[源代码]

表示一个用于 MLflow 跟踪的 Spark 数据集（例如，来自 Spark 表/文件目录或 Delta 表的数据）。

property df

Spark DataFrame 实例。

返回:

Spark DataFrame 实例。

property predictions: str | None

预测列的名称。如果在创建数据集时未指定预测列，则可能为 None。

property profile: Any | None

数据集的概况。如果没有可用的概况，则可能为 None。

property schema: Schema | None

Spark 数据集的 MLflow ColSpec 模式。

property source: SparkDatasetSource | DeltaDatasetSource

Spark 数据集来源信息。

返回:

一个 SparkDatasetSource 或 DeltaDatasetSource 的实例。

property targets: str | None

包含用于监督学习的Spark DataFrame列的名称（标签）。

返回:

包含目标的 Spark DataFrame 列的字符串名称。

to_dict() → Dict[str, str][源代码]

为数据集创建配置字典。

返回一个包含以下字段的字符串字典：名称、摘要、源、源类型、模式和配置文件。

Hugging Face

mlflow.data.huggingface_dataset.from_huggingface(ds, path: str | None = None, targets: str | None = None, data_dir: str | None = None, data_files: str | Sequence[str] | Mapping[str, str | Sequence[str]] | None = None, revision=None, name: str | None = None, digest: str | None = None, trust_remote_code: bool | None = None) → HuggingFaceDataset[源代码]

从 Hugging Face 数据集创建一个 mlflow.data.huggingface_dataset.HuggingFaceDataset。

参数:

  • ds – 一个 Hugging Face 数据集。必须是 datasets.Dataset 的实例。其他类型，如 datasets.DatasetDict，是不支持的。

  • path – 用于构建源的 Hugging Face 数据集的路径。这与 datasets.load_dataset() 函数中的 path 参数相同。为了能够通过 MLflow 重新加载数据集，path 必须与数据集在 hub 上的路径匹配，例如，“databricks/databricks-dolly-15k”。如果没有指定路径，将使用 CodeDatasetSource，它将从运行上下文中获取信息。

  • targets – Hugging Face dataset.Dataset 数据集中包含监督学习目标（标签）的列名。

  • data_dir – Hugging Face 数据集配置的 data_dir。该目录由 datasets.load_dataset() 函数使用，以便通过 HuggingFaceDataset.source.load() 重新加载数据集。

  • data_files – Hugging Face 数据集配置的源数据文件路径。这被 datasets.load_dataset() 函数使用，以便通过 HuggingFaceDataset.source.load() 重新加载数据集。

  • revision – 要加载的数据集脚本的版本。这是由 datasets.load_dataset() 函数使用的，用于通过 HuggingFaceDataset.source.load() 重新加载数据集。

  • name – 数据集的名称。例如：“wiki_train”。如果未指定，将自动生成一个名称。

  • digest – 数据集的摘要（哈希，指纹）。如果未指定，则会自动计算摘要。

  • trust_remote_code – 是否信任来自数据集仓库的远程代码。

class mlflow.data.huggingface_dataset.HuggingFaceDataset[源代码]

表示一个用于 MLflow 跟踪的 HuggingFace 数据集。

property ds: datasets.Dataset

Hugging Face 的 datasets.Dataset 实例。

返回:

Hugging Face 的 datasets.Dataset 实例。

property profile: Any | None

Hugging Face 数据集的汇总统计数据，包括行数、大小和字节大小。

property schema: Schema | None

Hugging Face 数据集的 MLflow ColSpec 模式。

property source: HuggingFaceDatasetSource

Hugging Face 数据集来源信息。

返回:

一个 mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource

property targets: str | None

Hugging Face 数据集列的名称，包含监督学习的目标（标签）。

返回:

包含目标的 Hugging Face 数据集列的字符串名称。

to_dict() → Dict[str, str][源代码]

为数据集创建配置字典。

返回一个包含以下字段的字符串字典：名称、摘要、源、源类型、模式和配置文件。

to_evaluation_dataset(path=None, feature_names=None) → EvaluationDataset[源代码]

将数据集转换为用于模型评估的 EvaluationDataset。使用 mlflow.evaluate() 时需要。

TensorFlow

mlflow.data.tensorflow_dataset.from_tensorflow(features, source: str | DatasetSource | None = None, targets=None, name: str | None = None, digest: str | None = None) → TensorFlowDataset[源代码]

从 TensorFlow 数据、可选目标和源构建一个 TensorFlowDataset 对象。

如果源路径类似于路径，那么这将根据源路径构造一个 DatasetSource 对象。否则，源路径被假定为一个 DatasetSource 对象。

参数:

  • features – 一个 TensorFlow 数据集或特征张量。

  • source – 数据来源，例如文件系统路径、S3 URI、HTTPS URL、带有版本的delta表名或spark表等。如果来源不是一个路径字符串，直接传入一个DatasetSource对象。如果没有指定来源，则使用CodeDatasetSource，它将从运行上下文中获取信息。

  • targets – 一个 TensorFlow 数据集或目标张量。可选。

  • name – 数据集的名称。如果未指定，则会生成一个名称。

  • digest – 数据集摘要（哈希）。如果未指定，则自动计算摘要。

class mlflow.data.tensorflow_dataset.TensorFlowDataset[源代码]

表示一个用于 MLflow 跟踪的 TensorFlow 数据集。

property data

底层 TensorFlow 数据。

property profile: Any | None

数据集的概况。如果没有可用的概况，则可能为 None。

property schema: TensorDatasetSchema | None

表示张量数据集的 MLflow TensorSpec 模式

property source: DatasetSource

数据集的来源。

property targets

数据集的目标。

to_dict() → Dict[str, str][源代码]

为数据集创建配置字典。

返回一个包含以下字段的字符串字典：名称、摘要、源、源类型、模式和配置文件。

to_evaluation_dataset(path=None, feature_names=None) → EvaluationDataset[源代码]

将数据集转换为用于模型评估的 EvaluationDataset。仅在数据集为张量时支持。使用 mlflow.evaluate() 时必需。

class mlflow.data.evaluation_dataset.EvaluationDataset[源代码]

用于模型评估的输入数据集。这是为与 mlflow.models.evaluate() API 一起使用而设计的。

NUM_SAMPLE_ROWS_FOR_HASH = 5
SPARK_DATAFRAME_LIMIT = 10000
property feature_names
property features_data

将特征数据作为 numpy 数组或 pandas DataFrame 返回。

property has_predictions

如果数据集有目标则返回 True，否则返回 False。

property has_targets

如果数据集有目标则返回 True，否则返回 False。

property hash

数据集哈希，包括前20行和后20行的哈希。

property labels_data

将标签数据作为 numpy 数组返回

property name

数据集名称，如果用户未指定名称，则为指定的数据集名称或数据集哈希。

property path

数据集路径

property predictions_data

将标签数据作为 numpy 数组返回

property predictions_name

返回预测名称

property targets_name

返回目标名称

数据集来源

class mlflow.data.filesystem_dataset_source.FileSystemDatasetSource[源代码]

表示存储在文件系统上的数据集的源，例如本地UNIX文件系统、S3等blob存储服务。

abstract classmethod from_dict(source_dict: Dict[Any, Any]) → FileSystemDatasetSource[源代码]
参数:

source_dict – FileSystemDatasetSource 的字典表示。

abstract load(dst_path=None) → str[源代码]

将数据集源下载到本地文件系统。

参数:

dst_path – 本地文件系统目标目录的路径，用于下载数据集源。如果目录不存在，则创建该目录。如果未指定，数据集源将下载到本地文件系统上的一个新创建的唯一命名目录，除非数据集源已经存在于本地文件系统中，在这种情况下，直接返回其本地路径。

返回:

下载的数据集源在本地文件系统中的路径。

abstract to_dict() → Dict[Any, Any][源代码]
返回:

FileSystemDatasetSource 的 JSON 兼容字典表示形式。

abstract property uri

指向数据集源文件系统位置的URI。

返回:

指向数据集源文件系统位置的URI，例如 “s3://mybucket/path/to/mydataset”，”/tmp/path/to/my/dataset” 等。

class mlflow.data.http_dataset_source.HTTPDatasetSource[源代码]

表示存储在网络位置并通过HTTP或HTTPS URL引用的数据集的源。

classmethod from_dict(source_dict: Dict[Any, Any]) → HTTPDatasetSource[源代码]
参数:

source_dict – HTTPDatasetSource 的字典表示。

load(dst_path=None) → str[源代码]

将数据集源下载到本地文件系统。

参数:

dst_path – 本地文件系统目标目录的路径，用于下载数据集源。如果该目录不存在，则创建它。如果未指定，数据集源将下载到本地文件系统上的一个新创建的唯一命名目录中。

返回:

下载的数据集源在本地文件系统中的路径。

to_dict() → Dict[Any, Any][源代码]
返回:

HTTPDatasetSource 的 JSON 兼容字典表示。

property url

指向数据集源位置的HTTP/S URL。

返回:

指向数据集源位置的HTTP/S URL。

class mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource[源代码]

表示在MLflow Tracking中使用的Hugging Face数据集的源。

classmethod from_dict(source_dict: Dict[Any, Any]) → HuggingFaceDatasetSource[源代码]

从字典表示构建 DatasetSource 的一个实例。

参数:

source_dict – 数据集源的字典表示。

返回:

一个 DatasetSource 实例。

load(**kwargs)[源代码]

基于 HuggingFaceDatasetSource 加载 Hugging Face 数据集。

参数:

kwargs – 用于通过 Hugging Face 的 datasets.load_dataset() 方法加载数据集的额外关键字参数。

返回:

datasets.Dataset 的一个实例。

to_dict() → Dict[Any, Any][源代码]

获取 DatasetSource 的 JSON 兼容字典表示。

返回:

DatasetSource 的 JSON 兼容字典表示形式。

class mlflow.data.delta_dataset_source.DeltaDatasetSource[源代码]

表示存储在 delta 表中的数据集的源。

property delta_table_id: str | None
property delta_table_name: str | None
property delta_table_version: int | None
classmethod from_dict(source_dict: Dict[Any, Any]) → DeltaDatasetSource[源代码]

从字典表示构建 DatasetSource 的一个实例。

参数:

source_dict – 数据集源的字典表示。

返回:

一个 DatasetSource 实例。

load(**kwargs)[源代码]

加载数据集源作为 Delta 数据集源。

返回:

pyspark.sql.DataFrame 的一个实例。

property path: str | None
to_dict() → Dict[Any, Any][源代码]

获取 DatasetSource 的 JSON 兼容字典表示。

返回:

DatasetSource 的 JSON 兼容字典表示形式。

class mlflow.data.spark_dataset_source.SparkDatasetSource[源代码]

表示存储在spark表中的数据集的源。

classmethod from_dict(source_dict: Dict[Any, Any]) → SparkDatasetSource[源代码]

从字典表示构建 DatasetSource 的一个实例。

参数:

source_dict – 数据集源的字典表示。

返回:

一个 DatasetSource 实例。

load(**kwargs)[源代码]

将数据集源加载为 Spark 数据集源。

返回:

pyspark.sql.DataFrame 的一个实例。

to_dict() → Dict[Any, Any][源代码]

获取 DatasetSource 的 JSON 兼容字典表示。

返回:

DatasetSource 的 JSON 兼容字典表示形式。