Package: datatrove
DataTroveDataTrove is a library to process, filter and deduplicate text data at a very large scale. It provides a set of prebuilt commonly used processing blocks with a framework to easily add custom functionality.DataTrove processing pipelines are platform-agnostic, running out of the box locally or on a slurm cluster. Its (relatively) low memory usage and multiple step design makes it ideal for large workloads, such as to process an LLM's training data.Local, remote and other file systems are supported throughfsspec.Table of contentsInstallationQuickstart examplesPipelineDataTrove DocumentTypes of pipeline blocksFull pipelineExecutorsLocalPipelineExecutorSlurmPipelineExecutorLoggingDataFolder / pathsPractical guidesReading dataExtracting textFiltering dataSaving dataDeduplicating dataCustom blocksSimple dataCustom functionCustom blockContributingInstallationpipinstalldatatrove[FLAVOUR]Available flavours (combine them with,i.e.[processing,s3]):allinstalls everything:pip install datatrove[all]iodependencies to readwarc/arc/wetfiles and arrow/parquet formats:pip install datatrove[io]processingdependencies for text extraction, filtering and tokenization:pip install datatrove[processing]s3s3 support:pip install datatrove[s3]clifor command line tools:pip install datatrove[cli]Quickstart examplesYou can check the followingexamples:process_common_crawl_dump.pyfull pipeline to read commoncrawl warc files, extract their text content, filters and save the resulting data to s3. Runs on slurmtokenize_c4.pyreads data directly from huggingface's hub to tokenize the english portion of the C4 dataset using thegpt2tokenizerminhash_deduplication.pyfull pipeline to run minhash deduplication of text datasentence_deduplication.pyexample to run sentence level exact deduplicationexact_substrings.pyexample to run ExactSubstr (requiresthis repo)PipelineDataTrove DocumentEach pipeline block processes data in the datatroveDocumentformat:textthe actual text content for each sampleida unique id (string) for this samplemetadataa dictionary where any additional info may be storedTypes of pipeline blocksEach pipeline block takes a generator ofDocumentas input and returns another generator ofDocument.readersread data from different formats and yieldDocumentwriterssaveDocumentto disk/cloud in different formatsextractorsextract text content from raw formats (such as webpage html)filtersfilter out (remove) someDocuments based on specific rules/criteriastatsblocks to collect statistics on the datasettokensblocks to tokenize data or count tokensdedupblocks for deduplicationFull pipelineA pipeline is defined as a list of pipeline blocks. As an example, the following pipeline would read data from disk, randomly filter (remove) some documents and write them back to disk:fromdatatrove.pipeline.readersimportCSVReaderfromdatatrove.pipeline.filtersimportSamplerFilterfromdatatrove.pipeline.writersimportJsonlWriterpipeline=[CSVReader(data_folder="/my/input/path"),SamplerFilter(rate=0.5),JsonlWriter(output_folder="/my/output/path")]ExecutorsPipelines are platform-agnostic, which means that the same pipeline can smoothly run on different execution environments without any changes to its steps. Each environment has its own PipelineExecutor.
Some options common to all executors:pipelinea list consisting of the pipeline steps that should be runlogging_dira datafolder where log files, statistics and more should be saved. Do not reuse folders for different pipelines/jobs as this will overwrite your stats, logs and completions.skip_completed(bool,Trueby default) datatrove keeps track of completed tasks so that when you relaunch a job they can be skipped. Set this toFalseto disable this behaviourCall an executor'srunmethod to execute its pipeline.[!TIP]
Datatrove keeps track of which tasks successfully completed by creating a marker (an empty file) in the${logging_dir}/completionsfolder. Once the job finishes, if some of its tasks have failed, you cansimply relaunch the exact same executorand datatrove will check and only run the tasks that were not previously completed.[!CAUTION]
If you relaunch a pipeline because some tasks failed,do not change the total number of tasksas this will affect the distribution of input files/sharding.LocalPipelineExecutorThis executor will launch a pipeline on a local machine.
Options:taskstotal number of tasks to runworkershow many tasks to run simultaneously. If-1, no limit. Anything> 1will use multiprocessing to execute the tasks.start_methodmethod to use to spawn a multiprocessing Pool. Ignored ifworkersis 1Example executorfromdatatrove.executorimportLocalPipelineExecutorexecutor=LocalPipelineExecutor(pipeline=[...],logging_dir="logs/",tasks=10,workers=5)executor.run()SlurmPipelineExecutorThis executor will launch a pipeline on a slurm cluster, using slurm job arrays to group and manage tasks.
Options:taskstotal number of tasks to run.requiredtimeslurm time limit string.requiredpartitionslurm partition.requiredworkershow many tasks to run simultaneously. If-1, no limit. Slurm will runworkerstasks at a time. (default:-1)job_nameslurm job name (default: "data_processing)dependsanother SlurmPipelineExecutor instance, which will be a dependency of this pipeline (current pipeline will only start executing after the depended on pipeline successfully completes)sbatch_argsdictionary with any other arguments you would like to pass to sbatchslurm_logs_folderwhere to save the slurm log files. If using a local path forlogging_dir, they will be saved onlogging_dir/slurm_logs. If not, they will be saved as a subdir of the current directory.Other optionscpus_per_taskhow many cpus to give each task (default:1)qosslurm qos (default: "normal")mem_per_cpu_gbmemory per cpu, in GB (default: 2)env_commandcustom command to activate a python environment, if neededcondaenvconda environment to activatevenv_pathpath to a python environment to activatemax_array_sizetheMaxArraySizevalue in$ scontrol show config. If number of tasks exceeds this number, it will split into multiple array jobs (default: 1001)max_array_launch_parallelif we need multiple jobs due to max_array_size, whether to launch them all in one go (parallel) or sequentially (default:False)stagger_max_array_jobswhen max_array_launch_parallel is True, this determines how many seconds to wait between launching each of the parallel jobs (default:0)run_on_dependency_failstart executing when a job we depend on finishes even if it has failed (default:False)randomize_startrandomize the start of each task in a job in a ~3 min window. Useful when heavily hitting an s3 bucket for example. (default:False)Example executorfromdatatrove.executorimportSlurmPipelineExecutorexecutor1=SlurmPipelineExecutor(pipeline=[...],job_name="my_cool_job1",logging_dir="logs/job1",tasks=500,workers=100,# omit to run all at oncetime="10:00:00",# 10 hourspartition="hopper-cpu")executor2=SlurmPipelineExecutor(pipeline=[...],job_name="my_cool_job2",logging_dir="logs/job2",tasks=1,time="5:00:00",# 5 hourspartition="hopper-cpu",depends=executor1# this pipeline will only be launched after executor1 successfuly completes)# executor1.run()executor2.run()# this will actually launch executor1, as it is a dependency, so no need to launch it explicitlyLoggingFor a pipeline withlogging_dirmylogspath/exp1, the following folder structure would be created:See folder structure└── mylogspath/exp1
    │── executor.json ⟵ json dump of the executor options and pipeline steps
    │── launch_script.slurm ⟵ the slurm config created and used to launch this job (if running on slurm)
    │── executor.pik ⟵ the slurm config created and used to launch this job (if running on slurm)
    │── ranks_to_run.json ⟵ list of tasks that are being run
    │── logs/
    │   └──[task_00000.log, task_00001.log, task_00002.log, ...] ⟵ individual logging files for each task
    │── completions/
    │   └──[00004, 00007, 00204, ...] ⟵ empty files marking a task as completed. Using when relaunching/resuming a job (only unfinished tasks will be run)
    │── stats/
    │   └──[00000.json, 00001.json, 00002.json, ...] ⟵ individual stats for each task (number of samples processed, filtered, removed, etc)
    └── stats.json ⟵ global stats from all tasksDataFolder / pathsDatatrove supports a wide variety of input/output sources throughfsspec.There are a few ways to provide a path to a datatrove block (forinput_folder,logging_dir,data_folderand so on arguments):str: the simplest way is to pass a single string. Example:/home/user/mydir,s3://mybucket/myinputdata,hf://datasets/allenai/c4/en/(str, fsspec filesystem instance): a string path and a fully initialized filesystem object. Example:("s3://mybucket/myinputdata", S3FileSystem(client_kwargs={"endpoint_url": endpoint_uri}))(str, dict): a string path and a dictionary with options to initialize a fs. Example (equivalent to the previous line):("s3://mybucket/myinputdata", {"client_kwargs": {"endpoint_url": endpoint_uri}})DataFolder: you can initialize aDataFolderobject directly and pass it as an argumentUnder the hood these argument combinations are parsed byget_datafolder.Practical guidesReading dataUsually, pipelines will start with aReaderblock.
Most readers take adata_folderargument — a path to a folder containing the data to be read.These files will be distributed across each task. If you haveNtasks, task with ranki(0-based) will process filesi, i+N, i+2N, i+3N,....Internally, each reader reads data and converts it into a dictionary before creating aDocumentobject.Some options common to most readers:text_keythe dictionary key containing the text content for each sample. Default:textid_keythe dictionary key containing the id for each sample. Default:iddefault_metadataa dictionary for any default metadata values you would like to add (such as their source, for example)recursivewhether to look for files recursively indata_folder's subdirectoriesglob_patternuse this field to match specific files. For instance,glob_pattern="*/warc/*.warc.gz"will match files with a.warc.gzfile extension on thewarc/folder of each of thedata_folder's subdirectoriesadapterthis function takes the raw dictionary obtained from the reader and returns a dictionary withDocument's field names. You may overwrite this function (_default_adapter) if you would like.limitread only a certain number of samples. Useful for testing/debuggingExtracting textYou can useextractorsto extract text content from raw html. The most commonly used extractor in datatrove isTrafilatura, which uses thetrafilaturalibrary.Filtering dataFiltersare some of the most important blocks of any data processing pipeline. Datatrove's filter blocks take aDocumentand return a boolean (Trueto keep a document,Falseto remove it). Removed samples do not continue to the next pipeline stage. You can also save the removed samples to disk by passing aWriterto theexcluded_writerparameter.Saving dataOnce you are done processing your data you will probably want to save it somewhere. For this you can use awriter.
Writers require anoutput_folder(the path where data should be saved). You can choose thecompressionto use (default:gzip) and the filename to save each file as.
For theoutput_filename, a template is applied using the following arguments:${rank}replaced with the current task's rank. Note that if this tag isn't present,different tasks may try to write to the same location${id}replaced with the sample idmetadata: any other${tag}will be replaced with the correspondingdocument.metadata['tag']valueAn example to separate samples by language based on theirlangmetadata field:JsonlWriter(
    f"{MAIN_OUTPUT_PATH}/non_english/",
    output_filename="${language}/" + DUMP + "/${rank}.jsonl.gz",  # folder structure: language/dump/file
)Deduplicating dataFor deduplication check the examplesminhash_deduplication.py,sentence_deduplication.pyandexact_substrings.py.Custom blocksSimple dataYou can pass an iterable ofDocumentdirectly as a pipeline block like so:fromdatatrove.dataimportDocumentfromdatatrove.pipeline.filtersimportSamplerFilterfromdatatrove.pipeline.writersimportJsonlWriterpipeline=[[Document(text="some data",id="0"),Document(text="some more data",id="1"),Document(text="even more data",id="2"),],SamplerFilter(rate=0.5),JsonlWriter(output_folder="/my/output/path")]Do note, however, that this iterable will not be sharded (if you launch more than 1 task they will all get the full iterable).
This is usually useful for small workloads/testing.Custom functionFor simple processing you can simply pass in a custom function with the following signature:fromdatatrove.dataimportDocumentsPipelinedefuppercase_everything(data:DocumentsPipeline,rank:int=0,world_size:int=1)->DocumentsPipeline:"""`data` is a generator of Document. You must also return a generator of Document (yield)You can optionally use `rank` and `world_size` for sharding"""fordocumentindata:document.text=document.text.upper()yielddocumentpipeline=[...,uppercase_everything,...][!TIP]
You might have some pickling issues due to the imports. If this happens, simply move whatever imports you need inside the function body.Custom blockYou can also define a full block inheriting fromPipelineStepor one of its subclasses:fromdatatrove.pipeline.baseimportPipelineStepfromdatatrove.dataimportDocumentsPipelinefromdatatrove.ioimportDataFolderLike,get_datafolderclassUppercaserBlock(PipelineStep):def__init__(self,some_folder:DataFolderLike,some_param:int=5):super().__init__()# you can take whatever parameters you need and save them hereself.some_param=some_param# to load datafolders use get_datafolder()self.some_folder=get_datafolder(some_folder)defrun(self,data:DocumentsPipeline,rank:int=0,world_size:int=1)->DocumentsPipeline:# you could also load data from the `some_folder`:forfilepathinself.some_folder.get_shard(rank,world_size):# it also accepts a glob pattern, among other thingswithself.some_folder.open(filepath,"rt")asf:# do something...yielddoc## OR process data from previous blocks (`data`)#fordocindata:withself.track_time():# you can wrap the main processing code in `track_time` to know how much each document took to processnr_uppercase_letters=sum(map(lambdac:c.isupper(),doc.text))# you can also keep track of stats per document using stat_updateself.stat_update("og_upper_letters",value=nr_uppercase_letters)doc.text=doc.text.upper()# make sure you keep the yield outside the track_time block, or it will affect the time calculationyielddoc## OR save data to disk#withself.some_folder.open("myoutput","wt")asf:fordocindata:f.write(doc...)pipeline=[...,UppercaserBlock("somepath"),...]You could also inherit fromBaseExtractor,BaseFilter,BaseReader/BaseDiskReader, orDiskWriter.Contributinggitclonegit@github.com:huggingface/datatrove.git&&cddatatrove
pipinstall-e".[dev]"Install pre-commit code style hooks:pre-commitinstallRun the tests:pytest-sv./tests/