DataTrove: 大规模文本数据处理利器

datatrove

DataTrove简介

DataTrove是Hugging Face开源的一个强大的文本数据处理库,专门用于大规模处理、过滤和去重文本数据。它提供了一系列预构建的常用处理模块,以及一个易于添加自定义功能的框架,使得处理海量文本数据变得简单高效。

DataTrove的处理流水线设计跨平台兼容,可以在本地机器或Slurm集群上无缝运行。其相对较低的内存占用和多步骤设计非常适合处理大规模工作负载,比如处理LLM的训练数据。通过fsspec,DataTrove还支持本地、远程和其他文件系统。

DataTrove架构图

主要特性

大规模处理能力: 专为处理TB级数据设计,可以高效地处理和过滤海量文本数据。
模块化设计: 提供了一系列预构建的处理模块,包括读取器、写入器、提取器、过滤器等,可以灵活组合。
跨平台兼容: 同一个处理流水线可以在本地机器或Slurm集群上无缝运行。
低内存占用: 采用流式处理和多步骤设计,即使在处理海量数据时也能保持较低的内存占用。
易于扩展: 提供了清晰的接口用于添加自定义功能模块。
多种文件系统支持: 通过fsspec支持本地、S3、HDFS等多种文件系统。
丰富的统计功能: 内置了多种统计模块,可以方便地收集数据集的各种统计信息。

安装方法

DataTrove的安装非常简单,只需使用pip即可:

pip install datatrove[FLAVOUR]

其中FLAVOUR可以是以下选项(可以用逗号组合):

all: 安装所有依赖
io: 安装读取warc/arc/wet文件和arrow/parquet格式的依赖
processing: 安装用于文本提取、过滤和分词的依赖
s3: 安装S3支持
cli: 安装命令行工具

例如,要安装所有功能:

pip install datatrove[all]

快速入门

DataTrove提供了多个示例来帮助用户快速上手:

fineweb.py: 完整重现FineWeb数据集的处理流程。
process_common_crawl_dump.py: 完整的流水线示例,用于读取CommonCrawl的warc文件,提取文本内容,过滤并将结果保存到S3。在Slurm上运行。
tokenize_c4.py: 直接从Hugging Face的hub读取数据,使用gpt2分词器对C4数据集的英文部分进行分词。
minhash_deduplication.py: 使用MinHash算法对文本数据进行去重的完整流水线示例。
sentence_deduplication.py: 在句子级别进行精确去重的示例。

这些示例涵盖了DataTrove的主要使用场景,可以帮助用户快速理解和上手使用DataTrove进行大规模文本数据处理。

核心概念

在深入了解DataTrove的使用之前,我们需要先理解一些核心概念:

pipeline: 要执行的处理步骤列表(读取数据、过滤、写入磁盘等)。
executor: 在特定执行环境(Slurm、多CPU机器等)上运行特定pipeline的执行器。
job: 在给定executor上执行pipeline的过程。
task: 一个job由多个task组成,用于并行执行。通常每个task处理一个数据分片。
file: 单个输入文件(.json、.csv等)。
shard: 输入数据的分组(通常是一组文件),会被分配给特定的task处理。
worker: 执行单个task的计算资源。

了解这些概念有助于我们更好地设计和优化数据处理流程。

处理流水线

DataTrove的核心是其处理流水线(pipeline)。每个pipeline由一系列处理模块(block)组成,每个模块处理Document格式的数据:

class Document:
    text: str  # 每个样本的实际文本内容
    id: str    # 该样本的唯一ID
    metadata: dict  # 存储任何额外信息的字典

流水线模块类型

DataTrove提供了多种类型的处理模块:

readers: 从不同格式读取数据并生成Document
writers: 将Document保存到磁盘/云端
extractors: 从原始格式(如网页HTML)中提取文本内容
filters: 基于特定规则/标准过滤(删除)某些Document
stats: 收集数据集统计信息的模块
tokens: 用于分词或计算token的模块
dedup: 用于去重的模块

示例流水线

下面是一个简单的流水线示例,它从磁盘读取数据,随机过滤一些文档,然后将结果写回磁盘:

from datatrove.pipeline.readers import CSVReader
from datatrove.pipeline.filters import SamplerFilter
from datatrove.pipeline.writers import JsonlWriter

pipeline = [
    CSVReader(
        data_folder="/my/input/path"
    ),
    SamplerFilter(rate=0.5),
    JsonlWriter(
        output_folder="/my/output/path"
    )
]

执行器

DataTrove的流水线设计是平台无关的,这意味着同一个流水线可以在不同的执行环境中无缝运行。每种环境都有自己的PipelineExecutor。

LocalPipelineExecutor

这个执行器会在本地机器上启动流水线。主要选项包括:

tasks: 要运行的总任务数
workers: 同时运行多少个任务
start_method: 用于生成多进程Pool的方法

示例:

from datatrove.executor import LocalPipelineExecutor

executor = LocalPipelineExecutor(
    pipeline=[
        ...
    ],
    logging_dir="logs/",
    tasks=10,
    workers=5
)
executor.run()

SlurmPipelineExecutor

这个执行器会在Slurm集群上启动流水线,使用Slurm作业数组来管理任务。主要选项包括:

tasks: 要运行的总任务数
time: Slurm时间限制字符串
partition: Slurm分区
workers: 同时运行多少个任务
job_name: Slurm作业名称

示例:

from datatrove.executor import SlurmPipelineExecutor

executor = SlurmPipelineExecutor(
    pipeline=[
        ...
    ],
    job_name="my_cool_job",
    logging_dir="logs/job1",
    tasks=500,
    workers=100,
    time="10:00:00",
    partition="hopper-cpu"
)
executor.run()