使用 TensorFlow Datasets 构建高性能数据管道-优快云博客

老铁们，今天我们来聊聊 TensorFlow Datasets，这个工具可以说是相当给力。它提供了一系列可以直接用于 TensorFlow 或其他 Python 机器学习框架（如 Jax）的数据集，能让我们轻松构建高性能的数据输入管道。要是你还没用过这个，那咱们今天就来一波实战演示。

技术背景介绍

TensorFlow Datasets（TFDS）是个专门为 TensorFlow 和其他机器学习框架提供标准化数据集的库。这些数据集都被暴露为 tf.data.Datasets 对象，方便你在训练模型时使用。

安装

要使用 TensorFlow Datasets，老铁们需要先安装 tensorflow 和 tensorflow-datasets 这两个 Python 包：

%pip install --upgrade --quiet tensorflow
%pip install --upgrade --quiet tensorflow-datasets

原理深度解析

今天我们以 mlqa/en 数据集为例，这个数据集是一个用于评估多语言问答性能的基准数据集，包含了 7 种语言。这里展示的是如何将 TFDS 数据集加载为可用于下游任务的 Document 格式。

实战代码演示

下面是如何加载并转换 mlqa/en 数据集的完整实例：

import tensorflow as tf
import tensorflow_datasets as tfds

# 加载数据集
ds = tfds.load("mlqa/en", split="test")
ds = ds.take(1)  # 只拿一个示例

# 定义自定义转换函数
from langchain_core.documents import Document

def decode_to_str(item: tf.Tensor) -> str:
    return item.numpy().decode("utf-8")

def mlqaen_example_to_document(example: dict) -> Document:
    return Document(
        page_content=decode_to_str(example["context"]),
        metadata={
            "id": decode_to_str(example["id"]),
            "title": decode_to_str(example["title"]),
            "question": decode_to_str(example["question"]),
            "answer": decode_to_str(example["answers"]["text"][0]),
        },
    )

# 测试转换函数
for example in ds:
    doc = mlqaen_example_to_document(example)
    print(doc)
    break

这里我们使用自定义函数将数据集示例转换为 Document 格式，context 字段作为 Document 的 page_content，其他字段作为 metadata。

优化建议分享

在实际操作中，你可能需要处理更大规模的数据集。此时，建议使用 TensorflowDatasetLoader，它可以帮助你设定加载文档的数量，还能自定义样本到 Document 的转换函数，确保数据处理的稳定性。

from langchain_community.document_loaders import TensorflowDatasetLoader

loader = TensorflowDatasetLoader(
    dataset_name="mlqa/en",
    split_name="test",
    load_max_docs=3,
    sample_to_document_function=mlqaen_example_to_document,
)

docs = loader.load()
print(len(docs))  # 输出: 3

补充说明和总结

说白了，TensorFlow Datasets 提供了标准化的数据集，可以极大减轻数据准备的工作量，而将其转换为 Document 格式后，在处理自然语言处理任务时更加方便。老铁们若是在开发中遇到问题，可以在评论区交流哦。

今天的技术分享就到这里，希望对大家有帮助。开发过程中遇到问题也可以在评论区交流~

—END—