Hugging Face datasets¶
pyroid.integrations.hf has functions to pass to Dataset.map and Dataset.filter with
batched=True: ftfy, fastText language ID, the Gopher, C4 and FineWeb quality filters, and the
whole FineWeb recipe with MinHash dedup. On Arrow-format datasets pyroid reads the text column in
place, without building a Python str per document. hf.fineweb(ds) keeps the same documents,
with the same text, as datatrove 0.10.1's FineWeb pipeline
(parity). Tested with datasets 3.6, 4.8 and 5.1.
$ pip install "pyroid[hf]" # datasets, pyarrow, numpy
Change one line¶
-ds = ds.map(lambda doc: {"text": ftfy.fix_text(doc["text"])})
+ds = ds.with_format("arrow").map(hf.fix_text(), batched=True)
The pieces of the FineWeb recipe, one call each:
import datasets
from pyroid.integrations import hf
ds = datasets.Dataset.from_parquet("sample.parquet").with_format("arrow")
ds = ds.map(hf.fix_text(), batched=True) # ftfy on the text column
ds = ds.map(hf.langid("lid.176.bin"), batched=True) # + language, language_score
ds = ds.filter(hf.quality_mask(), batched=True) # Gopher, C4 and FineWeb filters
ds.num_rows # 170 of the 200 documents
The functions¶
Each returns a callable for batched=True:
| Function | For | What the callable does |
|---|---|---|
fix_text(column="text", config=None, *, threads=None, **kwargs) |
map |
replaces column with ftfy.fix_text of it (pyroid.ftfy); config and **kwargs as for fix_text_batch |
langid(lid_model=None, column="text", *, k=1, threads=None, language_column="language", score_column="language_score") |
map |
adds fastText's top label (without __label__) and its probability; with k other than 1 (-1: all labels), list columns, best first |
quality_mask(column="text", filters=None, *, threads=None) |
filter |
keeps what passes the quality filters, run in order as datatrove does (pyroid.filters.apply); filters=None is FineWeb's four |
annotate(column="text", *, lid_model=None, threads=None) |
map |
the FineWeb chain but the dedup: sets the cleaned text and the columns keep, reason, language, language_score |
keep_mask(column="keep") |
filter |
keeps the rows whose boolean column is true (e.g. annotate's keep) |
Each document goes to fastText as one line (newlines become spaces), as in datatrove's
LanguageFilter. quality_mask uses only the filters' decision, so C4's line removal is not
applied to the text. annotate and fineweb give the text as it leaves the chain. threads is
pyroid's threads= (Threads), and None uses the whole thread pool for
each batch.
annotate's reason is None for a kept document, else "lid",
"<stage>:<datatrove reason>" or "null", as in
pyroid clean:
ds = datasets.Dataset.from_parquet("sample.parquet")
ann = ds.with_format("arrow").map(hf.annotate(lid_model="lid.176.bin"), batched=True)
ann.with_format(None).select_columns(["keep", "reason"])[:2]
# {'keep': [False, True], 'reason': ['gopher_quality:gopher_below_alpha_threshold', None]}
kept = ann.filter(hf.keep_mask(), batched=True)
kept.num_rows # 170
Arrow batches and Python batches¶
The callables take both kinds of batch that datasets passes:
- With Arrow batches (
ds.with_format("arrow"), the zero-copy path), the callable gets apyarrow.Tableand pyroid reads the text column in place. Measured on a 3 MB batch,fix_textpeaked at 932 bytes of Python heap, with no growth of pyarrow's memory pool. A map callable returns apyarrow.Tablewith every column of the batch, whichdatasetsuses in place of the batch. Soremove_columns=has no effect; call.remove_columns()on the dataset instead. A filter callable returns a booleanpyarrow.ChunkedArray. - With Python batches (the default format), the callable gets the dict of lists. It returns a
dict of the new or changed columns, which
datasetsmerges into the batch, or alist[bool]forfilter.
None texts stay None, and filters drop them (an Arrow null in a mask drops the row too). Don't
pass input_columns=: the callables read their column by name.
ds = datasets.Dataset.from_dict({"text": ["été 2024", None, "plain"]})
ds = ds.map(hf.fix_text(), batched=True) # a Python-format dataset
ds.to_dict() # {'text': ['été 2024', None, 'plain']}
The callables also work in a streaming IterableDataset's map and filter. An Arrow-format
IterableDataset.filter needs datasets 5; older versions reject it, whatever the function.
fineweb needs a Dataset, because MinHash reads the survivors back.
The whole recipe: fineweb(ds)¶
hf.fineweb runs FineWeb's recipe on a Dataset, MinHash included:
ds.map(annotate(...), batched=True)on Arrow batches: ftfy, language ID (English, probability above 0.65) and the four quality filters;- MinHash with datatrove's FineWeb settings over the survivors, read back in dataset order. Each
cluster of near-duplicates keeps one document, as in datatrove, and the others get
reason="minhash". This step runs in this process over the whole dataset, so its result does not depend onbatch_sizeornum_proc.
import datasets
from pyroid.integrations import hf
res = hf.fineweb(datasets.Dataset.from_parquet("sample.parquet"), lid_model="lid.176.bin")
res.stats.docs, res.stats.kept # (200, 170)
res.stats.dropped
# {'lid': 0, 'gopher_repetition': 0, 'gopher_quality': 27, 'c4_quality': 0,
# 'fineweb_quality': 3, 'minhash': 0}
res.stats.reasons
# {'gopher_quality:gopher_below_alpha_threshold': 24, 'fineweb_quality:line_punct_ratio': 3,
# 'gopher_quality:gopher_too_many_end_ellipsis': 3}
res.dataset.to_parquet("clean.parquet")
It returns a FineWebResult:
dataset: the kept documents, in order, with the cleaned text, the input's other columns andlanguage/language_score;annotated: every document, with the text as it left the chain and the columnskeep(the final decision) andreason("minhash"for a near-duplicate);stats: documents in and kept, the drops per stage (in chain order) and per reason, and wall times (annotate,dedup,total);print(res.stats)gives a summary.
Both datasets get the input's format back. **map_kwargs go to step 1's Dataset.map
(keep_in_memory, load_from_cache_file, cache_file_name, desc, ...). An empty dataset
raises ValueError. annotate and fineweb run the recipe object that pyroid clean uses,
pyroid.integrations.recipes.FineWeb, which you
can also call on Arrow batches yourself.
Caching and num_proc¶
The callables pickle by their settings. So datasets can fingerprint them, reusing a cached
result when you run the same map again, and ship them to num_proc worker processes, each of
which loads the fastText model itself. A loaded pyroid.fasttext.Model can't be pickled, so pass
the model's path. pyroid already runs every batch on all cores, so num_proc rarely helps. If you
use it, threads of about os.cpu_count() // num_proc avoids oversubscription.
lid_model=None finds the model as the command line does: $PYROID_LID_MODEL, else a cached
lid.176.bin, downloaded once and sha256-checked
(details).
The notebook¶
examples/notebooks/clean_fineweb_in_60_seconds.ipynb
streams the first 10,000 documents of FineWeb's sample-10BT, runs hf.fineweb on them, and
prints the counts per stage and per reason, the time it took and a few dropped documents per
stage. Then it writes the kept documents to Parquet. It is ready for Colab.
To run it on local files without network, set PYROID_NOTEBOOK_DATA (a FineWeb Parquet file),
PYROID_LID_MODEL, and optionally PYROID_NOTEBOOK_DOCS and PYROID_NOTEBOOK_OUT. pyroid's test
suite executes it that way, headless, and checks the counts against datatrove's pipeline. See the
notebook's README.
The Gradio app¶
examples/space/ is a minimal Gradio
app, ready for a Hugging Face Space but not deployed. On the Paste text tab you see what each
stage of the chain decides for one document: whether ftfy changed it, the language and its
probability, and which filter dropped it and why. On the Upload a file tab you run the whole
recipe with hf.fineweb on a Parquet or JSONL file with a text column, and download the
cleaned file. Every input is capped for a public demo (characters, upload size, rows). Run
locally, it serves on http://127.0.0.1:7860 only:
$ pip install -r requirements.txt
$ python app.py
Speed¶
The callables call pyroid's batch functions on each datasets batch, so they run at the
modules' speed. See pyroid.ftfy,
pyroid.fasttext, pyroid.filters and
pyroid.dedup. This page has no separate timings for them.