diff --git a/CHANGELOG_DEV.md b/CHANGELOG_DEV.md index 43d0c6e2d..f4cf88213 100644 --- a/CHANGELOG_DEV.md +++ b/CHANGELOG_DEV.md @@ -217,4 +217,791 @@ This PR improves training monitoring and logging across runs besides some other * Add tutorials on Einsum Transformer (Example model integration) and profiling **Breaking Changes** -* experiments_root_path is now exposed on an API level \ No newline at end of file +* experiments_root_path is now exposed on an API level + +## PR #XXX Quality-based document selection and up/downsampling + +This PR adds a way to build a training blend by filtering documents on quality signals +and choosing how heavily each dataset is sampled, plus a fast way to see the token +budget a given selection yields before committing to a tokenization run. + +**Motivation** + +Two kinds of quality signal exist in practice: metrics a corpus already carries in its +own records (`fw_edu_scores`, `proxy_score`, `finemath_scores`, `perplexity`, ...), and +external per-document annotations that have to be joined on. Neither could be used to +shape a blend, and the only way to change a dataset's share was to duplicate it on disk +-- which is what the `_epoch_1` / `_epoch_2` directory convention did. + +**General changes** + +* New `modalities quality` command group with one subcommand per stage: + `calibrate`, `build-sidecar`, `join-annotations`, `build-cube`, `preview`, `apply` + and `write-packing-configs`. +* New `src/modalities/dataloader/preprocessing/quality/` package: + * `registry` declares each dataset's source and how it joins to annotations. Four key + kinds are supported, covering corpora that store a plain id, an id wrapped in + ``, no id at all (keyed by a hash of the text), and a `/` + pointer into a separate source corpus. + * `tokens` measures per-dataset token estimators. Estimates are per document and + based on the text rather than the stored line, because quality correlates with + length and several corpora keep multiple renderings of a document in one record. + * `sidecar` streams each JSONL once and records one row per document: position, + length, estimated tokens, join key and native metrics. + * `annotation_join` joins annotations by bucketing both sides on a hash of the key, + so a split of billions of rows never needs a single hash table. Coverage, + duplicate keys and unmatched documents are reported rather than hidden. + * `cube` aggregates a sidecar into grouped document and token counts, which is what + makes `preview` return in microseconds. Thresholds landing on a bin edge are exact; + one landing inside a bin is reported as interpolated instead of silently guessed. + * `selection` evaluates a YAML selection over both annotation labels and native + metrics, with ordinal scales declared explicitly. + * `materialize` writes the selection out as filtered `.idx` files. +* New `WeightedCombinedDataset` (component `dataset`/`weighted_combined`), which takes a + float repeat factor per dataset. A ratio of 2.5 draws a dataset two and a half times + per epoch and 0.3 draws three tenths of it, without duplicating anything on disk. The + partial pass is chosen by a seeded affine permutation, so it is deterministic across + ranks and restarts and spreads across the whole dataset rather than taking a prefix. +* New `TokenizerInstantiationModel`, so a tool can reuse a packing config for its + tokenizer without also having to satisfy that config's `settings`. + +**Notes** + +Selection produces a filtered index rather than a filtered copy of the corpus: +`PackedDataGenerator` already tokenizes exactly the documents its index lists, so +`pack_encoded_data` consumes the output unchanged and no packing code was touched. An +ablation therefore costs megabytes of index rather than a second copy of the data, and +the source tree is never written to. + +**Breaking Changes** + +None. `CombinedDataset` and every existing config keep working as before. + + +## PR #XXX Quality selection: performance and sharding + +Follow-up to the quality-selection PR, from profiling the pipeline against the real +43 TB blend rather than a fixture. Three of the four stages were far slower than they +needed to be, and two of them could not be parallelised at all. + +**General changes** + +* Native metrics declared as a plain field path (`.fw_edu_scores`, + `.metadata.dclm_plus2."__label__1"`) are now read by direct dictionary lookup instead + of jq. `jq.compile(...).input_value(record)` re-serialises the whole document on every + call, which on a 21 KB record cost more than twenty times the rest of building a + sidecar row; the full pass measured **29 MB/s with jq against 374 MB/s without, 13x + end to end**. Patterns jq cannot reduce to a field chain still use jq, and + `build-sidecar` now warns when one does, since that pattern then dominates the stage. +* `build_cube` groups with Arrow's C++ kernels instead of a Python loop over documents, + batching row groups so cardinality saturates before each grouping pass: + **246k -> 2.39M rows/s, 9.9x**. For the full blend that is ~50 minutes rather than + ~8.5 hours, and the stage was not shardable, so it was a hard floor. +* `build-sidecar` takes `--shard_id`/`--num_shards`. Work is divided per JSONL file + across every selected dataset, so one array covers the whole blend however unevenly + the file counts fall. Previously the only split was per dataset, leaving a floor of + the slowest dataset -- ~71 h for `finepdfs-en`. +* Annotation bucketing moved out of `join-annotations` into its own shardable + `bucket-annotations` stage. Each task writes its own file per bucket and the join + reads all of them, so the result is identical to a single-task run. Bucketing 13.9 bn + annotation rows was ~32 h serial, with a ~8.7 h floor from the largest single split. +* `join-annotations` refuses to run against an incomplete bucketing run rather than + silently dropping the annotations a missing task was carrying, which would have looked + exactly like a corpus that was never annotated. + +**Notes** + +Measured on this cluster: sequential read from `/data` is ~282 MB/s per stream and +~3.8 GB/s aggregate. With the jq fix the sidecar pass is I/O-bound, so beyond ~16 +concurrent tasks the storage is the limit rather than the code. End to end the one-time +setup goes from ~88 h to ~4 h on two nodes; previewing a selection is unaffected at +~10 s for the whole blend, because it only ever reads the cubes. + +**Breaking Changes** + +* `modalities quality build-sidecar` no longer takes `--file_id`; use + `--shard_id`/`--num_shards`. +* `modalities quality join-annotations` no longer buckets. Run + `modalities quality bucket-annotations` first. Its `--num_buckets` and + `--rebuild_buckets` options moved to that command (`--rebuild_buckets` is now + `--force`). + + +## PR #XXX Fix: calibration read scaled with file count + +`quality calibrate` read a fixed 20,000 lines from *every* file of a dataset to collect a +2,000-document sample, so its cost scaled with the file count rather than the sample size. +Over the real blend that came to **30.3 TB and ~30 hours** -- `dolmino` alone, at 40,003 +files, accounted for 26 TB. Reported from a real run that was still going after an hour +having finished 9 of 19 datasets. + +**General changes** + +* The sampler now draws from at most `max_probe_files` files (default 32), spaced evenly + across the dataset, reading only enough lines from each to fill the sample. Full + calibration of all 19 datasets: **4 minutes**, about 1 GB read. The three worst + datasets together (`dolmino`, `finephrase`, `hplt-de`) now take 67 s. +* `calibration.yaml` is written after each dataset instead of once at the end, so + interrupting the stage keeps what it already measured. Previously an hour of work was + discarded on Ctrl-C. + +**Notes** + +Spread is preserved -- the probe files are spaced across the dataset, not taken from the +front -- and the sample is still trimmed with a seeded choice, so calibration stays +reproducible. Tests cover both: that the number of files opened stays bounded regardless +of dataset size, and that probe files reach both ends of the file list. + + +## PR #XXX Fix: bucket writer buffered the whole input when bucket count was high + +`bucket-annotations` OOM-killed all 64 tasks of a real run at 24 GB each. The writer +flushed a bucket once it held `flush_rows` (100,000) rows, which bounds nothing: spread +50 M rows over 1024 buckets and each holds ~49 k, so no bucket ever reaches the threshold +and the entire input accumulates as Python dicts until the writer closes. + +**General changes** + +* The cap is now on the **total** rows buffered across all buckets (`max_buffered_rows`, + default 500,000); reaching it flushes every bucket. Measured on the shard that caused + the failure -- 50 M rows at 1024 buckets -- peak RSS is **2.36 GB** against the + previous unbounded growth, at an unchanged 121k rows/s. +* `bucket_annotations` refuses to write into a directory holding output from a run with a + different `num_shards`. A sharded run cannot clear the directory, so leftovers would + otherwise be mixed in and a bucket would be read as rows from two incompatible runs. + +**Notes** + +The completeness guard added earlier did its job here: `join-annotations` refused to run +against the partial buckets the OOMed array left behind ("63 of 64 bucketing tasks +finished, missing shard id 0") rather than silently producing sidecars with no labels. +Without it the failure would have surfaced much later as an unexplained 0 % coverage. + +Tests cover the bound directly: 20,000 rows over 1024 buckets with a 1,000-row cap, with +the buffered total asserted after every row, plus that repeated flushes of one bucket +still yield one file with every row. + + +## PR #XXX Fix: metadata write/read race in bucket-annotations + +The guard added in the previous entry -- refusing to mix output from runs with different +array sizes -- read every `_meta.*.json` in a split directory with a bare `json.loads`. +The writer used `Path.write_text`, which truncates before writing, so the file is briefly +empty; sibling tasks of the same array read those files, and 12 of 64 tasks of a real run +died with `JSONDecodeError: Expecting value: line 1 column 1 (char 0)`. + +**General changes** + +* Metadata is written to a per-task `_meta..json.tmp` and renamed onto the final + name. Rename is atomic, so a concurrent reader sees either the old file or the complete + new one. +* Both readers -- the guard and `read_bucket_metadata` -- skip a file they cannot parse + rather than propagating the error, and both ignore `*.tmp`. Skipping is the safe + direction for `read_bucket_metadata`: an unread file leaves its shard id unseen, so the + run reports as incomplete instead of joining missing annotations. If no file at all can + be read it now raises a clear error rather than an `AttributeError`. + +**Notes** + +Verified on the real `finewiki` split: four sequential shards into one directory, 43.1 M +rows, metadata reading back as a complete run with a matching row total and no `.tmp` +left behind. Tests cover a truncated metadata file mid-run, an unreadable file making the +run report incomplete, `.tmp` being ignored, and a thread rewriting metadata while the +guard runs repeatedly. + + +## PR #XXX Fix: the join re-read the whole annotation split per sidecar part + +`join_annotations` looped over sidecar parts on the outside and annotation buckets on the +inside, so the entire bucketed split was read once per part. Because a part's documents +hash across every bucket, that is full re-reads, not partial ones. Measured against the +real blend after bucketing completed: + +``` +hplt-es 502 parts x 107.3 GB = 53.9 TB +nemotron-cc 5,319 parts x 23.1 GB = 122.7 TB +climbmix-en 6,543 parts x 24.6 GB = 161.3 TB +TOTAL 454.2 TB = 451 h of reading +``` + +**General changes** + +* Sidecar parts are processed in batches (`max_batch_keys`, default 20 M documents) and + each annotation bucket is read once per batch. Read amplification drops from the part + count to the batch count: **454 TB to 5.9 TB**, 451 h to 5.8 h serial. +* Within a batch each bucket is filtered with `pyarrow.compute.is_in` before anything is + materialised in Python, so memory is bounded by the batch rather than by the bucket -- a + bucket of a billion-row split holds millions of rows, of which one batch wants a few + thousand. +* Bucket file lists are globbed once and cached, saving 65,536 directory scans per batch on + a 1024-bucket split written by 64 tasks. +* Duplicate annotation keys are now counted once, among the keys the join actually wants. + The previous figure was inflated by the part count: `finewiki-it` reported 868,586 where + the real number is 36,747. +* New `3a_join_annotations.sbatch` runs one array task per annotated dataset, resolving the + dataset from the registry so the mapping cannot drift. Wall time becomes the slowest + dataset (~2 h) rather than the sum. `3b_build_cubes.sbatch` follows it. + +**Notes** + +Verified on the real `finewiki-it` sidecar: 1,799,759 of 1,799,759 documents annotated, +100 % coverage, in 48 s against 107 s for a quarter of the documents before. Tests assert +that the batch size does not change any document's label, that a bucket is never read +twice within a batch, and that a duplicate key is counted once rather than once per part. + + +## PR #XXX Fix: resumable join, and a cube stage that fails clearly + +A `nemotron-cc` join reached `5319/5319` parts and was then killed by the 12-hour limit in +`3a_join_annotations.sbatch` during its final write-back, leaving 10 of 5,319 parts without +labels. `build-cube` then died on that inconsistency with +`KeyError: Field "educational_value" does not exist in schema`, after building 9 cubes and +before attempting 6 datasets that were perfectly healthy. + +**General changes** + +* `join-annotations --resume` skips sidecar parts that already carry the label columns and + reports the count. Finishing the interrupted run took **9 minutes rather than 12 hours**, + skipping 5,309 parts. Off by default: resuming after re-bucketing the annotations would + silently keep the old labels, so continuing has to be asked for explicitly. + `3a_join_annotations.sbatch` passes it when `JOIN_RESUME` is set. +* `build_cube` reads every part's schema rather than assuming they match the first one's, + and raises `CubeError` naming the dataset, how many parts are missing which columns, and + the `--resume` command that finishes the job. +* `build_cubes` no longer aborts the stage on one bad dataset. It builds every healthy one, + logs the failures together and re-raises so the job still exits non-zero. +* Per-dataset `join_report/.json`, plus a merged `join_report.json`. Sixteen + parallel `--only` tasks had been overwriting one shared file, leaving only the last + dataset's coverage. +* `3a_join_annotations.sbatch` time limit 12 h -> 48 h, with the measured `nemotron-cc` + figure recorded next to it. + +**Notes** + +The join being CPU-bound in Python is why `nemotron-cc` took 12 h where bytes-read implied +2 h. Resolving in Arrow with `pc.index_in` and `Table.take` instead of per-key Python dicts +should be worth 10-50x and is worth doing, but it is a separate change and not needed to +produce a blend. + + +## PR #XXX Fix: preview refused to admit it was scanning 1.7 bn documents + +`evaluate_blend` caught the `SelectionError` a cube raises for a field it was not grouped +on and quietly scanned the per-document sidecar instead. A sidecar scan is exact but reads +every document, so a preview advertised as taking seconds ran for over ten minutes on the +real blend: `nemotron-cc` thresholded `commercial_bias`, which the join attaches but the +cube does not group, and `dolmino` thresholded `dclm_plus2`, which exists only under one of +its subdirectories and was null in all of 400 sampled sidecar parts. + +**General changes** + +* The fallback is now opt-in via `--allow_fallback`. Without it every unanswerable + predicate is collected and reported together, naming the dataset, the field, the + dimensions the cube does carry, and how many documents a scan would read. The real + selection now fails in **16 s** with both problems named, instead of hanging. +* `build-cube --label_dimension` (repeatable) chooses which annotation columns to group on, + so a field a selection needs can be added. The flag replaces the default seven rather + than extending it, and each field multiplies the cell count by its number of levels. +* The example selection no longer thresholds fields the default cubes lack, with a comment + at each site saying why and how to re-enable it. The registry records that dolmino's + `dclm_plus2` is confined to `stem-heavy-crawl`. + +**Notes** + +First full preview of the real blend: 13.7 s across 19 cubes, 3.04 T effective tokens +against a 400 B target. Coverage is 100 % for finewiki, HPLT, Nemotron-CC, ClimbMix and +KletterMix, and 92.5-99.7 % for FinePDFs. + + +## PR #XXX Fix: pin the source file list a sidecar was built against + +A sidecar row locates its document by `(file_id, byte_offset, byte_len)`, and `file_id` +was a position in a file list re-derived from the filesystem at every stage. That makes +every recorded offset depend on the source tree never changing, with nothing recording +what it looked like. + +An ongoing transfer then re-sharded four corpora after their sidecars were built. +`Nemotron-CC` went from 606 MB files to 137 MB files while keeping nearly the same file +count, so the count comparison that existed passed and the offsets pointed past end of +file. Eleven of nineteen datasets were unusable and only one -- whose file count fell to +zero -- failed loudly. The rest would have packed a blend of wrong byte ranges. + +**General changes** + +* `build-sidecar` records its file list in `sidecar//_files.json`, written + atomically because a sharded build has every task describing the same list. +* `apply` resolves ids through that manifest instead of re-globbing, so a file added to + the tree cannot renumber anything, and refuses to run if a recorded file changed size or + disappeared. Paths are relative, so moving or snapshotting a tree stays valid. +* New `quality verify-sidecar`: seeks to recorded offsets and compares the document found + against the recorded text length, catching a file rewritten at the same size. Rows at + offset zero are skipped -- the first document of any JSONL file parses, so they succeed + against a completely different file, which is how the broken sidecars looked healthy. + `--adopt` stamps a manifest onto a pre-existing sidecar, but only one that verifies. +* `build_sidecars` no longer defaults its index root to the source tree. `SidecarBuilder` + writes a `.idx` beside each JSONL when given no index root, which would modify a shared + read-only corpus; it now defaults to `work_dir/idx`. + +**Testing** + +* `slurm/make_smoke_snapshot.py` freezes ~1 GB of five corpora, chosen to cover all four + distinct join-key kinds plus the native-metrics-only path -- every branch of the join -- + and `smoke_registry.yaml` / `smoke_selection.yaml` run the full pipeline over it in + minutes rather than 15 hours. +* `slurm/check_smoke_run.py` compares packed token counts against the preview's estimates, + loads the output as a `WeightedCombinedDataset` including a fractional repeat factor, + and asserts nothing was written under the source root. +* `test_file_manifest.py` covers the re-shard that preserves the file count, a prepended + file that would renumber ids, a removed file, `apply` refusing a drifted tree, and that + adoption is refused for a sidecar that does not verify. + + +## PR #XXX Fix: a resumed join reported 0% coverage + +`--resume` counted only the parts it re-joined, so a run that skipped everything wrote a +report saying 0 documents and 0.0 coverage. On the smoke run it overwrote three datasets' +genuine coverage figures with zeros, which reads as a failed join rather than a skipped +one. Coverage is a property of the sidecar, not of the run, so skipped parts now +contribute their existing labels to the totals, and `n_parts_resumed` records how many +were not redone. + + +## PR #XXX Fix: the token estimator was 16-19% out, and looked stable while being wrong + +The end-to-end smoke run compared the preview's estimates against a real packing run for +the first time. Document counts matched exactly, but `finewiki-de` tokens were 10.7% out, +and chasing that found two independent defects in the calibration. + +**The sample was a prefix.** `_sample_documents` took the first N documents of each probe +file. Being deterministic, it produced identical ratios across every seed, which reads as +stability rather than as a sample that never moves. On the FineWiki snapshot the first +2,000 documents gave 3.531 bytes per token where the whole file gives 4.214 -- 16% out, +applied to every token figure downstream. I introduced this when bounding an earlier 30 TB +read: I capped the read by taking a prefix and never made the within-file sample spread. + +**One global ratio cannot describe a corpus.** FineWiki's ratio runs from 3.571 for +documents under a kilobyte to 34.648 for the nine documents above 256 KB -- and those nine +hold 5.7% of all bytes. A single sum ratio is therefore hostage to whether the sample +caught them, which is why the global estimator's error swung between -19.4% and +4.2% +across seeds. + +**General changes** + +* Documents are sampled at offsets spread evenly across each file, and the document + *containing* each offset is taken rather than the one following it. That makes selection + proportional to length, which is what a byte-weighted ratio needs: the top stratum went + from 2 sampled documents per 2,000 to 44. +* The ratio is measured per size stratum (log-spaced, six of them) and applied per document + from the length the sidecar already records exactly. A stratum reached by fewer than 20 + documents falls back to the corpus-wide ratio rather than becoming an estimator of its + own. +* The corpus-wide ratio now uses inverse-probability weights. Under length-proportional + sampling a plain sum ratio is weighted by the square of length and came out 62% low. +* Each document's ratio is measured on a slice of at most 64 KB, taken from a random + position inside it. Length-proportional sampling means the multi-megabyte documents do get + sampled, and tokenizing them in full took calibration from 1.5 minutes to over 10; a + document's ratio is far more uniform within itself than across the corpus, so a slice + measures it well. Calibration is now ~15 s per dataset. +* `--sample_size` default raised from 2000 to 4000, for margin. + +**Result on the FineWiki snapshot**, against the true token count of all 27,846 documents: + +| estimator | worst error over 5 seeds | +|---|---| +| global ratio, prefix sample (before) | 16.2% bias, invisible across seeds | +| global ratio, spread sample | 19.4% | +| stratified, spread sample | 0.8% | + +**Notes** + +Calibrations written before this change have no strata and fall back to the global ratio, +so they still load; they should be re-measured. Changing a calibration changes `est_tokens` +in the sidecars, which means rebuilding them -- worth folding into the rebuild the source +re-shard already forces. + + +## PR #XXX End-to-end validation of the estimator against a real packing run + +Re-ran the whole smoke pipeline after the estimator fix. Estimated against packed tokens, +with document counts as the exact control: + +| dataset | est tokens | packed | error | docs selected | docs packed | +|---|---|---|---|---|---| +| finewiki-de | 18,506,144 | 18,438,976 | -0.36% | 20,693 | 20,693 | +| finepdfs-es | 8,644,194 | 8,505,597 | -1.60% | 825 | 825 | +| climbmix-en | 54,527,783 | 54,493,399 | -0.06% | 60,834 | 60,834 | +| klettermix-de | 19,799,809 | 19,818,254 | +0.09% | 22,915 | 22,915 | +| dolmino | 48,320,010 | 48,508,512 | +0.39% | 8,106 | 8,106 | +| **total** | **149,797,940** | **149,764,738** | **-0.02%** | | | + +`finewiki-de` was -10.72% before the fix. Document counts match exactly, which is the +stronger check: the filtered index names exactly the selected documents, so any difference +would be a defect in materialize rather than estimator error. + +The blend also loads: 8 packed files combined through `WeightedCombinedDataset` with repeat +factors 0.5/1.0/1.5/2.0, length 74,033 against 74,032 expected, samples pulled at both +boundaries and the middle. The fractional factors exercise the partial-pass permutation, +which no test on real data had reached. Nothing was written under the source root. + + +## PR #XXX Perf: resolve the join in Arrow, and drop bucket routing + +The join built a `dict[key, list[(part, row)]]` over every document, a per-row dict of +labels for every document, and a Python list comprehension per label column -- a handful of +Python objects per document, at 1.7 bn documents for Nemotron-CC. Each part is now resolved +with one `index_in` against the batch's lookup table and one `take` per label column, and +the key column stays Arrow in the outer loop, where materialising it had been 1.7 bn Python +strings before any joining began. + +Profiling what remained showed the next cost was not per-row work at all: 22 s of reads and +14 s of a thousand separate `is_in` calls, out of 47 s, because a 554 MB split is +partitioned into 1024 files of roughly 540 KB. The routing those buckets exist for also +turned out to decide nothing -- a batch holds millions of keys, which hash across every +bucket, so the profile recorded all 1024 files being read anyway, after a blake2b call per +key in Python to choose them. Routing is gone, replaced by one `pyarrow.dataset` scan with +the key filter pushed into it. Arrow applies the filter per row group and reads in parallel, +so memory stays bounded by matching rows rather than by split size. + +**Measured** on `finewiki-en`, 6.6 M documents against the 43 M-row FineWiki split, both +implementations on separate copies of the same sidecar: + +| | elapsed | per document | +|---|---|---| +| before | 88.7 s | 13.41 us | +| after | 42.1 s | 6.36 us | + +Equivalence was verified twice rather than assumed: 79,375,860 label values identical across +12 columns on the benchmark, and 1,441,584 values identical across all four join-key kinds +on the smoke blend, with duplicate-key counts matching exactly (390 / 2 / 1,936 / 6,025). +That check mattered because "keep the first row seen" had to survive reimplementation as +`index_in` over unique values, which reports each value's first position. + +**What this does not fix, and it is the larger cost.** The annotation split is scanned once +per batch, and batch count scales with documents: + +| dataset | documents | batches @ 20 M | annotation rows | row scans | +|---|---|---|---|---| +| finewiki-en | 6.6 M | 1 | 43 M | 43 M | +| hplt-de | 176 M | 9 | 3.76 bn | 33.8 bn | +| climbmix-en | 553 M | 28 | 552 M | 15.5 bn | +| nemotron-cc | 1.70 bn | 85 | 747 M | 63.5 bn | + +So `nemotron-cc` reads its 22 GB split 85 times, about 1.9 TB, and that is what its twelve +hours were mostly spent on. The benchmark above has exactly one batch, so it measures the +overhead this commit addresses and none of the re-scanning, and it should not be read as a +prediction for `nemotron-cc`. + +Reducing the scans means letting batches hold far more keys, which today is bounded by the +batch holding a full `pa.Table` per part. Collecting only key columns for the scan and +re-reading the parts to write back would cut that by roughly an order of magnitude, at the +cost of reading the sidecar twice -- cheap against 85 scans of the split. Not attempted +here; it changes the memory profile of a stage that has already been OOM-killed once. + + +## PR #XXX Measured: what the join speedup actually is, on nemotron-cc + +Two proxy datasets failed to predict this one. Non-scan cost per document came out 6.36 us on +`finewiki-en` (15-character keys, 1,024 split files) and 37.59 us on `climbmix-en` +(64-character keys, 12,288 files), and the projection from the second, 21.9 h, exceeded the +12 h actually observed with the slower pre-vectorisation code. Two points cannot separate key +length from fragment count, so the dataset was measured directly. + +60 parts of the real `nemotron-cc` sidecar -- 18,748,230 documents with 36-character UUID +keys -- against the real 746,648,080-row split in 16,384 files, at 2.51 % density against the +2.68 % a real 20 M-key batch sees: + +| implementation | batches | elapsed | us/document | peak RSS | +|---|---|---|---|---| +| before | 1 | 772.6 s | 41.21 | 24.9 GiB | +| after | 1 | 505.0 s | 26.94 | 51.9 GiB | +| after | 2 | 757.5 s | 40.40 | 51.9 GiB | + +**1.53x on the dataset that matters**, and equivalence holds: 224,978,760 label values +identical across 12 columns. This also settles a specific worry -- replacing bucket routing +with one filtered dataset scan could have been a regression here, since the case that had +degraded was the many-file, large-key-set one, and this dataset is 16,384 files at 20 M keys +per batch. It is not a regression; it is where the gain is. + +**Memory doubled, and that needed acting on.** 51.9 GiB peak against the 64 G the join sbatch +requested is 19 % headroom on a stage that has already been OOM-killed once, so the request is +now 160 G. Holding the matched annotation rows in Arrow is what costs it. + +**Decomposition**, from the one- to two-batch delta: one scan of this split costs 252.5 s, and +non-scan work is 13.47 us/document. Extrapolated to all 1,696,565,570 documents: + +| batches | scan | per-document | total | +|---|---|---|---| +| 85 (today's 20 M default) | 6.0 h | 6.3 h | 12.3 h | +| 17 | 1.2 h | 6.3 h | 7.5 h | +| 9 | 0.6 h | 6.3 h | 7.0 h | + +Treat these as order-of-magnitude. The same method applied to the old code projects 19 h +where 12 h was observed, so the extrapolation carries roughly 60 % error; what it does +establish is the shape. Per-document work is the floor at about 6 h, so no amount of batching +gets `nemotron-cc` below that, and larger batches are worth roughly 1.8x rather than the 10x +the original vectorisation estimate implied. + +Reducing batch count needs the batch to stop holding a full table per part -- which would cut +peak memory as well, the two being the same constraint seen from different sides. + + +## PR #XXX Feature: quality-aware upsampling curves + +A `ratio` gives every surviving document of a dataset the same repeat factor, so a dataset +filtered to "content quality at least adequate" repeats its barely-adequate documents exactly +as often as its excellent ones. `upsampling` replaces the scalar with a curve whose factor +rises with quality. + +Method and functional form from Dolma 3 / Olmo 3 (arXiv:2512.13961 §3.4.4, appendix A.2.4), +which measured it against flat quality filtering on 1B models and found it better at every +matched repetition factor -- 0.740 against 0.843-0.870 bits-per-byte on their maths suite. +Quality is placed on a [0, 1] axis where a bucket's width is its share of the dataset's +tokens; the factor is `C * (x - a)**p` above the discard threshold, with the integral pinned +to the token target, no bucket above `max_factor`, and monotone. + +Reproducing their published example -- twenty vigintiles, discard the bottom 40%, cap 7x, +draw one pool's worth of tokens -- returns exactly their figure: bottom eight buckets dropped, +top at 7.00x, monotone between. + +**Deliberate departures from the paper** + +* Their family carries an extra `exp(lam * (x - a))`. Fixing `lam = 0` makes the solution + unique instead of a feasible region, and every integral analytic, so no quadrature is + needed. The remaining degree of freedom is spent pushing the top bucket to exactly + `max_factor`, which is where their own figure sits. +* The exponent is capped at 8. Steepest-admissible is right when the cap binds, but when the + target is a small fraction of the pool -- the regime of any blend that draws far fewer + tokens than it holds -- no exponent violates the cap, "steepest" is unbounded, and the + budget collapses onto the top bucket: the hard top-k filtering the curve exists to beat. + A first attempt did exactly that, solving to exponent 512 with one surviving bucket. +* Unannotated documents cannot be ordered, so they form an explicit `` bucket at + the bottom of the axis. Named in the report rather than folded in, because for a dataset + with partial coverage this decides the fate of the majority. + +**Downstream** + +`apply` writes one index tree per bucket, since the packer emits one file per source file and +factors only mean anything if documents carrying different ones are in different indexes. Each +bucket becomes a manifest row named `__` carrying `source_dataset`, so the +registry lookup in `write-packing-configs` still resolves, and lands as its own repeat factor +in `WeightedCombinedDataset`. The curve is re-solved from the exact counts `apply` observes +rather than from the cube's interpolated ones. + +**Limits** + +Ordinal fields only; a numeric axis needs the cube's quantile edges carried into +materialisation, and a non-ordinal `quality_field` is refused at config load rather than after +a preview has succeeded. And the axis is only as fine as the field: propella labels have about +five levels, unevenly filled -- `climbmix-en` holds 81.9% of its tokens in `moderate` -- so a +curve cannot express a finer preference than the levels allow, where Dolma 3 cuts twenty +vigintiles from a continuous score. + +Not ablated. The cluster was busy, so this is verified against the paper's published example +and by unit tests, not by a training run. + + +## PR #XXX Repetition accounting, and per-predicate attribution + +Two ideas taken from a colleague's TokenFabric design, which applies routing policy at +training time rather than offline. Most of that architecture does not transfer -- chunk-boundary +activation, shard servers, live telemetry solve problems an offline pipeline does not have -- +but two of its principles apply directly. + +**Repetition is per pass, and nothing was checking it.** `target_tokens` was documented as +"only used to report the gap", and that is all it did. A ratio is relative to one pass over the +blend, so if the run consumes more tokens than the blend yields, the loader wraps and every +factor is multiplied. The upsampling curves added in the previous change can assign 7x, which +two passes turn into 14x -- past the point where repetition pays, and silently. + +`target_tokens` now means what the run consumes. `max_total_exposure` caps how often anything +may be seen; a dataset with a curve uses its own `max_factor` instead. `preview` reports the +accounting, `apply` refuses, `--allow_overexposure` overrides. Caps apply to the requested +factor even with no target declared -- a ratio of 9 against a cap of 2 is a violation at any +number of passes -- and the report says the wrapping multiplier is unknown rather than assuming +one pass. + +On the smoke blend, declaring a 300 M run against its 151 M yield: `climbmix-en / high` asked +for 2.53x and would be seen **5.01x** against its own declared 4x cap. + +**`preview --explain` attributes retention to individual predicates.** The per-dataset figure +says nothing about which condition caused it. This reports, per predicate, what it matches +alone and its *marginal* effect -- how many more tokens the dataset would keep with that +predicate dropped and the others left in place -- plus a pairwise overlap matrix. It is +milliseconds: the per-cell weights are the whole computation. + +It found two redundant predicates on the smoke blend at once: `content_integrity at_least +mostly_complete` on `finepdfs-es` (21 k marginal out of 65 M) and `educational_value at_least +basic` on `klettermix-de` (matches 100.0% of tokens, 425 marginal). + +**Two bugs found while building it** + +* The overlap diagonal computed each predicate's factor times itself. For an interpolated + predicate the factors are fractional, so squaring them undercounts -- it read 8.19 M where + the predicate matched 9.75 M. The diagonal is now the match count. +* `_cube_weights` tested the cumulative exactness flag when recording which predicates were + interpolated, so once any predicate was interpolated every later one was reported as + interpolated too. Now per predicate. + + +## PR #XXX Fixes found by the first full production run + +Three failures on the real blend, two of them silent-failure bugs of the same shape as the +ones fixed earlier this week, and one schema drift in the delivered data. + +**Pointer resolution destroyed the pointers it could not resolve.** `Nemotron-ClimbMix`, the +English corpus KletterMix's `source_pointer` keys resolve into, was moved to +`/data/annealing_unused/`. `resolve_source_pointers` resolved nothing, wrote 225 M nulls over +the pointers, logged `resolved 0 of 4,463,748 pointers` as INFO, and carried on -- surfacing +two hours later as 0% join coverage, by which point the pointers were gone and only a full +sidecar rebuild could recover them. It now refuses to write a part that resolved nothing. + +**The join scan was not bounded by matches.** The comment claimed memory "stays bounded by +the rows that match"; `to_table` buffers the whole scan. On HPLT -- 65,536 bucket files, +3.76 bn annotation rows -- that reached 167 GB and was OOM-killed against a 160 G request, +though only ~10 M rows can match. Now streamed via `Scanner.to_batches()` with bounded +readahead. + +**The join report merge raced.** Each array task writes its own report then re-reads the +directory to rebuild a merged view. Non-atomic writes meant a sibling could read a +half-written file: `finewiki-es` died on `JSONDecodeError` *after* its own join had succeeded +at 100%. Same fix as the bucket metadata: `os.replace` plus a tolerant reader. + +**Delivered data changed shape.** Nemotron-CC renamed `warc_record_id` to `id` and re-chunked +from 1.70 bn documents of ~1 KB to 504 M of ~12 KB in the same 1.8 TB. Registry key updated; +a spot check matches 400/400 ids in both `high_actual` and `high_diverse_qa_pairs`. + + +## PR #XXX Fix: join memory scales with fragment count, not data volume + +The streamed scan in the previous change fixed the wrong layer. HPLT still exceeded the +160 GiB limit, and the measurements show why -- peak memory tracks how many bucket *files* a +split has, not how much data: + +| split | files | documents | peak RSS | +|---|---|---|---| +| nemotron-cc | 16,384 | 504 M | 81.6 GB | +| nemotron-climbmix | 12,288 | 226 M | 95.0 GB | +| hplt (each) | 65,536 | 10.7 M | OOM > 160 GB | + +Four times the files and one forty-seventh of the documents, and HPLT is the one that dies. +A dataset holds a fragment per file with that file's parquet metadata, so the scan machinery +scales with the fragment count. Batching the record batches never touched that. + +The split is now scanned in groups of `SCAN_FILE_GROUP` files, bounding the fragment +metadata held at once. All three HPLT splits then completed in 42-45 minutes, and `hplt-it` +-- which had succeeded before the change -- returns byte-identical coverage. + +Peak still sits close to the limit, so the deeper fix is the bucket layout: 1024 buckets per +shard-task was chosen so a bucket could be loaded whole, a constraint filter pushdown +removed. 65,536 files per split is now actively harmful, and re-bucketing at lower fanout +would help every future re-join. + + +## PR #XXX Pack many configs per process + +`modalities data pack_encoded_data` rebuilds its components, tokenizer included, from the +config on every call. Measured on a compute node that startup is 24.7 s, against ~3 s of +real work for a Dolmino file. The blend renders 54,738 configs -- one per source file, of +which 40,003 are Dolmino -- so driving the CLI per file spends roughly 375 core-hours +loading the tokenizer to do about 48 core-hours of tokenising. + +`slurm/pack_many.py` loads the tokenizer once per process and constructs a +`PackedDataGenerator` per config, which is what `pack_encoded_data` does internally. Inside +the driver the load costs 1.2 s and is paid once. Everything else still comes from the +rendered config, so the output is identical. + +Two details that matter at this scale: the slice is strided rather than contiguous, because +the config list is grouped by dataset and contiguous slices would hand one task all 40,003 +Dolmino files and another all 15 FineWiki ones; and a failing config logs and continues +rather than killing its 854 siblings, with `--skip_existing` making reruns resumable. + +Measured on the real blend: 250 GB/h of packed output across 10 concurrent 32-core tasks, +with `num_cpus` confirmed resolving to 32. That is about 54k tokens/s per core, well below +what a fast tokenizer manages, which points at the per-document seeks the filtered index +implies rather than at tokenisation. + + +## PR #XXX Verification for a real packed blend, and a stale-file trap + +The smoke checker assumes a handful of packed files and loads every document index at once. +At 54,738 files and 6.0 TB that does not hold, so two stage-appropriate checks now exist. + +`slurm/scan_pbins.py` reads every packed file's header and reports any whose data section is +empty or unreadable. On the first full run it found exactly one bad file in 54,738: +`finewiki-it/000_00000.pbin`, 151 MB on disk and reporting `data_len=0`. It was a leftover +from an earlier probe, and `pack_many.py --skip_existing` had skipped it because the file +existed -- existence taken for health. Had the final check counted files rather than reading +headers, that dataset would have gone to training 567 M tokens short. `--skip_existing` now +reads the header and requires a non-empty data section. + +`slurm/verify_blend.py` compares packed tokens against the manifest's estimates per dataset, +packed documents against documents selected, and packed files against index files, reading +one document index at a time so 54,738 are never resident together. + +**Result on the real blend**, after repacking the one bad file: + +| | | +|---|---| +| datasets | 18, all document counts exact | +| packed tokens | 1,637,755,654,948 against 1,646,919,435,199 estimated | +| total token error | **-0.56%** | +| worst dataset | klettermix-de at -4.39%, estimated from a rescaled native count | +| files written into the source tree | 0 | + +Document counts matching exactly for all 18 -- 499,676,886 selected and packed for +nemotron-cc -- is the stronger result: the filtered index names precisely the selected +documents, so any difference would be a defect rather than estimator error. + + +## PR #XXX Feature: export sampled JSONL instead of packed tokens + +The pipeline no longer tokenizes. Its final stage writes the selected documents out as JSONL, +one shard per source file under `out//`, and the training set is the concatenation of +those files. `write-packing-configs` and `pack` are gone from this pipeline; the core packer +(`pack_encoded_data`, `PackedDataGenerator`) and `WeightedCombinedDataset` are untouched, since +they are general modalities features used well outside it. + +**The ratios had to move into the bytes.** They used to be metadata: `mix_manifest.yaml` +carried them and `WeightedCombinedDataset` applied them at training time, fractional factors +included. A concatenation carries no weights, so `export-jsonl` materialises them -- a dataset +at 3.0 has each document written three times, one at 0.6 loses two of every five. Fractional +factors are resolved per document rather than by truncating a list: 1.2 means every document +once and a hash-chosen fifth of them twice, keyed on the selection's `seed` and the document's +position via blake2b, so the choice is identical across runs and machines. That reproducibility +is what makes the stage resumable. + +Copies of a document are adjacent, and the documents of a curve's quality buckets are merged +back into one output directory in source order rather than grouped by bucket. + +**The footgun this introduces, and how it is closed.** `mix_manifest.yaml` still says +`ratio: 3.0` after the repetition is already on disk; feeding that to a `weighted_combined` +config would train the data nine times. `export_manifest.yaml` therefore reports +`training_ratio: 1.0` and `repeat_factor_applied: true`, and says so in a `note` field. + +**Sizes**, measured from source bytes and per-dataset cube token totals rather than estimated: + +| | | +|---|---| +| source bytes over the 18 blended datasets | 13.70 TB | +| effective tokens (ratios applied) | 1.695 T | +| JSONL output | **~9.75 TB** | +| packed `.pbin` it replaces | 6.0 TB | + +Verified on the real `finewiki-it` at ratio 3.0: 5,399,277 lines written, exactly +3 x 1,799,759, 67 GB, with 2,670,579 lines replayed byte-for-byte against the corpus. + +**Lessons carried over rather than relearned.** A shard counts as complete only against a +recorded line and byte count, never against merely existing -- the distinction that let a +truncated `.pbin` reporting `data_len=0` survive into a blend. The per-dataset records are +written one file per dataset rather than into one shared manifest, because the stage runs as an +array and concurrent writers to one file lose each other's entries; the blend-wide manifest is +merged afterwards by `--finalize_only`. + +`slurm/verify_jsonl.py` replaces `verify_blend.py`, `scan_pbins.py` and `load_blend.py`: it +checks line counts against the records, the realised lines-per-document against the requested +ratio, replays sampled shards byte-for-byte against the corpus, and asserts nothing was written +into the source tree. `slurm/check_token_estimates.py` replaces `check_smoke_run.py` -- since +nothing tokenizes any more, it tokenizes a sample of the export to keep the calibration that +the whole token budget rests on under check. The tokenizer configs survive for exactly that +reason, reduced to their `tokenizer:` blocks and renamed `annealing_tokenizer.yaml` and +`smoke_tokenizer.yaml`. diff --git a/config_files/data_preparation/quality/README.md b/config_files/data_preparation/quality/README.md new file mode 100644 index 000000000..6a2e3dbc3 --- /dev/null +++ b/config_files/data_preparation/quality/README.md @@ -0,0 +1,383 @@ +# Quality-based selection and up/downsampling + +Builds a training blend by filtering documents on quality signals and choosing how +heavily each dataset is sampled. Two kinds of signal are addressed with the same syntax: +metrics a corpus already carries in its records, and external per-document annotations +that are joined on. + +The source data is never copied or modified. Selection produces a filtered `.idx` naming +exactly the documents that survived, so trying a different threshold costs megabytes of index +rather than a second copy of the corpus. Only the final export writes document bytes. + +## Two files describe a blend + +`annealing_registry.yaml` — where each dataset lives, how it joins to annotations, and +which native metrics to read out of it. Written once and changed rarely. + +`annealing_selection.yaml` — thresholds and an up/downsample ratio per dataset. This is +the file you edit per ablation. + +## Workflow + +```bash +REG=config_files/data_preparation/quality/annealing_registry.yaml +SEL=config_files/data_preparation/quality/annealing_selection.yaml +WORK=/path/to/scratch/blend_v1 + +# 1. Measure how each corpus's records map to our tokenizer's counts. +# Reuses a packing config so the calibration tokenizer cannot drift from the +# tokenizer the packing will actually use. +modalities quality calibrate --registry $REG --work_dir $WORK \ + --tokenizer_config config_files/data_preparation/packed_cc_en_2048.yaml + +# 2. One row per document: position, length, estimated tokens, join key, native metrics. +# The only stage that reads the raw data, so run it as an array. Work is divided by +# file across all datasets, so one array covers the whole blend. +modalities quality build-sidecar --registry $REG --work_dir $WORK --index_root $WORK/idx \ + --shard_id $SLURM_ARRAY_TASK_ID --num_shards 64 + +# 3a. Partition the annotation splits by a hash of their key. The expensive half of +# the join -- HPLT alone is ~12 bn rows -- so run it as an array. +modalities quality bucket-annotations --registry $REG --work_dir $WORK --num_buckets 1024 \ + --shard_id $SLURM_ARRAY_TASK_ID --num_shards 64 + +# 3b. Attach the labels to each dataset's sidecar and report coverage. Cheap; one task. +modalities quality join-annotations --registry $REG --work_dir $WORK + +# 4. Aggregate, so any threshold combination can be costed without reading the sidecars. +modalities quality build-cube --registry $REG --work_dir $WORK + +# 5. Edit $SEL and re-run this as often as you like. It reads only the cubes. +modalities quality preview --selection $SEL --work_dir $WORK + +# 6. Write the filtered indexes and a manifest recording what was selected. +modalities quality apply --selection $SEL --registry $REG --work_dir $WORK --output_dir $WORK/blend + +# 7. Export the selected documents as JSONL, with the sampling baked into the bytes. +modalities quality export-jsonl --manifest $WORK/blend/mix_manifest.yaml \ + --registry $REG --output_dir $WORK/out +``` + +The training set is `cat $WORK/out/*/*.jsonl`. Tokenization is left to whatever consumes it. + +Steps 1–4 are run once per blend. Step 5 is the loop you actually iterate in. + +## What it costs + +Measured on `/data/annealing` (43 TB across 19 datasets, ~7.6 bn documents): + +| Stage | Cost | How often | +|---|---|---| +| calibrate | minutes | once per blend | +| build-sidecar | ~7 h on 64 tasks (first run; ~3.5 h after) | once per blend | +| bucket + join annotations | ~0.5 h on 64 tasks | once per blend | +| build-cube | ~50 min, single task | once per blend | +| **preview** | **~10 s for the whole blend** | **every threshold you try** | +| apply | ~1 h | once you have settled | +| export-jsonl | proportional to what survived | once you have settled | + +Changing thresholds, ratios or the missing-annotation policy costs only a `preview`. +Adding a dataset or a native metric means rebuilding that dataset's sidecar and cube, +because the metric has to come out of the raw records. + +Two things dominate if you get them wrong, both measured: + +* **Keep native-metric patterns to plain field paths** (`.fw_edu_scores`, + `.metadata.dclm_plus2."__label__1"`). Those are evaluated by direct dictionary lookup. + Anything jq cannot reduce to a field chain -- pipes, filters, indexing -- falls back to + jq, which re-serialises the whole document per call and costs about 13x more for the + entire pass. `build-sidecar` warns when a pattern takes that route. +* **Sequential read from `/data` runs at ~282 MB/s per stream and ~3.8 GB/s aggregate.** + With plain paths the sidecar pass reaches ~374 MB/s per core, so it is I/O-bound and + more than ~16 concurrent tasks buys little. +* **`build-sidecar` reads every file twice on its first run**: once for `IndexGenerator` + to write the `.idx`, then once to build the sidecar. The floor is therefore + `2 x 43 TB / 3.8 GB/s`, about 6.3 h -- not the 3.2 h a single pass suggests. Measured + The real run took ~15 h rather than the ~7 h that floor implies, for a reason that is + about placement rather than throughput -- see the next point. + + The index pass is not throwaway work: the export reads those `.idx` files and every later + run reuses them, so a second `build-sidecar` over the same data -- after + adding a native metric, say -- is roughly twice as fast. +* **Per-node bandwidth binds before cluster aggregate does, so spread the tasks.** The + 3.8 GB/s above was measured on the login node and does not transfer to a compute node. + On the real run SLURM packed 12 of the 64 tasks onto one node and 8 onto another while + four nodes sat idle; each got ~95 MB/s against the 282 MB/s a single stream gets, and + the job spent a nine-hour tail at a third of its starting speed. + + Both array scripts therefore request 8 CPUs per task. A task is single-threaded and + needs one; the request exists to cap how many bandwidth-hungry tasks land on a node + (4 rather than 16), spreading an array of 64 across the cluster. Raise the array size + only after the tasks are spread -- more tasks on the same node buys nothing. + +## What the preview reports + +``` +dataset docs kept row% tokens kept tok% ratio effective share +------------------------------------------------------------------------------------ +hplt-de 1.42B 61.3% 238.10B 58.4% 0.60 142.86B 18.1% +finepdfs-en 0.31B 44.0% 97.40B 51.2% 1.40 136.36B 17.3% +------------------------------------------------------------------------------------ +TOTAL 789.20B 100.0% + +target 400.00B tokens -- 389.20B over (97.3%) +``` + +`row%` and `tok%` differ on purpose. Quality correlates with length, so a filter that +keeps the better documents keeps a larger share of the tokens than of the documents — +which is why the row retention alone cannot be used to predict a token budget. + +A `~` next to a row means a numeric threshold fell inside a cube bin rather than on its +edge, so that row was interpolated. Re-run with `--exact` to scan the per-document +sidecars instead. + +**How wrong can the interpolation be?** Measured on the smoke blend, comparing the cube +against the exact per-document figures `apply` writes into the manifest: + +| dataset | predicate kind | cube tokens | exact tokens | error | +|---|---|---|---|---| +| finewiki-de | ordinal only | 18.51M | 18.51M | 0.00% | +| climbmix-en | ordinal only | 54.53M | 54.53M | 0.00% | +| klettermix-de | ordinal + score | 19.73M | 19.80M | -0.35% | +| dolmino | score only | 46.58M | 48.32M | -3.60% | +| finepdfs-es | ordinal + score | 9.73M | 8.64M | +12.61% | +| **total** | | **150.80M** | **149.80M** | **+0.67%** | + +Ordinal predicates are exact -- levels are cube dimensions, so no interpolation happens. +Numeric thresholds are only exact when they land on a bin edge, and the error does not +shrink with dataset size, because the bin count is fixed rather than proportional. Use the +cube to explore, then `--exact` before committing to a budget, or raise +`--num_score_bins` when building the cube. `apply` always scans the sidecar, so the +manifest it writes carries exact figures regardless. + +### When a predicate is not in the cube + +The join attaches twelve annotation columns; the cube groups on seven (`audience_level`, +`commercial_bias`, `content_length`, `content_ratio` and `time_sensitivity` are attached but +not grouped), because grouping on all twelve would multiply the cell count by thousands. A +native metric that is null throughout a dataset is dropped from its cube too. + +Thresholding a field the cube does not carry is therefore expressible in a selection but +cannot be answered from the cube, and `preview` refuses rather than quietly reading every +document: + +``` +these predicates cannot be answered from the cubes: + nemotron-cc: cube for 'nemotron-cc' was not grouped on label 'commercial_bias' ... + (answering it from the sidecar means reading 1,696,565,570 documents) +``` + +Three ways out: drop or replace the predicate; rebuild that dataset's cube naming the field +(`build-cube --only nemotron-cc --label_dimension commercial_bias --label_dimension ...`, +listing the others too, since the flag replaces the default set); or accept the cost with +`--allow_fallback`. That last is what used to happen silently, and it turned a 13-second +preview into a job still running after ten minutes. + +## The ratio is applied by the export, not at training time + +`export-jsonl` materialises the sampling: a dataset at 2.0 has each of its documents written +twice, one at 0.6 has two of every five dropped. The training set is the concatenation of the +exported files, drawn once. + +This is a change from how the pipeline used to work, and the trap is worth stating plainly. +`mix_manifest.yaml` still records `ratio: 2.0`, because that is what was asked for. Feeding +that number to a `weighted_combined` dataset now would apply it a **second** time. Read +`export_manifest.yaml` instead: it reports `training_ratio: 1.0` and +`repeat_factor_applied: true`. + +Fractional factors are resolved per document rather than by truncating a list -- 1.2 means +every document once and a hash-chosen fifth of them twice -- keyed on the selection's `seed` +and the document's position, so it is reproducible across runs and machines. Copies of a +document are written adjacent to one another, so a sequential reader needs a shuffle buffer +larger than the run of copies to separate them. Nothing is +duplicated on disk, and changing the blend means changing a number rather than rebuilding +data. + +## Quality-aware upsampling curves + +A `ratio` treats every surviving document alike: `nemotron-cc` filtered to "content quality +at least adequate" and set to 1.2 repeats the barely-adequate documents exactly as often as +the excellent ones. A curve makes the repeat factor rise with quality instead, so the token +budget is spent where the quality signal says it should be. + +```yaml +- name: nemotron-cc + upsampling: + quality_field: content_quality # ordinal axis, worst to best + target_tokens: 900_000_000_000 # or target_ratio, as a multiple of what is available + max_factor: 7.0 # no bucket repeated more than this + discard_below_percentile: 40 # drop the weakest 40% of tokens + predicates: [...] # still applied first; the curve works on survivors +``` + +`ratio` and `upsampling` are mutually exclusive -- both express the same decision, so +setting both is refused rather than one being silently ignored. + +The method and the functional form come from Dolma 3 / Olmo 3 (arXiv:2512.13961 §3.4.4, +appendix A.2.4), which measured it against flat quality filtering on 1B models and found it +better at every matched repetition factor -- 0.740 against 0.843-0.870 bits-per-byte on their +maths suite. Quality goes on a [0, 1] axis where each bucket's width is its share of the +dataset's tokens, and the repeat factor is `C * (x - a)**p` above the discard threshold `a`, +with `p` chosen so the top bucket sits exactly at `max_factor`. + +Reproducing their published example -- twenty equal vigintiles, discard the bottom 40%, +repeat at most 7x, draw as many tokens as the pool holds -- gives exactly their figure: the +bottom eight buckets dropped, the top at 7.00x, monotone in between. + +**What `apply` does with it.** Documents with different factors have to live in different +indexes, so each bucket becomes its own index tree and its own manifest row, named +`__` with `source_dataset` naming the registry entry. The curve is re-solved +from the exact token counts found during `apply` rather than from the cube, since that stage +reads every document anyway. `export-jsonl` then merges the buckets back into one output +directory per input dataset, each bucket repeated by its own factor, with the documents +written in source order rather than grouped by quality level. + +### Two limits worth knowing + +**The axis is only as fine as the field.** Dolma 3 cuts twenty vigintiles because their +quality signal is a continuous classifier score. Propella labels are ordinal with about five +levels, and the levels are not evenly filled: on the smoke blend, `climbmix-en` has 81.9% of +its tokens in `moderate` and 17.9% in `high`. A curve over that axis can still do something +useful -- it drew 2.53x the `high` tokens and 6% of the `moderate` ones -- but it cannot +express a finer preference than the levels allow. A threshold falling inside a level applies +a fractional factor to the whole level rather than splitting it, and which documents survive +is then whatever the dataloader's even spread picks. + +**Only ordinal fields for now.** A native numeric metric would need the cube's quantile edges +carried into materialisation, which is not built. Setting a non-ordinal `quality_field` is +refused when the config loads, rather than after a preview has already succeeded. + +## Repetition is per pass, so declare what the run consumes + +A `ratio` of 2.0 draws a dataset twice *per pass over the blend*. If the run consumes more +tokens than the blend yields, the loader comes round again and every factor is multiplied by +the number of passes. A bucket set to 7x under a curve becomes 14x on a second pass, which is +well past the point where repetition pays for itself. + +`target_tokens` is what the run will actually consume, and it is what makes this checkable: + +```yaml +target_tokens: 400_000_000_000 # what the run consumes +max_total_exposure: 3.0 # refuse if anything is seen more often than this +``` + +A dataset with an upsampling curve uses its own `max_factor` as the cap instead, since that is +already declared. `preview` always reports the accounting; `apply` refuses, because that is the +point of commitment: + +``` +repetition, once wrapping is counted + run consumes 300,000,000 tokens from 150,800,110 effective -- wraps 1.99 passes + factor exposure cap + climbmix-en / high 2.53 5.04 4 OVER + klettermix-de 2.00 3.98 2.5 OVER +``` + +Pass `--allow_overexposure` to `apply` to proceed anyway. Without `target_tokens` the caps +still apply to the requested factors, but the multiplier from wrapping is unknown and the +report says so rather than assuming one pass. + +## Which predicate is actually doing the work + +`preview --explain` attributes a dataset's retention to its individual predicates and shows +how they overlap. It costs milliseconds, because the per-cell weights are the whole +computation and the cube is already loaded. + +``` + finepdfs-es: 9,727,716 of 65,540,305 tokens kept + predicate matches share marginal + fw_edu gte 1.5 9,748,842 14.9% 55,753,485 ~ + content_integrity at_least mostly_complete 65,481,201 99.9% 21,126 +``` + +`marginal` is the number to read: how many more tokens the dataset would keep if that +predicate were dropped and the others left alone. A marginal of zero means the predicate is +fully shadowed by its neighbours -- it changes nothing and only makes the selection harder to +read. On the smoke blend this immediately found two: `content_integrity` on `finepdfs-es` +(21 k marginal out of 65 M) and `educational_value at_least basic` on `klettermix-de`, which +matches 100.0% of tokens. + +The overlap matrix below the table gives tokens matching each pair. Its diagonal is each +predicate's own match count -- not the product with itself, which for an interpolated +predicate squares a fractional survival rate and undercounts. + +## The source tree must not move underneath a sidecar + +A sidecar row locates its document by `(file_id, byte_offset, byte_len)`, where `file_id` +is a position in the dataset's sorted file list. If the source tree changes between +building the sidecar and using it, the same id names a different file and every offset is +wrong. + +This is not hypothetical. A transfer re-sharded four corpora after their sidecars were +built -- `Nemotron-CC` went from 606 MB files to 137 MB files while keeping nearly the +same file *count* -- and the only check that existed compared counts. Eleven of nineteen +datasets had unusable sidecars, and just one of them failed loudly. + +So `build-sidecar` records the file list it used in `sidecar//_files.json`, and +`apply` resolves ids through that list and refuses to run if any recorded file has +changed size or vanished. Paths are stored relative to the root, so moving or snapshotting +a tree is fine; only the file set has to agree. + +Run the direct check after any transfer, and before `apply` on a tree that might have been +touched: + +```bash +$MQ -m modalities quality verify-sidecar --registry $REG --work_dir $WORK +``` + +It seeks to recorded offsets and compares the document it finds against the recorded text +length, so it catches a file that was rewritten at the same size. It skips rows at offset +zero on purpose: the first document of any JSONL file parses, so those rows succeed even +against a completely different file, and sampling them is exactly how a broken sidecar +looked healthy. `--adopt` stamps a manifest onto a sidecar built before manifests existed, +but only if it verifies. + +A drifted dataset needs its sidecar rebuilt, then re-joined and re-cubed. The annotation +buckets are unaffected -- they are built from the annotation cache, not the corpora -- so +the expensive bucketing stage does not repeat. + +## Testing the pipeline end to end + +`slurm/make_smoke_snapshot.py` freezes about 1 GB of five corpora into a snapshot, and +`smoke_registry.yaml` / `smoke_selection.yaml` run the whole pipeline over it in minutes. +The five datasets cover all four distinct join-key kinds plus the native-metrics-only +path, which is every branch of the join; HPLT is left out because it shares FineWiki's key +kind and would add 327 GB of bucket reads to exercise no new code. +`slurm/check_token_estimates.py` then tokenizes a sample of the exported JSONL and compares +it against the preview's estimates, and `slurm/verify_jsonl.py` replays sampled shards against +the corpus and asserts nothing was written into the source tree. + +Use it after any change to the sidecar, join, cube, or materialize stages. It is much +cheaper than discovering a bug 15 hours into a real build. + +## Two things to be careful about + +**Token counts are estimates, and a corpus has no single bytes-per-token ratio.** On the +German FineWiki snapshot the ratio runs from 3.571 for documents under a kilobyte to 34.648 +for the nine documents above 256 KB -- and those nine hold 5.7% of all bytes. So the +calibration measures a ratio per log-spaced size stratum and applies it per document from +the byte length the sidecar records exactly. Getting there needed the sample to be drawn at +offsets spread across each file, taking the document *containing* each offset so that +selection is proportional to length; sampling uniformly by document count found 2 of the +top-stratum documents per 2,000 draws, where length-proportional sampling finds 44. + +Measured against the true token count of all 27,846 documents, worst case over five seeds: +16.2% error for the original estimator (a fixed bias, identical across seeds, so it looked +stable), 19.4% for a global ratio on a spread sample, 0.8% for the stratified estimate. + +Datasets carrying their own token count in every record (FinePDFs, KletterMix, FinePhrase) +are estimated from that field rescaled to our tokenizer, per document, and need no +stratifying. + +Still validate on your own data: export one small dataset and run +`slurm/check_token_estimates.py`, which tokenizes a sample and reports the error per dataset. +Nothing in the pipeline tokenizes any more, so this is the only check on the model the whole +token budget rests on. Document counts, by contrast, are not estimates -- the filtered index +names precisely the selected documents -- and `verify_jsonl.py` requires them to match +exactly. + +**Decide what to do with unannotated documents.** `missing_annotation: keep` treats an +annotation predicate as satisfied for documents that have no label; `drop` treats it as +failed. On a partly downloaded annotation split the unannotated documents can be the +majority, and the two policies then give completely different blends. The join report +(`join_report.json` in the working directory) tells you the coverage per dataset. diff --git a/config_files/data_preparation/quality/annealing_registry.yaml b/config_files/data_preparation/quality/annealing_registry.yaml new file mode 100644 index 000000000..2779f1186 --- /dev/null +++ b/config_files/data_preparation/quality/annealing_registry.yaml @@ -0,0 +1,181 @@ +# Corpus registry for the annealing blend. +# +# Every join key below was verified by sampling real keys out of the JSONL and scanning +# the propella `id` columns for them. Coverage measured 2026-08-17 against the caches +# named in `annotation_root`: +# +# finewiki (all 5 languages) 100% +# hplt-4-unfiltered (de/fr/it/es) 100% +# nemotron-cc/high-actual 100% (also covers the synthetic +# high_diverse_qa_pairs subdirectory, +# which shares warc_record_id) +# nemotron-climbmix 100% (English, and German via the source) +# finepdfs (eng/deu/fra/spa) 10-34% only 1-5 shards fetched per split +# +# Datasets with no `annotation_split` have no propella corpus at all and can only be +# shaped by their native metrics. + +annotation_root: /data/michael.fromm/hf-cache/datasets--openeurollm--propella-annotations/snapshots/e80fc1407801a15b956f40c642d3709b528abbc9/data/propella-1-4b +extra_annotation_roots: + # The HPLT shards were fetched into a second cache; both are searched. + - /data/alex.jude/.cache/huggingface/datasets--openeurollm--propella-annotations/snapshots/9e9c5083f81dc4bbd2708b65816a4dc41f59b911/data/propella-1-4b + +datasets: + # ---------------------------------------------------------------- FineWiki + # Ids look like `dewiki/6851323` and are used verbatim on both sides. + - name: finewiki-en + jsonl_root: /data/annealing/english/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + native_metrics: + - {name: bytes_html, jq_pattern: .bytes_html} + - name: finewiki-de + jsonl_root: /data/annealing/german/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + - name: finewiki-fr + jsonl_root: /data/annealing/french/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + - name: finewiki-it + jsonl_root: /data/annealing/italian/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + - name: finewiki-es + jsonl_root: /data/annealing/spanish/Finewiki + annotation_split: finewiki + key: {kind: field, field: id} + + # ---------------------------------------------------------------- HPLT v4 + # 32-hex ids, used verbatim. No quality score of any kind in the records, so propella + # is the only quality signal available for these 14 TB. + - name: hplt-de + jsonl_root: /data/annealing/german/hpltv4 + annotation_split: hplt-4-unfiltered/deu_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + - name: hplt-fr + jsonl_root: /data/annealing/french/hpltv4 + annotation_split: hplt-4-unfiltered/fra_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + - name: hplt-it + jsonl_root: /data/annealing/italian/hpltv4 + annotation_split: hplt-4-unfiltered/ita_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + - name: hplt-es + jsonl_root: /data/annealing/spanish/hpltv4 + annotation_split: hplt-4-unfiltered/spa_Latn + key: {kind: field, field: id} + native_metrics: + - {name: lid_prob, jq_pattern: ".\"openlid-v3\".prob", aggregation: max} + - {name: cluster_size, jq_pattern: .cluster_size} + + # ---------------------------------------------------------------- Nemotron-CC + # No `id` field at all; `warc_record_id` is the annotation key. The annotation ids are + # not unique -- roughly 4% recur -- so the join keeps the first occurrence. This is not + # specific to this split: finewiki measured 868,586 duplicate keys in 43.1 M rows (2.0%) + # on a real run, so expect the join to report duplicates for most splits. + # The 2026-08 delivery renamed `warc_record_id` to `id`, keeping the same bare UUIDs. + # The old name silently produced a null join key for every document and 0% coverage. + # Both subdirectories -- high_actual and high_diverse_qa_pairs -- are covered by the + # high-actual annotation split; 400/400 sampled ids from each were found in it. + - name: nemotron-cc + jsonl_root: /data/annealing/english/Nemotron-CC + annotation_split: nemotron-cc/high-actual + key: {kind: field, field: id} + + # ---------------------------------------------------------------- ClimbMix + # No identifier of any kind; the annotation key is the SHA-256 of the exact text. + - name: climbmix-en + jsonl_root: /data/annealing/english/Climbmix + annotation_split: nemotron-climbmix + key: {kind: sha256_text} + + # German KletterMix is a translation of ClimbMix. Its own text is German, so hashing + # it matches nothing -- the annotation belongs to the English original, which the id + # points at as `/`, zero-indexed. + - name: klettermix-de + jsonl_root: /data/annealing/german/AIML-TUDA-KletterMix-filtered + glob: "*.jsonl" + annotation_split: nemotron-climbmix + key: + kind: source_pointer + field: id + source_root: /data/annealing_unused/Nemotron-ClimbMix + source_line_offset: 0 + native_metrics: + - {name: proxy_score, jq_pattern: .proxy_score} + - {name: token_count, jq_pattern: .token_count} + + # ---------------------------------------------------------------- FinePDFs + # Ids are UUIDs stored in two forms -- `` and bare -- mixed within + # single files on both sides, so both sides are normalised before comparing. + - name: finepdfs-en + jsonl_root: /data/annealing/english/Finepdfs + annotation_split: finepdfs/eng_Latn + key: {kind: urn_uuid_field, field: id} + note: "annotation coverage was 34% on 2026-08-17; finish the fetch before relying on a propella predicate" + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: dclm, jq_pattern: .dclm_scores, aggregation: max} + - {name: ocr_quality, jq_pattern: .ocr_quality_scores, aggregation: min} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + - {name: duplicate_count, jq_pattern: .duplicate_count} + - name: finepdfs-de + jsonl_root: /data/annealing/german/Finepdfs + annotation_split: finepdfs/deu_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + - name: finepdfs-fr + jsonl_root: /data/annealing/french/Finepdfs + annotation_split: finepdfs/fra_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + - name: finepdfs-es + jsonl_root: /data/annealing/spanish/Finepdfs + annotation_split: finepdfs/spa_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + + # `italian/Finepdfs` holds English data, byte-identical to `english/Finepdfs` in 578 + # of its 580 files. Enabling it would count the same 5.4 TB of English documents + # twice, so it stays declared and disabled until that is resolved. + - name: finepdfs-it + jsonl_root: /data/annealing/italian/Finepdfs + annotation_split: finepdfs/eng_Latn + key: {kind: urn_uuid_field, field: id} + enabled: false + note: "English data in the Italian folder; duplicates finepdfs-en" + + # -------------------------------------------- no propella corpus, native only + # The largest subset of the blend. Every candidate annotation key returns 0%, so it + # can only be shaped by its category directory and native fields. + - name: nemotron-cc-v2 + jsonl_root: /data/annealing/english/Nemotron-CC-v2 + native_metrics: [] + - name: finephrase + jsonl_root: /data/annealing/english/Finephrase + native_metrics: + - {name: fw_edu, jq_pattern: .score} + - {name: language_score, jq_pattern: .language_score} + - {name: token_count, jq_pattern: .token_count} + - name: dolmino + jsonl_root: /data/annealing/english/Dolmino + note: "dclm_plus2 exists only under stem-heavy-crawl; it was null in all of 400 sampled sidecar parts, so the cube drops it" + native_metrics: + - {name: dclm_plus2, jq_pattern: '.metadata.dclm_plus2."__label__1"'} + - {name: len_cl100k_base, jq_pattern: .metadata.len_cl100k_base} diff --git a/config_files/data_preparation/quality/annealing_selection.yaml b/config_files/data_preparation/quality/annealing_selection.yaml new file mode 100644 index 000000000..7d9ff71af --- /dev/null +++ b/config_files/data_preparation/quality/annealing_selection.yaml @@ -0,0 +1,145 @@ +# Example selection for the annealing blend. +# +# Each dataset states which documents to keep and how heavily to sample what survives. +# `modalities quality preview` costs this in documents and tokens in seconds; nothing +# is read or written until `modalities quality apply`. +# +# Predicates within a dataset are combined with AND. Ordinal levels come from the +# declared scales in `selection.py`; note that `information_density` orders +# `moderate` *below* `adequate`, which is easy to get backwards. + +# What to do with documents that carry no annotation. `keep` treats a propella +# predicate as satisfied for them, which is the safe default while FinePDFs is only +# partly downloaded -- `drop` there would silently discard two thirds of the corpus for +# having no label rather than for failing the filter. +missing_annotation: keep + +# What the training run will consume. It adjusts no ratio, but it is what makes the ratios +# checkable: a ratio is per pass over the blend, so if the blend yields less than this the +# loader wraps and every factor is multiplied by the number of passes. +target_tokens: 400_000_000_000 + +# Refuse to materialise if any dataset would be seen more often than this once wrapping is +# counted. Set above the largest ratio below (3.0) so it does not fire on the intended +# design, only on an accident -- a blend trimmed too far, or a ratio raised without +# re-checking the yield. A dataset with an `upsampling` curve uses its own `max_factor` +# instead, since that is already a declared cap. Pass --allow_overexposure to override. +max_total_exposure: 4.0 + +datasets: + # Web text carries the most junk, so it gets the strictest filter and is downsampled. + - name: hplt-de + ratio: 0.6 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + - {field: content_safety, op: at_least, value: mild_concerns} + - name: hplt-fr + ratio: 0.6 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + - {field: content_safety, op: at_least, value: mild_concerns} + - name: hplt-it + ratio: 0.8 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + - name: hplt-es + ratio: 0.6 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: content_integrity, op: at_least, value: mostly_complete} + + # Already filtered upstream to a high-quality subset, so a light touch and a mild + # upsample. + # + # This is the natural place to try a quality-aware curve instead of the flat ratio: a + # single number repeats the barely-adequate documents exactly as often as the excellent + # ones. Replace `ratio` with the block below to make the repeat factor rise with quality. + # `preview` prints the resulting per-bucket factors, and `apply` then writes one index + # tree per bucket so each can carry its own factor into training. + # + # upsampling: + # quality_field: content_quality # ordinal axis, worst to best + # target_tokens: 900_000_000_000 # what this dataset should contribute + # max_factor: 7.0 # no bucket repeated more than this + # discard_below_percentile: 40 # drop the weakest 40% of tokens + # + # Note that `ratio` and `upsampling` are mutually exclusive: the curve already determines how much + # is drawn, so setting both is refused rather than silently ignoring one. + - name: nemotron-cc + ratio: 1.2 + predicates: + - {field: content_quality, op: at_least, value: adequate} + # `commercial_bias` is joined onto the sidecar but is not one of the seven fields + # the cube groups on by default, so thresholding it here would force the preview to + # read all 1.7 bn documents. To use it, rebuild this dataset's cube with + # build-cube --only nemotron-cc --label_dimension commercial_bias ... + # naming the other seven too, since the flag replaces the default set. + + - name: climbmix-en + ratio: 1.0 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: information_density, op: at_least, value: moderate} + + # The German pool is small relative to the English one, so what survives is upsampled. + # Combines a propella predicate with the corpus's own proxy score. + - name: klettermix-de + ratio: 2.0 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: proxy_score, op: gte, value: 0.65} + + # Reference text: keep all of it, upsample the smaller languages. + - name: finewiki-en + ratio: 1.0 + - name: finewiki-de + ratio: 2.0 + - name: finewiki-fr + ratio: 2.0 + - name: finewiki-it + ratio: 3.0 + - name: finewiki-es + ratio: 3.0 + + # Only 10-34% annotated. `missing_annotation: keep` above means the propella + # predicate applies to the annotated minority and the rest passes on the native + # score alone -- deliberate, and worth revisiting once the fetch is finished. + - name: finepdfs-en + ratio: 1.0 + predicates: + - {field: fw_edu, op: gte, value: 2.0} + - {field: content_integrity, op: at_least, value: mostly_complete} + - name: finepdfs-de + ratio: 1.5 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + - name: finepdfs-fr + ratio: 1.5 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + - name: finepdfs-es + ratio: 1.5 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + + # No propella corpus exists for these, so they are shaped by native metrics only. + - name: finephrase + ratio: 0.5 + predicates: + - {field: fw_edu, op: gte, value: 2.5} + - {field: language_score, op: gte, value: 0.9} + - name: dolmino + ratio: 1.0 + predicates: + # `dclm_plus2` only exists in the stem-heavy-crawl subdirectory -- it was null in + # every one of 400 sampled sidecar parts -- so it is dropped from the cube and + # cannot be thresholded. Length is the signal that is actually present here. + - {field: len_cl100k_base, op: gte, value: 200} + + # Declared but excluded, so the reason is recorded rather than implied by absence. + - name: nemotron-cc-v2 + ratio: 1.0 + enabled: false diff --git a/config_files/data_preparation/quality/annealing_tokenizer.yaml b/config_files/data_preparation/quality/annealing_tokenizer.yaml new file mode 100644 index 000000000..7a6c4f027 --- /dev/null +++ b/config_files/data_preparation/quality/annealing_tokenizer.yaml @@ -0,0 +1,19 @@ +# The tokenizer the token estimates are measured with. +# +# Used by `modalities quality calibrate --tokenizer_config `, which measures +# how many tokens each dataset's bytes are worth. Nothing here tokenizes the corpus: the +# pipeline now exports JSONL and leaves tokenization to whatever consumes it. The estimate +# is what makes `preview` able to cost a selection in tokens, so it should still name the +# tokenizer the run actually trains with -- getting it wrong does not fail loudly, it +# produces a plausible token budget for the wrong tokenizer. + +tokenizer: + component_key: tokenizer + variant_key: pretrained_hf_tokenizer + config: + # The tokenizer the long-context pipeline uses. The token audit under + # /data/michael.fromm used the Super-120B variant instead -- confirm which this run + # trains with before trusting any token figure. + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + padding: false + truncation: false diff --git a/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw b/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw new file mode 100644 index 000000000..3bb6560be --- /dev/null +++ b/config_files/data_preparation/quality/data_preprocessing_pipeline.excalidraw @@ -0,0 +1,4426 @@ +{ + "type": "excalidraw", + "version": 2, + "source": "https://excalidraw.com", + "elements": [ + { + "id": "el0001", + "type": "text", + "x": 60, + "y": 46, + "width": 386.1, + "height": 32.5, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 107919, + "version": 1, + "versionNonce": 304729, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Data Preprocessing Pipeline", + "originalText": "Data Preprocessing Pipeline", + "fontSize": 26, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0002", + "type": "text", + "x": 60, + "y": 82, + "width": 336.6, + "height": 22.5, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 115838, + "version": 1, + "versionNonce": 409458, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "==================================", + "originalText": "==================================", + "fontSize": 18, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0003", + "type": "rectangle", + "x": 60, + "y": 130, + "width": 275, + "height": 320, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 123757, + "version": 1, + "versionNonce": 514187, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0004", + "type": "text", + "x": 74, + "y": 140, + "width": 42.900000000000006, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 131676, + "version": 1, + "versionNonce": 618916, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Inputs", + "originalText": "Inputs", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0005", + "type": "rectangle", + "x": 82, + "y": 172, + "width": 230, + "height": 54, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 139595, + "version": 1, + "versionNonce": 723645, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0006" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0006", + "type": "text", + "x": 88, + "y": 184.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 147514, + "version": 1, + "versionNonce": 828374, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Source Corpora\n/data/annealing (19 subsets)", + "originalText": "Source Corpora\n/data/annealing (19 subsets)", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0005", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0007", + "type": "ellipse", + "x": 301, + "y": 161, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 155433, + "version": 1, + "versionNonce": 933103, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0008" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0008", + "type": "text", + "x": 303, + "y": 165, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 163352, + "version": 1, + "versionNonce": 1037832, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0007", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0009", + "type": "rectangle", + "x": 82, + "y": 240, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 171271, + "version": 1, + "versionNonce": 1142561, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0010" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0010", + "type": "text", + "x": 88, + "y": 250.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 179190, + "version": 1, + "versionNonce": 1247290, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Corpus Registry\nannealing_registry.yaml", + "originalText": "Corpus Registry\nannealing_registry.yaml", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0009", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0011", + "type": "ellipse", + "x": 301, + "y": 229, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 187109, + "version": 1, + "versionNonce": 1352019, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0012" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0012", + "type": "text", + "x": 303, + "y": 233, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 195028, + "version": 1, + "versionNonce": 1456748, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0011", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0013", + "type": "rectangle", + "x": 82, + "y": 304, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 202947, + "version": 1, + "versionNonce": 1561477, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0014" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0014", + "type": "text", + "x": 88, + "y": 314.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 210866, + "version": 1, + "versionNonce": 1666206, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Propella Annotations\n(external parquet cache)", + "originalText": "Propella Annotations\n(external parquet cache)", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0013", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0015", + "type": "ellipse", + "x": 301, + "y": 293, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 218785, + "version": 1, + "versionNonce": 1770935, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0016" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0016", + "type": "text", + "x": 303, + "y": 297, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 226704, + "version": 1, + "versionNonce": 1875664, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0015", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0017", + "type": "rectangle", + "x": 82, + "y": 368, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 234623, + "version": 1, + "versionNonce": 1980393, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0018" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0018", + "type": "text", + "x": 88, + "y": 378.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 242542, + "version": 1, + "versionNonce": 2085122, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Tokenizer Config\n(for token estimates only)", + "originalText": "Tokenizer Config\n(for token estimates only)", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0017", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0019", + "type": "ellipse", + "x": 301, + "y": 357, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 250461, + "version": 1, + "versionNonce": 2189851, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0020" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0020", + "type": "text", + "x": 303, + "y": 361, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 258380, + "version": 1, + "versionNonce": 2294580, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0019", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0021", + "type": "rectangle", + "x": 400, + "y": 168, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 266299, + "version": 1, + "versionNonce": 2399309, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0022" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0022", + "type": "text", + "x": 406, + "y": 189.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 274218, + "version": 1, + "versionNonce": 2504038, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "1. calibrate", + "originalText": "1. calibrate", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0021", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0023", + "type": "ellipse", + "x": 541, + "y": 153, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 282137, + "version": 1, + "versionNonce": 2608767, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0024" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0024", + "type": "text", + "x": 543, + "y": 157, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 290056, + "version": 1, + "versionNonce": 2713496, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0023", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0025", + "type": "text", + "x": 402, + "y": 232, + "width": 82.5, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 297975, + "version": 1, + "versionNonce": 2818225, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~15 s / dataset", + "originalText": "~15 s / dataset", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0026", + "type": "rectangle", + "x": 400, + "y": 262, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 305894, + "version": 1, + "versionNonce": 2922954, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0027" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0027", + "type": "text", + "x": 406, + "y": 271.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 313813, + "version": 1, + "versionNonce": 3027683, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "calibration.yaml\nstratified ratios", + "originalText": "calibration.yaml\nstratified ratios", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0026", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0028", + "type": "rectangle", + "x": 590, + "y": 168, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 321732, + "version": 1, + "versionNonce": 3132412, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0029" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0029", + "type": "text", + "x": 596, + "y": 189.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 329651, + "version": 1, + "versionNonce": 3237141, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "2. build-sidecar", + "originalText": "2. build-sidecar", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0028", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0030", + "type": "ellipse", + "x": 731, + "y": 153, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 337570, + "version": 1, + "versionNonce": 3341870, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0031" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0031", + "type": "text", + "x": 733, + "y": 157, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 345489, + "version": 1, + "versionNonce": 3446599, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0030", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0032", + "type": "text", + "x": 592, + "y": 232, + "width": 110.00000000000001, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 353408, + "version": 1, + "versionNonce": 3551328, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~15 h (SLURM array)", + "originalText": "~15 h (SLURM array)", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0033", + "type": "rectangle", + "x": 590, + "y": 262, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 361327, + "version": 1, + "versionNonce": 3656057, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0034" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0034", + "type": "text", + "x": 596, + "y": 271.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 369246, + "version": 1, + "versionNonce": 3760786, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "sidecar/*.parquet\n+ _files.json", + "originalText": "sidecar/*.parquet\n+ _files.json", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0033", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0035", + "type": "rectangle", + "x": 590, + "y": 340, + "width": 170, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 377165, + "version": 1, + "versionNonce": 3865515, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0036" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0036", + "type": "text", + "x": 596, + "y": 361.5, + "width": 158, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 385084, + "version": 1, + "versionNonce": 3970244, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "3. bucket-annotations", + "originalText": "3. bucket-annotations", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0035", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0037", + "type": "ellipse", + "x": 741, + "y": 325, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 393003, + "version": 1, + "versionNonce": 4074973, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0038" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0038", + "type": "text", + "x": 743, + "y": 329, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 400922, + "version": 1, + "versionNonce": 4179702, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0037", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0039", + "type": "text", + "x": 592, + "y": 404, + "width": 115.50000000000001, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 408841, + "version": 1, + "versionNonce": 4284431, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~4.5 h (SLURM array)", + "originalText": "~4.5 h (SLURM array)", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0040", + "type": "rectangle", + "x": 590, + "y": 418, + "width": 170, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 416760, + "version": 1, + "versionNonce": 4389160, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0041" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0041", + "type": "text", + "x": 596, + "y": 434.125, + "width": 158, + "height": 13.75, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 424679, + "version": 1, + "versionNonce": 4493889, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "buckets/ (key-hashed)", + "originalText": "buckets/ (key-hashed)", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0040", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0042", + "type": "rectangle", + "x": 860, + "y": 250, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 432598, + "version": 1, + "versionNonce": 4598618, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0043" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0043", + "type": "text", + "x": 866, + "y": 271.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 440517, + "version": 1, + "versionNonce": 4703347, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "4. join-annotations", + "originalText": "4. join-annotations", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0042", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0044", + "type": "ellipse", + "x": 1001, + "y": 235, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 448436, + "version": 1, + "versionNonce": 4808076, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0045" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0045", + "type": "text", + "x": 1003, + "y": 239, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 456355, + "version": 1, + "versionNonce": 4912805, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0044", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0046", + "type": "text", + "x": 862, + "y": 314, + "width": 88.0, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 464274, + "version": 1, + "versionNonce": 5017534, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "hours; --resume", + "originalText": "hours; --resume", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0047", + "type": "rectangle", + "x": 860, + "y": 330, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 472193, + "version": 1, + "versionNonce": 5122263, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0048" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0048", + "type": "text", + "x": 866, + "y": 339.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 480112, + "version": 1, + "versionNonce": 5226992, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "labelled sidecar\n100% coverage", + "originalText": "labelled sidecar\n100% coverage", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0047", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0049", + "type": "rectangle", + "x": 1080, + "y": 250, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 488031, + "version": 1, + "versionNonce": 5331721, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0050" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0050", + "type": "text", + "x": 1086, + "y": 271.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 495950, + "version": 1, + "versionNonce": 5436450, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "5. build-cube", + "originalText": "5. build-cube", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0049", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0051", + "type": "ellipse", + "x": 1221, + "y": 235, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 503869, + "version": 1, + "versionNonce": 5541179, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0052" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0052", + "type": "text", + "x": 1223, + "y": 239, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 511788, + "version": 1, + "versionNonce": 5645908, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0051", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0053", + "type": "text", + "x": 1082, + "y": 314, + "width": 38.5, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 519707, + "version": 1, + "versionNonce": 5750637, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~20 min", + "originalText": "~20 min", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0054", + "type": "rectangle", + "x": 1080, + "y": 330, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 527626, + "version": 1, + "versionNonce": 5855366, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0055" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0055", + "type": "text", + "x": 1086, + "y": 339.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 535545, + "version": 1, + "versionNonce": 5960095, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "cube/*.parquet\ncontingency table", + "originalText": "cube/*.parquet\ncontingency table", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0054", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0056", + "type": "rectangle", + "x": 400, + "y": 500, + "width": 420, + "height": 150, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 543464, + "version": 1, + "versionNonce": 6064824, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0057", + "type": "text", + "x": 414, + "y": 510, + "width": 300.3, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 551383, + "version": 1, + "versionNonce": 6169553, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Selection Loop (cheap -- iterate freely)", + "originalText": "Selection Loop (cheap -- iterate freely)", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0058", + "type": "rectangle", + "x": 424, + "y": 542, + "width": 175, + "height": 56, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 559302, + "version": 1, + "versionNonce": 6274282, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0059" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0059", + "type": "text", + "x": 430, + "y": 555.0, + "width": 163, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 567221, + "version": 1, + "versionNonce": 6379011, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "annealing_selection.yaml\nthresholds + ratios", + "originalText": "annealing_selection.yaml\nthresholds + ratios", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0058", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0060", + "type": "ellipse", + "x": 588, + "y": 527, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 575140, + "version": 1, + "versionNonce": 6483740, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0061" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0061", + "type": "text", + "x": 590, + "y": 531, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 583059, + "version": 1, + "versionNonce": 6588469, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0060", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0062", + "type": "rectangle", + "x": 640, + "y": 542, + "width": 160, + "height": 56, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 590978, + "version": 1, + "versionNonce": 6693198, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0063" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0063", + "type": "text", + "x": 646, + "y": 555.0, + "width": 148, + "height": 30.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 598897, + "version": 1, + "versionNonce": 6797927, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "6. preview\n--exact / --allow_fallback", + "originalText": "6. preview\n--exact / --allow_fallback", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0062", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0064", + "type": "ellipse", + "x": 789, + "y": 527, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 606816, + "version": 1, + "versionNonce": 6902656, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0065" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0065", + "type": "text", + "x": 791, + "y": 531, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 614735, + "version": 1, + "versionNonce": 7007385, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0064", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0066", + "type": "text", + "x": 424, + "y": 612, + "width": 181.50000000000003, + "height": 12.5, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 622654, + "version": 1, + "versionNonce": 7112114, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "~14 s per iteration over 19 cubes", + "originalText": "~14 s per iteration over 19 cubes", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0067", + "type": "arrow", + "x": 601, + "y": 558, + "width": 37, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 630573, + "version": 1, + "versionNonce": 7216843, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 37, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0068", + "type": "arrow", + "x": 638, + "y": 584, + "width": 37, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 638492, + "version": 1, + "versionNonce": 7321572, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + -37, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0069", + "type": "rectangle", + "x": 880, + "y": 530, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 646411, + "version": 1, + "versionNonce": 7426301, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0070" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0070", + "type": "text", + "x": 886, + "y": 551.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 654330, + "version": 1, + "versionNonce": 7531030, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "7. apply", + "originalText": "7. apply", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0069", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0071", + "type": "ellipse", + "x": 1021, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 662249, + "version": 1, + "versionNonce": 7635759, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0072" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0072", + "type": "text", + "x": 1023, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 670168, + "version": 1, + "versionNonce": 7740488, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0071", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0073", + "type": "rectangle", + "x": 880, + "y": 608, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 678087, + "version": 1, + "versionNonce": 7845217, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0074" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0074", + "type": "text", + "x": 886, + "y": 617.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 686006, + "version": 1, + "versionNonce": 7949946, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "filtered *.idx\n+ mix_manifest.yaml", + "originalText": "filtered *.idx\n+ mix_manifest.yaml", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0073", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0075", + "type": "rectangle", + "x": 1070, + "y": 530, + "width": 160, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 693925, + "version": 1, + "versionNonce": 8054675, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0076" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0076", + "type": "text", + "x": 1076, + "y": 551.5, + "width": 148, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 701844, + "version": 1, + "versionNonce": 8159404, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "8. export-jsonl", + "originalText": "8. export-jsonl", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0075", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0077", + "type": "ellipse", + "x": 1211, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 709763, + "version": 1, + "versionNonce": 8264133, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0078" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0078", + "type": "text", + "x": 1213, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 717682, + "version": 1, + "versionNonce": 8368862, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0077", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0079", + "type": "rectangle", + "x": 1070, + "y": 608, + "width": 160, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 725601, + "version": 1, + "versionNonce": 8473591, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0080" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0080", + "type": "text", + "x": 1076, + "y": 617.25, + "width": 148, + "height": 27.5, + "angle": 0, + "strokeColor": "#0c8599", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 733520, + "version": 1, + "versionNonce": 8578320, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "*.jsonl per dataset\n(sampling in the bytes)", + "originalText": "*.jsonl per dataset\n(sampling in the bytes)", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0079", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0081", + "type": "rectangle", + "x": 1260, + "y": 530, + "width": 175, + "height": 58, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 741439, + "version": 1, + "versionNonce": 8683049, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0082" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0082", + "type": "text", + "x": 1266, + "y": 551.5, + "width": 163, + "height": 15.0, + "angle": 0, + "strokeColor": "#e8590c", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 749358, + "version": 1, + "versionNonce": 8787778, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "cat out/*/*.jsonl", + "originalText": "cat out/*/*.jsonl", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0081", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0083", + "type": "ellipse", + "x": 1421, + "y": 515, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 757277, + "version": 1, + "versionNonce": 8892507, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0084" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0084", + "type": "text", + "x": 1423, + "y": 519, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 765196, + "version": 1, + "versionNonce": 8997236, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0083", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0085", + "type": "rectangle", + "x": 1260, + "y": 608, + "width": 175, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "dashed", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 773115, + "version": 1, + "versionNonce": 9101965, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0086" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0086", + "type": "text", + "x": 1266, + "y": 617.25, + "width": 163, + "height": 27.5, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 781034, + "version": 1, + "versionNonce": 9206694, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "-> Trainings Pipeline\n(Trainings Loop)", + "originalText": "-> Trainings Pipeline\n(Trainings Loop)", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0085", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0087", + "type": "rectangle", + "x": 60, + "y": 560, + "width": 275, + "height": 265, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 788953, + "version": 1, + "versionNonce": 9311423, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0088", + "type": "text", + "x": 74, + "y": 570, + "width": 135.85000000000002, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 796872, + "version": 1, + "versionNonce": 9416152, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Validation & Guards", + "originalText": "Validation & Guards", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0089", + "type": "rectangle", + "x": 82, + "y": 602, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 804791, + "version": 1, + "versionNonce": 9520881, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0090" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0090", + "type": "text", + "x": 88, + "y": 612.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 812710, + "version": 1, + "versionNonce": 9625610, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "verify-sidecar\nbyte offsets vs source", + "originalText": "verify-sidecar\nbyte offsets vs source", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0089", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0091", + "type": "ellipse", + "x": 301, + "y": 591, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 820629, + "version": 1, + "versionNonce": 9730339, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0092" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0092", + "type": "text", + "x": 303, + "y": 595, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 828548, + "version": 1, + "versionNonce": 9835068, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0091", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0093", + "type": "rectangle", + "x": 82, + "y": 666, + "width": 230, + "height": 46, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 836467, + "version": 1, + "versionNonce": 9939797, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0094" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0094", + "type": "text", + "x": 88, + "y": 674.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 844386, + "version": 1, + "versionNonce": 10044526, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "join coverage report\nper dataset", + "originalText": "join coverage report\nper dataset", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0093", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0095", + "type": "ellipse", + "x": 301, + "y": 655, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 852305, + "version": 1, + "versionNonce": 10149255, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0096" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0096", + "type": "text", + "x": 303, + "y": 659, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 860224, + "version": 1, + "versionNonce": 10253984, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0095", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0097", + "type": "rectangle", + "x": 82, + "y": 726, + "width": 230, + "height": 50, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 868143, + "version": 1, + "versionNonce": 10358713, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0098" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0098", + "type": "text", + "x": 88, + "y": 736.0, + "width": 218, + "height": 30.0, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 876062, + "version": 1, + "versionNonce": 10463442, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "smoke snapshot\n+ check_smoke_run", + "originalText": "smoke snapshot\n+ check_smoke_run", + "fontSize": 12, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0097", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0099", + "type": "ellipse", + "x": 301, + "y": 715, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 883981, + "version": 1, + "versionNonce": 10568171, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0100" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0100", + "type": "text", + "x": 303, + "y": 719, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 891900, + "version": 1, + "versionNonce": 10672900, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0099", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0101", + "type": "text", + "x": 82, + "y": 786, + "width": 126.50000000000001, + "height": 25.0, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 899819, + "version": 1, + "versionNonce": 10777629, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "run after any transfer,\nand before apply", + "originalText": "run after any transfer,\nand before apply", + "fontSize": 10, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0102", + "type": "arrow", + "x": 314, + "y": 198, + "width": 82, + "height": 6, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 907738, + "version": 1, + "versionNonce": 10882358, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 82, + -6 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0103", + "type": "arrow", + "x": 562, + "y": 197, + "width": 24, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 915657, + "version": 1, + "versionNonce": 10987087, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 24, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0104", + "type": "arrow", + "x": 315, + "y": 330, + "width": 271, + "height": 36, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 923576, + "version": 1, + "versionNonce": 11091816, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 271, + 36 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0105", + "type": "arrow", + "x": 752, + "y": 200, + "width": 104, + "height": 62, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 931495, + "version": 1, + "versionNonce": 11196545, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 104, + 62 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0106", + "type": "arrow", + "x": 762, + "y": 366, + "width": 94, + "height": 68, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 939414, + "version": 1, + "versionNonce": 11301274, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 94, + -68 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0107", + "type": "arrow", + "x": 1022, + "y": 279, + "width": 54, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 947333, + "version": 1, + "versionNonce": 11406003, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 54, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0108", + "type": "arrow", + "x": 1160, + "y": 378, + "width": 360, + "height": 120, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 955252, + "version": 1, + "versionNonce": 11510732, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 0, + 98 + ], + [ + -360, + 98 + ], + [ + -360, + 120 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0109", + "type": "arrow", + "x": 802, + "y": 570, + "width": 74, + "height": 14, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 963171, + "version": 1, + "versionNonce": 11615461, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 74, + -14 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0110", + "type": "arrow", + "x": 1042, + "y": 559, + "width": 26, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 971090, + "version": 1, + "versionNonce": 11720190, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 26, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0111", + "type": "arrow", + "x": 1232, + "y": 559, + "width": 26, + "height": 0, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 2 + }, + "seed": 979009, + "version": 1, + "versionNonce": 11824919, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "points": [ + [ + 0, + 0 + ], + [ + 26, + 0 + ] + ], + "lastCommittedPoint": null, + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow", + "elbowed": false + }, + { + "id": "el0112", + "type": "rectangle", + "x": 1680, + "y": 130, + "width": 120, + "height": 200, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": { + "type": 3 + }, + "seed": 986928, + "version": 1, + "versionNonce": 11929648, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0113", + "type": "text", + "x": 1694, + "y": 140, + "width": 35.75, + "height": 16.25, + "angle": 0, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 994847, + "version": 1, + "versionNonce": 12034377, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "Owner", + "originalText": "Owner", + "fontSize": 13, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + }, + { + "id": "el0114", + "type": "ellipse", + "x": 1701, + "y": 167, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "#a5d8ff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1002766, + "version": 1, + "versionNonce": 12139106, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0115" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0115", + "type": "text", + "x": 1703, + "y": 171, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1010685, + "version": 1, + "versionNonce": 12243835, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "R", + "originalText": "R", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0114", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0116", + "type": "ellipse", + "x": 1701, + "y": 197, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "#b2f2bb", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1018604, + "version": 1, + "versionNonce": 12348564, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0117" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0117", + "type": "text", + "x": 1703, + "y": 201, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1026523, + "version": 1, + "versionNonce": 12453293, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "T", + "originalText": "T", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0116", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0118", + "type": "ellipse", + "x": 1701, + "y": 227, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#f08c00", + "backgroundColor": "#ffec99", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1034442, + "version": 1, + "versionNonce": 12558022, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0119" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0119", + "type": "text", + "x": 1703, + "y": 231, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#f08c00", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1042361, + "version": 1, + "versionNonce": 12662751, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "S", + "originalText": "S", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0118", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0120", + "type": "ellipse", + "x": 1701, + "y": 257, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#e03131", + "backgroundColor": "#ffc9c9", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1050280, + "version": 1, + "versionNonce": 12767480, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0121" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0121", + "type": "text", + "x": 1703, + "y": 261, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#e03131", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1058199, + "version": 1, + "versionNonce": 12872209, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "H", + "originalText": "H", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0120", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0122", + "type": "ellipse", + "x": 1701, + "y": 287, + "width": 22, + "height": 22, + "angle": 0, + "strokeColor": "#099268", + "backgroundColor": "#96f2d7", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1066118, + "version": 1, + "versionNonce": 12976938, + "isDeleted": false, + "boundElements": [ + { + "type": "text", + "id": "el0123" + } + ], + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "el0123", + "type": "text", + "x": 1703, + "y": 291, + "width": 18, + "height": 14, + "angle": 0, + "strokeColor": "#099268", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1074037, + "version": 1, + "versionNonce": 13081667, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "M", + "originalText": "M", + "fontSize": 11, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "middle", + "containerId": "el0122", + "lineHeight": 1.25, + "autoResize": false + }, + { + "id": "el0124", + "type": "text", + "x": 1680, + "y": 340, + "width": 64.35000000000001, + "height": 11.25, + "angle": 0, + "strokeColor": "#868e96", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": null, + "roundness": null, + "seed": 1081956, + "version": 1, + "versionNonce": 13186396, + "isDeleted": false, + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "text": "initials only", + "originalText": "initials only", + "fontSize": 9, + "fontFamily": 1, + "textAlign": "left", + "verticalAlign": "top", + "containerId": null, + "lineHeight": 1.25, + "autoResize": true + } + ], + "appState": { + "gridSize": null, + "viewBackgroundColor": "#ffffff" + }, + "files": {} +} \ No newline at end of file diff --git a/config_files/data_preparation/quality/make_pipeline_diagram.py b/config_files/data_preparation/quality/make_pipeline_diagram.py new file mode 100644 index 000000000..899ec245b --- /dev/null +++ b/config_files/data_preparation/quality/make_pipeline_diagram.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +"""Generates the Data Preprocessing Pipeline diagram as an Excalidraw file. + +Matches the visual language of the existing training/pruning diagram: hand-drawn rounded +boxes, orange stage labels, coloured owner badges, grouped containers with titles, and a +section header underlined with '=' characters. + +Layout is two converging lanes -- documents along the top, annotations below -- meeting at +the join. That is not decoration: a single row needs arrows that skip past boxes, and those +arrows cross other boxes wherever they go. + +Kept as a generator rather than hand-written JSON so the layout can be nudged in one place. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +OUT = Path( + "/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality/" + "data_preprocessing_pipeline.excalidraw" +) + +INK = "#1e1e1e" +STAGE_TEXT = "#e8590c" +ARTIFACT_TEXT = "#0c8599" +NOTE_TEXT = "#868e96" +GUARD_TEXT = "#2f9e44" + +# Owner badge palette, keyed by the initial used in the existing diagram. +BADGES = { + "R": ("#1971c2", "#a5d8ff"), + "T": ("#2f9e44", "#b2f2bb"), + "S": ("#f08c00", "#ffec99"), + "H": ("#e03131", "#ffc9c9"), + "M": ("#099268", "#96f2d7"), +} + +elements: list[dict] = [] +_counter = [0] + + +def _next_id() -> str: + _counter[0] += 1 + return f"el{_counter[0]:04d}" + + +def _base(kind: str, x: float, y: float, w: float, h: float, **over) -> dict: + element = { + "id": _next_id(), + "type": kind, + "x": x, + "y": y, + "width": w, + "height": h, + "angle": 0, + "strokeColor": INK, + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 1, + "opacity": 100, + "groupIds": [], + "frameId": None, + "roundness": {"type": 3}, + "seed": 100_000 + _counter[0] * 7919, + "version": 1, + "versionNonce": 200_000 + _counter[0] * 104_729, + "isDeleted": False, + "boundElements": None, + "updated": 1, + "link": None, + "locked": False, + } + element.update(over) + return element + + +def text(content: str, x: float, y: float, size: int = 14, color: str = INK) -> dict: + lines = content.split("\n") + element = _base( + "text", x, y, + max(len(line) for line in lines) * size * 0.55, len(lines) * size * 1.25, + strokeColor=color, roundness=None, + ) + element.update({ + "text": content, "originalText": content, "fontSize": size, "fontFamily": 1, + "textAlign": "left", "verticalAlign": "top", "containerId": None, + "lineHeight": 1.25, "autoResize": True, + }) + elements.append(element) + return element + + +def box(label: str, x: float, y: float, w: float, h: float, color: str = STAGE_TEXT, + dashed: bool = False, size: int = 12) -> dict: + """A rounded rectangle with centred bound text.""" + rect = _base("rectangle", x, y, w, h, strokeStyle="dashed" if dashed else "solid") + lines = label.split("\n") + label_element = _base( + "text", x + 6, y + (h - len(lines) * size * 1.25) / 2, w - 12, len(lines) * size * 1.25, + strokeColor=color, roundness=None, + ) + label_element.update({ + "text": label, "originalText": label, "fontSize": size, "fontFamily": 1, + "textAlign": "center", "verticalAlign": "middle", "containerId": rect["id"], + "lineHeight": 1.25, "autoResize": False, + }) + rect["boundElements"] = [{"type": "text", "id": label_element["id"]}] + elements.append(rect) + elements.append(label_element) + return rect + + +def container(title: str, x: float, y: float, w: float, h: float) -> dict: + rect = _base("rectangle", x, y, w, h) + elements.append(rect) + text(title, x + 14, y + 10, size=13) + return rect + + +def badge(letter: str, cx: float, cy: float, r: float = 11) -> None: + stroke, fill = BADGES[letter] + ellipse = _base("ellipse", cx - r, cy - r, r * 2, r * 2, + strokeColor=stroke, backgroundColor=fill, roundness=None) + elements.append(ellipse) + label = _base("text", cx - r + 2, cy - 7, r * 2 - 4, 14, strokeColor=stroke, roundness=None) + label.update({ + "text": letter, "originalText": letter, "fontSize": 11, "fontFamily": 1, + "textAlign": "center", "verticalAlign": "middle", "containerId": ellipse["id"], + "lineHeight": 1.25, "autoResize": False, + }) + ellipse["boundElements"] = [{"type": "text", "id": label["id"]}] + elements.append(label) + + +def arrow(*waypoints: tuple[float, float], dashed: bool = False, color: str = INK) -> None: + """An arrow through a sequence of absolute points, so routes can dodge boxes.""" + x0, y0 = waypoints[0] + points = [[x - x0, y - y0] for x, y in waypoints] + element = _base( + "arrow", x0, y0, + max(abs(p[0]) for p in points), max(abs(p[1]) for p in points), + strokeColor=color, strokeStyle="dashed" if dashed else "solid", + roundness={"type": 2}, + ) + element.update({ + "points": points, "lastCommittedPoint": None, "startBinding": None, + "endBinding": None, "startArrowhead": None, "endArrowhead": "arrow", + "elbowed": False, + }) + elements.append(element) + + +# --------------------------------------------------------------------------- header +text("Data Preprocessing Pipeline", 60, 46, size=26) +text("=" * 34, 60, 82, size=18) + +# --------------------------------------------------------------------------- inputs +container("Inputs", 60, 130, 275, 320) +box("Source Corpora\n/data/annealing (19 subsets)", 82, 172, 230, 54) +badge("R", 312, 172) +box("Corpus Registry\nannealing_registry.yaml", 82, 240, 230, 50) +badge("R", 312, 240) +box("Propella Annotations\n(external parquet cache)", 82, 304, 230, 50) +badge("R", 312, 304) +box("Tokenizer Config\n(for token estimates only)", 82, 368, 230, 50) +badge("R", 312, 368) + +# ------------------------------------------------- lane 1: documents (top), y = 168 +box("1. calibrate", 400, 168, 160, 58) +badge("R", 552, 164) +text("~15 s / dataset", 402, 232, size=10, color=NOTE_TEXT) +box("calibration.yaml\nstratified ratios", 400, 262, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +box("2. build-sidecar", 590, 168, 160, 58) +badge("R", 742, 164) +text("~15 h (SLURM array)", 592, 232, size=10, color=NOTE_TEXT) +box("sidecar/*.parquet\n+ _files.json", 590, 262, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +# ------------------------------------------------- lane 2: annotations (below), y = 340 +box("3. bucket-annotations", 590, 340, 170, 58) +badge("R", 752, 336) +text("~4.5 h (SLURM array)", 592, 404, size=10, color=NOTE_TEXT) +box("buckets/ (key-hashed)", 590, 418, 170, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +# ------------------------------------------------- lanes converge +box("4. join-annotations", 860, 250, 160, 58) +badge("R", 1012, 246) +text("hours; --resume", 862, 314, size=10, color=NOTE_TEXT) +box("labelled sidecar\n100% coverage", 860, 330, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +box("5. build-cube", 1080, 250, 160, 58) +badge("R", 1232, 246) +text("~20 min", 1082, 314, size=10, color=NOTE_TEXT) +box("cube/*.parquet\ncontingency table", 1080, 330, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +# ------------------------------------------------- selection loop +container("Selection Loop (cheap -- iterate freely)", 400, 500, 420, 150) +box("annealing_selection.yaml\nthresholds + ratios", 424, 542, 175, 56) +badge("R", 599, 538) +box("6. preview\n--exact / --allow_fallback", 640, 542, 160, 56) +badge("R", 800, 538) +text("~14 s per iteration over 19 cubes", 424, 612, size=10, color=NOTE_TEXT) +arrow((601, 558), (638, 558)) +arrow((638, 584), (601, 584)) + +# ------------------------------------------------- tail +for x, name, artifact in ( + (880, "7. apply", "filtered *.idx\n+ mix_manifest.yaml"), + (1070, "8. export-jsonl", "*.jsonl per dataset\n(sampling in the bytes)"), +): + box(name, x, 530, 160, 58) + badge("R", x + 152, 526) + box(artifact, x, 608, 160, 46, color=ARTIFACT_TEXT, dashed=True, size=11) + +box("cat out/*/*.jsonl", 1260, 530, 175, 58) +badge("R", 1432, 526) +box("-> Trainings Pipeline\n(Trainings Loop)", 1260, 608, 175, 46, color=GUARD_TEXT, dashed=True, size=11) + +# ------------------------------------------------- validation +container("Validation & Guards", 60, 560, 275, 265) +box("verify-sidecar\nbyte offsets vs source", 82, 602, 230, 50, color=GUARD_TEXT) +badge("R", 312, 602) +box("join coverage report\nper dataset", 82, 666, 230, 46, color=GUARD_TEXT) +badge("R", 312, 666) +box("smoke snapshot\n+ check_token_estimates", 82, 726, 230, 50, color=GUARD_TEXT) +badge("R", 312, 726) +text("run after any transfer,\nand before apply", 82, 786, size=10, color=NOTE_TEXT) + +# ------------------------------------------------- flow arrows (routed to avoid boxes) +arrow((314, 198), (396, 192)) # corpora -> calibrate +arrow((562, 197), (586, 197)) # calibrate -> build-sidecar +arrow((315, 330), (586, 366)) # propella -> bucket-annotations +arrow((752, 200), (856, 262)) # sidecar -> join +arrow((762, 366), (856, 298)) # buckets -> join +arrow((1022, 279), (1076, 279)) # join -> cube +arrow((1160, 378), (1160, 476), (800, 476), (800, 498)) # cube -> selection loop +arrow((802, 570), (876, 556)) # preview -> apply +arrow((1042, 559), (1068, 559)) # apply -> export-jsonl +arrow((1232, 559), (1258, 559)) # export-jsonl -> concatenation + +# ------------------------------------------------- legend +container("Owner", 1680, 130, 120, 200) +for row, letter in enumerate("RTSHM"): + badge(letter, 1712, 178 + row * 30) +text("initials only", 1680, 340, size=9, color=NOTE_TEXT) + +document = { + "type": "excalidraw", + "version": 2, + "source": "https://excalidraw.com", + "elements": elements, + "appState": {"gridSize": None, "viewBackgroundColor": "#ffffff"}, + "files": {}, +} + +OUT.parent.mkdir(parents=True, exist_ok=True) +OUT.write_text(json.dumps(document, indent=2)) +print(f"wrote {OUT} ({len(elements)} elements)") diff --git a/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch new file mode 100755 index 000000000..c4468de6a --- /dev/null +++ b/config_files/data_preparation/quality/slurm/1_build_sidecar.sbatch @@ -0,0 +1,43 @@ +#!/bin/bash +# Per-document sidecar over /data/annealing. The only stage that reads the raw data. +# Work is divided per JSONL file across every enabled dataset, so this one array covers +# the whole blend regardless of how unevenly file counts fall. +#SBATCH --job-name=q_sidecar +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +# 8 CPUs is not what a task needs -- it is single-threaded. It is what stops SLURM +# packing 16 of them onto one node. The binding constraint is per-node bandwidth to +# /data, not CPU: a real run put 12 tasks on one node and each got 95 MB/s against the +# 282 MB/s a single stream gets, while four nodes sat idle. At 8 CPUs only 4 tasks fit +# per node, so an array of 64 spreads over 16 nodes instead of piling onto 4. +#SBATCH --cpus-per-task=8 +#SBATCH --mem=16G +#SBATCH --time=24:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/1_sidecar_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/1_sidecar_%A_%a.err +#SBATCH --array=0-63 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" + +# Each task is single-threaded and I/O-bound at ~374 MB/s where bandwidth allows. The +# 3.8 GB/s aggregate measured on the login node is not what a compute node gets, so +# spreading tasks across nodes matters more than raising the array size. +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" +echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" + +srun "$MQ" -m modalities quality build-sidecar \ + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" \ + --work_dir "$WORK" \ + --index_root "$WORK/idx" \ + --shard_id "$SLURM_ARRAY_TASK_ID" \ + --num_shards "$NUM_SHARDS" + +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch new file mode 100755 index 000000000..8cbdfd4c3 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/2_bucket_annotations.sbatch @@ -0,0 +1,43 @@ +#!/bin/bash +# Partition the propella splits by a hash of the join key. The expensive half of the +# join: HPLT alone is ~12 bn annotation rows. Each task writes its own file per bucket. +#SBATCH --job-name=q_bucket +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +# As in the sidecar script: the CPU request is there to limit how many bandwidth-hungry +# tasks SLURM packs onto one node, not because a task needs 8 cores. +#SBATCH --cpus-per-task=8 +# Peak RSS measured at 2.4 GB for a 50 M-row shard at 1024 buckets, so 24 GB is ample. +# It was not always: the writer used to buffer per bucket, which bounded nothing, and a +# real run OOM-killed all 64 tasks at this same limit. +#SBATCH --mem=24G +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/2_bucket_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/2_bucket_%A_%a.err +#SBATCH --array=0-63 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +NUM_BUCKETS="${NUM_BUCKETS:-1024}" + +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +NUM_SHARDS="${SLURM_ARRAY_TASK_COUNT:-64}" +echo "START $(date) shard ${SLURM_ARRAY_TASK_ID}/${NUM_SHARDS}" + +# Every task in the array must finish. join-annotations refuses to run against an +# incomplete bucketing run rather than silently dropping a missing task's annotations, +# so a failed task is a hard stop, not a quiet gap. Re-run just that shard id. +srun "$MQ" -m modalities quality bucket-annotations \ + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" \ + --work_dir "$WORK" \ + --num_buckets "$NUM_BUCKETS" \ + --shard_id "$SLURM_ARRAY_TASK_ID" \ + --num_shards "$NUM_SHARDS" + +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch b/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch new file mode 100755 index 000000000..ac620209e --- /dev/null +++ b/config_files/data_preparation/quality/slurm/3_join_and_cube.sbatch @@ -0,0 +1,31 @@ +#!/bin/bash +# Attach the annotations to the sidecars, then aggregate into cubes. Needs the sidecar +# and bucketing arrays finished. Single task; the cube build is the longer half (~50 min +# for the full blend) and holds a few GB while grouping. +#SBATCH --job-name=q_join_cube +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/3_join_cube_%j.out +#SBATCH --error=/home/richard.rutmann/logs/quality/3_join_cube_%j.err + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true +echo "START $(date)" + +srun "$MQ" -m modalities quality join-annotations \ + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" --work_dir "$WORK" + +srun "$MQ" -m modalities quality build-cube \ + --registry "${REGISTRY:-$QDIR/annealing_registry.yaml}" --work_dir "$WORK" + +echo "Coverage per dataset is in $WORK/join_report.json -- read it before trusting a selection." +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch new file mode 100755 index 000000000..e0e464ad9 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/3a_join_annotations.sbatch @@ -0,0 +1,57 @@ +#!/bin/bash +# Attach the bucketed annotations to each dataset's sidecar. One array task per annotated +# dataset: the joins are independent, so wall time is the slowest dataset (~2 h for +# nemotron-cc) rather than the sum (~6 h). +# +# Set the array upper bound to (number of annotated datasets) - 1. Print the list with: +# $MQ -c "from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry; \ +# r=CorpusRegistry.from_yaml('$QDIR/annealing_registry.yaml'); \ +# print([d.name for d in r.enabled_datasets() if d.annotation_split])" +#SBATCH --job-name=q_join +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +# 8 CPUs to cap tasks per node: this is bandwidth-bound, not CPU-bound. See the README. +#SBATCH --cpus-per-task=8 +# Measured 51.9 GiB peak on an 18.7 M-document nemotron-cc batch, against a real batch of +# 20 M keys. Resolving in Arrow roughly doubled peak memory over the old per-row code (24.9 +# GiB), because the batch holds the matched annotation rows and every part's table at once. +# 64 G left only 19% headroom on a stage that has already been OOM-killed once. +#SBATCH --mem=160G +# nemotron-cc measured ~12 h at 1.7 bn documents and was killed by a 12 h limit at 99.8% +# complete. The other 15 datasets each finished inside 6.5 h. +#SBATCH --time=48:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/3a_join_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/3a_join_%A_%a.err +#SBATCH --array=0-15 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +# Resolve this task's dataset from the registry, so the mapping cannot drift from the +# config the way a hand-maintained list would. +DATASET=$("$MQ" - "$REGISTRY" "$SLURM_ARRAY_TASK_ID" <<'PY' +import sys +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +reg = CorpusRegistry.from_yaml(sys.argv[1]) +names = [d.name for d in reg.enabled_datasets() if d.annotation_split] +idx = int(sys.argv[2]) +print(names[idx] if idx < len(names) else "") +PY +) + +if [[ -z "$DATASET" ]]; then + echo "array index $SLURM_ARRAY_TASK_ID is past the last annotated dataset; nothing to do" + exit 0 +fi + +echo "START $(date) dataset=$DATASET" +srun "$MQ" -m modalities quality join-annotations \ + --registry "$REGISTRY" --work_dir "$WORK" --only "$DATASET" ${JOIN_RESUME:+--resume} +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch b/config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch new file mode 100755 index 000000000..f2deb9dc1 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/3b_build_cubes.sbatch @@ -0,0 +1,26 @@ +#!/bin/bash +# Aggregate the joined sidecars into cubes. Must run after every join task has finished. +# Single task; ~50 min for the whole blend, and it holds a few GB while grouping. +#SBATCH --job-name=q_cubes +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=64G +#SBATCH --time=8:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/3b_cubes_%j.out +#SBATCH --error=/home/richard.rutmann/logs/quality/3b_cubes_%j.err + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +echo "START $(date)" +srun "$MQ" -m modalities quality build-cube --registry "$REGISTRY" --work_dir "$WORK" +echo "Coverage per dataset: $WORK/join_report.json -- read it before trusting a selection." +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch b/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch new file mode 100755 index 000000000..15ac5f759 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/4_export_jsonl.sbatch @@ -0,0 +1,75 @@ +#!/bin/bash +# Write the selected documents out as JSONL, one shard per source file, with the up- and +# downsampling materialised in the bytes. One array task per dataset: the datasets are +# independent, so wall time is the largest one rather than the sum. +# +# NUM_FILE_SHARDS splits each dataset's source files across several tasks, striding rather +# than slicing contiguously. One task per dataset is fine until one dataset holds far more +# files than the rest: dolmino has 40,003 against a median near 500, and at ~1.1 s per file +# that is 12.2 h against 20 minutes -- past this script's own wall limit. With it set, the +# array is (datasets x NUM_FILE_SHARDS) tasks and the task index is decomposed below. +# +# Set the array upper bound to (number of datasets in the mix manifest) - 1. Print the list: +# $MQ -c "import yaml,sys; m=yaml.safe_load(open('$WORK/mix/mix_manifest.yaml')); \ +# print(sorted({d.get('source_dataset') or d['name'] for d in m['datasets']}))" +# +# Each task writes its own dataset record; the blend-wide export_manifest.yaml is merged +# afterwards by the --finalize_only step below, because eighteen tasks writing one shared +# manifest would race and the last writer would erase the rest. +#SBATCH --job-name=q_export +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +# Pure byte copying: read the selected lines, write them out. Bandwidth-bound, not CPU-bound, +# so 8 CPUs is here to cap tasks per node rather than because they are used. +#SBATCH --cpus-per-task=8 +# One source file's index at a time plus an mmap of that file. Nothing accumulates across +# files, so this is far below the join or the apply stage. +#SBATCH --mem=32G +# Estimated 9.75 TB written in total, dominated by nemotron-cc at 2.14 TB and klettermix-de +# at 1.68 TB -- both inflated by their ratios. Packing 6.0 TB took 1 h 40 m with tokenisation; +# this does less work per byte but writes more of them. +#SBATCH --time=12:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/4_export_%A_%a.out +#SBATCH --error=/home/richard.rutmann/logs/quality/4_export_%A_%a.err +#SBATCH --array=0-17 + +set -euo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" +MANIFEST="${MANIFEST:-$WORK/mix/mix_manifest.yaml}" +OUT="${OUT:-$WORK/out}" +NUM_FILE_SHARDS="${NUM_FILE_SHARDS:-1}" + +unset SLURM_MEM_PER_CPU || true +unset SLURM_MEM_PER_GPU || true + +# Resolved from the manifest rather than a hand-maintained list, so the mapping cannot drift +# from what was actually materialised. Curve buckets collapse to their source dataset. +DATASET=$("$MQ" - "$MANIFEST" "$SLURM_ARRAY_TASK_ID" "$NUM_FILE_SHARDS" <<'PY' +import sys, yaml +m = yaml.safe_load(open(sys.argv[1])) +names = sorted({d.get("source_dataset") or d["name"] for d in m["datasets"]}) +# With file sharding the array covers datasets x shards, so the flat task index splits into +# a dataset and a shard within it. +idx, n_shards = int(sys.argv[2]), int(sys.argv[3]) +dataset_idx, shard_idx = divmod(idx, n_shards) +print(f"{names[dataset_idx]} {shard_idx}" if dataset_idx < len(names) else "") +PY +) + +if [[ -z "$DATASET" ]]; then + echo "array index $SLURM_ARRAY_TASK_ID is past the last dataset; nothing to do" + exit 0 +fi +SHARD_ID="${DATASET##* }" +DATASET="${DATASET%% *}" + +echo "START $(date) dataset=$DATASET shard=$SHARD_ID/$NUM_FILE_SHARDS" +srun "$MQ" -m modalities quality export-jsonl \ + --manifest "$MANIFEST" --registry "$REGISTRY" --output_dir "$OUT" \ + --only "$DATASET" --no_finalize \ + --shard_id "$SHARD_ID" --num_shards "$NUM_FILE_SHARDS" +echo "END $(date)" diff --git a/config_files/data_preparation/quality/slurm/5_verify.sbatch b/config_files/data_preparation/quality/slurm/5_verify.sbatch new file mode 100755 index 000000000..6f732a7de --- /dev/null +++ b/config_files/data_preparation/quality/slurm/5_verify.sbatch @@ -0,0 +1,35 @@ +#!/bin/bash +# Check what is actually on disk, not what the pipeline believes it wrote. +# +# One stage now that the output is JSONL rather than packed tokens: +# verify_jsonl -- every dataset's line count against its record, the realised +# lines-per-document against the ratio that was asked for, a byte-for-byte +# replay of sampled shards against the corpus, and a check that nothing was +# written into the source tree. +# +# The replay is the one that would catch a real problem: it recomputes which documents the +# export should have written and how many times, and compares that to the bytes on disk. +#SBATCH --job-name=q_verify +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=8 +# Streams shards and source files a line at a time; nothing accumulates. +#SBATCH --mem=32G +#SBATCH --time=08:00:00 +#SBATCH --output=/home/richard.rutmann/logs/quality/verify_%j.out +#SBATCH --error=/home/richard.rutmann/logs/quality/verify_%j.err + +set -uo pipefail + +MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +QDIR="${QDIR:-/home/richard.rutmann/repos/modalities/config_files/data_preparation/quality}" +WORK="${WORK:?WORK is not set}" +SOURCE_ROOT="${SOURCE_ROOT:-/data/annealing}" +# Re-counting every line reads the whole 9.75 TB export; off by default, worth it once. +COUNT_LINES="${COUNT_LINES:-}" + +srun "$MQ" "$QDIR/slurm/verify_jsonl.py" \ + --manifest "$WORK/out/export_manifest.yaml" \ + --mix_manifest "$WORK/mix/mix_manifest.yaml" \ + --source_root "$SOURCE_ROOT" \ + --shards 8 ${COUNT_LINES:+--count_lines} diff --git a/config_files/data_preparation/quality/slurm/README.md b/config_files/data_preparation/quality/slurm/README.md new file mode 100644 index 000000000..c8f039abb --- /dev/null +++ b/config_files/data_preparation/quality/slurm/README.md @@ -0,0 +1,356 @@ +# Running the quality pipeline on /data/annealing + +`/home/richard.rutmann/data` is a symlink to `/data`, so `/home/richard.rutmann/data/annealing` +and `/data/annealing` are the same tree (same inode). The registry already uses the +`/data/...` form; nothing needs changing. + +## Environment + +The `nemo_25_11.sif` container cannot run this branch: it ships torch 2.9 and modalities +needs `>=2.10` (`ScheduleDualPipeV`). These stages are CPU-only, so a plain venv is +simpler than rebuilding the container: + +```bash +/opt/conda/bin/python3 -m venv /data/user/richard.rutmann/venvs/modalities-quality +V=/data/user/richard.rutmann/venvs/modalities-quality +$V/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu +$V/bin/pip install -e /home/richard.rutmann/repos/modalities +``` + +Already built and verified at that path. Every command below uses `$MQ`: + +```bash +cd /home/richard.rutmann/repos/modalities +source config_files/data_preparation/quality/slurm/env.sh +``` + +`env.sh` sets `MQ`, `REPO`, `QDIR`, `WORK`, `HF_HOME`, `HF_TOKEN` and `EXPORTS`, only +filling in what is unset, and refuses to continue if any of them ends up empty or points +at something missing. Source it in every new shell. Setting them by hand works too, but a +shell where `QDIR` is empty turns `$QDIR/slurm/x.sbatch` into `/slurm/x.sbatch`, and one +where `WORK` is empty turns `rm -rf $WORK/buckets` into `rm -rf /buckets`. Both have +happened. + +`calibrate` downloads the tokenizer, so it needs HF auth. A token is already at +`~/.config/huggingface/token`: + +```bash +export HF_TOKEN="$(tr -d '\r\n' < ~/.config/huggingface/token)" +export HF_HOME=/data/cache/hf_cache +``` + +## Before you start + +Confirm the tokenizer in `annealing_tokenizer.yaml`. It is set to +`nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, the one the long-context pipeline uses; +the token audit under `/data/michael.fromm` used the Super-120B variant. Every token +figure downstream depends on this, and a wrong choice fails silently. + +Two registry entries are deliberately disabled or worth a second look: +`finepdfs-it` is disabled (it is English data duplicating `finepdfs-en`), and +`nemotron-cc-v2` is enabled at 9.86 TB but has no annotations at all. Drop it from the +registry if the blend does not need it -- it is ~16 h of sidecar work. + +## Running everything with timings + +`run_all_timed.sh` runs the steps in order and prints how long each took, plus a summary. +Array jobs go through `sbatch --wait`, so the timing is the job's real duration rather +than how long submission took. + +```bash +bash $QDIR/slurm/run_all_timed.sh # all steps +bash $QDIR/slurm/run_all_timed.sh 1 2 3 4 # only the once-per-blend stages +bash $QDIR/slurm/run_all_timed.sh 5 # just re-preview +``` + +It stops at the first failing step and still prints the summary, and it honours +`REGISTRY`, `SELECTION`, `TOKENIZER_CONFIG`, `WORK`, `SIDECAR_TASKS`, `BUCKET_TASKS`, +`NUM_BUCKETS`, `PACK_TASKS`, `SAMPLE_SIZE` and `BLEND_NAME` as environment overrides. + +The individual commands are below if you would rather drive them by hand. + +## Steps 1-4: once per blend + +Nothing is written into `/data/annealing`. Indexes go to `$WORK/idx`. + +```bash +# 1. Token calibration. ~4 min for all 19 datasets, one task, ~1 GB read. +# Writes calibration.yaml after each dataset, so an interruption keeps what it +# measured; re-run with --only to fill in the rest. +$MQ -m modalities quality calibrate \ + --registry $QDIR/annealing_registry.yaml --work_dir $WORK \ + --tokenizer_config $QDIR/annealing_tokenizer.yaml --sample_size 2000 + +# 2. Sidecar. The only stage that reads the raw data. Array of 64. +sbatch $QDIR/slurm/1_build_sidecar.sbatch + +# 3. Bucket the annotation splits. Array of 64. Wait for step 2 only if you +# like -- it touches different data, so the two can run concurrently. +sbatch $QDIR/slurm/2_bucket_annotations.sbatch + +# 4a. Join, one array task per annotated dataset. The joins are independent, so wall +# time is the slowest dataset (~2 h for nemotron-cc) not the sum (~6 h). +sbatch --export=$EXPORTS --array=0-15 $QDIR/slurm/3a_join_annotations.sbatch + +# 4b. Aggregate into cubes. After every join task has finished. ~50 min, one task. +sbatch --export=$EXPORTS $QDIR/slurm/3b_build_cubes.sbatch +``` + +`3_join_and_cube.sbatch` still exists and does both in one task if you would rather not +manage two submissions; it takes ~6 h instead of ~3 h. + +Check coverage before trusting anything: `$WORK/join_report.json` gives the annotated +fraction per dataset. FinePDFs was 10-34% at last measurement. + +## Step 5: the loop you actually iterate + +```bash +# Edit thresholds and ratios, then: +$MQ -m modalities quality preview \ + --selection $QDIR/annealing_selection.yaml --work_dir $WORK +``` + +~10 s for the whole blend, on the login node. Reads only the cubes. Repeat as often as +you like. `--exact` scans the per-document sidecars instead, for a threshold that fell +inside a cube bin. + +## Steps 6-8: once you have settled + +```bash +# 6. Filtered indexes plus a manifest of what was selected. +$MQ -m modalities quality apply \ + --selection $QDIR/annealing_selection.yaml \ + --registry $QDIR/annealing_registry.yaml \ + --work_dir $WORK --output_dir $WORK/blend_v1 + +# 7. Export the sampled documents as JSONL, one shard per source file. +$MQ -m modalities quality export-jsonl \ + --manifest $WORK/blend_v1/mix_manifest.yaml \ + --registry $QDIR/annealing_registry.yaml \ + --output_dir $WORK/out +``` + +The output is `$WORK/out//.jsonl`, and the training set is the +concatenation of all of it: + +```bash +cat $WORK/out/*/*.jsonl > training.jsonl # or feed the shards in directly +``` + +**The ratios are already in the bytes.** This is the one thing to get right about this stage. +A dataset at 3.0 has each of its documents written three times; one at 0.6 has two of every +five dropped. `mix_manifest.yaml` still records `ratio: 3.0`, because that is what was asked +for -- but `export_manifest.yaml` records `training_ratio: 1.0` and +`repeat_factor_applied: true`, and that is the number a training config must use. Carrying +the mix manifest's ratio into a `weighted_combined` dataset after this stage would apply it a +second time, training that data nine times rather than three. + +Fractional factors are resolved per document, not by truncating a list: 1.2 means every +document once and a hash-chosen fifth of them twice. The choice depends on the selection's +`seed` and the document's position, so it is identical on every run and every machine, which +is what makes the stage resumable. + +Copies of a document are written **adjacent** to each other. A sequential reader will see the +same document several times in a row unless the training shuffle buffer is larger than the run +of copies. + +Reruns skip shards that are already complete. Completeness is checked against a recorded line +and byte count, not against the file merely existing -- the same distinction that let a +truncated `.pbin` survive into a blend once. + +No `weighted_combined` dataset is needed any more: the repetition is on disk, so every +exported file is drawn exactly once. + +## Validate the token estimate before trusting a large budget + +Predicted token counts are estimates. Confirm them on one small dataset by comparing the +preview against what packing actually produced: + +```bash +$MQ -m modalities quality preview --selection $QDIR/annealing_selection.yaml \ + --work_dir $WORK 2>&1 | grep finewiki-it +$MQ $QDIR/slurm/check_token_estimates.py \ + --manifest $WORK/out/export_manifest.yaml \ + --mix_manifest $WORK/mix/mix_manifest.yaml \ + --tokenizer_config $QDIR/annealing_tokenizer.yaml +``` + +Nothing in the pipeline tokenizes any more, so this is the only thing that checks the +calibration the whole token budget rests on. It tokenizes a sample of the exported JSONL and +reports the error per dataset. A few percent is expected; tens of percent means the +calibration is modelling something other than what the export writes. + +## Resuming an interrupted join + +The join writes labels into each sidecar part as it goes, so an interrupted run can be +continued rather than redone: + +```bash +JOIN_RESUME=1 sbatch --wait --export=$EXPORTS,JOIN_RESUME=1 \ + --array= $QDIR/slurm/3a_join_annotations.sbatch +``` + +`--resume` skips parts that already carry the label columns and reports how many it +skipped. A real interrupted `nemotron-cc` join finished in **9 minutes instead of 12 +hours** this way, having skipped 5,309 of 5,319 parts. + +Leave it off after re-bucketing the annotations: resuming would keep the labels from the +previous bucketing rather than picking up the new ones. + +## Clearing the bucketed annotations + +A sharded bucketing run cannot clear its own output directory -- sibling tasks are writing +into it -- so clearing it is a separate, deliberate step: + +```bash +bash $QDIR/slurm/reset_buckets.sh +``` + +It refuses to run without `WORK` set, and refuses obviously wrong paths. Sidecars are left +alone; only `$WORK/buckets` goes. + +## Re-running a failed shard + +Every stage is idempotent per shard, so a failed array task is re-run on its own: + +```bash +# sidecar: rebuilds only that task's parquet parts +sbatch --array=17 $QDIR/slurm/1_build_sidecar.sbatch + +# bucketing: --num_shards must match the original array, or the buckets will not line up +sbatch --array=41 $QDIR/slurm/2_bucket_annotations.sbatch +``` + +`bucket-annotations` skips a split whose output is already complete, so re-submitting the +whole array is safe but pointless. Use `--force` to genuinely re-bucket. + +## What the stages leave behind + +``` +$WORK/calibration.yaml bytes-per-token per dataset +$WORK/idx// .idx per source file (reusable; packing needs these too) +$WORK/sidecar// one parquet part per source file +$WORK/buckets// partitioned annotations +$WORK/cube/.parquet what preview reads; a few MB each +$WORK/join_report.json annotated fraction per dataset <- read this +$WORK/blend_v1/ filtered indexes + mix_manifest.yaml +$WORK/out// exported .jsonl shards + export_manifest.yaml +``` + +Only `$WORK` is written. `/data/annealing` is read-only throughout -- verified on a real +run of `finewiki-it`, which produced 452,714 sidecar rows and added no file to the source +tree. + +## Measured on a real run of finewiki-it + +| Stage | Measured | +|---|---| +| calibrate, all 19 datasets | 4 min 5 s (2000 documents sampled each) | +| build-sidecar, 1 of 4 files | 55 s for 6 GB / 452,714 documents, incl. index creation | +| bucket-annotations, finewiki split | 6 m 36 s for 43,097,138 rows (108.8k rows/s) | +| join-annotations | 1 m 47 s, **100 % coverage**, 868,586 duplicate annotation keys (2.0 %) | +| build-cube | 12 s -> 456 cells over 452,714 documents, 1.20 B estimated tokens | +| preview | 8 s, of which ~6 s is interpreter startup | + +Two things worth carrying forward from that run: + +* **Duplicate annotation keys are normal, not a Nemotron-CC quirk.** finewiki reported + 868,586 of them. The join keeps the first occurrence and reports the count; check it in + `join_report.json` rather than assuming a clean one-to-one join. +* **`preview` pays ~6 s of import overhead** before it does any work, because the CLI + imports the component registry and therefore torch. The cube evaluation itself is + milliseconds to a couple of seconds, so budget roughly 15-20 s for a full-blend preview + rather than the 10 s the cube maths alone suggests. + +On that run token retention (62.6 %) exceeded row retention (51.5 %) on real data, which +is the length correlation the design exists to account for. +## Full run, as measured on the annealing blend + +Timings are from the complete run over `/data/annealing` (20.21 TB, 19 datasets, August +2026). Everything long goes through SLURM: the login node has usage limits and killed two +sessions during this run. + +```bash +cd /home/richard.rutmann/repos/modalities +source config_files/data_preparation/quality/slurm/env.sh +REG=$QDIR/annealing_registry.yaml +SEL=$QDIR/annealing_selection.yaml +TOK=$QDIR/annealing_tokenizer.yaml +``` + +### 0. Gate + +Do not start while the corpus is still being written. A transfer that re-shards a corpus +after its sidecar is built invalidates every byte offset, and the only loud symptom is a +dataset whose file count fell to zero. + +```bash +find /data/annealing -name '*.jsonl' -mmin -180 | wc -l # must be 0 +$MQ -c " +from pathlib import Path +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +r = CorpusRegistry.from_yaml(Path('$REG')) +[print('MISSING', d.name, d.jsonl_root) for d in r.enabled_datasets() if not d.jsonl_root.exists()]" +``` + +### 1-9. The pipeline + +```bash +# 1. Calibrate. 48 min for 19 datasets. +$MQ -m modalities quality calibrate --registry $REG --work_dir $WORK --tokenizer_config $TOK + +# 2. Sidecars. 2 h 09 m, 64 tasks, the only stage that reads all 20 TB. +sbatch --wait --export=$EXPORTS $QDIR/slurm/1_build_sidecar.sbatch + +# 3. Verify the offsets before spending anything more. 143 s. +$MQ -m modalities quality verify-sidecar --registry $REG --work_dir $WORK + +# 4. Buckets. 4.5 h -- SKIP if $WORK/buckets is intact and the annotations have not moved. +# sbatch --wait --export=$EXPORTS $QDIR/slurm/2_bucket_annotations.sbatch + +# 5. Join. ~6 h; nemotron-cc is the long pole at 6 h alone. Add JOIN_RESUME=1 to continue +# an interrupted run -- never after re-bucketing, which would keep stale labels. +sbatch --wait --export=$EXPORTS $QDIR/slurm/3a_join_annotations.sbatch + +# 6. Cubes. 12 min. +sbatch --wait --export=$EXPORTS $QDIR/slurm/3b_build_cubes.sbatch + +# 7. Preview. 14 s -- edit $SEL and repeat as often as you like. +$MQ -m modalities quality preview --selection $SEL --work_dir $WORK +$MQ -m modalities quality preview --selection $SEL --work_dir $WORK --explain # which predicate binds +$MQ -m modalities quality preview --selection $SEL --work_dir $WORK --exact # before committing + +# 8. Apply. 42 min, peaks near 160 GB: it holds every kept document's offsets. +sbatch --wait --job-name=q_apply --nodes=1 --cpus-per-task=8 --mem=220G --time=12:00:00 \ + --output=$HOME/logs/quality/apply_%j.out --error=$HOME/logs/quality/apply_%j.err \ + --export=$EXPORTS --wrap="srun $MQ -m modalities quality apply --selection $SEL \ + --registry $REG --work_dir $WORK --output_dir $WORK/mix" + +# 9. Export as JSONL. One array task per dataset; ~9.75 TB written. +sbatch --wait --export=$EXPORTS,MANIFEST=$WORK/mix/mix_manifest.yaml,OUT=$WORK/out \ + --array=0-17 $QDIR/slurm/4_export_jsonl.sbatch +$MQ -m modalities quality export-jsonl --manifest $WORK/mix/mix_manifest.yaml \ + --registry $REG --output_dir $WORK/out --finalize_only + +# 10. Verify: line counts, realised ratios, a byte-for-byte replay, source tree untouched. +sbatch --wait --export=$EXPORTS,SOURCE_ROOT=/data/annealing $QDIR/slurm/5_verify.sbatch +``` + +About 12 hours end to end, with the join as the long pole. The export replaces what was +1 h 40 m of packing; it writes more bytes but does no tokenization. + +### What each check catches + +`verify-sidecar` reads the source bytes at recorded offsets. A re-sharded corpus once left +11 of 19 datasets with unusable sidecars and only one failed loudly. + +`5_verify.sbatch` replays sampled shards against the corpus: it recomputes which documents the +export should have written and how many copies of each, then compares that to the bytes on +disk. Counting files proves nothing -- one `.pbin` in 54,738 was once 151 MB on disk reporting +`data_len=0`, and counting would have shipped that dataset 567 M tokens short. The same trap +exists here, which is why a shard counts as complete only against a recorded line and byte +count. + +Document counts in the verification must match **exactly**. They are not estimates -- the +filtered index names precisely the selected documents -- so any difference is a defect in +materialize or the index, not estimator error. Token estimates landed within -0.56% overall. diff --git a/config_files/data_preparation/quality/slurm/check_token_estimates.py b/config_files/data_preparation/quality/slurm/check_token_estimates.py new file mode 100755 index 000000000..fec00dded --- /dev/null +++ b/config_files/data_preparation/quality/slurm/check_token_estimates.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +"""Compares the pipeline's token estimates against a real tokenization of the output. + +Every figure `preview` reports -- the blend's yield, each dataset's share, whether the run +wraps -- is estimated from text bytes and a per-dataset calibration. Nothing in the pipeline +tokenizes any more, so nothing checks that model unless something like this does. + +This tokenizes a sample of the exported JSONL and reports the error per dataset. It is a +sample rather than a full pass because the point is to catch a calibration that is wrong by +tens of percent, not to measure the last percent: the calibration itself was built from +64 KB slices, and a few thousand documents is already far more than that. + +An error of a few percent is expected and fine. An error of tens of percent means the +calibration is measuring something other than what the export writes -- a changed tokenizer, +a text field that is not the one being counted, or a dataset whose records shifted shape. + +**Trust the blend row; treat a single dataset's row as indicative.** Sampling in proportion +to length over-weights very large documents, and on a corpus whose bytes-per-token ratio +spans an order of magnitude with the tail holding much of the mass -- FineWiki runs 3.571 for +documents under a kilobyte to 34.648 for the nine above 256 KB -- the measured tokens-per-byte +comes out low and the estimate looks inflated. Measured on the real blend: finewiki-es +reported +51.1% on one seed and +8.9% on another, while the blend total moved only from -0.8% +to -1.1%. Re-run with a different --seed before believing any single dataset's number. + +Reads only. +""" + +from __future__ import annotations + +import argparse +import json +import random +import sys +from pathlib import Path + +import yaml + +from modalities.config.config import load_app_config_dict +from modalities.tokenization.tokenizer_wrapper import PreTrainedHFTokenizer + + +def sample_lines(shard_paths: list[Path], n: int, rng: random.Random) -> list[bytes]: + """Samples lines at byte offsets spread across a dataset's shards. + + Seeks to a random offset, discards the partial line it lands in, and takes the next whole + one. That selects documents in proportion to their length, which is what makes the + measured tokens-per-byte ratio comparable to the estimate -- the calibration is itself a + bytes-to-tokens ratio. + + Reading a prefix instead, which this did at first, is the bias the calibration work + already had to fix once: the opening lines of a shard are not a random draw from it, and + on the real blend that produced errors from -65% to +106% with no consistent sign. + + Args: + shard_paths (list[Path]): The dataset's shards. + n (int): How many lines to aim for. + rng (random.Random): Chooses shards and offsets. + + Returns: + list[bytes]: The sampled lines, without terminators. + """ + chosen = rng.sample(shard_paths, min(8, len(shard_paths))) + per_shard = max(1, n // len(chosen)) + lines: list[bytes] = [] + for shard in chosen: + size = shard.stat().st_size + if size < 2: + continue + with shard.open("rb") as f: + for _ in range(per_shard): + f.seek(rng.randrange(size)) + f.readline() # discard the partial line this offset fell inside + line = f.readline() # the document containing the next boundary + if not line: + f.seek(0) + line = f.readline() + if line.strip(): + lines.append(line.rstrip(b"\n")) + return lines + + +def main() -> int: + """Reports estimated against measured tokens. + + Returns: + int: 0 always; this is a report, and what counts as too much error is a judgement + the reader makes with the table in front of them. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True, help="export_manifest.yaml") + parser.add_argument("--mix_manifest", type=Path, required=True, help="mix_manifest.yaml, for the estimates.") + parser.add_argument("--tokenizer_config", type=Path, required=True, help="Config holding the tokenizer section.") + parser.add_argument("--sample", type=int, default=2000, help="Lines to tokenize per dataset.") + parser.add_argument("--text_field", default="text", help="The JSON field holding the document text.") + parser.add_argument("--seed", type=int, default=0, help="Seed for choosing shards.") + args = parser.parse_args() + + with args.manifest.open() as f: + manifest = yaml.safe_load(f) + with args.mix_manifest.open() as f: + mix = yaml.safe_load(f) + root = args.manifest.parent + + estimated_tokens: dict[str, float] = {} + estimated_documents: dict[str, int] = {} + for row in mix["datasets"]: + name = row.get("source_dataset") or row["name"] + estimated_tokens[name] = estimated_tokens.get(name, 0.0) + row.get("est_effective_tokens", 0) + estimated_documents[name] = estimated_documents.get(name, 0) + row["n_documents_kept"] + + section = load_app_config_dict(args.tokenizer_config)["tokenizer"]["config"] + tokenizer = PreTrainedHFTokenizer( + pretrained_model_name_or_path=section["pretrained_model_name_or_path"], + padding=section.get("padding", False), + truncation=section.get("truncation", False), + ) + + rng = random.Random(args.seed) + print(f"{'dataset':<18} {'sampled':>9} {'est tok/kB':>11} {'real tok/kB':>12} {'est tokens':>13} " + f"{'implied':>13} {'error':>8}") + print("-" * 84) + total_estimated = total_measured = 0.0 + for dataset in manifest["datasets"]: + name = dataset["name"] + shards = sorted((root / name).rglob("*.jsonl")) + if not shards: + continue + lines = sample_lines(shards, args.sample, rng) + if not lines: + continue + sampled_tokens = sampled_bytes = 0 + for line in lines: + try: + text = json.loads(line).get(args.text_field, "") + except ValueError: + continue + sampled_tokens += len(tokenizer.tokenize(text)) if isinstance(text, str) else 0 + sampled_bytes += len(line) + 1 # the newline the export writes + + if not sampled_bytes: + continue + # Both sides are byte-weighted, which is what makes them comparable: the estimate is + # a bytes-to-tokens ratio, and the sample is drawn in proportion to length. + measured_ratio = sampled_tokens / sampled_bytes + estimated_ratio = estimated_tokens[name] / dataset["n_bytes"] if dataset["n_bytes"] else 0.0 + implied = measured_ratio * dataset["n_bytes"] + error = (estimated_tokens[name] - implied) / implied if implied else 0.0 + total_estimated += estimated_tokens[name] + total_measured += implied + print( + f"{name:<18} {len(lines):>9,} {estimated_ratio * 1000:>11,.1f} {measured_ratio * 1000:>12,.1f} " + f"{estimated_tokens[name] / 1e9:>12,.1f}B {implied / 1e9:>12,.1f}B {error:>7.1%}" + ) + print("-" * 84) + overall = (total_estimated - total_measured) / total_measured if total_measured else 0.0 + print(f"{'BLEND':<18} {'':>9} {'':>11} {'':>12} {total_estimated / 1e9:>12,.1f}B " + f"{total_measured / 1e9:>12,.1f}B {overall:>7.1%}") + print() + print(" A few percent is expected. Tens of percent means the calibration is modelling") + print(" something other than what the export writes -- check the tokenizer and text field.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/slurm/env.sh b/config_files/data_preparation/quality/slurm/env.sh new file mode 100644 index 000000000..a7d35c3b5 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/env.sh @@ -0,0 +1,44 @@ +# Source this before running any stage by hand: +# +# source config_files/data_preparation/quality/slurm/env.sh +# +# Every variable can be overridden by setting it first; this only fills in what is unset. +# Sourcing rather than exporting by hand is the point: a stage invoked with an empty $QDIR +# silently becomes "/slurm/...", and an empty $WORK turns "rm -rf $WORK/buckets" into +# "rm -rf /buckets". Both have happened. + +export MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +export REPO="${REPO:-/home/richard.rutmann/repos/modalities}" +export QDIR="${QDIR:-$REPO/config_files/data_preparation/quality}" +export WORK="${WORK:-/data/user/richard.rutmann/annealing_blend}" +export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" +export NUM_BUCKETS="${NUM_BUCKETS:-1024}" + +if [[ -f "$HOME/.config/huggingface/token" ]]; then + export HF_TOKEN="$(tr -d '\r\n' < "$HOME/.config/huggingface/token")" +fi + +export EXPORTS="ALL,MQ=$MQ,QDIR=$QDIR,WORK=$WORK,NUM_BUCKETS=$NUM_BUCKETS,HF_HOME=$HF_HOME" + +mkdir -p "$WORK" "$HOME/logs/quality" + +# Fail loudly here rather than as a confusing error from sbatch or, worse, an rm against +# the wrong path. +for _v in MQ REPO QDIR WORK; do + if [[ -z "${!_v:-}" ]]; then + echo "quality env: $_v is empty -- refusing to continue" >&2 + return 1 2>/dev/null || exit 1 + fi +done +for _p in "$MQ" "$QDIR/annealing_registry.yaml" "$QDIR/slurm/2_bucket_annotations.sbatch"; do + if [[ ! -e "$_p" ]]; then + echo "quality env: expected to find $_p but it is missing" >&2 + return 1 2>/dev/null || exit 1 + fi +done +unset _v _p + +echo "quality env ready:" +echo " MQ $MQ" +echo " QDIR $QDIR" +echo " WORK $WORK" diff --git a/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py b/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py new file mode 100755 index 000000000..aa7456300 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/make_smoke_snapshot.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Freezes a tiny slice of the annealing corpora so the pipeline can be tested end to end. + +Why a snapshot rather than pointing at the corpora directly: the source tree is still +being transferred, and a corpus that is re-sharded halfway through a run invalidates every +byte offset recorded against it. A frozen copy makes the test repeatable and immune to +that. It also keeps the test honest about scale -- 1.5 GB runs in minutes, so a broken +stage is found in minutes. + +The five datasets are chosen to cover all four distinct join-key kinds plus the +native-metrics-only path, which is full code coverage of the join. HPLT is deliberately +absent: it uses the same ``field`` key kind as FineWiki and would add 327 GB of annotation +bucket reads to learn nothing new. + +Reads only. Nothing is written under the source root. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import sys +from pathlib import Path + +SOURCE_ROOT = Path("/data/annealing") + +# (name, relative source dir, glob, byte budget, pointer filter) +# `pointer_part` restricts KletterMix to documents whose source pointer lands in a single +# part of the 1.6 TB original corpus, so resolving them reads one 16 GB file instead of +# scattering across many. +DATASETS = [ + ("finewiki-de", "german/Finewiki", "*.jsonl", 256 << 20, None), + ("finepdfs-es", "spanish/Finepdfs", "*.jsonl", 256 << 20, None), + ("climbmix-en", "english/Climbmix", "*.jsonl", 256 << 20, None), + ("klettermix-de", "german/AIML-TUDA-KletterMix-filtered", "*.jsonl", 3 << 30, "part_65.detokenized.jsonl"), + ("dolmino", "english/Dolmino", "**/*.jsonl", 64 << 20, None), +] + + +def copy_prefix(src: Path, dst: Path, budget: int) -> tuple[int, int]: + """Copies a line-aligned prefix of a JSONL file. + + Args: + src (Path): Source JSONL file, opened read-only. + dst (Path): Destination path. Parents are created. + budget (int): Approximate byte budget. The last, partial line is dropped, since a + truncated JSON document would fail to parse and look like a pipeline bug. + + Returns: + tuple[int, int]: Documents and bytes written. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + n_docs = 0 + n_bytes = 0 + with src.open("rb") as fin, dst.open("wb") as fout: + for line in fin: + if not line.endswith(b"\n"): + break + fout.write(line) + n_bytes += len(line) + n_docs += 1 + if n_bytes >= budget: + break + return n_docs, n_bytes + + +def copy_filtered(src: Path, dst: Path, budget: int, pointer_part: str) -> tuple[int, int]: + """Copies documents whose ``id`` pointer names one particular source part. + + Args: + src (Path): Source JSONL file, opened read-only. + dst (Path): Destination path. Parents are created. + budget (int): How many source bytes to scan before stopping. + pointer_part (str): Keep only documents whose ``id`` starts with this. + + Returns: + tuple[int, int]: Documents and bytes written. + """ + dst.parent.mkdir(parents=True, exist_ok=True) + n_docs = 0 + n_bytes = 0 + scanned = 0 + with src.open("rb") as fin, dst.open("wb") as fout: + for line in fin: + scanned += len(line) + if not line.endswith(b"\n"): + break + try: + record = json.loads(line) + except ValueError: + continue + if str(record.get("id", "")).startswith(pointer_part): + fout.write(line) + n_bytes += len(line) + n_docs += 1 + if scanned >= budget: + break + return n_docs, n_bytes + + +def main() -> int: + """Builds the snapshot. + + Returns: + int: Process exit status. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="Snapshot root. Must not be under the source root.") + parser.add_argument("--force", action="store_true", help="Replace an existing snapshot.") + args = parser.parse_args() + + out = args.out.resolve() + if SOURCE_ROOT in out.parents or out == SOURCE_ROOT: + print(f"refusing to write inside the source tree {SOURCE_ROOT}", file=sys.stderr) + return 2 + if out.exists(): + if not args.force: + print(f"{out} exists; pass --force to replace it", file=sys.stderr) + return 2 + shutil.rmtree(out) + + total_docs = 0 + total_bytes = 0 + for name, rel, pattern, budget, pointer_part in DATASETS: + src_dir = SOURCE_ROOT / rel + if not src_dir.exists(): + print(f" {name:<14} SKIPPED, {src_dir} does not exist") + continue + sources = sorted(src_dir.glob(pattern)) + if not sources: + print(f" {name:<14} SKIPPED, no files match {pattern}") + continue + # Dolmino's files are small and its metrics live under a subdirectory, so take a + # handful of whole files rather than a prefix of one. + picks = sources[:4] if name == "dolmino" else sources[:1] + docs = written = 0 + for src in picks: + dst = out / rel / src.relative_to(src_dir) + if pointer_part: + d, b = copy_filtered(src, dst, budget, pointer_part) + else: + d, b = copy_prefix(src, dst, budget) + docs += d + written += b + print(f" {name:<14} {docs:>9,} docs {written / (1 << 20):>8.1f} MiB from {len(picks)} file(s)") + total_docs += docs + total_bytes += written + + print(f"\n snapshot at {out}: {total_docs:,} docs, {total_bytes / (1 << 30):.2f} GiB") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/slurm/reset_buckets.sh b/config_files/data_preparation/quality/slurm/reset_buckets.sh new file mode 100755 index 000000000..9f87a7099 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/reset_buckets.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Deletes the bucketed annotations so a split can be re-bucketed from scratch. A sharded +# bucketing run cannot clear the directory itself -- sibling tasks are writing into it -- +# so clearing it is a deliberate separate step. +# +# source .../env.sh && bash .../reset_buckets.sh +set -euo pipefail + +# ":?" is the point of this script existing: an unset WORK would otherwise make this +# "rm -rf /buckets". +BUCKETS="${WORK:?WORK is not set -- source env.sh first}/buckets" + +case "$BUCKETS" in + /buckets|/|"") echo "refusing to delete $BUCKETS" >&2; exit 1 ;; +esac + +if [[ ! -d "$BUCKETS" ]]; then + echo "nothing to remove: $BUCKETS does not exist" + exit 0 +fi + +echo "removing $BUCKETS" +echo " split dirs: $(ls "$BUCKETS" | wc -l), parquet: $(find "$BUCKETS" -name '*.parquet' | wc -l)" +rm -rf "$BUCKETS" +echo "done. Sidecars under $WORK/sidecar are untouched." diff --git a/config_files/data_preparation/quality/slurm/run_all_timed.sh b/config_files/data_preparation/quality/slurm/run_all_timed.sh new file mode 100755 index 000000000..408daa119 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/run_all_timed.sh @@ -0,0 +1,141 @@ +#!/usr/bin/env bash +# Runs the quality pipeline end to end against /data/annealing, printing how long each +# step took and a summary at the end. +# +# bash config_files/data_preparation/quality/slurm/run_all_timed.sh +# bash .../run_all_timed.sh 1 2 3 4 # only these steps +# +# Array jobs are submitted with `sbatch --wait`, so each timing is the job's real +# duration rather than how long submission took. That also makes the steps run in order, +# which steps 4 and 6 onwards require. + +set -uo pipefail + +# ----------------------------------------------------------------------- settings +export MQ="${MQ:-/data/user/richard.rutmann/venvs/modalities-quality/bin/python}" +export REPO="${REPO:-/home/richard.rutmann/repos/modalities}" +export QDIR="${QDIR:-$REPO/config_files/data_preparation/quality}" +export WORK="${WORK:-/data/user/richard.rutmann/annealing_blend}" +export HF_HOME="${HF_HOME:-/data/cache/hf_cache}" + +# Overridable so a blend variant can be run without editing the shared configs. +REGISTRY="${REGISTRY:-$QDIR/annealing_registry.yaml}" +SELECTION="${SELECTION:-$QDIR/annealing_selection.yaml}" +TOKENIZER_CONFIG="${TOKENIZER_CONFIG:-$QDIR/annealing_tokenizer.yaml}" + +SIDECAR_TASKS="${SIDECAR_TASKS:-64}" +BUCKET_TASKS="${BUCKET_TASKS:-64}" +NUM_BUCKETS="${NUM_BUCKETS:-1024}" +PACK_TASKS="${PACK_TASKS:-64}" +SAMPLE_SIZE="${SAMPLE_SIZE:-2000}" +BLEND_NAME="${BLEND_NAME:-blend_v1}" + +if [[ -f "$HOME/.config/huggingface/token" ]]; then + export HF_TOKEN="$(tr -d '\r\n' < "$HOME/.config/huggingface/token")" +fi +mkdir -p "$WORK" "$HOME/logs/quality" +cd "$REPO" + +WANTED=("$@") +want() { + [[ ${#WANTED[@]} -eq 0 ]] && return 0 + local s + for s in "${WANTED[@]}"; do [[ "$s" == "$1" ]] && return 0; done + return 1 +} + +# ------------------------------------------------------------------------ timing +declare -a SUMMARY=() +OVERALL_START=$(date +%s) + +fmt() { + local s=$1 + printf '%dh %02dm %02ds' $((s / 3600)) $(((s % 3600) / 60)) $((s % 60)) +} + +step() { + local id="$1" label="$2"; shift 2 + if ! want "$id"; then + SUMMARY+=("$(printf '%-3s %-34s %14s' "$id" "$label" 'skipped')") + return 0 + fi + printf '\n===== step %s: %s\n===== started %s\n' "$id" "$label" "$(date '+%F %T')" + local t0 rc dt + t0=$(date +%s) + "$@" + rc=$? + dt=$(( $(date +%s) - t0 )) + printf '===== step %s took %s (exit %d)\n' "$id" "$(fmt "$dt")" "$rc" + SUMMARY+=("$(printf '%-3s %-34s %14s exit %d' "$id" "$label" "$(fmt "$dt")" "$rc")") + if [[ $rc -ne 0 ]]; then + echo "!!!!! step $id failed; stopping. Fix it and re-run with: bash $0 $id ..." >&2 + exit "$rc" + fi +} + +print_summary() { + printf '\n%s\n' "==========================================================================" + printf '%-3s %-34s %14s\n' "id" "step" "elapsed" + printf '%s\n' "--------------------------------------------------------------------------" + local line + for line in "${SUMMARY[@]}"; do printf '%s\n' "$line"; done + printf '%s\n' "--------------------------------------------------------------------------" + printf '%-38s %14s\n' "TOTAL" "$(fmt $(( $(date +%s) - OVERALL_START )))" + printf '%s\n' "==========================================================================" +} +trap print_summary EXIT + +EXPORTS="ALL,MQ=$MQ,QDIR=$QDIR,WORK=$WORK,NUM_BUCKETS=$NUM_BUCKETS,HF_HOME=$HF_HOME,REGISTRY=$REGISTRY" + +# ============================================================ once per blend +step 1 "calibrate tokens" \ + "$MQ" -m modalities quality calibrate \ + --registry "$REGISTRY" --work_dir "$WORK" \ + --tokenizer_config "$TOKENIZER_CONFIG" --sample_size "$SAMPLE_SIZE" + +step 2 "build sidecar (array)" \ + sbatch --wait --export="$EXPORTS" --array="0-$((SIDECAR_TASKS - 1))" \ + "$QDIR/slurm/1_build_sidecar.sbatch" + +step 3 "bucket annotations (array)" \ + sbatch --wait --export="$EXPORTS" --array="0-$((BUCKET_TASKS - 1))" \ + "$QDIR/slurm/2_bucket_annotations.sbatch" + +step 4 "join annotations + build cubes" \ + sbatch --wait --export="$EXPORTS" "$QDIR/slurm/3_join_and_cube.sbatch" + +# ============================================================ per ablation +step 5 "preview selection" \ + "$MQ" -m modalities quality preview \ + --selection "$SELECTION" --work_dir "$WORK" + +step 6 "apply selection (filtered indexes)" \ + "$MQ" -m modalities quality apply \ + --selection "$SELECTION" \ + --registry "$REGISTRY" \ + --work_dir "$WORK" --output_dir "$WORK/$BLEND_NAME" + +export_jsonl() { + # One array task per dataset in the mix manifest. Each writes its own record; the + # blend-wide manifest is merged afterwards, because concurrent tasks writing one shared + # file would race and the last writer would erase the rest. + n=$("$MQ" -c "import yaml,sys; m=yaml.safe_load(open('$WORK/mix/mix_manifest.yaml')); \ + print(len({d.get('source_dataset') or d['name'] for d in m['datasets']}))") + if [[ "$n" -eq 0 ]]; then + echo "no datasets in $WORK/mix/mix_manifest.yaml" >&2 + return 1 + fi + echo "exporting $n dataset(s) as $n task(s)" + sbatch --wait --export="$EXPORTS,MANIFEST=$WORK/mix/mix_manifest.yaml,OUT=$WORK/out" \ + --array="0-$((n - 1))" "$QDIR/slurm/4_export_jsonl.sbatch" + "$MQ" -m modalities quality export-jsonl \ + --manifest "$WORK/mix/mix_manifest.yaml" \ + --registry "$REG" --output_dir "$WORK/out" --finalize_only +} + +step 7 "export sampled documents as jsonl (array)" export_jsonl + +echo +echo "Coverage per dataset: $WORK/join_report.json" +echo "Blend manifest: $WORK/$BLEND_NAME/mix_manifest.yaml" +echo "Take the per-dataset 'ratio' values into a weighted_combined dataset in the training config." diff --git a/config_files/data_preparation/quality/slurm/verify_jsonl.py b/config_files/data_preparation/quality/slurm/verify_jsonl.py new file mode 100755 index 000000000..3296138c4 --- /dev/null +++ b/config_files/data_preparation/quality/slurm/verify_jsonl.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +"""Checks the exported JSONL against the manifest, and against the corpus it came from. + +Three questions, in increasing order of how much they would cost to get wrong: + +1. Does every shard hold the number of lines its record claims? A short shard means a job + died mid-write and the resume logic accepted it, which would silently shrink the training + set by however much it lost. +2. Do the line counts realise the ratios that were asked for? Up- and downsampling is now in + the bytes, so a dataset at 3.0 that emitted 1x is not a reporting error, it is the wrong + training set. +3. Are the exported lines the documents they claim to be? Sampled shards are read back and + compared against the source file at the recorded offset. Nothing is re-serialised by the + export, so this is a byte comparison, and any mismatch means the corpus moved under us. + +Reads only; writes nothing anywhere. +""" + +from __future__ import annotations + +import argparse +import pickle +import random +import sys +from pathlib import Path + +import yaml + +from modalities.dataloader.preprocessing.quality.export import copies_for + + +def load(path: Path) -> dict: + """Reads a YAML file. + + Args: + path (Path): The file. + + Returns: + dict: Its contents. + """ + with Path(path).open() as f: + return yaml.safe_load(f) + + +def count_lines(path: Path) -> int: + """Counts newlines in a file without holding it in memory. + + Args: + path (Path): The file. + + Returns: + int: Number of lines. + """ + total = 0 + with path.open("rb") as f: + while chunk := f.read(1 << 24): + total += chunk.count(b"\n") + return total + + +def main() -> int: + """Verifies an export. + + Returns: + int: 0 if everything checks out, else 1. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True, help="export_manifest.yaml") + parser.add_argument("--mix_manifest", type=Path, default=None, help="mix_manifest.yaml, to check the ratios.") + parser.add_argument("--shards", type=int, default=5, help="Shards to replay against the corpus in full.") + parser.add_argument("--count_lines", action="store_true", help="Re-count every shard. Reads the whole export.") + parser.add_argument("--seed", type=int, default=0, help="Seed for choosing which shards to sample.") + parser.add_argument( + "--source_root", + type=Path, + default=None, + help="If given, assert that no pipeline artifact was written into the source tree.", + ) + args = parser.parse_args() + + manifest = load(args.manifest) + mix = load(args.mix_manifest) if args.mix_manifest else None + root = args.manifest.parent + problems: list[str] = [] + + print(f"export at {root}") + print(f" {manifest['n_lines']:,} lines, {manifest['n_documents']:,} documents, " + f"{manifest['n_bytes'] / 1e12:,.2f} TB") + if not manifest.get("repeat_factor_applied"): + problems.append("the manifest does not record that the repeat factors were applied") + print() + + print(f"{'dataset':<18} {'shards':>8} {'documents':>14} {'lines':>15} {'lines/doc':>10} {'factors':>22}") + print("-" * 92) + for dataset in manifest["datasets"]: + factors = dataset["factors_applied"] + ratio = dataset["n_lines"] / dataset["n_documents"] if dataset["n_documents"] else 0.0 + described = ",".join(f"{v:g}" for v in sorted(factors.values())) + print( + f"{dataset['name']:<18} {dataset['n_shards']:>8,} {dataset['n_documents']:>14,} " + f"{dataset['n_lines']:>15,} {ratio:>10.2f} {described:>22}" + ) + if dataset["ratio"] != 1.0: + problems.append(f"{dataset['name']} reports a training ratio of {dataset['ratio']}, not 1.0") + # With a single flat factor the realised lines-per-document must land on it. A curve + # has several factors, so only the range is checkable. + if len(factors) == 1: + expected = next(iter(factors.values())) + if dataset["n_documents"] and abs(ratio - expected) > max(0.05, 0.05 * expected): + problems.append( + f"{dataset['name']} realised {ratio:.3f} lines per document against {expected:g} requested" + ) + elif factors and dataset["n_documents"] and not (min(factors.values()) <= ratio <= max(factors.values())): + problems.append( + f"{dataset['name']} realised {ratio:.3f} lines per document, outside its curve's " + f"{min(factors.values()):g}-{max(factors.values()):g}" + ) + print() + + if args.count_lines: + print("re-counting every shard") + print("-" * 92) + for dataset in manifest["datasets"]: + counted = sum(count_lines(p) for p in sorted((root / dataset["name"]).rglob("*.jsonl"))) + status = "ok" if counted == dataset["n_lines"] else "MISMATCH" + print(f" {dataset['name']:<18} {counted:>15,} {status}") + if counted != dataset["n_lines"]: + problems.append( + f"{dataset['name']} holds {counted:,} lines, manifest says {dataset['n_lines']:,}" + ) + print() + + print(f"replaying {args.shards} shard(s) against the corpus, line by line") + print("-" * 92) + if mix is None: + print(" skipped: pass --mix_manifest to locate the source files and their indexes") + else: + rng = random.Random(args.seed) + contributions: dict[str, dict[str, list]] = {} + for row in mix["datasets"]: + name = row.get("source_dataset") or row["name"] + for source_path, index_path in row["index_files"].items(): + contributions.setdefault(name, {}).setdefault(source_path, []).append( + (index_path, float(row["ratio"])) + ) + seed = mix.get("seed", 42) + + candidates = [ + (name, source) for name, sources in contributions.items() for source in sources + ] + checked_lines = 0 + for name, source_path in rng.sample(candidates, min(args.shards, len(candidates))): + entries = [] + for index_path, factor in contributions[name][source_path]: + for offset, length in pickle.loads(Path(index_path).read_bytes()): + entries.append((offset, length, factor)) + entries.sort() + + shard = root / name / Path(source_path).name.replace(".jsonl", ".jsonl") + matches = sorted((root / name).rglob(Path(source_path).name)) + shard = matches[0] if matches else shard + if not shard.is_file(): + problems.append(f"{name}: no shard for {source_path}") + continue + + # Replays exactly what the export should have written -- same offsets, same copy + # counts -- and compares it to what is on disk. This validates content, ordering + # and repetition together, which counting lines cannot. + mismatch = None + with shard.open("rb") as out, Path(source_path).open("rb") as src: + for offset, length, factor in entries: + copies = copies_for(factor, seed, str(source_path), offset) + if copies == 0: + continue + src.seek(offset) + expected = src.read(length) + b"\n" + for _ in range(copies): + actual = out.readline() + checked_lines += 1 + if actual != expected: + mismatch = (offset, expected[:80], actual[:80]) + break + if mismatch: + break + trailing = out.readline() if not mismatch else b"" + if mismatch: + offset, expected, actual = mismatch + problems.append(f"{name} {shard.name}: line at offset {offset} differs") + print(f" MISMATCH {name}/{shard.name} at offset {offset}") + print(f" expected {expected!r}") + print(f" actual {actual!r}") + elif trailing: + problems.append(f"{name} {shard.name}: has more lines than the index accounts for") + else: + print(f" ok {name}/{shard.name}") + print(f" {checked_lines:,} lines replayed byte-for-byte") + print() + + if args.source_root is not None: + print(f"checking that nothing was written into {args.source_root}") + print("-" * 92) + # Only artifacts this pipeline could have produced. Flagging every non-.jsonl file + # instead reported the corpus's own 44 .gitattributes files, delivered with the data + # eleven days before the export -- a check that cries wolf gets ignored, and then it + # is worth nothing on the day something really is written here. + artifacts = (".idx", ".pbin", ".meta.json", ".partial", ".fingerprint") + stray = [ + str(p) + for p in args.source_root.rglob("*") + if p.is_file() and (p.suffix in artifacts or p.name.endswith(".jsonl.meta.json")) + ] + others = sum( + 1 for p in args.source_root.rglob("*") if p.is_file() and p.suffix != ".jsonl" + ) - len(stray) + if stray: + problems.append(f"{len(stray):,} pipeline artifact(s) written into the source root") + for path in stray[:10]: + print(f" STRAY {path}") + else: + print(" clean: no pipeline artifact under the source root") + if others: + print(f" ({others:,} other non-.jsonl file(s) present, e.g. .gitattributes -- " + f"delivered with the corpus, not written by us)") + print() + + if problems: + print("PROBLEMS") + for problem in problems: + print(f" {problem}") + return 1 + print("OK: the export matches its manifest and the ratios it was asked for.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/config_files/data_preparation/quality/smoke_registry.yaml b/config_files/data_preparation/quality/smoke_registry.yaml new file mode 100644 index 000000000..ecb68643f --- /dev/null +++ b/config_files/data_preparation/quality/smoke_registry.yaml @@ -0,0 +1,65 @@ +# Registry for the end-to-end smoke test. +# +# Points at the frozen snapshot built by `slurm/make_smoke_snapshot.py`, not at the live +# corpora. The corpora are still being transferred, and a re-shard mid-run invalidates +# every byte offset recorded against them -- which is exactly what happened to the first +# full run. The snapshot makes this test repeatable. +# +# The five datasets cover all four distinct join-key kinds plus the native-metrics-only +# path, which is every branch of the join: +# +# field finewiki-de the plain case +# urn_uuid_field finepdfs-es ids stored both wrapped and bare, mixed within a file +# sha256_text climbmix-en no identifier at all; keyed by a hash of its own text +# source_pointer klettermix-de translated corpus, keyed via the English original +# (none) dolmino native metrics only, no annotations +# +# HPLT is absent on purpose: same `field` kind as FineWiki, and its annotation buckets are +# 327 GB, so including it would cost an hour of reads to exercise no new code. + +annotation_root: /data/michael.fromm/hf-cache/datasets--openeurollm--propella-annotations/snapshots/e80fc1407801a15b956f40c642d3709b528abbc9/data/propella-1-4b +extra_annotation_roots: + - /data/alex.jude/.cache/huggingface/datasets--openeurollm--propella-annotations/snapshots/9e9c5083f81dc4bbd2708b65816a4dc41f59b911/data/propella-1-4b + +datasets: + - name: finewiki-de + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/german/Finewiki + glob: "*.jsonl" + annotation_split: finewiki + key: {kind: field, field: id} + + - name: finepdfs-es + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/spanish/Finepdfs + glob: "*.jsonl" + annotation_split: finepdfs/spa_Latn + key: {kind: urn_uuid_field, field: id} + native_metrics: + - {name: fw_edu, jq_pattern: .fw_edu_scores, aggregation: max} + - {name: full_doc_lid_score, jq_pattern: .full_doc_lid_score} + + - name: climbmix-en + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/english/Climbmix + glob: "*.jsonl" + annotation_split: nemotron-climbmix + key: {kind: sha256_text} + + # The snapshot keeps only documents pointing into `part_65`, so resolving these pointers + # reads one 16 GB file of the original corpus rather than scattering across many. The + # source corpus is the real one -- it is read, never written. + - name: klettermix-de + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/german/AIML-TUDA-KletterMix-filtered + glob: "*.jsonl" + annotation_split: nemotron-climbmix + key: + kind: source_pointer + field: id + source_root: /data/annealing/Nemotron-ClimbMix + source_line_offset: 0 + native_metrics: + - {name: proxy_score, jq_pattern: .proxy_score} + - {name: token_count, jq_pattern: .token_count} + + - name: dolmino + jsonl_root: /data/user/richard.rutmann/annealing_smoke_data/english/Dolmino + native_metrics: + - {name: len_cl100k_base, jq_pattern: .metadata.len_cl100k_base} diff --git a/config_files/data_preparation/quality/smoke_selection.yaml b/config_files/data_preparation/quality/smoke_selection.yaml new file mode 100644 index 000000000..a181a0c80 --- /dev/null +++ b/config_files/data_preparation/quality/smoke_selection.yaml @@ -0,0 +1,46 @@ +# Selection for the end-to-end smoke test. +# +# Chosen to exercise the machinery rather than to be a sensible blend: an ordinal +# predicate, a native-score threshold, a dataset with both, a dataset with neither, a +# downsample, an upsample, and a fractional ratio. The fractional one matters -- a repeat +# factor of 1.5 is what drives `WeightedCombinedDataset` through its partial-pass +# permutation, which no test on real data has ever reached. + +missing_annotation: keep + +# Small, so the report shows an achievable target rather than a 600 % overshoot. +target_tokens: 40_000_000 + +datasets: + # Ordinal predicate only, and a fractional upsample. + - name: finewiki-de + ratio: 1.5 + predicates: + - {field: educational_value, op: at_least, value: basic} + + # Ordinal plus native score, on the key kind that mixes wrapped and bare UUIDs. + - name: finepdfs-es + ratio: 1.0 + predicates: + - {field: fw_edu, op: gte, value: 1.5} + - {field: content_integrity, op: at_least, value: mostly_complete} + + # Two ordinals, on the hash-keyed dataset. + - name: climbmix-en + ratio: 0.5 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: information_density, op: at_least, value: moderate} + + # Ordinal plus the corpus's own proxy score, on the pointer-keyed dataset. + - name: klettermix-de + ratio: 2.0 + predicates: + - {field: educational_value, op: at_least, value: basic} + - {field: proxy_score, op: gte, value: 0.65} + + # No annotations at all: native metric only, unfiltered ratio. + - name: dolmino + ratio: 1.0 + predicates: + - {field: len_cl100k_base, op: gte, value: 200} diff --git a/config_files/data_preparation/quality/smoke_tokenizer.yaml b/config_files/data_preparation/quality/smoke_tokenizer.yaml new file mode 100644 index 000000000..2713d5731 --- /dev/null +++ b/config_files/data_preparation/quality/smoke_tokenizer.yaml @@ -0,0 +1,16 @@ +# The tokenizer the token estimates are measured with. +# +# Used by `modalities quality calibrate --tokenizer_config `, which measures +# how many tokens each dataset's bytes are worth. Nothing here tokenizes the corpus: the +# pipeline now exports JSONL and leaves tokenization to whatever consumes it. The estimate +# is what makes `preview` able to cost a selection in tokens, so it should still name the +# tokenizer the run actually trains with -- getting it wrong does not fail loudly, it +# produces a plausible token budget for the wrong tokenizer. + +tokenizer: + component_key: tokenizer + variant_key: pretrained_hf_tokenizer + config: + pretrained_model_name_or_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 + padding: false + truncation: false diff --git a/src/modalities/__main__.py b/src/modalities/__main__.py index bb29ce2fe..c25e813d2 100644 --- a/src/modalities/__main__.py +++ b/src/modalities/__main__.py @@ -28,6 +28,10 @@ from modalities.config.config import ProcessGroupBackendType, load_app_config_dict from modalities.config.instantiation_models import TrainingComponentsInstantiationModel from modalities.dataloader.create_instruction_tuning_data import create_instruction_tuning_data +from modalities.dataloader.preprocessing.quality import export as quality_export +from modalities.dataloader.preprocessing.quality import pipeline as quality_pipeline +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +from modalities.dataloader.preprocessing.quality.verify import format_verify_report from modalities.main import Main from modalities.models.huggingface_adapters.hf_adapter import HFModelAdapter from modalities.running_env.cuda_env import CudaEnv @@ -723,6 +727,569 @@ def CMD_entry_point_run_train_step_profiler( ) +@main.group(name="quality") +def quality() -> None: + """ + Quality-based document selection and up/downsampling of a training blend. + + The stages are meant to be run in order: `calibrate` measures how records map to + token counts, `build-sidecar` records one row per document, `join-annotations` + attaches external labels, and `build-cube` aggregates the result. After that, + `preview` costs a selection in seconds and `apply` writes filtered index files that + `modalities data pack_encoded_data` consumes unchanged. + """ + pass + + +@quality.command(name="calibrate") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option( + "--tokenizer_config", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to a packing config whose tokenizer section is used for calibration.", +) +@click.option( + "--sample_size", + type=int, + default=4000, + show_default=True, + help="Documents tokenized per dataset to measure the estimator. Sampled in proportion to " + "document length and grouped into size strata, so the rare very long documents that " + "dominate a corpus's byte count are represented.", +) +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +def CMD_quality_calibrate( + registry_path: Path, work_dir: Path, tokenizer_config: Path, sample_size: int, only: tuple[str, ...] +) -> None: + """Measures how each dataset's records relate to the training tokenizer's counts. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + tokenizer_config (Path): Packing config supplying the tokenizer. + sample_size (int): Documents tokenized per dataset. + only (tuple[str, ...]): Restrict to these dataset names. + """ + from modalities.config.component_factory import ComponentFactory + from modalities.config.instantiation_models import TokenizerInstantiationModel + from modalities.registry.components import COMPONENTS + from modalities.registry.registry import Registry + + config_dict = load_app_config_dict(tokenizer_config) + factory = ComponentFactory(registry=Registry(COMPONENTS)) + tokenizer = factory.build_components( + config_dict=config_dict, components_model_type=TokenizerInstantiationModel + ).tokenizer + tokenizer_name = str(config_dict["tokenizer"]["config"].get("pretrained_model_name_or_path", "unknown")) + + quality_pipeline.calibrate_blend( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + tokenizer=tokenizer, + tokenizer_name=tokenizer_name, + sample_size=sample_size, + only=list(only) or None, + ) + + +@quality.command(name="build-sidecar") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--index_root", + type=Path, + default=None, + help="Where JSONL index files live or should be created. Use this when the source tree is read-only.", +) +@click.option( + "--shard_id", + type=int, + default=0, + show_default=True, + help="This task's index. Set from SLURM_ARRAY_TASK_ID to run the build as an array.", +) +@click.option( + "--num_shards", + type=int, + default=1, + show_default=True, + help="How many tasks share the work. Files are divided across all selected datasets, " + "so one array covers the whole blend.", +) +def CMD_quality_build_sidecar( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + index_root: Optional[Path], + shard_id: int, + num_shards: int, +) -> None: + """Records one row per document: position, estimated tokens, key and native metrics. + + The only stage that reads the raw data, so the one worth running as an array. Each + task writes its own parquet parts, so tasks never contend. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to these dataset names. + index_root (Optional[Path]): Where JSONL index files live or should be created. + shard_id (int): This task's index in [0, num_shards). + num_shards (int): Total number of tasks sharing the work. + """ + written = quality_pipeline.build_sidecars( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + index_root=index_root, + shard_id=shard_id, + num_shards=num_shards, + ) + for name, n_documents in written.items(): + print_rank_0(f"{name}: {n_documents:,} documents") + + +@quality.command(name="bucket-annotations") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to the splits these datasets need (repeatable).") +@click.option( + "--num_buckets", + type=int, + default=1024, + show_default=True, + help="Partitions per annotation split. Each is loaded whole during the join, so raise this for large splits.", +) +@click.option( + "--shard_id", + type=int, + default=0, + show_default=True, + help="This task's index. Set from SLURM_ARRAY_TASK_ID to bucket as an array.", +) +@click.option("--num_shards", type=int, default=1, show_default=True, help="How many tasks bucket each split.") +@click.option("--force", is_flag=True, default=False, help="Re-bucket a split even if its output is complete.") +def CMD_quality_bucket_annotations( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + num_buckets: int, + shard_id: int, + num_shards: int, + force: bool, +) -> None: + """Partitions the annotation splits by a hash of their key, ready for joining. + + The expensive half of the join, since a split can run to billions of rows. Shardable, + and splits shared by several datasets are bucketed only once. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to the splits these datasets need. + num_buckets (int): Partitions per split. + shard_id (int): This task's index in [0, num_shards). + num_shards (int): How many tasks bucket each split. + force (bool): Re-bucket even if the output is complete. + """ + written = quality_pipeline.bucket_blend_annotations( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_buckets=num_buckets, + shard_id=shard_id, + num_shards=num_shards, + force=force, + ) + for split, n_rows in written.items(): + print_rank_0(f"{split}: {n_rows:,} rows bucketed by this task") + + +@quality.command(name="join-annotations") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--resume", + is_flag=True, + default=False, + help="Skip sidecar parts that already carry labels, to continue an interrupted run. " + "Omit it after re-bucketing the annotations, or the old labels are kept.", +) +def CMD_quality_join_annotations(registry_path: Path, work_dir: Path, only: tuple[str, ...], resume: bool) -> None: + """Attaches the bucketed annotations to each dataset's sidecar and reports coverage. + + Run `bucket-annotations` first. Read the reported coverage before trusting a + selection: on a partly downloaded split most documents may carry no label at all. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to these dataset names. + resume (bool): Skip parts that already carry labels. + """ + reports = quality_pipeline.join_blend_annotations( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + resume=resume, + ) + for report in reports: + print_rank_0(report.summary()) + + +@quality.command(name="verify-sidecar") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory holding the sidecars.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--num_parts", + type=int, + default=8, + show_default=True, + help="Sidecar parts to sample per dataset.", +) +@click.option( + "--num_rows_per_part", + type=int, + default=4, + show_default=True, + help="Documents to probe per sampled part. Rows at offset 0 are skipped, since the first " + "document of any JSONL file parses and so proves nothing.", +) +@click.option( + "--adopt", + is_flag=True, + default=False, + help="Write a source file manifest for verified sidecars that lack one, so later stages can " + "detect drift cheaply. Only sidecars that pass verification are stamped.", +) +def CMD_quality_verify_sidecar( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + num_parts: int, + num_rows_per_part: int, + adopt: bool, +) -> None: + """Checks that the sidecars' byte offsets still describe the current source files. + + Run this after any data transfer, and before apply on a blend whose source tree may + have been touched. A corpus that was re-sharded after its sidecar was built yields a + blend of wrong byte ranges, and this is the only check that reads the source bytes. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory holding the sidecars. + only (tuple[str, ...]): Restrict to these dataset names. + num_parts (int): Sidecar parts to sample per dataset. + num_rows_per_part (int): Documents to probe per sampled part. + adopt (bool): Stamp a manifest onto verified sidecars that lack one. + """ + reports = quality_pipeline.verify_sidecars( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_parts=num_parts, + n_rows_per_part=num_rows_per_part, + adopt=adopt, + ) + print_rank_0(format_verify_report(reports)) + if any(not r.ok for r in reports): + raise SystemExit(1) + + +@quality.command(name="build-cube") +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory for the blend's intermediates.") +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--num_score_bins", + type=int, + default=10, + show_default=True, + help="Quantile bins per native metric. A threshold on a bin edge stays exact.", +) +@click.option( + "--label_dimension", + "label_dimensions", + multiple=True, + help="Annotation column to group on, repeatable. Defaults to the seven ordinal fields; " + "name a field here if a selection thresholds on it, or the preview must scan the sidecar. " + "Each added field multiplies the cell count by its number of levels.", +) +def CMD_quality_build_cube( + registry_path: Path, + work_dir: Path, + only: tuple[str, ...], + num_score_bins: int, + label_dimensions: tuple[str, ...], +) -> None: + """Aggregates the sidecars so a selection can be costed without reading them again. + + Args: + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory for the blend's intermediates. + only (tuple[str, ...]): Restrict to these dataset names. + num_score_bins (int): Quantile bins per native metric. + label_dimensions (tuple[str, ...]): Annotation columns to group on. + """ + quality_pipeline.build_cubes( + registry=CorpusRegistry.from_yaml(registry_path), + work_dir=work_dir, + only=list(only) or None, + n_score_bins=num_score_bins, + label_dimensions=list(label_dimensions) or None, + ) + + +@quality.command(name="preview") +@click.option( + "--selection", + "selection_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the selection YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory holding the cubes and sidecars.") +@click.option( + "--exact", + is_flag=True, + default=False, + help="Scan the per-document sidecars instead of the cubes. Slower, but exact for any threshold.", +) +@click.option( + "--allow_fallback", + is_flag=True, + default=False, + help="Let a dataset whose cube cannot answer a predicate be scanned from its sidecar. " + "That reads every document, so it costs minutes to hours rather than seconds.", +) +@click.option( + "--explain", + is_flag=True, + default=False, + help="Attribute each dataset's retention to its individual predicates, and show how they " + "overlap. Answers which condition binds and which is redundant.", +) +def CMD_quality_preview( + selection_path: Path, work_dir: Path, exact: bool, allow_fallback: bool, explain: bool +) -> None: + """Reports how many documents and tokens a selection yields, per dataset and in total. + + Args: + selection_path (Path): Path to the selection YAML. + work_dir (Path): Working directory holding the cubes and sidecars. + exact (bool): Scan the sidecars instead of the cubes. + allow_fallback (bool): Permit per-dataset sidecar scans where a cube falls short. + explain (bool): Attribute retention to individual predicates. + """ + _, report = quality_pipeline.preview_selection( + selection_path=selection_path, + work_dir=work_dir, + force_exact=exact, + allow_sidecar_fallback=allow_fallback, + explain=explain, + ) + print_rank_0(report) + + +@quality.command(name="apply") +@click.option( + "--selection", + "selection_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the selection YAML.", +) +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--work_dir", type=Path, required=True, help="Working directory holding the sidecars.") +@click.option( + "--output_dir", type=Path, required=True, help="Directory receiving the filtered index files and the mix manifest." +) +@click.option( + "--allow_overexposure", + is_flag=True, + default=False, + help="Materialise even when the run would repeat data past a declared cap. Off by default: " + "ratios are per pass, so a run that wraps multiplies every one of them.", +) +def CMD_quality_apply( + selection_path: Path, + registry_path: Path, + work_dir: Path, + output_dir: Path, + allow_overexposure: bool, +) -> None: + """Writes a selection out as filtered index files plus a manifest. + + The source data is not copied or modified. Point `pack_encoded_data` at a written + index to tokenize only the selected documents. + + Args: + selection_path (Path): Path to the selection YAML. + registry_path (Path): Path to the corpus registry YAML. + work_dir (Path): Working directory holding the sidecars. + output_dir (Path): Directory receiving the index files and manifest. + """ + manifest_path = quality_pipeline.apply_selection( + allow_overexposure=allow_overexposure, + selection_path=selection_path, + registry_path=registry_path, + work_dir=work_dir, + output_dir=output_dir, + ) + print_rank_0(f"Manifest written to {manifest_path}") + + +@quality.command(name="export-jsonl") +@click.option( + "--manifest", + "manifest_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the mix_manifest.yaml written by 'apply'.", +) +@click.option( + "--registry", + "registry_path", + type=click_pathlib.Path(exists=True), + required=True, + help="Path to the corpus registry YAML.", +) +@click.option("--output_dir", type=Path, required=True, help="Directory receiving one subdirectory per dataset.") +@click.option( + "--seed", + type=int, + default=None, + help="Overrides the seed the selection was applied with, which decides which documents " + "receive the extra copy of a fractional repeat factor.", +) +@click.option("--only", multiple=True, help="Restrict to these dataset names (repeatable).") +@click.option( + "--resume/--no_resume", + default=True, + help="Leave shards that are already complete alone (default: resume).", +) +@click.option( + "--finalize/--no_finalize", + default=True, + help="Merge the per-dataset records into export_manifest.yaml. Pass --no_finalize in an " + "array task, then run 'export-jsonl --finalize_only' once the array has finished.", +) +@click.option( + "--finalize_only", + is_flag=True, + help="Write export_manifest.yaml from the per-dataset records already on disk, exporting nothing.", +) +@click.option("--shard_id", type=int, default=0, show_default=True, help="This task's index in [0, num_shards).") +@click.option( + "--num_shards", + type=int, + default=1, + show_default=True, + help="Split each dataset's source files across this many tasks. One task per dataset is fine " + "until one dataset holds far more files than the rest.", +) +def CMD_quality_export_jsonl( + manifest_path: Path, + registry_path: Path, + output_dir: Path, + seed: Optional[int], + only: tuple[str, ...], + resume: bool, + finalize: bool, + finalize_only: bool, + shard_id: int, + num_shards: int, +) -> None: + """Writes the selected documents out as JSONL, with the sampling baked into the bytes. + + Up- and downsampling is materialised here: a dataset at 3.0 has each of its documents + written three times, and one at 0.6 loses two of every five. The training set is the + concatenation of the resulting files, so their ratios must not be applied again. + + Args: + manifest_path (Path): Path to the mix manifest. + registry_path (Path): Path to the corpus registry YAML. + output_dir (Path): Directory receiving the exported JSONL. + seed (Optional[int]): Overrides the selection's seed. + only (tuple[str, ...]): Restrict to these dataset names. + resume (bool): Leave complete shards alone. + finalize (bool): Merge the per-dataset records afterwards. + finalize_only (bool): Only merge the records; export nothing. + shard_id (int): This task's index. + num_shards (int): Tasks splitting each dataset's files. + """ + if finalize_only: + print_rank_0(f"Export manifest written to {quality_export.finalize_export(output_dir)}") + return + + exports = quality_pipeline.export_jsonl( + manifest_path=manifest_path, + registry_path=registry_path, + output_dir=output_dir, + seed=seed, + only=list(only) or None, + resume=resume, + finalize=finalize, + shard_id=shard_id, + num_shards=num_shards, + ) + n_lines = sum(e.n_lines for e in exports) + n_bytes = sum(e.n_bytes for e in exports) + skipped = sum(1 for e in exports for s in e.shards if s.skipped) + print_rank_0( + f"Exported {len(exports)} dataset(s): {n_lines:,} lines, {n_bytes / 1e12:,.2f} TB" + + (f", {skipped:,} shard(s) already complete" if skipped else "") + ) + + def _format_exception_as_json(e: Exception, environment: dict[str, Any]) -> str: # Format an exception into a structured JSON string with error message, type, and stack trace. error = { diff --git a/src/modalities/config/config.py b/src/modalities/config/config.py index 45d8eead3..9571cd084 100644 --- a/src/modalities/config/config.py +++ b/src/modalities/config/config.py @@ -475,6 +475,12 @@ class CombinedDatasetConfig(BaseModel): datasets: list[PydanticDatasetIFType] +class WeightedCombinedDatasetConfig(BaseModel): + datasets: list[PydanticDatasetIFType] + repeat_factors: list[Annotated[float, Field(strict=False, ge=0)]] + seed: Annotated[int, Field(strict=True, ge=0)] = 42 + + class BatchSamplerConfig(BaseModel): sampler: PydanticSamplerIFType batch_size: Annotated[int, Field(strict=True, gt=0)] diff --git a/src/modalities/config/instantiation_models.py b/src/modalities/config/instantiation_models.py index fd7fd3b78..5e0172e50 100644 --- a/src/modalities/config/instantiation_models.py +++ b/src/modalities/config/instantiation_models.py @@ -223,6 +223,18 @@ class PackedDatasetSettings(BaseModel): settings: PackedDatasetSettings +class TokenizerInstantiationModel(BaseModel): + """Builds only the tokenizer of a config. + + Lets a tool reuse a packing config for its tokenizer without also having to satisfy + that config's settings, which name a specific source file the tool has no interest + in. Building from the packing config is what guarantees the same tokenizer is used + for measuring token estimates and for the packing those estimates predict. + """ + + tokenizer: PydanticTokenizerIFType + + class TextGenerationInstantiationModel(BaseModel): class TextGenerationSettings(BaseModel): model_path: FilePath diff --git a/src/modalities/dataloader/dataset.py b/src/modalities/dataloader/dataset.py index 0ef4d9076..794bdcc73 100644 --- a/src/modalities/dataloader/dataset.py +++ b/src/modalities/dataloader/dataset.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math from enum import Enum from pathlib import Path from typing import Optional @@ -462,3 +463,127 @@ def __getitem__(self, idx: int) -> dict: local_idx = idx - (self.cumulative_sizes[dataset_idx - 1] if dataset_idx > 0 else 0) return self.datasets[dataset_idx][local_idx] + + +class WeightedCombinedDataset(Dataset): + """Combines multiple datasets at runtime, each contributing a chosen number of epochs. + + `CombinedDataset` concatenates its datasets once each, so the only way to change a + dataset's share of the blend is to change how much of it is on disk. This class + takes a repeat factor per dataset instead: 2.5 draws a dataset two and a half times + per epoch, 0.3 draws three tenths of it. Nothing is duplicated on disk, fractional + factors work, and the blend becomes a config value. + + The fractional part is realised without storing an index map, so memory stays + constant no matter how large the datasets are. The documents making up a partial + pass are picked by a seeded affine permutation of the dataset's indices, which + spreads them evenly across the whole dataset rather than taking a prefix, and gives + the same selection on every rank and every restart. + + Note: + The partial-pass selection is evenly spread rather than statistically random. + That is what makes it O(1), and it is a good property here -- a prefix would + over-sample whatever the corpus happens to be ordered by -- but it is not a + substitute for shuffling, which the sampler still does. + """ + + def __init__(self, datasets: list[Dataset], repeat_factors: list[float], seed: int = 42): + """Initializes the WeightedCombinedDataset. + + Args: + datasets (list[Dataset]): The datasets to combine. + repeat_factors (list[float]): How many times to draw each dataset per + epoch. Must be non-negative and align one-to-one with `datasets`. + A factor of 0 excludes a dataset while keeping it declared. + seed (int): Seed for the partial-pass selection. + + Raises: + ValueError: If the lengths disagree or a factor is negative. + """ + if len(datasets) != len(repeat_factors): + raise ValueError( + f"got {len(datasets)} datasets but {len(repeat_factors)} repeat factors; they must correspond" + ) + if any(factor < 0 for factor in repeat_factors): + raise ValueError(f"repeat factors must be non-negative, got {repeat_factors}") + + self.datasets = datasets + self.repeat_factors = list(repeat_factors) + self.seed = seed + + self._full_passes: list[int] = [] + self._num_partial: list[int] = [] + self._permutation_params: list[tuple[int, int]] = [] + virtual_lengths: list[int] = [] + + for dataset_idx, (dataset, factor) in enumerate(zip(datasets, repeat_factors)): + num_samples = len(dataset) + full_passes = int(factor) + num_partial = int(round((factor - full_passes) * num_samples)) + # Rounding up to a whole extra pass is expressed as one more full pass, so + # `num_partial` never equals `num_samples` and the permutation stays a + # strict subset. + if num_partial >= num_samples > 0: + full_passes += 1 + num_partial = 0 + self._full_passes.append(full_passes) + self._num_partial.append(num_partial) + self._permutation_params.append(self._affine_permutation_params(num_samples, seed, dataset_idx)) + virtual_lengths.append(full_passes * num_samples + num_partial) + + self.cumulative_sizes = np.cumsum(virtual_lengths, dtype=np.int64) + + @staticmethod + def _affine_permutation_params(num_samples: int, seed: int, dataset_idx: int) -> tuple[int, int]: + # `index -> (multiplier * index + offset) % num_samples` is a bijection exactly + # when the multiplier is coprime with num_samples, which is what makes the + # partial pass a subset with no repeats. + if num_samples <= 1: + return 1, 0 + rng = np.random.default_rng([seed, dataset_idx]) + multiplier = 1 + for _ in range(1000): + candidate = int(rng.integers(1, num_samples)) + if math.gcd(candidate, num_samples) == 1: + multiplier = candidate + break + return multiplier, int(rng.integers(0, num_samples)) + + def __len__(self) -> int: + """Returns the number of samples one epoch of the blend yields. + + Returns: + int: Sum over datasets of `repeat_factor * len(dataset)`, rounded per + dataset. + """ + return int(self.cumulative_sizes[-1]) if len(self.cumulative_sizes) else 0 + + def __getitem__(self, idx: int) -> dict: + """Retrieves a sample from the blend. + + Args: + idx (int): Index into the blend. + + Returns: + dict: The sample from whichever dataset the index falls in. + + Raises: + IndexError: If `idx` is outside the blend. + """ + if idx < 0: + idx += len(self) + if not 0 <= idx < len(self): + raise IndexError(f"index {idx} is out of range for a blend of {len(self)} samples") + + dataset_idx = int(np.searchsorted(self.cumulative_sizes, idx, side="right")) + local_idx = idx - (self.cumulative_sizes[dataset_idx - 1] if dataset_idx > 0 else 0) + + num_samples = len(self.datasets[dataset_idx]) + num_in_full_passes = self._full_passes[dataset_idx] * num_samples + if local_idx < num_in_full_passes: + sample_idx = local_idx % num_samples + else: + multiplier, offset = self._permutation_params[dataset_idx] + sample_idx = (multiplier * (local_idx - num_in_full_passes) + offset) % num_samples + + return self.datasets[dataset_idx][int(sample_idx)] diff --git a/src/modalities/dataloader/dataset_factory.py b/src/modalities/dataloader/dataset_factory.py index 1eab9e328..00e6ef1f3 100644 --- a/src/modalities/dataloader/dataset_factory.py +++ b/src/modalities/dataloader/dataset_factory.py @@ -12,6 +12,7 @@ MemMapDataset, PackedMemMapDatasetContinuous, PackedMemMapDatasetMegatron, + WeightedCombinedDataset, ) @@ -126,3 +127,19 @@ def get_combined_dataset(datasets: list[Dataset]) -> Dataset: Dataset: CombinedDataset object. """ return CombinedDataset(datasets=datasets) + + @staticmethod + def get_weighted_combined_dataset(datasets: list[Dataset], repeat_factors: list[float], seed: int = 42) -> Dataset: + """Factory method for creating a combined dataset with per-dataset epoch counts. + + Args: + datasets (list[Dataset]): List of datasets to combine. + repeat_factors (list[float]): How many times each dataset is drawn per + epoch. Fractional values are supported, so a dataset can contribute a + partial pass without being duplicated on disk. + seed (int): Seed for selecting the documents of a partial pass. + + Returns: + Dataset: WeightedCombinedDataset object. + """ + return WeightedCombinedDataset(datasets=datasets, repeat_factors=repeat_factors, seed=seed) diff --git a/src/modalities/dataloader/preprocessing/quality/__init__.py b/src/modalities/dataloader/preprocessing/quality/__init__.py new file mode 100644 index 000000000..9243afc3f --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/__init__.py @@ -0,0 +1,14 @@ +"""Quality-based document selection and up/downsampling for pretraining blends. + +The package turns per-document quality signals into a training mix: + +1. ``registry`` declares where each dataset lives and how it joins to annotations. +2. ``sidecar`` streams the JSONL once and records per-document position, length, + estimated token count and native quality metrics. +3. ``propella_join`` attaches external propella annotations to those documents. +4. ``cube`` aggregates the result so any threshold combination can be costed + without touching the data again. +5. ``selection`` evaluates a YAML selection against the cube. +6. ``materialize`` writes a filtered ``.idx`` that ``pack_encoded_data`` consumes + unchanged, so only selected documents are ever tokenized. +""" diff --git a/src/modalities/dataloader/preprocessing/quality/annotation_join.py b/src/modalities/dataloader/preprocessing/quality/annotation_join.py new file mode 100644 index 000000000..65de258d5 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/annotation_join.py @@ -0,0 +1,680 @@ +"""Attaches external per-document annotations to a dataset's sidecar. + +The two sides of this join are both far too large to hold in memory -- one annotation +split alone runs to billions of rows -- so neither can be turned into a hash table. Both +sides are instead partitioned by a hash of the join key into buckets small enough to +join one at a time. Documents whose key appears in no annotation shard keep null +labels; that is the normal outcome for a split that has only been partly downloaded, +and the join reports how often it happens rather than hiding it. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.dataset as ds +import pyarrow.parquet as pq +from tqdm import tqdm + +from modalities.utils.logger_utils import get_logger + +# Annotation columns worth carrying into the sidecar. The annotation corpora also hold +# free-text and list-valued columns; those are read on demand rather than copied onto +# every document, because they cannot be aggregated into the cube anyway. +DEFAULT_LABEL_COLUMNS: tuple[str, ...] = ( + "content_integrity", + "content_quality", + "information_density", + "reasoning_indicators", + "educational_value", + "content_safety", + "pii_presence", + "audience_level", + "commercial_bias", + "time_sensitivity", + "content_ratio", + "content_length", +) + +KEY_COLUMN = "id" + + +# Files per dataset scan. This bounds how much parquet fragment metadata is held at once, +# which is what actually drives the join's peak memory -- see the note in ``flush``. +SCAN_FILE_GROUP = 2048 + + +class AnnotationJoinError(RuntimeError): + """Raised when a join cannot be carried out as specified.""" + + +@dataclass +class JoinReport: + """What the join did, in numbers worth acting on. + + Attributes: + dataset (str): Dataset that was joined. + split (str): Annotation split used. + n_documents (int): Documents in the sidecar. + n_matched (int): Documents that received labels. + n_annotation_rows (int): Annotation rows read. + n_duplicate_keys (int): Annotation keys seen more than once. Duplicates are + real in at least one published split, so they are counted rather than + assumed away. + n_missing_key (int): Documents whose sidecar row had no join key at all. + n_parts_resumed (int): Parts that already carried labels and were not re-joined. + Their documents still count towards the totals above, read back from the + labels already on disk: coverage describes the sidecar, not the run, and a + resumed join that reported 0% would read as a failed one. + label_columns (list[str]): Columns actually copied across. + """ + + dataset: str + split: str + n_documents: int = 0 + n_matched: int = 0 + n_annotation_rows: int = 0 + n_duplicate_keys: int = 0 + n_missing_key: int = 0 + n_parts_resumed: int = 0 + label_columns: list[str] = field(default_factory=list) + + @property + def coverage(self) -> float: + """Share of documents that received labels. + + Returns: + float: Matched documents over total documents; 0.0 for an empty sidecar. + """ + return self.n_matched / self.n_documents if self.n_documents else 0.0 + + def to_dict(self) -> dict: + """Renders the report as a plain dictionary. + + Returns: + dict: Report fields plus the derived coverage. + """ + return { + "dataset": self.dataset, + "split": self.split, + "n_documents": self.n_documents, + "n_matched": self.n_matched, + "coverage": round(self.coverage, 6), + "n_annotation_rows": self.n_annotation_rows, + "n_duplicate_keys": self.n_duplicate_keys, + "n_missing_key": self.n_missing_key, + "n_parts_resumed": self.n_parts_resumed, + "label_columns": self.label_columns, + } + + def summary(self) -> str: + """One-line human-readable summary. + + Returns: + str: Dataset, coverage and the counts a reader should notice. + """ + return ( + f"{self.dataset}: {self.n_matched:,}/{self.n_documents:,} documents annotated " + f"({self.coverage:.1%}) from {self.n_annotation_rows:,} annotation rows" + + (f", {self.n_duplicate_keys:,} duplicate keys" if self.n_duplicate_keys else "") + + (f", {self.n_missing_key:,} without a key" if self.n_missing_key else "") + + (f", {self.n_parts_resumed:,} parts already joined" if self.n_parts_resumed else "") + ) + + +def _metadata_paths(annotation_bucket_dir: Path) -> list[Path]: + # Excludes the `.tmp` files an interrupted write can leave behind, so a half-written + # file is never mistaken for a task's metadata. + return sorted(p for p in Path(annotation_bucket_dir).glob("_meta.*.json") if p.suffix == ".json") + + +def _read_metadata(path: Path) -> Optional[dict]: + """Reads one bucketing metadata file, returning None if it cannot be used. + + Args: + path (Path): Path to a ``_meta..json`` file. + + Returns: + Optional[dict]: The parsed metadata, or None if the file is absent or unparseable. + + Note: + Metadata is written atomically, so an unparseable file should not occur. It is + tolerated anyway because the alternative is a fifteen-hour pipeline dying on one + unreadable sidecar file -- which is exactly what happened when this was a bare + ``json.loads``: tasks read a sibling's file mid-write and 12 of 64 crashed. + Skipping is also the safe direction for :func:`read_bucket_metadata`, where a + missing entry makes the run look incomplete rather than joinable. + """ + try: + return json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return None + + +def _write_metadata(path: Path, payload: dict) -> None: + """Writes bucketing metadata so a concurrent reader never sees a partial file. + + Args: + path (Path): Final path of the metadata file. + payload (dict): Content to write. + + Note: + ``Path.write_text`` truncates before writing, leaving a window in which the file + is empty. Sibling tasks of the same array read these files, so that window is a + real race. Writing to a per-task temporary name and renaming makes the swap + atomic: a reader sees either the old file or the complete new one. + """ + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload)) + os.replace(tmp, path) + + +def bucket_of(key: str, n_buckets: int) -> int: + """Assigns a join key to a bucket. + + Args: + key (str): The join key. + n_buckets (int): Number of buckets. + + Returns: + int: Bucket index in ``[0, n_buckets)``. + + Note: + Uses blake2b rather than the built-in ``hash``, whose string seed varies per + process. Both sides of the join are bucketed in separate runs, so an unstable + hash would silently send matching keys to different buckets. + """ + digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, "little") % n_buckets + + +class _BucketWriter: + """Streams rows out to one parquet file per bucket, with bounded memory. + + The shard suffix in the filename lets many tasks bucket one split at once: each + writes its own file per bucket, and the join reads every file belonging to a bucket. + + Memory is capped on the **total** rows held across all buckets, not per bucket. A + per-bucket threshold does not bound anything: with 1024 buckets a 50 M-row input + puts only ~49 k rows in each, so no bucket ever reaches a 100 k threshold and the + whole input ends up buffered as Python dicts. That is how this OOM-killed all 64 + tasks of a real run at 24 GB each. + """ + + def __init__( + self, + out_dir: Path, + schema: pa.Schema, + n_buckets: int, + shard_id: int = 0, + max_buffered_rows: int = 500_000, + ): + """ + Args: + out_dir (Path): Directory receiving the bucket files. + schema (pa.Schema): Schema of the rows being written. + n_buckets (int): Number of buckets. + shard_id (int): This task's index, used to name its files. + max_buffered_rows (int): Rows held across all buckets before everything is + flushed. Bounds memory regardless of the bucket count. + """ + self._out_dir = Path(out_dir) + self._out_dir.mkdir(parents=True, exist_ok=True) + self._schema = schema + self._n_buckets = n_buckets + self._shard_id = shard_id + self._max_buffered_rows = max_buffered_rows + self._writers: dict[int, pq.ParquetWriter] = {} + self._buffers: dict[int, list[dict]] = {} + self._buffered_rows = 0 + + def add(self, bucket: int, row: dict) -> None: + """Queues one row for its bucket, flushing everything if memory is up. + + Args: + bucket (int): Bucket the row belongs to. + row (dict): The row, matching the writer's schema. + """ + self._buffers.setdefault(bucket, []).append(row) + self._buffered_rows += 1 + if self._buffered_rows >= self._max_buffered_rows: + self.flush_all() + + def flush_all(self) -> None: + """Writes every buffered row out and releases the memory.""" + for bucket in list(self._buffers): + self._flush(bucket) + self._buffered_rows = 0 + + def _flush(self, bucket: int) -> None: + buffer = self._buffers.get(bucket) + if not buffer: + return + if bucket not in self._writers: + path = self._out_dir / f"bucket-{bucket:04d}.{self._shard_id:04d}.parquet" + self._writers[bucket] = pq.ParquetWriter(path, self._schema, compression="zstd") + self._writers[bucket].write_table(pa.Table.from_pylist(buffer, schema=self._schema)) + self._buffers[bucket] = [] + + def close(self) -> None: + """Flushes what is left and closes every open parquet writer.""" + self.flush_all() + for writer in self._writers.values(): + writer.close() + self._writers.clear() + + +def bucket_annotations( + shard_paths: list[Path], + out_dir: Path, + n_buckets: int, + label_columns: Optional[list[str]] = None, + key_column: str = KEY_COLUMN, + normalize_key: Optional[str] = None, + shard_id: int = 0, + num_shards: int = 1, + max_buffered_rows: int = 500_000, + show_progress: bool = True, +) -> tuple[int, list[str]]: + """Partitions annotation shards by a hash of their key. + + This is the expensive half of the join, since a split can run to billions of rows, + so it is shardable: run it as an array of ``num_shards`` tasks, each taking a subset + of the input shards. The result is identical to a single-task run. + + Args: + shard_paths (list[Path]): Annotation parquet shards of one split. + out_dir (Path): Directory receiving the bucket files. Cleared first, so a + partial previous run cannot contribute stale rows. + n_buckets (int): Number of buckets. Higher means smaller working set per join + step; a split of billions of rows wants at least 1024. + label_columns (Optional[list[str]]): Columns to carry. Defaults to + ``DEFAULT_LABEL_COLUMNS``, intersected with what the shards actually have. + key_column (str): Column holding the annotation key. + normalize_key (Optional[str]): Set to ``"urn_uuid"`` to strip + ```` wrappers, which occur mixed with bare UUIDs on both + sides of some joins. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): How many tasks are bucketing this split. + max_buffered_rows (int): Rows held in memory across all buckets before being + flushed. Raise it for throughput, lower it under a tight memory limit. + show_progress (bool): Whether to show a progress bar. + + Returns: + tuple[int, list[str]]: Rows written by this task, and the label columns carried. + + Raises: + AnnotationJoinError: If no shards are given, the key column is absent, or the + shard selection is out of range. + """ + if not shard_paths: + raise AnnotationJoinError("no annotation shards to bucket") + if not 0 <= shard_id < num_shards: + raise AnnotationJoinError(f"shard_id {shard_id} is not in [0, {num_shards})") + + available = set(pq.ParquetFile(shard_paths[0]).schema_arrow.names) + if key_column not in available: + raise AnnotationJoinError(f"annotation shards have no {key_column!r} column; found {sorted(available)}") + wanted = list(label_columns) if label_columns is not None else list(DEFAULT_LABEL_COLUMNS) + carried = [c for c in wanted if c in available] + if not carried: + raise AnnotationJoinError( + f"none of the requested label columns exist in the shards; available: {sorted(available)}" + ) + + out_dir = Path(out_dir) + # Only a single-task run may clear the directory; sibling tasks are writing into it. + if num_shards == 1 and out_dir.exists(): + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # A sharded run cannot clear the directory, so leftovers from a previous run with a + # different array size would be silently mixed in and the join would read a bucket + # made of rows from two incompatible runs. + for stale in _metadata_paths(out_dir): + meta = _read_metadata(stale) + if meta is None: + continue + previous = meta.get("num_shards", 1) + if previous != num_shards: + raise AnnotationJoinError( + f"{out_dir} holds output from a run with num_shards={previous}, but this task has " + f"num_shards={num_shards}. Delete the directory and re-bucket the split from scratch." + ) + # Strided rather than contiguous, so tasks stay balanced when shard sizes trend + # across a split. + my_shards = [path for i, path in enumerate(sorted(shard_paths)) if i % num_shards == shard_id] + schema = pa.schema([pa.field("key", pa.large_string())] + [pa.field(c, pa.large_string()) for c in carried]) + writer = _BucketWriter(out_dir, schema, n_buckets, shard_id=shard_id, max_buffered_rows=max_buffered_rows) + + from modalities.dataloader.preprocessing.quality.registry import strip_urn_uuid + + n_rows = 0 + try: + for shard in tqdm(my_shards, desc="bucketing annotations", disable=not show_progress): + parquet_file = pq.ParquetFile(shard) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx, columns=[key_column] + carried) + keys = table.column(key_column).to_pylist() + columns = {c: table.column(c).to_pylist() for c in carried} + for i, key in enumerate(keys): + if key is None: + continue + key = str(key) + if normalize_key == "urn_uuid": + key = strip_urn_uuid(key) + row = {"key": key} + for c in carried: + value = columns[c][i] + row[c] = None if value is None else str(value) + writer.add(bucket_of(key, n_buckets), row) + n_rows += 1 + finally: + writer.close() + + # One metadata file per task, so concurrent tasks never overwrite each other's, and + # written atomically so a sibling reading the directory cannot catch it half-written. + _write_metadata( + out_dir / f"_meta.{shard_id:04d}.json", + { + "n_buckets": n_buckets, + "label_columns": carried, + "n_rows": n_rows, + "shard_id": shard_id, + "num_shards": num_shards, + "n_input_shards": len(my_shards), + }, + ) + return n_rows, carried + + +def read_bucket_metadata(annotation_bucket_dir: Path) -> dict: + """Merges the metadata a bucketing run left behind, checking it is complete. + + Args: + annotation_bucket_dir (Path): Directory written by :func:`bucket_annotations`. + + Returns: + dict: ``n_buckets``, ``label_columns`` and the total ``n_rows`` bucketed. + + Raises: + AnnotationJoinError: If no metadata is present, the tasks disagree on the bucket + count or label columns, or some announced task never finished. Joining an + incomplete run would silently drop the annotations that task was carrying, + which looks exactly like a corpus that was never annotated. + """ + metadata_paths = _metadata_paths(annotation_bucket_dir) + if not metadata_paths: + raise AnnotationJoinError( + f"{annotation_bucket_dir} holds no bucketing metadata; run 'modalities quality bucket-annotations' first" + ) + merged: Optional[dict] = None + total_rows = 0 + seen_shards: set[int] = set() + for path in metadata_paths: + meta = _read_metadata(path) + if meta is None: + # Leaving the shard id unseen makes the run report as incomplete, which is the + # safe outcome: better to refuse the join than to join missing annotations. + continue + if merged is None: + merged = meta + elif (meta["n_buckets"], meta["label_columns"]) != (merged["n_buckets"], merged["label_columns"]): + raise AnnotationJoinError( + f"{annotation_bucket_dir} mixes incompatible bucketing runs: " + f"{meta['n_buckets']} buckets / {meta['label_columns']} vs " + f"{merged['n_buckets']} / {merged['label_columns']}. Re-bucket the split from scratch." + ) + total_rows += meta.get("n_rows", 0) + seen_shards.add(meta.get("shard_id", 0)) + + if merged is None: + raise AnnotationJoinError( + f"{annotation_bucket_dir} has {len(metadata_paths)} metadata file(s) but none could be read; " + "the bucketing run did not complete. Delete the directory and re-bucket the split." + ) + + expected = merged.get("num_shards", 1) + if len(seen_shards) != expected: + missing = sorted(set(range(expected)) - seen_shards) + raise AnnotationJoinError( + f"{annotation_bucket_dir} is incomplete: {len(seen_shards)} of {expected} bucketing task(s) " + f"finished, missing shard id(s) {missing}. Joining now would lose their annotations." + ) + return {"n_buckets": merged["n_buckets"], "label_columns": merged["label_columns"], "n_rows": total_rows} + + +def _part_has_labels(part: Path, label_columns: list[str]) -> bool: + """Whether a sidecar part has already been written back by a join. + + Args: + part (Path): The sidecar part. + label_columns (list[str]): Columns the join adds. + + Returns: + bool: True if every label column is present. Reads only the parquet footer, so + this is cheap enough to check for every part of a large dataset. + """ + try: + names = set(pq.ParquetFile(part).schema_arrow.names) + except OSError: + return False + return all(column in names for column in label_columns) + + +def _iter_sidecar_parts(sidecar_dir: Path) -> list[Path]: + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise AnnotationJoinError(f"no sidecar parts found in {sidecar_dir}") + return parts + + +def join_annotations( + sidecar_dir: Path, + annotation_bucket_dir: Path, + dataset_name: str, + split_name: str, + duplicate_policy: str = "first", + max_batch_keys: int = 20_000_000, + resume: bool = False, + show_progress: bool = True, +) -> JoinReport: + """Copies annotation labels onto a dataset's sidecar, in place. + + Sidecar parts are processed in batches, and each annotation bucket is read once per + batch rather than once per part. That distinction decides whether this finishes: + reading the bucketed split per part meant 454 TB of reads over the real blend -- + 5,319 parts for Nemotron-CC against a 23 GB split -- because every part's documents + hash across every bucket. Batching turns the read amplification from the part count + into the far smaller number of batches. + + Within a batch each bucket is filtered to the keys the batch actually wants before + anything is materialised in Python, so memory is bounded by the batch rather than by + the bucket -- a bucket of a billion-row split holds millions of rows, of which a batch + typically wants a few thousand. + + Args: + sidecar_dir (Path): Directory of sidecar parts to enrich. + annotation_bucket_dir (Path): Output of :func:`bucket_annotations`. + dataset_name (str): Dataset name, for the report. + split_name (str): Annotation split name, for the report. + duplicate_policy (str): What to do when one key carries several annotation rows. + ``"first"`` keeps the first row seen; ``"error"`` refuses to join. + max_batch_keys (int): Documents to hold per batch. Larger batches read the + annotation side fewer times but hold more keys in memory. + resume (bool): Skip parts that already carry the label columns, to continue an + interrupted run. The write-back adds columns and values together, so a part + having them means it was processed. Off by default: re-bucketing the + annotations and then resuming would silently keep the old labels, so + continuing has to be asked for. + show_progress (bool): Whether to show progress bars. + + Returns: + JoinReport: Coverage and the counts needed to judge whether a selection built on + these labels is meaningful. + + Raises: + AnnotationJoinError: If the bucket directory is unusable, or duplicates are found + under ``duplicate_policy="error"``. + """ + annotation_bucket_dir = Path(annotation_bucket_dir) + # Read for its own sake as well as for the label columns: it refuses an incomplete + # bucketing run rather than letting the join silently drop a missing task's annotations. + meta = read_bucket_metadata(annotation_bucket_dir) + label_columns: list[str] = meta["label_columns"] + + parts = _iter_sidecar_parts(sidecar_dir) + report = JoinReport(dataset=dataset_name, split=split_name, label_columns=label_columns) + report.n_annotation_rows = meta.get("n_rows", 0) + + # Globbed once for the whole split rather than per bucket. Routing keys to buckets is + # gone: a batch holds millions of keys, which hash across every bucket, so every bucket + # file was read anyway -- the profile confirmed 1024 of 1024 -- and the routing cost a + # blake2b call per key in Python to decide nothing. + all_bucket_files = sorted(annotation_bucket_dir.glob("bucket-*.parquet")) + if not all_bucket_files: + raise AnnotationJoinError( + f"no bucket files in {annotation_bucket_dir}; run 'quality bucket-annotations' first" + ) + + def flush(batch: list[tuple[Path, pa.Table]]) -> None: + """Resolves one batch of parts and writes their label columns back. + + Resolution stays in Arrow from end to end. The obvious implementation -- + materialise the keys, build a dict from key to row, look every document up, emit a + list per label column -- costs a handful of Python objects per document, and at + 1.7 bn documents for Nemotron-CC that was twelve hours. Here each part's labels come + from one ``index_in`` against the batch's lookup table followed by one ``take`` per + label column, so the per-document work happens in Arrow's kernels. + + The single remaining Python loop is over the batch's *unique* keys, to route them to + buckets. That one cannot be vectorised: both sides of the join are bucketed in + separate runs, so the hash has to be stable across processes, and blake2b is not + something Arrow's compute layer can do. It is per unique key rather than per + document, and it is cheap relative to reading the buckets. + """ + if not batch: + return + + chunks: list[pa.Array] = [] + for _, table in batch: + for chunk in table.column("join_key").chunks: + if len(chunk) > 0: + chunks.append(chunk.cast(pa.large_string())) + wanted_keys = ( + pc.drop_null(pc.unique(pa.chunked_array(chunks, type=pa.large_string()))) + if chunks + else pa.array([], type=pa.large_string()) + ) + + # Scanned in groups of files rather than as one dataset over the whole split. + # + # A dataset holds a fragment per file, with that file's parquet metadata, and the + # scan machinery scales with the fragment count rather than with the data. The + # measurements say so plainly: nemotron-cc, 16,384 files and 504 M documents, peaked + # at 81.6 GB; HPLT, 65,536 files but only 10.7 M documents, exceeded 160 GB and was + # OOM-killed. Four times the files, one fiftieth of the data, and it is the one that + # dies. Streaming the batches was not enough, because the fragments themselves are + # what costs. + pieces: list[pa.Table] = [] + if len(wanted_keys) > 0: + batches: list[pa.RecordBatch] = [] + for start in range(0, len(all_bucket_files), SCAN_FILE_GROUP): + group = all_bucket_files[start : start + SCAN_FILE_GROUP] + scanner = ds.dataset(group, format="parquet").scanner( + columns=["key"] + label_columns, + filter=ds.field("key").isin(wanted_keys), + batch_size=131_072, + batch_readahead=4, + fragment_readahead=2, + ) + batches.extend(batch for batch in scanner.to_batches() if batch.num_rows) + if batches: + pieces.append(pa.Table.from_batches(batches)) + + lookup_keys: Optional[pa.Array] = None + lookup: Optional[pa.Table] = None + if pieces: + lookup = pa.concat_tables(pieces, promote_options="permissive") + keys_column = lookup.column("key").cast(pa.large_string()).combine_chunks() + unique_keys = pc.unique(keys_column) + n_duplicates = len(keys_column) - len(unique_keys) + if n_duplicates: + report.n_duplicate_keys += n_duplicates + if duplicate_policy == "error": + counts = pc.value_counts(keys_column) + repeated = counts.field("values").filter(pc.greater(counts.field("counts"), 1)) + example = repeated[0].as_py() if len(repeated) else "" + raise AnnotationJoinError( + f"annotation key {example!r} appears more than once in split " + f"{split_name!r}; choose duplicate_policy='first' to keep the first" + ) + # index_in reports the first position of each value, which is exactly the + # "keep the first row seen" policy, done in one pass instead of a loop. + lookup = lookup.take(pc.index_in(unique_keys, value_set=keys_column)) + lookup_keys = unique_keys + else: + lookup_keys = keys_column + + for part, table in batch: + keys = table.column("join_key").cast(pa.large_string()) + if lookup is None or lookup_keys is None or len(lookup_keys) == 0: + indices = None + else: + indices = pc.index_in(keys, value_set=lookup_keys) + report.n_matched += len(indices) - indices.null_count + + for column in label_columns: + if indices is None: + array = pa.nulls(table.num_rows, type=pa.large_string()) + else: + # A null index yields a null label, so unmatched documents fall out + # correctly without being special-cased. + array = pc.take(lookup.column(column), indices).cast(pa.large_string()) + existing = table.schema.get_field_index(column) + if existing >= 0: + table = table.set_column(existing, pa.field(column, pa.large_string()), array) + else: + table = table.append_column(pa.field(column, pa.large_string()), array) + pq.write_table(table, part, compression="zstd") + + batch: list[tuple[Path, pa.Table]] = [] + batch_keys = 0 + n_skipped = 0 + for part in tqdm(parts, desc=f"join {dataset_name}", disable=not show_progress): + if resume and label_columns and _part_has_labels(part, label_columns): + # Count what this part already holds instead of ignoring it, so a resumed run + # reports the sidecar's real coverage rather than only what it happened to do. + n_skipped += 1 + existing = pq.read_table(part, columns=["join_key", label_columns[0]]) + report.n_documents += existing.num_rows + report.n_matched += existing.num_rows - existing.column(label_columns[0]).null_count + report.n_missing_key += existing.column("join_key").null_count + continue + table = pq.read_table(part) + # Kept as Arrow: materialising this column was 1.7 bn Python strings for the + # largest dataset, before any joining had happened at all. + report.n_documents += table.num_rows + report.n_missing_key += table.column("join_key").null_count + batch.append((part, table)) + batch_keys += table.num_rows + if batch_keys >= max_batch_keys: + flush(batch) + batch, batch_keys = [], 0 + flush(batch) + + report.n_parts_resumed = n_skipped + if n_skipped: + get_logger(name="main").info( + f"{dataset_name}: resumed, skipped {n_skipped:,} of {len(parts):,} parts that already carried labels" + ) + get_logger(name="main").info(report.summary()) + return report diff --git a/src/modalities/dataloader/preprocessing/quality/cube.py b/src/modalities/dataloader/preprocessing/quality/cube.py new file mode 100644 index 000000000..cd9bd9f78 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/cube.py @@ -0,0 +1,365 @@ +"""Aggregates a sidecar so any threshold combination can be costed instantly. + +Tuning a blend means trying many threshold combinations and asking what each one costs +in tokens. Answering that from the per-document table would mean re-reading billions of +rows for every edit. Answering it from a cube does not: documents are grouped once by +the fields a selection may threshold on, and each group records how many documents and +tokens it holds. Any conjunction of predicates over those fields is then a sum over the +groups that satisfy it -- microseconds, and exact rather than sampled. + +The cube is exact for the dimensions it was built over. Predicates over anything else +cannot be answered from it, and the selection engine says so and falls back to the +sidecar rather than quietly approximating. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Optional + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + +# The propella label columns worth grouping on. Ordinal scales first, then the two +# binary-ish safety fields; together these are what selections actually threshold. +DEFAULT_LABEL_DIMENSIONS: tuple[str, ...] = ( + "educational_value", + "content_quality", + "information_density", + "reasoning_indicators", + "content_integrity", + "content_safety", + "pii_presence", +) + +# Marker for documents that carry no annotation. Kept as an explicit dimension value +# rather than dropped, because on a partly downloaded split the unannotated documents +# can be the majority and the policy for them changes the answer completely. +MISSING = "__missing__" + +N_SCORE_BINS = 10 + + +class CubeError(RuntimeError): + """Raised when a cube cannot be built or used as requested.""" + + +@dataclass(frozen=True) +class ScoreBinning: + """Bin edges turning a continuous native metric into a cube dimension. + + Attributes: + column (str): Sidecar column the edges were computed from. + edges (tuple[float, ...]): Ascending bin edges. Bin ``i`` covers + ``[edges[i], edges[i + 1])``, with the last bin closed at the top. + """ + + column: str + edges: tuple[float, ...] + + def bin_index(self, values: np.ndarray) -> np.ndarray: + """Assigns values to bins. + + Args: + values (np.ndarray): Metric values, possibly containing NaN. + + Returns: + np.ndarray: Bin index per value; ``-1`` where the value is missing. + """ + edges = np.asarray(self.edges, dtype=np.float64) + index = np.searchsorted(edges[1:-1], values, side="right").astype(np.int64) + return np.where(np.isnan(values), -1, index) + + def lower_bound_of(self, bin_index: int) -> float: + """Smallest value that can fall in a bin. + + Args: + bin_index (int): The bin. + + Returns: + float: The bin's lower edge. + """ + return self.edges[bin_index] + + def upper_bound_of(self, bin_index: int) -> float: + """Smallest value above a bin. + + Args: + bin_index (int): The bin. + + Returns: + float: The bin's upper edge. + """ + return self.edges[bin_index + 1] + + +@dataclass +class Cube: + """Grouped document and token counts for one dataset. + + Attributes: + dataset (str): Dataset the cube describes. + label_dimensions (list[str]): Annotation columns grouped on. + score_binnings (dict[str, ScoreBinning]): Native metrics grouped on, by name. + table (pa.Table): One row per non-empty group: the dimension values, plus + ``n_documents`` and ``n_tokens``. + n_documents (int): Documents represented. + n_tokens (int): Estimated tokens represented. + """ + + dataset: str + label_dimensions: list[str] + score_binnings: dict[str, ScoreBinning] + table: pa.Table + n_documents: int + n_tokens: int + + @property + def dimensions(self) -> list[str]: + """All groupable dimension names. + + Returns: + list[str]: Label dimensions followed by binned score dimensions. + """ + return list(self.label_dimensions) + [f"native_{name}" for name in self.score_binnings] + + def write(self, path: Path) -> None: + """Writes the cube to a parquet file with its metadata embedded. + + Args: + path (Path): Destination path. Parents are created. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + meta = { + b"quality_cube": json.dumps( + { + "dataset": self.dataset, + "label_dimensions": self.label_dimensions, + "score_binnings": {k: list(v.edges) for k, v in self.score_binnings.items()}, + "n_documents": self.n_documents, + "n_tokens": self.n_tokens, + } + ).encode() + } + table = self.table.replace_schema_metadata({**(self.table.schema.metadata or {}), **meta}) + pq.write_table(table, path, compression="zstd") + + @classmethod + def read(cls, path: Path) -> "Cube": + """Reads a cube written by :meth:`write`. + + Args: + path (Path): The cube parquet file. + + Returns: + Cube: The loaded cube. + + Raises: + CubeError: If the file carries no cube metadata. + """ + table = pq.read_table(path) + raw = (table.schema.metadata or {}).get(b"quality_cube") + if raw is None: + raise CubeError(f"{path} is not a quality cube (no metadata)") + meta = json.loads(raw) + return cls( + dataset=meta["dataset"], + label_dimensions=meta["label_dimensions"], + score_binnings={ + k: ScoreBinning(column=f"native_{k}", edges=tuple(v)) for k, v in meta["score_binnings"].items() + }, + table=table, + n_documents=meta["n_documents"], + n_tokens=meta["n_tokens"], + ) + + +def _quantile_edges(values: np.ndarray, n_bins: int) -> Optional[tuple[float, ...]]: + # Quantile edges keep every bin populated, which matters because native scores are + # heavily skewed: fixed-width bins would leave most of the range nearly empty and + # crowd almost all documents into one or two cells. + finite = values[np.isfinite(values)] + if finite.size == 0: + return None + quantiles = np.linspace(0.0, 1.0, n_bins + 1) + edges = np.unique(np.quantile(finite, quantiles)) + if edges.size < 2: + # A metric with a single distinct value cannot be binned, but must still be + # thresholdable, so it gets one bin wide enough to hold it. + return (float(edges[0]), float(np.nextafter(edges[0], np.inf))) + edges[0] = -np.inf + edges[-1] = np.inf + return tuple(float(e) for e in edges) + + +def _aggregate_cells(table: pa.Table, dimension_names: list[str]) -> pa.Table: + # Sums the two count columns per distinct combination of dimension values. Arrow + # suffixes aggregated columns, and does not promise where it puts them relative to + # the group keys, so the result is reassembled by name. + aggregated = table.group_by(dimension_names).aggregate([("n_documents", "sum"), ("n_tokens", "sum")]) + return pa.table( + { + **{name: aggregated.column(name) for name in dimension_names}, + "n_documents": aggregated.column("n_documents_sum"), + "n_tokens": aggregated.column("n_tokens_sum"), + } + ) + + +def _sidecar_parts(sidecar_dir: Path) -> list[Path]: + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise CubeError(f"no sidecar parts found in {sidecar_dir}") + return parts + + +def _sample_scores(parts: list[Path], column: str, max_rows: int) -> np.ndarray: + collected: list[np.ndarray] = [] + total = 0 + for part in parts: + parquet_file = pq.ParquetFile(part) + if column not in parquet_file.schema_arrow.names: + continue + for group_idx in range(parquet_file.metadata.num_row_groups): + chunk = ( + parquet_file.read_row_group(group_idx, columns=[column]).column(column).to_numpy(zero_copy_only=False) + ) + chunk = chunk.astype(np.float64) + collected.append(chunk) + total += chunk.size + if total >= max_rows: + return np.concatenate(collected)[:max_rows] + return np.concatenate(collected) if collected else np.empty(0, dtype=np.float64) + + +def build_cube( + sidecar_dir: Path, + dataset_name: str, + label_dimensions: Iterable[str] = DEFAULT_LABEL_DIMENSIONS, + score_columns: Optional[Iterable[str]] = None, + n_score_bins: int = N_SCORE_BINS, + binning_sample_rows: int = 2_000_000, + aggregate_batch_rows: int = 8_000_000, +) -> Cube: + """Groups a sidecar into a cube. + + Args: + sidecar_dir (Path): Directory of sidecar parts, after the annotation join. + dataset_name (str): Dataset name recorded in the cube. + label_dimensions (Iterable[str]): Annotation columns to group on. Columns + absent from the sidecar are skipped, so an unannotated dataset still gets + a usable cube over its native metrics. + score_columns (Optional[Iterable[str]]): Native metric columns to bin and group + on, named without the ``native_`` prefix. None uses every native column + present. + n_score_bins (int): Quantile bins per native metric. + binning_sample_rows (int): Rows sampled to compute bin edges. + aggregate_batch_rows (int): How many rows to accumulate before running a + grouping pass. Larger batches group more efficiently but hold more rows in + memory; 8 million costs roughly a gigabyte for a typical dimension set. + + Returns: + Cube: The aggregated cube. + + Raises: + CubeError: If the sidecar directory holds no parts. + """ + parts = _sidecar_parts(sidecar_dir) + # Every part's schema, not just the first one's. Assuming they agree turned an + # interrupted join into an opaque `KeyError: Field "educational_value" does not exist` + # from inside a dict comprehension, 10 unlabelled parts out of 5,319. + schemas = [set(pq.ParquetFile(p).schema_arrow.names) for p in parts] + available = set.intersection(*schemas) + widest = set.union(*schemas) + + inconsistent = sorted(widest - available) + if inconsistent: + n_missing = sum(1 for names in schemas if not widest <= names) + raise CubeError( + f"dataset {dataset_name!r}: {n_missing} of {len(parts)} sidecar parts are missing " + f"column(s) {inconsistent}. The join did not finish; re-run " + f"'modalities quality join-annotations --only {dataset_name} --resume' before building a cube." + ) + + used_labels = [c for c in label_dimensions if c in available] + if score_columns is None: + used_scores = sorted(c[len("native_") :] for c in available if c.startswith("native_")) + else: + used_scores = [c for c in score_columns if f"native_{c}" in available] + + binnings: dict[str, ScoreBinning] = {} + for name in used_scores: + edges = _quantile_edges(_sample_scores(parts, f"native_{name}", binning_sample_rows), n_score_bins) + if edges is not None: + binnings[name] = ScoreBinning(column=f"native_{name}", edges=edges) + + columns = used_labels + [f"native_{n}" for n in binnings] + ["est_tokens"] + dimension_names = used_labels + [f"native_{n}" for n in binnings] + schema = pa.schema( + [pa.field(c, pa.large_string()) for c in used_labels] + + [pa.field(f"native_{n}", pa.int16()) for n in binnings] + + [pa.field("n_documents", pa.int64()), pa.field("n_tokens", pa.int64())] + ) + + # Grouping stays inside Arrow's C++ kernels rather than a Python loop over documents, + # which is what makes this feasible over a whole blend. + # + # Row groups are batched before being aggregated. Aggregating each one alone barely + # compresses -- at these cardinalities a row group of a million rows yields nearly a + # million cells -- so the work would be done twice for no gain. Batching lets the + # cardinality saturate first, which is the whole reason a cube is small. + aggregates: list[pa.Table] = [] + pending: list[pa.Table] = [] + pending_rows = 0 + n_documents = 0 + n_tokens = 0 + + def flush() -> None: + nonlocal pending, pending_rows + if pending: + aggregates.append(_aggregate_cells(pa.concat_tables(pending), dimension_names)) + pending, pending_rows = [], 0 + + for part in parts: + parquet_file = pq.ParquetFile(part) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx, columns=columns) + if table.num_rows == 0: + continue + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False).astype(np.int64) + + grouped: dict[str, Any] = {c: pc.fill_null(table.column(c), MISSING) for c in used_labels} + for name, binning in binnings.items(): + bins = binning.bin_index( + table.column(f"native_{name}").to_numpy(zero_copy_only=False).astype(np.float64) + ) + grouped[f"native_{name}"] = pa.array(bins, type=pa.int16()) + grouped["n_documents"] = pa.array(np.ones(table.num_rows, dtype=np.int64)) + grouped["n_tokens"] = pa.array(tokens) + + pending.append(pa.table(grouped)) + pending_rows += table.num_rows + n_documents += table.num_rows + n_tokens += int(tokens.sum()) + if pending_rows >= aggregate_batch_rows: + flush() + flush() + + if aggregates: + table = _aggregate_cells(pa.concat_tables(aggregates), dimension_names).cast(schema) + else: + table = schema.empty_table() + + return Cube( + dataset=dataset_name, + label_dimensions=used_labels, + score_binnings=binnings, + table=table, + n_documents=n_documents, + n_tokens=n_tokens, + ) diff --git a/src/modalities/dataloader/preprocessing/quality/export.py b/src/modalities/dataloader/preprocessing/quality/export.py new file mode 100644 index 000000000..243e2dd0d --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/export.py @@ -0,0 +1,492 @@ +"""Writes a selection out as JSONL, with the up/downsampling baked into the bytes. + +The packing stage this replaces left the ratios as metadata: it tokenized the selected +documents once, and ``WeightedCombinedDataset`` applied the repeat factors at training time, +fractional ones included. The output here is meant to be *concatenated* into a training set, +and concatenation carries no weights, so the factors have to become bytes. A dataset at 3.0 +has each of its documents written three times; one at 0.6 loses two of every five. + +Fractional factors are resolved per document rather than by truncating a list, so a factor of +1.2 means every document once and a hash-chosen fifth of them twice. The choice is a function +of the document's position and the blend's seed, so it is identical on every run and on every +machine, and it does not depend on the order files happen to be processed in. + +Output lines are copied verbatim. Nothing is parsed or re-serialised, so each line is +byte-identical to the source line it came from -- which also makes the export a byte-for-byte +auditable operation rather than a transformation. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import mmap +import os +import pickle +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import yaml +from tqdm import tqdm + +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry +from modalities.utils.logger_utils import get_logger + +# The name of the manifest this stage writes, alongside the shards it describes. +EXPORT_MANIFEST = "export_manifest.yaml" + +# Resolution of the fractional part of a repeat factor. A factor of 1.2 is realised as "one +# copy always, plus a second for the 200,000 millionths of documents whose hash falls below +# the threshold", so this bounds how precisely a factor can be expressed. +FRACTION_MODULUS = 1_000_000 + + +class ExportError(RuntimeError): + """Raised when a selection cannot be exported as JSONL.""" + + +def fraction_hash(seed: int, source_path: str, byte_offset: int) -> int: + """Places one document on the axis that decides its fractional copy. + + Args: + seed (int): The blend's seed. + source_path (str): The source file the document lives in. + byte_offset (int): Its offset within that file. + + Returns: + int: A value in ``[0, FRACTION_MODULUS)``. + + Note: + blake2b rather than the built-in ``hash``, whose string seed varies per process -- + the same reasoning as :func:`~...annotation_join.bucket_of`. Position is used as the + identity because this stage reads indexes rather than sidecars and has no join key + in hand; it is stable because :mod:`file_manifest` pins every source file's size and + mtime and ``verify-sidecar`` refuses one that has drifted. + """ + digest = hashlib.blake2b(f"{seed}\x00{source_path}\x00{byte_offset}".encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, "little") % FRACTION_MODULUS + + +def copies_for(factor: float, seed: int, source_path: str, byte_offset: int) -> int: + """How many times one document is written. + + Args: + factor (float): The dataset's repeat factor. + seed (int): The blend's seed. + source_path (str): The source file the document lives in. + byte_offset (int): Its offset within that file. + + Returns: + int: Whole copies, at least zero. + """ + whole = math.floor(factor) + remainder = factor - whole + if remainder > 0 and fraction_hash(seed, source_path, byte_offset) < round(remainder * FRACTION_MODULUS): + return whole + 1 + return whole + + +@dataclass +class ShardReport: + """What one output shard holds. + + Attributes: + output_path (Path): The written shard. + source_path (Path): The JSONL file it was drawn from. + n_documents (int): Distinct documents drawn. + n_lines (int): Lines written, counting repeats. + n_bytes (int): Bytes written. + skipped (bool): Whether a complete shard was left alone. + """ + + output_path: Path + source_path: Path + n_documents: int = 0 + n_lines: int = 0 + n_bytes: int = 0 + skipped: bool = False + + +@dataclass +class DatasetExport: + """What one dataset contributed. + + Attributes: + name (str): The dataset. + factors (dict[str, float]): Repeat factor per manifest row that fed it. A curve + contributes several rows, one per quality bucket, each with its own factor. + shards (list[ShardReport]): The shards written. + """ + + name: str + factors: dict[str, float] = field(default_factory=dict) + shards: list[ShardReport] = field(default_factory=list) + + @property + def n_lines(self) -> int: + """Lines across every shard. + + Returns: + int: Total lines written for this dataset. + """ + return sum(s.n_lines for s in self.shards) + + @property + def n_documents(self) -> int: + """Distinct documents across every shard. + + Returns: + int: Total documents drawn for this dataset. + """ + return sum(s.n_documents for s in self.shards) + + @property + def n_bytes(self) -> int: + """Bytes across every shard. + + Returns: + int: Total bytes written for this dataset. + """ + return sum(s.n_bytes for s in self.shards) + + +def _meta_path(output_path: Path) -> Path: + """Where a shard's completion record lives. + + Args: + output_path (Path): The shard. + + Returns: + Path: ``.meta.json``. + """ + return output_path.with_name(output_path.name + ".meta.json") + + +def _completed(output_path: Path) -> Optional[dict]: + """The record of a finished shard, if it finished. + + A shard is complete only when its record exists *and* the file on disk is the size the + record claims. Existence alone is not health: a killed job leaves a plausible-looking + partial file, and taking that as done is exactly how a truncated ``.pbin`` once survived + into a blend. + + Args: + output_path (Path): The shard. + + Returns: + Optional[dict]: The record, or None if the shard is missing or unfinished. + """ + try: + record = json.loads(_meta_path(output_path).read_text()) + return record if output_path.stat().st_size == record["n_bytes"] else None + except (OSError, ValueError, KeyError): + return None + + +def export_file( + source_path: Path, + contributions: list[tuple[Path, float]], + output_path: Path, + seed: int, + resume: bool = True, +) -> ShardReport: + """Writes one source file's selected documents to one JSONL shard. + + Args: + source_path (Path): The source JSONL file. + contributions (list[tuple[Path, float]]): Index file and repeat factor, one pair per + manifest row drawing from this source file. A quality curve produces several. + output_path (Path): The shard to write. + seed (int): The blend's seed. + resume (bool): Leave a complete shard alone. + + Returns: + ShardReport: What was written, or what was already there. + + Raises: + ExportError: If an index names a document the source file cannot supply. + """ + if resume: + record = _completed(output_path) + if record is not None: + return ShardReport( + output_path=output_path, + source_path=source_path, + n_documents=record["n_documents"], + n_lines=record["n_lines"], + n_bytes=record["n_bytes"], + skipped=True, + ) + + # Every contributing index is merged and re-sorted by offset, so the shard is written in + # source order in a single forward pass. Writing bucket by bucket instead would group the + # output by quality level, which is a strong ordering artifact to hand to a trainer. + entries: list[tuple[int, int, float]] = [] + for index_path, factor in contributions: + for offset, length in pickle.loads(Path(index_path).read_bytes()): + entries.append((offset, length, factor)) + entries.sort() + + output_path.parent.mkdir(parents=True, exist_ok=True) + scratch = output_path.with_name(output_path.name + ".partial") + n_documents = n_lines = n_bytes = 0 + + # Mapped directly rather than through LargeFileLinesReader: that reader addresses a file + # by line number and insists on an index describing the whole file, whereas this stage + # already holds the selection's byte offsets and needs nothing but the bytes at them. + with source_path.open("rb") as raw: + data = mmap.mmap(raw.fileno(), 0, access=mmap.ACCESS_READ) + try: + with scratch.open("wb") as out: + for offset, length, factor in entries: + copies = copies_for(factor, seed, str(source_path), offset) + n_documents += 1 + if copies == 0: + continue + line = data[offset : offset + length] + if len(line) != length: + raise ExportError( + f"{source_path}: index names {length} bytes at offset {offset} but the " + f"file supplied {len(line)}. The source has changed since the sidecar " + f"was built; re-run 'modalities quality verify-sidecar'." + ) + # The index records the line without its terminator -- checked against the + # real corpus, where entries end on '}' with a one-byte gap to the next + # offset -- so the newline is added here. + record = line + b"\n" + for _ in range(copies): + out.write(record) + n_lines += copies + n_bytes += copies * len(record) + except BaseException: + scratch.unlink(missing_ok=True) + raise + finally: + data.close() + + os.replace(scratch, output_path) + # Written only after the shard is in place, so an interrupted export leaves no record and + # is redone rather than trusted. + _meta_path(output_path).write_text( + json.dumps({"n_documents": n_documents, "n_lines": n_lines, "n_bytes": n_bytes}) + ) + return ShardReport( + output_path=output_path, + source_path=source_path, + n_documents=n_documents, + n_lines=n_lines, + n_bytes=n_bytes, + ) + + +# The per-dataset record each array task writes. The blend-wide manifest is merged from these +# afterwards rather than written by every task: eighteen tasks rewriting one shared file would +# race, and the last writer would erase the rest. +DATASET_RECORD = "_export.yaml" + + +def _contributions_by_dataset(manifest: dict, registry: CorpusRegistry) -> dict[str, dict[Path, list]]: + """Groups the manifest's rows into work, keyed by dataset and source file. + + A quality curve emits one row per bucket, named ``__`` and each carrying + its own factor. They describe one input dataset and must land in one output directory, + so they are merged here, and a source file drawn from by several buckets yields one shard + fed by all of them. + + Args: + manifest (dict): The parsed mix manifest. + registry (CorpusRegistry): Resolves a dataset to its source root. + + Returns: + dict[str, dict[Path, list]]: Dataset name to source path to + ``[(index_path, factor), ...]``. + """ + grouped: dict[str, dict[Path, list]] = {} + for row in manifest["datasets"]: + name = row.get("source_dataset") or row["name"] + registry.get(name) # fails loudly here rather than midway through a terabyte + for source_path, index_path in row["index_files"].items(): + grouped.setdefault(name, {}).setdefault(Path(source_path), []).append( + (Path(index_path), float(row["ratio"])) + ) + return grouped + + +def export_blend( + manifest_path: Path, + registry_path: Path, + output_root: Path, + seed: Optional[int] = None, + only: Optional[list[str]] = None, + resume: bool = True, + show_progress: bool = True, + shard_id: int = 0, + num_shards: int = 1, +) -> list[DatasetExport]: + """Writes every dataset of a materialised selection out as JSONL. + + Args: + manifest_path (Path): The ``mix_manifest.yaml`` written by the apply stage. + registry_path (Path): The corpus registry YAML. + output_root (Path): Directory receiving one subdirectory per dataset. + seed (Optional[int]): Decides which documents get a fractional extra copy. Defaults + to the seed the selection was applied with, which the mix manifest records. + only (Optional[list[str]]): Restrict to these dataset names. + resume (bool): Leave complete shards alone. + show_progress (bool): Whether to show progress bars. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Total tasks splitting each dataset's source files between them. + A dataset is one array task by default, which is fine until one of them holds + far more files than the rest: `dolmino` has 40,003 against a median of about + 500, and at a second per file that is twelve hours against twenty minutes. + + Returns: + list[DatasetExport]: What each dataset contributed. + + Raises: + ExportError: If the manifest names an index that is not on disk. + ValueError: If ``shard_id`` is out of range. + """ + if not 0 <= shard_id < num_shards: + raise ValueError(f"shard_id {shard_id} is not in [0, {num_shards})") + with Path(manifest_path).open() as f: + manifest = yaml.safe_load(f) + registry = CorpusRegistry.from_yaml(registry_path) + output_root = Path(output_root) + seed = manifest.get("seed", 42) if seed is None else seed + + grouped = _contributions_by_dataset(manifest, registry) + factors = {} + for row in manifest["datasets"]: + factors.setdefault(row.get("source_dataset") or row["name"], {})[row["name"]] = float(row["ratio"]) + + exports: list[DatasetExport] = [] + for name in sorted(grouped): + if only and name not in only: + continue + entry = registry.get(name) + export = DatasetExport(name=name, factors=factors[name]) + # Strided rather than contiguous: a dataset's files are ordered by shard name and + # sizes run in streaks, so contiguous slices hand one task all the large files. + sources = sorted(grouped[name])[shard_id::num_shards] + if not sources: + continue + for source_path in tqdm(sources, desc=f"export {name}", disable=not show_progress): + relative = source_path.relative_to(entry.jsonl_root).with_suffix(".jsonl") + export.shards.append( + export_file( + source_path=source_path, + contributions=grouped[name][source_path], + output_path=output_root / name / relative, + seed=seed, + resume=resume, + ) + ) + _write_dataset_record(output_root, export, seed, shard_id, num_shards) + exports.append(export) + get_logger(name="main").info( + f"{name}: {export.n_lines:,} lines from {export.n_documents:,} documents " + f"({export.n_bytes / 1e9:,.1f} GB) over {len(export.shards):,} shards" + ) + return exports + + +def _write_dataset_record( + output_root: Path, export: DatasetExport, seed: int, shard_id: int = 0, num_shards: int = 1 +) -> Path: + """Records what one task produced for one dataset. + + Args: + output_root (Path): The export root. + export (DatasetExport): The dataset's result. + seed (int): The seed used. + shard_id (int): This task's index. + num_shards (int): Total tasks splitting the dataset. + + Returns: + Path: The written record. + """ + directory = output_root / export.name + directory.mkdir(parents=True, exist_ok=True) + # One record per task when a dataset is split, merged by finalize_export. Writing to a + # shared name instead would have concurrent tasks overwrite each other's counts, which + # is the same race the per-dataset split already avoids at the blend level. + path = directory / (DATASET_RECORD if num_shards == 1 else f"_export.{shard_id:05d}.yaml") + scratch = path.with_suffix(f".yaml.{os.getpid()}.tmp") + scratch.write_text( + yaml.safe_dump( + { + "name": export.name, + "seed": seed, + "n_documents": export.n_documents, + "n_lines": export.n_lines, + "n_bytes": export.n_bytes, + "n_shards": len(export.shards), + "factors_applied": export.factors, + # Baked into the bytes, so a training config must not weight this again. + "repeat_factor_applied": True, + "ratio": 1.0, + }, + sort_keys=False, + ) + ) + os.replace(scratch, path) + return path + + +def finalize_export(output_root: Path) -> Path: + """Merges the per-dataset records into one manifest for the whole export. + + Args: + output_root (Path): The export root. + + Returns: + Path: The written ``export_manifest.yaml``. + + Raises: + ExportError: If no dataset records are present. + """ + output_root = Path(output_root) + records = sorted(output_root.glob("*/_export*.yaml")) + if not records: + raise ExportError(f"no dataset records under {output_root}; nothing has been exported yet") + + # A dataset split across tasks leaves one record per task; they describe disjoint sets of + # source files, so merging is a sum. + merged: dict[str, dict] = {} + for path in records: + record = yaml.safe_load(path.read_text()) + existing = merged.get(record["name"]) + if existing is None: + merged[record["name"]] = record + continue + for key in ("n_documents", "n_lines", "n_bytes", "n_shards"): + existing[key] += record[key] + existing["factors_applied"].update(record["factors_applied"]) + datasets = [merged[name] for name in sorted(merged)] + manifest = { + # The ratios are already in the bytes. Anyone carrying the mix manifest's ratio into a + # 'weighted_combined' training config after this stage would apply it a second time -- + # 3.0 becoming 9.0 -- so this manifest states the training-time factor explicitly. + "repeat_factor_applied": True, + "training_ratio": 1.0, + "note": ( + "Up- and downsampling is materialised in these files. The training set is the " + "concatenation of every dataset's shards; do not apply the mix manifest's ratios again." + ), + "n_lines": sum(d["n_lines"] for d in datasets), + "n_documents": sum(d["n_documents"] for d in datasets), + "n_bytes": sum(d["n_bytes"] for d in datasets), + "datasets": datasets, + } + path = output_root / EXPORT_MANIFEST + scratch = path.with_suffix(f".yaml.{os.getpid()}.tmp") + scratch.write_text(yaml.safe_dump(manifest, sort_keys=False)) + os.replace(scratch, path) + get_logger(name="main").info( + f"Export manifest: {len(datasets)} dataset(s), {manifest['n_lines']:,} lines, " + f"{manifest['n_bytes'] / 1e12:,.2f} TB at {path}" + ) + return path diff --git a/src/modalities/dataloader/preprocessing/quality/file_manifest.py b/src/modalities/dataloader/preprocessing/quality/file_manifest.py new file mode 100644 index 000000000..b74e7bab8 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/file_manifest.py @@ -0,0 +1,257 @@ +"""Pins the source file list a sidecar was built against. + +A sidecar row locates its document by ``(file_id, byte_offset, byte_len)``, where +``file_id`` indexes the dataset's sorted file list. That list used to be re-derived from +the filesystem at every stage, which makes it a silent dependency on the source tree +never changing: rename a directory, re-shard a corpus, or let a transfer rewrite it, and +the same ``file_id`` names a different file. Every offset then points into the wrong +place, and nothing downstream can tell. + +We learned this the expensive way. A transfer re-sharded four corpora after their +sidecars were built; ``Nemotron-CC`` went from 606 MB files to 137 MB files while keeping +almost the same file count, so the count check that existed passed and the byte offsets +pointed past the end of the file. Only a dataset whose file count dropped to zero failed +loudly. + +So the file list is now recorded next to the sidecar and verified wherever it is used. +Drift becomes an error naming the file that moved, rather than a blend of garbage byte +ranges. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Optional + +from pydantic import BaseModel, Field + +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry + +MANIFEST_NAME = "_files.json" + + +class ManifestError(RuntimeError): + """Raised when a manifest is missing, unreadable, or no longer matches the source.""" + + +class SourceFile(BaseModel): + """One source file, as it was when the sidecar was built. + + Attributes: + file_id (int): Index of this file in the dataset's sorted file list. This is the + value stored in the sidecar's ``file_id`` column. + path (str): Path relative to the dataset's ``jsonl_root``. Relative so that a + whole tree can be moved or a snapshot taken without invalidating the + manifest, while the mapping from id to file stays exact. + size (int): Size in bytes. Checked because it is what actually catches a + re-sharded or partially transferred corpus, and unlike mtime it survives a + copy that preserves content. + """ + + file_id: int + path: str + size: int + + +class FileManifest(BaseModel): + """The source file list a sidecar was built against. + + Attributes: + dataset (str): Dataset name, so a manifest cannot be read for the wrong dataset. + jsonl_root (str): The root the paths are relative to, as recorded at build time. + Informational: verification resolves against the *current* registry root, so + moving the tree is fine and only the file set has to agree. + glob (str): The pattern that produced the list, recorded for diagnosis. + files (list[SourceFile]): One entry per file, ordered by ``file_id``. + """ + + dataset: str + jsonl_root: str + glob: str + files: list[SourceFile] = Field(default_factory=list) + + @classmethod + def from_entry(cls, entry: DatasetEntry) -> "FileManifest": + """Records the dataset's current file list. + + Args: + entry (DatasetEntry): The dataset to describe. + + Returns: + FileManifest: A manifest of the files matching the entry right now. + + Raises: + ManifestError: If the entry matches no files, which means a wrong root or + glob rather than an empty corpus. + """ + files = entry.iter_files() + if not files: + raise ManifestError(f"dataset {entry.name!r}: no files matched {entry.glob!r} under {entry.jsonl_root}") + return cls( + dataset=entry.name, + jsonl_root=str(entry.jsonl_root), + glob=entry.glob, + files=[ + SourceFile(file_id=i, path=str(p.relative_to(entry.jsonl_root)), size=p.stat().st_size) + for i, p in enumerate(files) + ], + ) + + @staticmethod + def path_for(sidecar_dir: Path) -> Path: + """Names the manifest belonging to a sidecar directory. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + + Returns: + Path: Where the manifest lives. + """ + return Path(sidecar_dir) / MANIFEST_NAME + + def write(self, sidecar_dir: Path) -> Path: + """Writes the manifest beside the sidecar parts. + + Written atomically via a uniquely named temporary file, because the sidecar build + is sharded across tasks that all describe the same file list and would otherwise + race: a reader could observe a half-written manifest. Each task writes identical + content, so last-writer-wins is correct. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. Created if absent. + + Returns: + Path: The manifest path. + """ + sidecar_dir = Path(sidecar_dir) + sidecar_dir.mkdir(parents=True, exist_ok=True) + target = self.path_for(sidecar_dir) + tmp = target.with_suffix(f".tmp.{os.getpid()}") + tmp.write_text(self.model_dump_json(indent=2)) + os.replace(tmp, target) + return target + + @classmethod + def read(cls, sidecar_dir: Path) -> "FileManifest": + """Loads the manifest belonging to a sidecar directory. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + + Returns: + FileManifest: The recorded file list. + + Raises: + ManifestError: If the manifest is absent or unparseable. + """ + path = cls.path_for(sidecar_dir) + if not path.exists(): + raise ManifestError( + f"no source file manifest at {path}. This sidecar was built before file lists were " + f"recorded, so its byte offsets cannot be checked against the source tree. Run " + f"'modalities quality verify-sidecar --adopt' to verify it and stamp a manifest, or " + f"rebuild the sidecar." + ) + try: + return cls.model_validate(json.loads(path.read_text())) + except Exception as e: + raise ManifestError(f"cannot read source file manifest {path}: {e}") from e + + def resolve(self, entry: DatasetEntry) -> list[Path]: + """Maps recorded file ids to paths under the entry's current root. + + Resolution is by recorded *path*, not by re-globbing, so a file added to or + removed from the source tree cannot shift the mapping. + + Args: + entry (DatasetEntry): The dataset, supplying the current root. + + Returns: + list[Path]: Absolute paths indexed by file id. + """ + root = Path(entry.jsonl_root) + return [root / f.path for f in sorted(self.files, key=lambda f: f.file_id)] + + def drift(self, entry: DatasetEntry, check_sizes: bool = True) -> list[str]: + """Describes how the source tree differs from what was recorded. + + Args: + entry (DatasetEntry): The dataset to check, supplying the current root. + check_sizes (bool): Whether to compare file sizes. Sizes are what catch a + re-sharded or half-transferred corpus, so this is on by default. + + Returns: + list[str]: One human-readable line per problem, empty if the tree agrees. + Truncated to the first 20 problems plus a count, since a re-shard makes + every file differ and a wall of text helps nobody. + """ + problems: list[str] = [] + root = Path(entry.jsonl_root) + if not root.exists(): + return [f"source root {root} does not exist"] + + for f in sorted(self.files, key=lambda f: f.file_id): + path = root / f.path + if not path.exists(): + problems.append(f"file_id {f.file_id}: {f.path} is gone") + continue + if check_sizes: + size = path.stat().st_size + if size != f.size: + problems.append( + f"file_id {f.file_id}: {f.path} is {size:,} bytes, was {f.size:,} when the sidecar was built" + ) + + recorded = {f.path for f in self.files} + current = {str(p.relative_to(root)) for p in entry.iter_files()} + added = sorted(current - recorded) + if added: + problems.append( + f"{len(added)} file(s) added to the source tree since the sidecar was built, e.g. " + f"{added[:3]} -- they hold no sidecar rows and will not take part in the blend" + ) + + if len(problems) > 20: + return problems[:20] + [f"... and {len(problems) - 20} further problems"] + return problems + + def require_current(self, entry: DatasetEntry, check_sizes: bool = True) -> list[Path]: + """Resolves file ids to paths, refusing if the source tree has drifted. + + Args: + entry (DatasetEntry): The dataset to check and resolve against. + check_sizes (bool): Whether to compare file sizes. + + Returns: + list[Path]: Absolute paths indexed by file id. + + Raises: + ManifestError: If the source tree no longer matches the manifest. Byte + offsets recorded against the old tree do not describe the new one, so + continuing would produce a blend of wrong documents. + """ + problems = self.drift(entry, check_sizes=check_sizes) + if problems: + listed = "\n".join(f" {p}" for p in problems) + raise ManifestError( + f"dataset {entry.name!r}: the source tree changed since the sidecar was built, so " + f"its byte offsets no longer describe these files:\n{listed}\n" + f"Rebuild this dataset's sidecar (and re-join and re-cube it) before using it." + ) + return self.resolve(entry) + + +def load_manifest(sidecar_dir: Path) -> Optional[FileManifest]: + """Reads a manifest if one is present. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + + Returns: + Optional[FileManifest]: The manifest, or None if there is none. + """ + if not FileManifest.path_for(sidecar_dir).exists(): + return None + return FileManifest.read(sidecar_dir) diff --git a/src/modalities/dataloader/preprocessing/quality/materialize.py b/src/modalities/dataloader/preprocessing/quality/materialize.py new file mode 100644 index 000000000..63508a665 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/materialize.py @@ -0,0 +1,591 @@ +"""Writes a selection out as index files the existing packer already understands. + +A modalities index is a pickled ``list[(byte_offset, byte_len)]`` naming the documents +of a JSONL file, and ``PackedDataGenerator`` tokenizes exactly the documents its index +lists. So a selection does not need a filtered copy of the corpus and does not need any +change to the packer: writing an index that lists only the surviving documents is +enough, and packing then reads only those. + +The practical consequence is that an ablation costs megabytes rather than terabytes. +The source tree is never written to, and several selections can coexist as several +index directories over the same untouched data. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import pickle +import shutil +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Optional + +import pyarrow.parquet as pq +import yaml +from tqdm import tqdm + +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest, ManifestError +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, DatasetEntry +from modalities.dataloader.preprocessing.quality.selection import ( + DatasetSelection, + MissingPolicy, + SelectionConfig, + document_mask, + exposure_report, + ordered_quality_levels, +) +from modalities.dataloader.preprocessing.quality.upsampling import ( + UNANNOTATED_BUCKET, + QualityBucket, + UpsamplingError, + solve_curve, +) +from modalities.utils.logger_utils import get_logger + + +class MaterializationError(RuntimeError): + """Raised when a selection cannot be written out.""" + + +@dataclass +class MaterializedDataset: + """Where one dataset's filtered indexes ended up. + + Attributes: + name (str): Dataset name. + ratio (float): Up/downsample factor recorded for training. + n_documents_total (int): Documents before filtering. + n_documents_kept (int): Documents listed in the written indexes. + tokens_kept (int): Estimated tokens of the kept documents. + index_files (dict[str, str]): Source JSONL path to written index path. + source_dataset (Optional[str]): The registry dataset this came from, when ``name`` + carries a quality bucket suffix and so no longer matches the registry. + quality_bucket (Optional[str]): The quality level this row covers, when the dataset + was split by a quality curve. + """ + + name: str + ratio: float + n_documents_total: int + n_documents_kept: int + tokens_kept: int + index_files: dict[str, str] + source_dataset: Optional[str] = None + quality_bucket: Optional[str] = None + + def to_dict(self) -> dict: + """Renders the record for the manifest. + + Returns: + dict: Plain-data form of this dataset's outcome. + """ + return { + "name": self.name, + "ratio": self.ratio, + "n_documents_total": self.n_documents_total, + "n_documents_kept": self.n_documents_kept, + "row_retention": round(self.n_documents_kept / self.n_documents_total, 6) + if self.n_documents_total + else 0.0, + "est_tokens_kept": self.tokens_kept, + "source_dataset": self.source_dataset or self.name, + "quality_bucket": self.quality_bucket, + "index_files": self.index_files, + } + + +def config_fingerprint(config: SelectionConfig) -> str: + """Fingerprints a selection so a manifest can be traced back to it. + + Args: + config (SelectionConfig): The selection. + + Returns: + str: Short stable digest of the selection's content. + """ + payload = json.dumps(config.model_dump(mode="json"), sort_keys=True).encode() + return hashlib.blake2b(payload, digest_size=8).hexdigest() + + +def materialize_dataset( + sidecar_dir: Path, + dataset_entry: DatasetEntry, + dataset_selection: DatasetSelection, + missing_policy: MissingPolicy, + output_dir: Path, + show_progress: bool = True, +) -> MaterializedDataset: + """Writes filtered index files for one dataset. + + Args: + sidecar_dir (Path): Directory of that dataset's sidecar parts. + dataset_entry (DatasetEntry): Registry entry, used to map file ids back to + source paths. + dataset_selection (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + output_dir (Path): Directory receiving the index files. The source tree's + directory structure is mirrored below it. + show_progress (bool): Whether to show a progress bar. + + Returns: + MaterializedDataset: Counts and the written index paths. + + Raises: + MaterializationError: If the sidecar is missing, or if the source tree changed + since the sidecar was built -- the byte offsets then describe documents that + are no longer at those positions, so the blend would be silently wrong. + """ + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise MaterializationError(f"no sidecar parts found in {sidecar_dir}") + + # File ids are positions in a file list, so they only mean anything against the list + # the sidecar was built from. Resolve through the recorded manifest and refuse if the + # tree has moved underneath us. + try: + source_files = FileManifest.read(sidecar_dir).require_current(dataset_entry) + except ManifestError as e: + raise MaterializationError(str(e)) from e + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Accumulated per source file, because one file's sidecar may span several parts + # and an index must list its documents in file order. + per_file: dict[int, list[tuple[int, int]]] = {} + n_total = 0 + n_kept = 0 + tokens_kept = 0 + + for part in tqdm(parts, desc=f"select {dataset_selection.name}", disable=not show_progress): + parquet_file = pq.ParquetFile(part) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx) + n_total += table.num_rows + mask = document_mask(table, dataset_selection, missing_policy) + if not mask.any(): + continue + file_ids = table.column("file_id").to_numpy(zero_copy_only=False)[mask] + offsets = table.column("byte_offset").to_numpy(zero_copy_only=False)[mask] + lengths = table.column("byte_len").to_numpy(zero_copy_only=False)[mask] + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False)[mask] + n_kept += int(mask.sum()) + tokens_kept += int(tokens.sum()) + for file_id, offset, length in zip(file_ids, offsets, lengths): + per_file.setdefault(int(file_id), []).append((int(offset), int(length))) + + index_files: dict[str, str] = {} + for file_id, entries in sorted(per_file.items()): + if file_id >= len(source_files): + raise MaterializationError( + f"dataset {dataset_selection.name!r}: sidecar references file id {file_id} but its manifest " + f"records only {len(source_files)} files. The sidecar is internally inconsistent; rebuild it." + ) + source_path = source_files[file_id] + # Index entries must be ordered by position, as a freshly generated index is. + entries.sort() + relative = source_path.relative_to(dataset_entry.jsonl_root).with_suffix(".idx") + index_path = output_dir / relative + index_path.parent.mkdir(parents=True, exist_ok=True) + index_path.write_bytes(pickle.dumps(entries)) + index_files[str(source_path)] = str(index_path) + + return MaterializedDataset( + name=dataset_selection.name, + ratio=dataset_selection.ratio, + n_documents_total=n_total, + n_documents_kept=n_kept, + tokens_kept=tokens_kept, + index_files=index_files, + ) + + +def materialize_dataset_buckets( + sidecar_dir: Path, + dataset_entry: DatasetEntry, + dataset_selection: DatasetSelection, + missing_policy: MissingPolicy, + output_dir: Path, + show_progress: bool = True, +) -> list[MaterializedDataset]: + """Writes one index tree per quality bucket, each with its own repeat factor. + + A quality curve gives a different repeat factor to each quality level, and the packer + emits one file per source file, so documents of different levels have to end up in + different indexes for their factors to mean anything. Each bucket therefore becomes its + own row of the manifest, its own packed output, and its own entry in the training + blend's repeat factors. + + The curve is solved from the token counts found here rather than from the cube, because + this stage reads every document anyway: the numbers are exact, so the curve hits its + token target exactly rather than to within the cube's interpolation error. + + Args: + sidecar_dir (Path): Directory of that dataset's sidecar parts. + dataset_entry (DatasetEntry): Registry entry, for mapping file ids to paths. + dataset_selection (DatasetSelection): The rule to apply. Must carry an + ``upsampling`` spec. + missing_policy (MissingPolicy): Policy for unannotated documents. + output_dir (Path): Directory receiving one subdirectory per bucket. + show_progress (bool): Whether to show a progress bar. + + Returns: + list[MaterializedDataset]: One entry per bucket that survived with a non-zero + factor, worst quality first. + + Raises: + MaterializationError: If the sidecar is missing, the source tree has drifted, the + quality field is absent, or the curve cannot be solved. + """ + spec = dataset_selection.upsampling + if spec is None: + raise MaterializationError(f"dataset {dataset_selection.name!r} has no upsampling spec") + + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise MaterializationError(f"no sidecar parts found in {sidecar_dir}") + try: + source_files = FileManifest.read(sidecar_dir).require_current(dataset_entry) + except ManifestError as e: + raise MaterializationError(str(e)) from e + + levels = list(ordered_quality_levels(spec.quality_field)) + # Unannotated documents cannot be ordered, so they form the bottom bucket. + bucket_labels = [UNANNOTATED_BUCKET] + levels + per_bucket: dict[str, dict[int, list[tuple[int, int]]]] = {label: {} for label in bucket_labels} + tokens_of: dict[str, int] = dict.fromkeys(bucket_labels, 0) + documents_of: dict[str, int] = dict.fromkeys(bucket_labels, 0) + n_total = 0 + + for part in tqdm(parts, desc=f"select {dataset_selection.name}", disable=not show_progress): + parquet_file = pq.ParquetFile(part) + if spec.quality_field not in parquet_file.schema_arrow.names: + raise MaterializationError( + f"dataset {dataset_selection.name!r}: sidecar has no column " + f"{spec.quality_field!r} to order quality by; join the annotations first" + ) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx) + n_total += table.num_rows + mask = document_mask(table, dataset_selection, missing_policy) + if not mask.any(): + continue + file_ids = table.column("file_id").to_numpy(zero_copy_only=False)[mask] + offsets = table.column("byte_offset").to_numpy(zero_copy_only=False)[mask] + lengths = table.column("byte_len").to_numpy(zero_copy_only=False)[mask] + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False)[mask] + quality = table.column(spec.quality_field).to_pylist() + kept_quality = [q for q, keep in zip(quality, mask) if keep] + + for file_id, offset, length, token, level in zip( + file_ids, offsets, lengths, tokens, kept_quality + ): + label = UNANNOTATED_BUCKET if level is None or level not in tokens_of else level + per_bucket[label].setdefault(int(file_id), []).append((int(offset), int(length))) + tokens_of[label] += int(token) + documents_of[label] += 1 + + buckets = [ + QualityBucket( + label=label, + n_documents=documents_of[label], + n_tokens=tokens_of[label], + unannotated=label == UNANNOTATED_BUCKET, + ) + for label in bucket_labels + if tokens_of[label] > 0 + ] + try: + plan = solve_curve(buckets, spec) + except UpsamplingError as e: + raise MaterializationError(f"dataset {dataset_selection.name!r}: {e}") from e + + results: list[MaterializedDataset] = [] + for bucket_plan in plan.buckets: + label = bucket_plan.bucket.label + if bucket_plan.factor <= 0: + continue + slug = label.strip("<>").replace(" ", "_") + bucket_dir = Path(output_dir) / slug + index_files: dict[str, str] = {} + for file_id, entries in sorted(per_bucket[label].items()): + if file_id >= len(source_files): + raise MaterializationError( + f"dataset {dataset_selection.name!r}: sidecar references file id {file_id} but its " + f"manifest records only {len(source_files)} files. Rebuild it." + ) + source_path = source_files[file_id] + entries.sort() + relative = source_path.relative_to(dataset_entry.jsonl_root).with_suffix(".idx") + index_path = bucket_dir / relative + index_path.parent.mkdir(parents=True, exist_ok=True) + index_path.write_bytes(pickle.dumps(entries)) + index_files[str(source_path)] = str(index_path) + + results.append( + MaterializedDataset( + name=f"{dataset_selection.name}__{slug}", + ratio=bucket_plan.factor, + n_documents_total=n_total, + n_documents_kept=bucket_plan.bucket.n_documents, + tokens_kept=bucket_plan.bucket.n_tokens, + index_files=index_files, + source_dataset=dataset_selection.name, + quality_bucket=label, + ) + ) + + get_logger(name="main").info( + f"{dataset_selection.name}: quality curve on {spec.quality_field} over " + f"{len(plan.buckets)} bucket(s), {len(results)} kept, exponent {plan.exponent:.2f}, " + f"drawing {plan.tokens_drawn:,.0f} of {plan.tokens_available:,} available tokens" + ) + return results + + +def _cap_for(config: SelectionConfig, materialized: MaterializedDataset) -> Optional[float]: + """Finds the repetition cap that applies to one materialised row. + + Args: + config (SelectionConfig): The selection. + materialized (MaterializedDataset): The row, which may be one bucket of a dataset. + + Returns: + Optional[float]: The dataset's own curve cap when it has one, otherwise the + blend-wide cap, or None if neither was declared. + """ + source = materialized.source_dataset or materialized.name + for dataset in config.datasets: + if dataset.name == source and dataset.upsampling is not None: + return dataset.upsampling.max_factor + return config.max_total_exposure + + +def materialize_blend( + config: SelectionConfig, + registry: CorpusRegistry, + sidecar_root: Path, + output_root: Path, + show_progress: bool = True, + allow_overexposure: bool = False, +) -> Path: + """Writes filtered indexes and a manifest for a whole selection. + + Args: + config (SelectionConfig): The blend specification. + registry (CorpusRegistry): Registry resolving dataset names to source files. + sidecar_root (Path): Directory holding one subdirectory of sidecar parts per + dataset. + output_root (Path): Directory receiving per-dataset index trees and the + manifest. + show_progress (bool): Whether to show progress bars. + + Returns: + Path: Path to the written ``mix_manifest.yaml``. + + Raises: + MaterializationError: If a selected dataset has no sidecar, or if the run would + repeat data past its declared cap and ``allow_overexposure`` is not set. + """ + output_root = Path(output_root) + output_root.parent.mkdir(parents=True, exist_ok=True) + + # Everything is built in a sibling directory and moved into place only once the + # exposure check has passed and the manifest is written. Writing into the destination + # directly meant a rejected apply left fresh index trees next to the previous run's + # manifest -- a directory that still looks complete but whose manifest no longer + # describes the indexes beside it. A sibling keeps the move a rename on one filesystem. + staging = output_root.parent / f".{output_root.name}.staging.{os.getpid()}" + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + try: + manifest_path = _materialize_into( + staging_root=staging, + published_root=output_root, + config=config, + registry=registry, + sidecar_root=sidecar_root, + show_progress=show_progress, + allow_overexposure=allow_overexposure, + ) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + # Two renames rather than one: os.replace refuses a non-empty destination directory, + # so the previous run steps aside first and is only deleted once the new tree is live. + superseded = output_root.parent / f".{output_root.name}.superseded.{os.getpid()}" + if output_root.exists(): + os.replace(output_root, superseded) + try: + os.replace(staging, output_root) + except BaseException: + if superseded.exists(): + os.replace(superseded, output_root) + shutil.rmtree(staging, ignore_errors=True) + raise + shutil.rmtree(superseded, ignore_errors=True) + + final_path = output_root / manifest_path.name + get_logger(name="main").info(f"Published the blend at {output_root}; manifest {final_path}.") + return final_path + + +def _materialize_into( + staging_root: Path, + published_root: Path, + config: SelectionConfig, + registry: CorpusRegistry, + sidecar_root: Path, + show_progress: bool, + allow_overexposure: bool, +) -> Path: + """Builds a complete blend inside ``staging_root``. + + Split out from ``materialize_blend`` so that every failure path -- a missing sidecar, + a bad curve, the exposure cap -- leaves the caller with a directory to discard rather + than a half-updated destination. + + Args: + staging_root (Path): Empty directory receiving the index trees and manifest. + published_root (Path): Where the staged tree will end up. Index paths are recorded + relative to this, not to the staging directory, which stops existing. + config (SelectionConfig): The blend specification. + registry (CorpusRegistry): Registry resolving dataset names to source files. + sidecar_root (Path): Directory holding one subdirectory of sidecar parts per dataset. + show_progress (bool): Whether to show progress bars. + allow_overexposure (bool): Whether to proceed despite exceeded repetition caps. + + Returns: + Path: Path to the manifest inside ``staging_root``. + + Raises: + MaterializationError: If a selected dataset has no sidecar, or if the run would + repeat data past its declared cap and ``allow_overexposure`` is not set. + """ + output_root = staging_root + materialized: list[MaterializedDataset] = [] + + for dataset_selection in config.enabled_datasets(): + entry = registry.get(dataset_selection.name) + sidecar_dir = Path(sidecar_root) / dataset_selection.name + if not sidecar_dir.is_dir(): + raise MaterializationError( + f"dataset {dataset_selection.name!r} has no sidecar at {sidecar_dir}; " + "run 'modalities data quality build-sidecar' for it first" + ) + arguments = dict( + sidecar_dir=sidecar_dir, + dataset_entry=entry, + dataset_selection=dataset_selection, + missing_policy=config.policy_for(dataset_selection), + output_dir=output_root / dataset_selection.name, + show_progress=show_progress, + ) + # A curve splits one dataset into several rows, one per quality bucket, so that each + # can carry its own repeat factor through packing and into the training blend. + if dataset_selection.upsampling is not None: + materialized.extend(materialize_dataset_buckets(**arguments)) + else: + materialized.append(materialize_dataset(**arguments)) + + # The index writers record the path they physically wrote to, which is inside the + # staging directory. That directory is renamed away on publication, so a manifest + # holding those paths would name files that no longer exist and every packing config + # built from it would point at nothing. + materialized = [_rebase_index_files(d, staging_root, published_root) for d in materialized] + + total_effective = sum(d.tokens_kept * d.ratio for d in materialized) + + # Ratios are per pass. If the run consumes more than one pass, every factor is + # multiplied, so a bucket set to 7x becomes 14x on a second pass -- well past where + # repetition pays. Checked here rather than at preview because this is the point of + # commitment: preview reports it, apply refuses. + exposure = exposure_report( + entries=[ + ( + d.name, + d.ratio, + _cap_for(config, d), + ) + for d in materialized + ], + effective_tokens=total_effective, + training_tokens=config.target_tokens, + ) + if exposure.exceeded and not allow_overexposure: + offenders = "\n".join( + f" {row.label}: {row.factor:.2f}x requested, seen {row.exposure:.2f}x over " + f"{exposure.passes:.2f} passes, cap {row.cap:g}x" + for row in exposure.exceeded + ) + raise MaterializationError( + f"the run would repeat data past its declared cap:\n{offenders}\n" + f"The blend yields {total_effective:,.0f} effective tokens and the run consumes " + f"{config.target_tokens:,.0f}, so it wraps {exposure.passes:.2f} times. Raise the " + f"blend's yield, lower target_tokens, relax the cap, or pass allow_overexposure " + f"to proceed anyway." + ) + manifest = { + "selection_fingerprint": config_fingerprint(config), + "missing_annotation": config.missing_annotation.value, + # Carried so the export stage can resolve fractional repeat factors identically + # without being handed the selection again. + "seed": config.seed, + "target_tokens": config.target_tokens, + "est_total_effective_tokens": int(total_effective), + "datasets": [ + { + **d.to_dict(), + "est_effective_tokens": int(d.tokens_kept * d.ratio), + "blend_share": round(d.tokens_kept * d.ratio / total_effective, 6) if total_effective else 0.0, + "predicates": [ + p.describe() + for p in next(s for s in config.datasets if s.name == (d.source_dataset or d.name)).predicates + ], + } + for d in materialized + ], + } + manifest_path = output_root / "mix_manifest.yaml" + with manifest_path.open("w") as f: + yaml.safe_dump(manifest, f, sort_keys=False) + + get_logger(name="main").info( + f"Staged {len(materialized)} filtered index tree(s); " + f"estimated {_humanise_tokens(total_effective)} effective tokens." + ) + return manifest_path + + +def _rebase_index_files( + materialized: MaterializedDataset, staging_root: Path, published_root: Path +) -> MaterializedDataset: + """Rewrites a row's index paths from where they were written to where they will live. + + Args: + materialized (MaterializedDataset): A row carrying staging-relative index paths. + staging_root (Path): The directory the indexes were written under. + published_root (Path): The directory they will be renamed into. + + Returns: + MaterializedDataset: The same row with index paths under ``published_root``. + """ + return replace( + materialized, + index_files={ + source: str(published_root / Path(index).relative_to(staging_root)) + for source, index in materialized.index_files.items() + }, + ) + + +def _humanise_tokens(n: float) -> str: + # Blend totals span from a few thousand tokens in a test to hundreds of billions in + # a real run, so a fixed unit renders one of those two cases uselessly. + for unit, size in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("k", 1e3)): + if abs(n) >= size: + return f"{n / size:.2f}{unit}" + return f"{n:.0f}" diff --git a/src/modalities/dataloader/preprocessing/quality/pipeline.py b/src/modalities/dataloader/preprocessing/quality/pipeline.py new file mode 100644 index 000000000..d8352cded --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/pipeline.py @@ -0,0 +1,667 @@ +"""Whole-blend orchestration of the quality selection stages. + +Each function here drives one stage across every dataset of a registry and leaves its +output in a fixed place under a working directory, so the stages can be run +independently, re-run for a single dataset, and resumed after a failure: + +``/calibration.yaml`` token estimator constants, one entry per dataset +``/sidecar//`` per-document parquet parts +``/buckets//`` annotation shards partitioned for joining +``/cube/.parquet`` aggregated counts, read by the preview +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Optional + + +from modalities.dataloader.preprocessing.quality.annotation_join import ( + JoinReport, + bucket_annotations, + join_annotations, + read_bucket_metadata, +) +from modalities.dataloader.preprocessing.quality.cube import Cube, build_cube +from modalities.dataloader.preprocessing.quality.export import export_blend, finalize_export +from modalities.dataloader.preprocessing.quality.materialize import materialize_blend +from modalities.dataloader.preprocessing.quality.registry import CorpusRegistry, KeyKind +from modalities.dataloader.preprocessing.quality.selection import ( + BlendResult, + SelectionConfig, + evaluate_blend, + SelectionError, + format_blend_report, + predicate_breakdown, +) +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder, resolve_source_pointers +from modalities.dataloader.preprocessing.quality.tokens import CalibrationSet, calibrate_dataset +from modalities.dataloader.preprocessing.quality.verify import VerifyReport, adopt_manifest, verify_sidecar +from modalities.utils.logger_utils import get_logger + + +def calibration_path(work_dir: Path) -> Path: + """Location of the calibration file within a working directory. + + Args: + work_dir (Path): The working directory. + + Returns: + Path: Path to ``calibration.yaml``. + """ + return Path(work_dir) / "calibration.yaml" + + +def sidecar_dir(work_dir: Path, dataset_name: str) -> Path: + """Location of one dataset's sidecar parts. + + Args: + work_dir (Path): The working directory. + dataset_name (str): The dataset name. + + Returns: + Path: Directory holding that dataset's parquet parts. + """ + return Path(work_dir) / "sidecar" / dataset_name + + +def cube_path(work_dir: Path, dataset_name: str) -> Path: + """Location of one dataset's cube. + + Args: + work_dir (Path): The working directory. + dataset_name (str): The dataset name. + + Returns: + Path: Path to that dataset's cube parquet. + """ + return Path(work_dir) / "cube" / f"{dataset_name}.parquet" + + +def bucket_dir(work_dir: Path, split: str) -> Path: + """Location of one annotation split's buckets. + + Args: + work_dir (Path): The working directory. + split (str): The annotation split path. + + Returns: + Path: Directory holding that split's bucket files. + """ + return Path(work_dir) / "buckets" / split.replace("/", "__") + + +def calibrate_blend( + registry: CorpusRegistry, + work_dir: Path, + tokenizer, + tokenizer_name: str, + sample_size: int = 2000, + only: Optional[list[str]] = None, +) -> CalibrationSet: + """Measures token-estimation constants for every dataset. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory receiving ``calibration.yaml``. + tokenizer: The tokenizer training will use. + tokenizer_name (str): Identifier recorded with each measurement. + sample_size (int): Documents to tokenize per dataset. + only (Optional[list[str]]): Restrict to these dataset names, merging the result + into any existing calibration file. + + Returns: + CalibrationSet: The calibrations, also written to ``calibration.yaml``. + """ + path = calibration_path(work_dir) + existing = CalibrationSet.from_yaml(path) if path.is_file() else CalibrationSet(tokenizer=tokenizer_name) + if existing.tokenizer != tokenizer_name: + get_logger(name="main").warning( + f"existing calibration was measured with {existing.tokenizer!r}, now measuring with " + f"{tokenizer_name!r}; entries for other datasets are stale and should be re-measured" + ) + existing.tokenizer = tokenizer_name + + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + calibration = calibrate_dataset( + dataset_name=dataset.name, + file_paths=dataset.iter_files(), + tokenizer=tokenizer, + tokenizer_name=tokenizer_name, + text_field=dataset.text_field, + sample_size=sample_size, + ) + existing.calibrations[dataset.name] = calibration + # Written after every dataset, not once at the end, so interrupting the stage + # keeps what it already measured. Re-running with `--only` then fills the rest. + existing.to_yaml(path) + get_logger(name="main").info( + f"{dataset.name}: {calibration.bytes_per_token:.3f} bytes/token" + + ( + f", using native field {calibration.native_field!r} scaled by {calibration.native_scale:.4f}" + if calibration.uses_native_field() + else "" + ) + ) + return existing + + +def plan_sidecar_work( + registry: CorpusRegistry, + only: Optional[list[str]] = None, + shard_id: int = 0, + num_shards: int = 1, +) -> dict[str, list[int]]: + """Assigns each task its share of the per-file sidecar work. + + The unit of work is one JSONL file, and the work list is flattened across every + selected dataset before being divided, so one array covers the whole blend however + unevenly the file counts fall -- and they fall very unevenly, from four files to + forty thousand. + + Args: + registry (CorpusRegistry): The blend's datasets. + only (Optional[list[str]]): Restrict to these dataset names. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Total number of tasks. + + Returns: + dict[str, list[int]]: File ids this task should build, per dataset. Datasets with + nothing for this task are absent. + + Raises: + ValueError: If the shard selection is out of range. + """ + if not 0 <= shard_id < num_shards: + raise ValueError(f"shard_id {shard_id} is not in [0, {num_shards})") + + work: list[tuple[str, int]] = [] + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + work.extend((dataset.name, file_id) for file_id in range(len(dataset.iter_files()))) + + assigned: dict[str, list[int]] = {} + # Strided, so no task ends up holding only the largest dataset's files. + for position, (name, file_id) in enumerate(work): + if position % num_shards == shard_id: + assigned.setdefault(name, []).append(file_id) + return assigned + + +def build_sidecars( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + index_root: Optional[Path] = None, + file_ids: Optional[list[int]] = None, + shard_id: int = 0, + num_shards: int = 1, + show_progress: bool = True, +) -> dict[str, int]: + """Builds the per-document table for every dataset. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory receiving ``sidecar//``. + only (Optional[list[str]]): Restrict to these dataset names. + index_root (Optional[Path]): Where JSONL index files live or should be created. + Defaults to ``work_dir/idx``, never the source tree. Source corpora are + typically shared and read-only, and an index written beside a JSONL file is a + modification of somebody else's data; pass an explicit path to override. + file_ids (Optional[list[int]]): Restrict to these file ids explicitly. Applies to + every selected dataset and cannot be combined with sharding. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Total number of tasks sharing the work. + show_progress (bool): Whether to show progress bars. + + Returns: + dict[str, int]: Documents written by this task, per dataset. + + Raises: + ValueError: If explicit file ids are combined with a shard selection, since the + two express the same thing and the outcome would depend on which won. + """ + if file_ids is not None and num_shards != 1: + raise ValueError("pass either explicit file_ids or a shard selection, not both") + + # Never default to writing indexes beside the source JSONL, which is what + # SidecarBuilder does when given no index root. + index_root = Path(index_root) if index_root is not None else Path(work_dir) / "idx" + + calibrations = CalibrationSet.from_yaml(calibration_path(work_dir)) + selected = [d for d in registry.enabled_datasets() if not only or d.name in only] + if file_ids is not None: + assignment: dict[str, Optional[list[int]]] = {d.name: file_ids for d in selected} + elif num_shards == 1: + assignment = {d.name: None for d in selected} + else: + assignment = plan_sidecar_work(registry, only=only, shard_id=shard_id, num_shards=num_shards) + get_logger(name="main").info( + f"shard {shard_id}/{num_shards} builds " + + (", ".join(f"{name}:{len(ids)} file(s)" for name, ids in sorted(assignment.items())) or "nothing") + ) + + written: dict[str, int] = {} + for dataset in selected: + if dataset.name not in assignment: + continue + builder = SidecarBuilder( + dataset=dataset, + calibration=calibrations.get(dataset.name), + index_root=index_root / dataset.name, + ) + parts = builder.build( + sidecar_dir(work_dir, dataset.name), + file_ids=assignment[dataset.name], + show_progress=show_progress, + ) + written[dataset.name] = sum(parts.values()) + + # Safe to run per task: each task owns the parts it just wrote. + if dataset.key is not None and dataset.key.kind == KeyKind.SOURCE_POINTER: + n_resolved = resolve_source_pointers( + sidecar_dir(work_dir, dataset.name), dataset, only_parts=assignment[dataset.name] + ) + get_logger(name="main").info( + f"{dataset.name}: resolved {n_resolved:,} of {written[dataset.name]:,} pointers into source-corpus keys" + ) + return written + + +def bucket_blend_annotations( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_buckets: int = 1024, + shard_id: int = 0, + num_shards: int = 1, + force: bool = False, + show_progress: bool = True, +) -> dict[str, int]: + """Partitions every annotation split the blend needs, ready for joining. + + The expensive stage of the join and the one worth parallelising. Splits shared by + several datasets are bucketed once. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory receiving ``buckets//``. + only (Optional[list[str]]): Restrict to the splits these datasets need. + n_buckets (int): Partitions per split. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): How many tasks bucket each split. + force (bool): Re-bucket a split whose output is already complete. + show_progress (bool): Whether to show progress bars. + + Returns: + dict[str, int]: Rows written by this task, per split. + """ + splits: dict[str, Optional[str]] = {} + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + if dataset.annotation_split and dataset.annotation_split not in splits: + # Whether keys need normalising is a property of the split's key space, so + # the first dataset naming a split settles it for every other user of it. + splits[dataset.annotation_split] = "urn_uuid" if dataset.key.kind == KeyKind.URN_UUID_FIELD else None + + written: dict[str, int] = {} + for split, normalize in splits.items(): + shards = registry.annotation_shards(split) + if not shards: + get_logger(name="main").warning(f"split {split!r}: no shards on disk, nothing to bucket") + continue + out_dir = bucket_dir(work_dir, split) + if not force: + try: + meta = read_bucket_metadata(out_dir) + except Exception: + pass + else: + get_logger(name="main").info( + f"split {split}: already bucketed ({meta['n_rows']:,} rows, {meta['n_buckets']} buckets), skipping" + ) + continue + n_rows, columns = bucket_annotations( + shard_paths=shards, + out_dir=out_dir, + n_buckets=n_buckets, + normalize_key=normalize, + shard_id=shard_id, + num_shards=num_shards, + show_progress=show_progress, + ) + written[split] = n_rows + get_logger(name="main").info( + f"split {split}: shard {shard_id}/{num_shards} wrote {n_rows:,} rows " + f"from {len(shards)} input shard(s), columns {columns}" + ) + return written + + +def join_blend_annotations( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + resume: bool = False, + show_progress: bool = True, +) -> list[JoinReport]: + """Attaches the bucketed annotations to every annotated dataset's sidecar. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory holding the sidecars and the buckets. + only (Optional[list[str]]): Restrict to these dataset names. + resume (bool): Skip sidecar parts that already carry labels, to continue an + interrupted run. + show_progress (bool): Whether to show progress bars. + + Returns: + list[JoinReport]: One report per joined dataset. + """ + reports: list[JoinReport] = [] + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + if not dataset.annotation_split: + continue + + buckets = bucket_dir(work_dir, dataset.annotation_split) + if not buckets.is_dir(): + get_logger(name="main").warning( + f"{dataset.name}: split {dataset.annotation_split!r} has not been bucketed; its documents stay " + "unannotated and any predicate on them falls back to the missing-annotation policy" + ) + continue + + reports.append( + join_annotations( + sidecar_dir=sidecar_dir(work_dir, dataset.name), + annotation_bucket_dir=buckets, + dataset_name=dataset.name, + split_name=dataset.annotation_split, + resume=resume, + show_progress=show_progress, + ) + ) + + # One file per dataset. A single shared join_report.json meant 16 parallel `--only` + # tasks overwrote each other and only the last one's coverage survived. + report_dir = Path(work_dir) / "join_report" + report_dir.mkdir(parents=True, exist_ok=True) + for r in reports: + # Atomic, because the merge below globs this directory while sibling array tasks are + # writing into it. A plain write_text let another task read a half-written file: + # finewiki-es died on JSONDecodeError after its own join had already succeeded. + target = report_dir / f"{r.dataset}.json" + tmp = target.with_suffix(f".tmp.{os.getpid()}") + tmp.write_text(json.dumps(r.to_dict(), indent=1)) + os.replace(tmp, target) + + # Merged view, rebuilt from whatever per-dataset files exist so far. Tolerant of a + # sibling mid-write for the same reason, and of the .tmp files that implies. + merged = [] + for path in sorted(report_dir.glob("*.json")): + try: + merged.append(json.loads(path.read_text())) + except (OSError, json.JSONDecodeError): + continue + (Path(work_dir) / "join_report.json").write_text(json.dumps(merged, indent=1)) + return reports + + +def verify_sidecars( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_parts: int = 8, + n_rows_per_part: int = 4, + adopt: bool = False, +) -> list[VerifyReport]: + """Checks every dataset's sidecar against its source files. + + Worth running before ``apply`` on any blend whose source tree might have been touched + since the sidecars were built, and after any data transfer. A sidecar whose corpus was + re-sharded underneath it produces a blend of wrong byte ranges, and this is the only + stage that looks at the source bytes to find out. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory holding the sidecars. + only (Optional[list[str]]): Restrict to these dataset names. + n_parts (int): Sidecar parts to sample per dataset. + n_rows_per_part (int): Documents to probe per sampled part. + adopt (bool): Write a file manifest for verified sidecars that lack one, so later + stages can check for drift cheaply. Only verified sidecars are stamped. + + Returns: + list[VerifyReport]: One report per dataset, in registry order. + """ + datasets = [d for d in registry.enabled_datasets() if only is None or d.name in set(only)] + reports: list[VerifyReport] = [] + for dataset in datasets: + directory = sidecar_dir(work_dir, dataset.name) + report = verify_sidecar(directory, dataset, n_parts=n_parts, n_rows_per_part=n_rows_per_part) + if adopt and report.ok and not report.has_manifest: + adopt_manifest(directory, dataset, report) + report.has_manifest = True + report.notes.append("manifest adopted after verification") + reports.append(report) + return reports + + +def build_cubes( + registry: CorpusRegistry, + work_dir: Path, + only: Optional[list[str]] = None, + n_score_bins: int = 10, + label_dimensions: Optional[list[str]] = None, +) -> dict[str, Cube]: + """Aggregates every dataset's sidecar into a cube. + + Args: + registry (CorpusRegistry): The blend's datasets. + work_dir (Path): Working directory holding sidecars and receiving cubes. + only (Optional[list[str]]): Restrict to these dataset names. + n_score_bins (int): Quantile bins per native metric. + label_dimensions (Optional[list[str]]): Annotation columns to group on. Defaults to + :data:`~...cube.DEFAULT_LABEL_DIMENSIONS`, which is a subset of the columns the + join attaches -- grouping on all twelve would multiply the cell count by + thousands. Name a field here when a selection needs to threshold on it, or the + preview will have to scan the sidecar instead. + + Returns: + dict[str, Cube]: The cubes, also written under ``cube/``. + """ + cubes: dict[str, Cube] = {} + failures: list[tuple[str, Exception]] = [] + for dataset in registry.enabled_datasets(): + if only and dataset.name not in only: + continue + directory = sidecar_dir(work_dir, dataset.name) + if not directory.is_dir(): + get_logger(name="main").warning(f"{dataset.name}: no sidecar at {directory}, skipping cube") + continue + # One unbuildable dataset must not cost the others their cubes. A single failure + # used to abort the stage: nine cubes were written and six perfectly healthy + # datasets were never attempted. + try: + cube = build_cube( + directory, + dataset.name, + n_score_bins=n_score_bins, + **({"label_dimensions": label_dimensions} if label_dimensions else {}), + ) + except Exception as e: # noqa: BLE001 - reported together at the end and re-raised + get_logger(name="main").error(f"{dataset.name}: cube failed: {e}") + failures.append((dataset.name, e)) + continue + cube.write(cube_path(work_dir, dataset.name)) + cubes[dataset.name] = cube + get_logger(name="main").info( + f"{dataset.name}: cube has {cube.table.num_rows:,} cells over {cube.n_documents:,} documents " + f"({cube.n_tokens / 1e9:.2f}B estimated tokens); dimensions {cube.dimensions}" + ) + + if failures: + get_logger(name="main").error( + f"built {len(cubes)} cube(s); {len(failures)} dataset(s) failed: " + ", ".join(name for name, _ in failures) + ) + raise failures[0][1] + return cubes + + +def load_cubes(work_dir: Path, names: Optional[list[str]] = None) -> dict[str, Cube]: + """Loads previously built cubes. + + Args: + work_dir (Path): Working directory holding ``cube/``. + names (Optional[list[str]]): Restrict to these dataset names. + + Returns: + dict[str, Cube]: Cubes by dataset name; datasets without a cube are absent. + """ + cubes: dict[str, Cube] = {} + directory = Path(work_dir) / "cube" + if not directory.is_dir(): + return cubes + for path in sorted(directory.glob("*.parquet")): + if names and path.stem not in names: + continue + cubes[path.stem] = Cube.read(path) + return cubes + + +def preview_selection( + selection_path: Path, + work_dir: Path, + force_exact: bool = False, + allow_sidecar_fallback: bool = False, + explain: bool = False, +) -> tuple[BlendResult, str]: + """Costs a selection in documents and tokens. + + Args: + selection_path (Path): The selection YAML. + work_dir (Path): Working directory holding the cubes and sidecars. + force_exact (bool): Scan the per-document sidecars instead of the cubes. + allow_sidecar_fallback (bool): Permit a sidecar scan for datasets whose cube cannot + answer a predicate, instead of reporting them. + + Returns: + tuple[BlendResult, str]: The evaluated blend and its rendered table. + """ + config = SelectionConfig.from_yaml(selection_path) + names = [d.name for d in config.enabled_datasets()] + cubes = load_cubes(work_dir, names) + sidecars = {name: sidecar_dir(work_dir, name) for name in names} + result = evaluate_blend( + config, + cubes, + sidecar_dirs=sidecars, + force_exact=force_exact, + allow_sidecar_fallback=allow_sidecar_fallback, + ) + report = format_blend_report(result, datasets_in_order=names, config=config) + + if explain: + # Attribution needs the cubes, so it is only available on the cube path; an exact + # sidecar scan does not produce per-cell weights to slice. + sections = ["", "per-predicate attribution"] + for dataset in config.enabled_datasets(): + cube = cubes.get(dataset.name) + if cube is None or not dataset.predicates: + continue + try: + sections.append(predicate_breakdown(cube, dataset, config.policy_for(dataset)).describe()) + except SelectionError as e: + sections.append(f" {dataset.name}: {e}") + report = report + "\n".join(sections) + + return result, report + + +def apply_selection( + selection_path: Path, + registry_path: Path, + work_dir: Path, + output_dir: Path, + show_progress: bool = True, + allow_overexposure: bool = False, +) -> Path: + """Writes a selection out as filtered index files plus a manifest. + + Args: + selection_path (Path): The selection YAML. + registry_path (Path): The corpus registry YAML. + work_dir (Path): Working directory holding the sidecars. + output_dir (Path): Directory receiving the index trees and manifest. + show_progress (bool): Whether to show progress bars. + allow_overexposure (bool): Proceed even when the run would repeat data past a + declared cap. + + Returns: + Path: Path to the written manifest. + """ + config = SelectionConfig.from_yaml(selection_path) + registry = CorpusRegistry.from_yaml(registry_path) + return materialize_blend( + config=config, + registry=registry, + sidecar_root=Path(work_dir) / "sidecar", + output_root=output_dir, + show_progress=show_progress, + allow_overexposure=allow_overexposure, + ) + + +def export_jsonl( + manifest_path: Path, + registry_path: Path, + output_dir: Path, + seed: Optional[int] = None, + only: Optional[list[str]] = None, + resume: bool = True, + show_progress: bool = True, + finalize: bool = True, + shard_id: int = 0, + num_shards: int = 1, +) -> list: + """Writes a materialised selection out as JSONL, sampling baked in. + + Args: + manifest_path (Path): The ``mix_manifest.yaml`` written by the apply stage. + registry_path (Path): The corpus registry YAML. + output_dir (Path): Directory receiving one subdirectory per dataset. + seed (Optional[int]): Overrides the seed the selection was applied with. + only (Optional[list[str]]): Restrict to these dataset names. + resume (bool): Leave complete shards alone. + show_progress (bool): Whether to show progress bars. + finalize (bool): Merge the per-dataset records into ``export_manifest.yaml``. An + array task exporting a single dataset should not, since the other datasets are + still being written; the merge is a separate step afterwards. + shard_id (int): This task's index in ``[0, num_shards)``. + num_shards (int): Tasks splitting each dataset's source files between them. + + Returns: + list: One :class:`~...export.DatasetExport` per dataset written. + """ + exports = export_blend( + manifest_path=manifest_path, + registry_path=registry_path, + output_root=output_dir, + seed=seed, + only=only, + resume=resume, + show_progress=show_progress, + shard_id=shard_id, + num_shards=num_shards, + ) + if finalize: + finalize_export(output_dir) + return exports diff --git a/src/modalities/dataloader/preprocessing/quality/registry.py b/src/modalities/dataloader/preprocessing/quality/registry.py new file mode 100644 index 000000000..a66ea1480 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/registry.py @@ -0,0 +1,356 @@ +"""Declares the datasets of a blend and how each one joins to external annotations. + +A blend mixes corpora that were built by different people at different times, so the +document identifier is different in almost every one of them. Some carry a plain +``id``, some wrap the same UUID in ````, one carries the identifier +under a different name entirely, and two carry no identifier at all and have to be +keyed by a hash of their own text. The registry is where those differences are +written down once, so the rest of the package can treat every dataset the same way. +""" + +from __future__ import annotations + +import hashlib +import json +from enum import Enum +from pathlib import Path +from typing import Annotated, Any, Optional + +import yaml +from pydantic import BaseModel, Field, model_validator + + +class KeyKind(str, Enum): + """How a document's annotation key is obtained. + + Attributes: + FIELD: The key is a top-level JSON field, used verbatim. + URN_UUID_FIELD: The key is a JSON field holding a UUID that may or may not be + wrapped in ````. Both forms occur within a single file, on + both sides of the join, so the wrapper is stripped before comparing. + SHA256_TEXT: No identifier is stored; the key is the SHA-256 hex digest of the + document text, taken over the exact UTF-8 bytes with no normalisation. + SOURCE_POINTER: The identifier is a ``/`` pointer into a separate + source corpus. The key is resolved by reading that line and hashing its + text, which is how a translated corpus inherits the annotations of the + original it was translated from. + """ + + FIELD = "field" + URN_UUID_FIELD = "urn_uuid_field" + SHA256_TEXT = "sha256_text" + SOURCE_POINTER = "source_pointer" + + +def strip_urn_uuid(value: str) -> str: + """Reduces a possibly ````-wrapped identifier to the bare UUID. + + Args: + value (str): The identifier as stored, wrapped or bare. + + Returns: + str: The bare identifier. Values that are not wrapped are returned unchanged. + """ + if value.startswith(""): + return value[len(" str: + """Hashes document text the way the annotation corpora key their rows. + + Args: + text (str): The document text, exactly as stored. + + Returns: + str: Lowercase hex SHA-256 digest of the UTF-8 encoding of ``text``. + + Note: + The digest is taken over the unmodified string. Stripping whitespace or + appending a newline both produce keys that match nothing. + """ + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +class KeySpec(BaseModel): + """Describes how to derive the annotation key for one dataset. + + Attributes: + kind (KeyKind): Which derivation to apply. + field (Optional[str]): The JSON field holding the identifier or pointer. + Required for every kind except ``SHA256_TEXT``. + text_field (str): The JSON field holding the document text. Used by + ``SHA256_TEXT``, and by ``SOURCE_POINTER`` when reading the source corpus. + source_root (Optional[Path]): Directory holding the source corpus that + ``SOURCE_POINTER`` resolves into. + source_line_offset (int): Index of the first line of a source file. The + pointers we have seen are zero-indexed; an off-by-one here shifts every + annotation by one document and still looks like a working join. + """ + + kind: KeyKind + field: Optional[str] = None + text_field: str = "text" + source_root: Optional[Path] = None + source_line_offset: int = 0 + + @model_validator(mode="after") + def _check_required_parts(self) -> "KeySpec": + if self.kind in (KeyKind.FIELD, KeyKind.URN_UUID_FIELD, KeyKind.SOURCE_POINTER) and not self.field: + raise ValueError(f"key kind {self.kind.value} requires 'field'") + if self.kind == KeyKind.SOURCE_POINTER and self.source_root is None: + raise ValueError("key kind source_pointer requires 'source_root'") + return self + + def derive(self, record: dict[str, Any]) -> Optional[str]: + """Derives the annotation key for a single decoded JSON record. + + Args: + record (dict[str, Any]): One decoded JSONL line. + + Returns: + Optional[str]: The key, or None if the record lacks the needed field. + + Note: + ``SOURCE_POINTER`` is not resolved here, because resolving it requires + reading another corpus. This returns the raw pointer; use + ``SourcePointerResolver`` to turn a batch of pointers into keys. + """ + if self.kind == KeyKind.SHA256_TEXT: + text = record.get(self.text_field) + return sha256_text(text) if isinstance(text, str) else None + + raw = record.get(self.field) + if raw is None: + return None + if self.kind == KeyKind.URN_UUID_FIELD: + return strip_urn_uuid(str(raw)) + return str(raw) + + +class SourcePointerResolver: + """Turns ``/`` pointers into text hashes of a separate source corpus. + + A translated corpus keeps a pointer back to the line of the original it came from. + Its own text is in another language, so hashing it matches nothing; the annotation + belongs to the original. Resolving means reading those specific lines of the + source corpus. + + Pointers are grouped by source file and each file is read once in a single + forward pass, because the source files are tens of gigabytes each and seeking per + pointer would be far slower than streaming. + """ + + def __init__(self, source_root: Path, text_field: str = "text", line_offset: int = 0): + """ + Args: + source_root (Path): Directory holding the source corpus files. + text_field (str): The JSON field holding the source document text. + line_offset (int): Index of the first line of a source file. + """ + self._source_root = Path(source_root) + self._text_field = text_field + self._line_offset = line_offset + + @staticmethod + def split_pointer(pointer: str) -> tuple[str, int]: + """Splits a ``/`` pointer into its parts. + + Args: + pointer (str): The pointer as stored, e.g. ``part_29.jsonl/5279753``. + + Returns: + tuple[str, int]: The source file name and the line number. + + Raises: + ValueError: If the pointer has no ``/`` or a non-integer line number. + """ + file_name, _, line_str = pointer.rpartition("/") + if not file_name: + raise ValueError(f"pointer {pointer!r} is not of the form /") + try: + return file_name, int(line_str) + except ValueError as e: + raise ValueError(f"pointer {pointer!r} has a non-integer line number") from e + + def resolve(self, pointers: list[str]) -> dict[str, str]: + """Resolves pointers to annotation keys. + + Args: + pointers (list[str]): Pointers of the form ``/``. + + Returns: + dict[str, str]: Maps each resolvable pointer to the SHA-256 digest of the + source document's text. Pointers whose source file is missing, or + whose line number is past the end of the file, are absent from the + result rather than mapped to a wrong key. + """ + wanted: dict[str, dict[int, str]] = {} + for pointer in pointers: + file_name, line_no = self.split_pointer(pointer) + wanted.setdefault(file_name, {})[line_no - self._line_offset] = pointer + + resolved: dict[str, str] = {} + for file_name, lines_wanted in wanted.items(): + source_path = self._source_root / file_name + if not source_path.is_file(): + continue + last_wanted = max(lines_wanted) + with source_path.open(errors="replace") as f: + for i, line in enumerate(f): + if i in lines_wanted: + try: + text = json.loads(line).get(self._text_field) + except json.JSONDecodeError: + continue + if isinstance(text, str): + resolved[lines_wanted[i]] = sha256_text(text) + if i >= last_wanted: + break + return resolved + + +class NativeMetric(BaseModel): + """A quality signal already present in the dataset's own records. + + Attributes: + name (str): Name used for this metric in selection configs and the cube. + jq_pattern (str): jq expression evaluated against each record, following the + same convention as ``PackedDataGenerator``'s ``jq_pattern``. + aggregation (Optional[str]): How to reduce a list-valued result to one number. + One of ``first``, ``min``, ``max``, ``mean``. Several corpora store + per-page score arrays rather than a single document score. + """ + + name: str + jq_pattern: str + aggregation: Optional[str] = None + + +class DatasetEntry(BaseModel): + """One dataset of the blend. + + Attributes: + name (str): Identifier used in selection configs and reports. + jsonl_root (Path): Directory containing the dataset's ``.jsonl`` files. + glob (str): Pattern matching the dataset's files below ``jsonl_root``. + annotation_split (Optional[str]): Path of the matching annotation split, + relative to the annotation root. None means the dataset has no external + annotations and can only be shaped by its native metrics. + key (Optional[KeySpec]): How to derive the annotation key. Required whenever + ``annotation_split`` is set. + native_metrics (list[NativeMetric]): Quality signals to read out of the + dataset's own records. + text_field (str): The JSON field holding the document text. + enabled (bool): Whether the dataset takes part in the blend. Kept so a dataset + can be registered, and the reason for excluding it recorded, without + deleting its entry. + note (Optional[str]): Free-text remark carried into reports. + """ + + name: str + jsonl_root: Path + glob: str = "**/*.jsonl" + annotation_split: Optional[str] = None + key: Optional[KeySpec] = None + native_metrics: list[NativeMetric] = Field(default_factory=list) + text_field: str = "text" + enabled: bool = True + note: Optional[str] = None + + @model_validator(mode="after") + def _check_key_present_for_annotated(self) -> "DatasetEntry": + if self.annotation_split and self.key is None: + raise ValueError(f"dataset {self.name!r} has an annotation_split but no key spec") + return self + + def iter_files(self) -> list[Path]: + """Lists the dataset's JSONL files in a stable order. + + Returns: + list[Path]: Sorted matching files. Sorted so that file ids assigned during + the sidecar build stay the same across runs. + """ + return sorted(self.jsonl_root.glob(self.glob)) + + +class CorpusRegistry(BaseModel): + """The set of datasets a blend is built from. + + Attributes: + annotation_root (Optional[Path]): Directory holding the annotation parquet + splits, i.e. the directory whose children are the split paths named by + ``DatasetEntry.annotation_split``. + extra_annotation_roots (list[Path]): Further directories searched for splits. + Annotation shards are often spread over more than one cache. + datasets (list[DatasetEntry]): The datasets themselves. + """ + + annotation_root: Optional[Path] = None + extra_annotation_roots: list[Path] = Field(default_factory=list) + datasets: Annotated[list[DatasetEntry], Field(min_length=1)] + + @model_validator(mode="after") + def _check_unique_names(self) -> "CorpusRegistry": + seen = set() + for dataset in self.datasets: + if dataset.name in seen: + raise ValueError(f"duplicate dataset name {dataset.name!r} in registry") + seen.add(dataset.name) + return self + + @classmethod + def from_yaml(cls, path: Path) -> "CorpusRegistry": + """Loads a registry from a YAML file. + + Args: + path (Path): Path to the registry YAML. + + Returns: + CorpusRegistry: The parsed registry. + """ + with Path(path).open() as f: + return cls.model_validate(yaml.safe_load(f)) + + def get(self, name: str) -> DatasetEntry: + """Looks a dataset up by name. + + Args: + name (str): The dataset name. + + Returns: + DatasetEntry: The matching entry. + + Raises: + KeyError: If no dataset of that name is registered. + """ + for dataset in self.datasets: + if dataset.name == name: + return dataset + raise KeyError(f"no dataset named {name!r} in registry; known: {[d.name for d in self.datasets]}") + + def enabled_datasets(self) -> list[DatasetEntry]: + """Lists the datasets taking part in the blend. + + Returns: + list[DatasetEntry]: Entries whose ``enabled`` flag is set. + """ + return [d for d in self.datasets if d.enabled] + + def annotation_shards(self, split: str) -> list[Path]: + """Finds the parquet shards of an annotation split across all roots. + + Args: + split (str): Split path relative to an annotation root. + + Returns: + list[Path]: Sorted, de-duplicated shard paths. Empty if the split has not + been downloaded, which is not an error: a split can be registered + before its shards are fetched. + """ + shards: set[Path] = set() + for root in [self.annotation_root, *self.extra_annotation_roots]: + if root is None: + continue + shards.update(Path(root).glob(f"{split}/*.parquet")) + return sorted(shards) diff --git a/src/modalities/dataloader/preprocessing/quality/selection.py b/src/modalities/dataloader/preprocessing/quality/selection.py new file mode 100644 index 000000000..47b9741df --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/selection.py @@ -0,0 +1,1332 @@ +"""Turns a YAML selection into kept-document and kept-token figures. + +A selection states, per dataset, which documents to keep and how heavily to sample what +remains. Predicates address two kinds of signal with the same syntax: ordinal +annotation labels such as ``educational_value``, and continuous native metrics such as +``fw_edu_scores``. + +Every predicate can be evaluated two ways. Against a :class:`~...cube.Cube` it answers +in microseconds, which is what makes threshold tuning interactive. Against the +per-document sidecar it answers exactly but has to read the table. The two agree except +where a numeric threshold falls inside a cube bin rather than on its edge; the cube +evaluation detects that case and reports the result as approximate instead of pretending +otherwise. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Iterable, Optional + +import numpy as np +import pyarrow.parquet as pq +import yaml +from pydantic import BaseModel, Field, model_validator + +from modalities.dataloader.preprocessing.quality.cube import MISSING, Cube +from modalities.dataloader.preprocessing.quality.upsampling import ( + UNANNOTATED_BUCKET, + QualityBucket, + UpsamplingError, + UpsamplingPlan, + UpsamplingSpec, + solve_curve, +) + +# Ordinal scales, worst value first. Ordering is what gives `at_least` its meaning, so +# these are stated explicitly rather than inferred: `information_density` in particular +# orders `moderate` below `adequate`, which no alphabetical or intuitive rule reproduces. +ORDINAL_SCALES: dict[str, tuple[str, ...]] = { + "educational_value": ("none", "minimal", "basic", "moderate", "high"), + "content_quality": ("unacceptable", "poor", "adequate", "good", "excellent"), + "information_density": ("empty", "thin", "moderate", "adequate", "dense"), + "reasoning_indicators": ("none", "minimal", "basic_reasoning", "explanatory", "analytical"), + "content_integrity": ("severely_degraded", "fragment", "mostly_complete", "complete"), + "content_safety": ("illegal", "harmful", "nsfw", "mild_concerns", "safe"), + "pii_presence": ("contains_pii", "no_pii"), + "commercial_bias": ("pure_marketing", "heavy", "moderate", "minimal", "none"), + "content_ratio": ("minimal_content", "mostly_navigation", "mixed_content", "mostly_content", "complete_content"), + "content_length": ("minimal", "brief", "moderate", "substantial"), +} + + +class Op(str, Enum): + """Comparison a predicate applies. + + Attributes: + AT_LEAST: Ordinal label at or above a level. + AT_MOST: Ordinal label at or below a level. + IN: Label is one of a set. + NOT_IN: Label is none of a set. + GTE: Numeric metric greater than or equal to a value. + LTE: Numeric metric less than or equal to a value. + BETWEEN: Numeric metric within an inclusive range. + """ + + AT_LEAST = "at_least" + AT_MOST = "at_most" + IN = "in" + NOT_IN = "not_in" + GTE = "gte" + LTE = "lte" + BETWEEN = "between" + + +class MissingPolicy(str, Enum): + """What to do with documents a predicate cannot be evaluated on. + + Attributes: + KEEP: Treat the predicate as satisfied. Right when annotations are only + partly downloaded and dropping the unannotated majority would silently + shrink the dataset. + DROP: Treat the predicate as failed. Right when the filter is a hard + requirement and an unannotated document cannot be shown to meet it. + """ + + KEEP = "keep" + DROP = "drop" + + +class SelectionError(RuntimeError): + """Raised when a selection is malformed or cannot be evaluated.""" + + +class Predicate(BaseModel): + """One condition a document must satisfy. + + Attributes: + field (str): Annotation label name, or native metric name. Native metrics are + named without the ``native_`` prefix the sidecar stores them under. + op (Op): The comparison. + value (Optional[Any]): Right-hand side for the single-valued comparisons. + values (Optional[list[Any]]): Right-hand side for ``IN``/``NOT_IN``, or the + two bounds of ``BETWEEN``. + missing (Optional[MissingPolicy]): Overrides the selection-wide policy for + this predicate. + """ + + field: str + op: Op + value: Optional[Any] = None + values: Optional[list[Any]] = None + missing: Optional[MissingPolicy] = None + + @model_validator(mode="after") + def _check_operands(self) -> "Predicate": + if self.op in (Op.AT_LEAST, Op.AT_MOST): + if self.value is None: + raise ValueError(f"{self.op.value} on {self.field!r} needs a 'value'") + scale = ORDINAL_SCALES.get(self.field) + if scale is None: + raise ValueError( + f"{self.op.value} needs an ordinal field; {self.field!r} has no declared scale. " + f"Ordinal fields: {sorted(ORDINAL_SCALES)}" + ) + if self.value not in scale: + raise ValueError(f"{self.value!r} is not a level of {self.field!r}; levels are {list(scale)}") + elif self.op in (Op.IN, Op.NOT_IN): + if not self.values: + raise ValueError(f"{self.op.value} on {self.field!r} needs a non-empty 'values'") + elif self.op == Op.BETWEEN: + if not self.values or len(self.values) != 2: + raise ValueError(f"between on {self.field!r} needs 'values: [low, high]'") + elif self.value is None: + raise ValueError(f"{self.op.value} on {self.field!r} needs a 'value'") + return self + + @property + def is_numeric(self) -> bool: + """Whether this predicate compares a continuous metric. + + Returns: + bool: True for ``GTE``/``LTE``/``BETWEEN``. + """ + return self.op in (Op.GTE, Op.LTE, Op.BETWEEN) + + def allowed_levels(self) -> set[str]: + """The label values that satisfy this predicate. + + Returns: + set[str]: Satisfying values for a categorical or ordinal predicate. + + Raises: + SelectionError: If called on a numeric predicate. + """ + if self.is_numeric: + raise SelectionError(f"predicate on {self.field!r} is numeric and has no level set") + if self.op == Op.AT_LEAST: + scale = ORDINAL_SCALES[self.field] + return set(scale[scale.index(self.value) :]) + if self.op == Op.AT_MOST: + scale = ORDINAL_SCALES[self.field] + return set(scale[: scale.index(self.value) + 1]) + if self.op == Op.IN: + return {str(v) for v in self.values} + scale = ORDINAL_SCALES.get(self.field) + if scale is None: + raise SelectionError( + f"not_in on {self.field!r} needs a declared value set; use 'in' with the values to keep instead" + ) + return set(scale) - {str(v) for v in self.values} + + def matches_value(self, value: Any, missing_policy: MissingPolicy) -> bool: + """Evaluates the predicate against one document's value. + + Args: + value (Any): The document's value for ``field``; None if absent. + missing_policy (MissingPolicy): Fallback policy for this selection. + + Returns: + bool: Whether the document satisfies the predicate. + """ + policy = self.missing or missing_policy + if value is None or value == MISSING or (isinstance(value, float) and np.isnan(value)): + return policy == MissingPolicy.KEEP + if self.is_numeric: + numeric = float(value) + if self.op == Op.GTE: + return numeric >= float(self.value) + if self.op == Op.LTE: + return numeric <= float(self.value) + return float(self.values[0]) <= numeric <= float(self.values[1]) + return str(value) in self.allowed_levels() + + def describe(self) -> str: + """Renders the predicate for reports. + + Returns: + str: A compact, readable form of the condition. + """ + if self.op == Op.BETWEEN: + return f"{self.field} in [{self.values[0]}, {self.values[1]}]" + if self.op in (Op.IN, Op.NOT_IN): + return f"{self.field} {self.op.value} {{{', '.join(str(v) for v in self.values)}}}" + return f"{self.field} {self.op.value} {self.value}" + + +class DatasetSelection(BaseModel): + """The rule applied to one dataset. + + Attributes: + name (str): Dataset name, matching the corpus registry. + ratio (float): Up/downsample factor applied after filtering. 1.0 uses the + surviving documents once; 2.5 draws them two and a half times; 0.3 keeps + three tenths of them. + predicates (list[Predicate]): Conditions combined with AND. An empty list + keeps every document, which is how an unannotated dataset participates. + missing_annotation (Optional[MissingPolicy]): Overrides the config-wide policy. + upsampling (Optional[UpsamplingSpec]): Replaces ``ratio`` with a quality-aware + curve, so the repeat factor rises with quality instead of being one number for + every surviving document. Mutually exclusive with a non-default ``ratio``. + enabled (bool): Whether this dataset takes part. + """ + + name: str + ratio: float = Field(default=1.0, ge=0.0) + predicates: list[Predicate] = Field(default_factory=list) + missing_annotation: Optional[MissingPolicy] = None + upsampling: Optional[UpsamplingSpec] = None + enabled: bool = True + + @model_validator(mode="after") + def _check_ratio_or_curve(self) -> "DatasetSelection": + # Silently ignoring one of them would make a config mean something other than it + # reads, and the two express the same decision. + if self.upsampling is not None and self.ratio != 1.0: + raise ValueError( + f"dataset {self.name!r} sets both 'ratio: {self.ratio}' and 'upsampling'; the curve " + f"already determines how much is drawn. Remove the ratio." + ) + if self.upsampling is not None and self.upsampling.quality_field not in ORDINAL_SCALES: + # Checked at config load rather than mid-run: a numeric axis would need quantile + # edges carried from the cube into materialisation, which is not built yet, and + # discovering that after the preview succeeded would be worse than refusing now. + raise ValueError( + f"dataset {self.name!r}: upsampling needs an ordinal quality_field, and " + f"{self.upsampling.quality_field!r} is not one. Available: {sorted(ORDINAL_SCALES)}" + ) + return self + + +class SelectionConfig(BaseModel): + """A complete blend specification. + + Attributes: + missing_annotation (MissingPolicy): Default policy for documents that carry no + annotation. + target_tokens (Optional[float]): Tokens the training run will consume. It does not + change any ratio, but it is what makes a ratio mean something: if the blend + yields fewer effective tokens than this, the loader wraps and every document is + seen more often than its ratio says. + seed (int): Decides which documents receive the extra copy of a fractional repeat + factor when the blend is exported. Changing it redraws that choice, so a blend + exported twice with different seeds differs in which documents were doubled. + max_total_exposure (Optional[float]): Refuse to materialise if any dataset -- or any + quality bucket of one -- would be seen more times than this once wrapping is + counted. A dataset with an upsampling curve uses its own ``max_factor`` instead, + since that is already a declared cap. + datasets (list[DatasetSelection]): Per-dataset rules. + """ + + missing_annotation: MissingPolicy = MissingPolicy.KEEP + target_tokens: Optional[float] = None + max_total_exposure: Optional[float] = Field(default=None, gt=0) + seed: int = 42 + datasets: list[DatasetSelection] + + @model_validator(mode="after") + def _check_unique(self) -> "SelectionConfig": + seen = set() + for dataset in self.datasets: + if dataset.name in seen: + raise ValueError(f"dataset {dataset.name!r} appears twice in the selection") + seen.add(dataset.name) + return self + + @classmethod + def from_yaml(cls, path: Path) -> "SelectionConfig": + """Loads a selection from YAML. + + Args: + path (Path): Path to the selection file. + + Returns: + SelectionConfig: The parsed selection. + """ + with Path(path).open() as f: + return cls.model_validate(yaml.safe_load(f)) + + def policy_for(self, dataset: DatasetSelection) -> MissingPolicy: + """Resolves the missing-annotation policy for one dataset. + + Args: + dataset (DatasetSelection): The dataset rule. + + Returns: + MissingPolicy: The dataset's own policy if set, else the config default. + """ + return dataset.missing_annotation or self.missing_annotation + + def enabled_datasets(self) -> list[DatasetSelection]: + """Lists participating datasets. + + Returns: + list[DatasetSelection]: Rules whose ``enabled`` flag is set. + """ + return [d for d in self.datasets if d.enabled] + + +@dataclass +class DatasetResult: + """What a selection costs for one dataset. + + Attributes: + name (str): Dataset name. + n_documents_total (int): Documents before filtering. + n_documents_kept (int): Documents surviving the predicates. + tokens_total (int): Estimated tokens before filtering. + tokens_kept (int): Estimated tokens surviving the predicates. + ratio (float): Up/downsample factor applied afterwards. + exact (bool): Whether the figures are exact. False when a numeric threshold + fell inside a cube bin, so the count had to be interpolated. + approximations (list[str]): Predicates that forced interpolation. + plan (Optional[UpsamplingPlan]): The solved quality curve, when the dataset uses one + instead of a flat ratio. + """ + + name: str + n_documents_total: int + n_documents_kept: int + tokens_total: int + tokens_kept: int + ratio: float + exact: bool = True + approximations: list[str] = field(default_factory=list) + plan: Optional[UpsamplingPlan] = None + + @property + def effective_tokens(self) -> float: + """Tokens the blend draws from this dataset. + + Returns: + float: Kept tokens scaled by the ratio, or what the curve draws. + """ + if self.plan is not None: + return self.plan.tokens_drawn + return self.tokens_kept * self.ratio + + @property + def ratio_label(self) -> str: + """How the up/downsampling is described in a report. + + Returns: + str: The flat ratio, or the curve's range of factors. + """ + if self.plan is None: + return f"{self.ratio:.2f}" + factors = [b.factor for b in self.plan.buckets if b.factor > 0] + if not factors: + return "curve" + return f"{min(factors):.2f}-{max(factors):.2f}x" + + @property + def row_retention(self) -> float: + """Share of documents kept. + + Returns: + float: Kept over total documents; 0.0 for an empty dataset. + """ + return self.n_documents_kept / self.n_documents_total if self.n_documents_total else 0.0 + + @property + def token_retention(self) -> float: + """Share of tokens kept. + + Returns: + float: Kept over total tokens; 0.0 for an empty dataset. + + Note: + This is normally higher than :attr:`row_retention`, because quality + correlates with length and quality filters keep the longer documents. + """ + return self.tokens_kept / self.tokens_total if self.tokens_total else 0.0 + + +@dataclass +class BlendResult: + """What a selection costs across the whole blend. + + Attributes: + datasets (list[DatasetResult]): Per-dataset outcomes. + target_tokens (Optional[float]): Budget the blend aimed at, if any. + """ + + datasets: list[DatasetResult] + target_tokens: Optional[float] = None + + @property + def total_effective_tokens(self) -> float: + """Tokens the blend yields in total. + + Returns: + float: Sum of the per-dataset effective token counts. + """ + return sum(d.effective_tokens for d in self.datasets) + + @property + def exact(self) -> bool: + """Whether every dataset's figures are exact. + + Returns: + bool: True if no dataset needed interpolation. + """ + return all(d.exact for d in self.datasets) + + def share_of(self, dataset: DatasetResult) -> float: + """Share of the blend one dataset contributes. + + Args: + dataset (DatasetResult): The dataset outcome. + + Returns: + float: Its effective tokens over the blend total; 0.0 for an empty blend. + """ + total = self.total_effective_tokens + return dataset.effective_tokens / total if total else 0.0 + + +def _bin_fraction_above(binning, threshold: float, bin_index: int) -> float: + # Fraction of a bin lying at or above a threshold, assuming values are spread + # evenly inside the bin. Only reached when the threshold splits a bin; on a bin + # edge the caller takes the whole bin or none of it and stays exact. + low = binning.lower_bound_of(bin_index) + high = binning.upper_bound_of(bin_index) + if not np.isfinite(low) or not np.isfinite(high) or high <= low: + return 0.5 + return float(np.clip((high - threshold) / (high - low), 0.0, 1.0)) + + +def quality_buckets_from_cube( + cube: Cube, quality_field: str, weight: np.ndarray +) -> list[QualityBucket]: + """Groups a cube's surviving cells into buckets ordered worst to best quality. + + Unannotated documents cannot be placed on a quality axis, so they form their own bucket + at the bottom of the order. Putting them there is a choice worth knowing about: it means + a dataset whose annotation coverage is partial will see its unannotated majority treated + as lowest quality, and discarded first. The report names the bucket explicitly so this is + visible rather than implied. + + Args: + cube (Cube): The dataset's cube. + quality_field (str): Ordinal label or native metric to order by. + weight (np.ndarray): Per-cell survival weight from the predicates. + + Returns: + list[QualityBucket]: Non-empty buckets, worst quality first. + + Raises: + SelectionError: If the cube was not grouped on the field. + """ + table = cube.table + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) * weight + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) * weight + + if quality_field not in cube.label_dimensions: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on {quality_field!r}, so it cannot order " + f"documents by it. Grouped labels: {cube.label_dimensions}. Rebuild the cube with " + f"--label_dimension {quality_field}." + ) + scale = ordered_quality_levels(quality_field) + rank = {level: i for i, level in enumerate(scale)} + values = table.column(quality_field).to_pylist() + keys = np.array( + [-1 if (v is None or v == MISSING) else rank.get(v, -1) for v in values], dtype=np.int64 + ) + order = sorted({int(k) for k in keys}) + labels = {k: (UNANNOTATED_BUCKET if k < 0 else scale[k]) for k in order} + + buckets: list[QualityBucket] = [] + for key in order: + mask = keys == key + n_tokens = int(round(float(tokens[mask].sum()))) + if n_tokens <= 0: + continue + buckets.append( + QualityBucket( + label=labels[key], + n_documents=int(round(float(documents[mask].sum()))), + n_tokens=n_tokens, + unannotated=key < 0, + ) + ) + return buckets + + +def ordered_bucket_labels(quality_field: str) -> list[str]: + """Bucket labels for a quality field, worst first, unannotated at the bottom. + + Shared by the cube path, the exact sidecar path and materialisation: three places that + must agree on the axis or a curve would mean different things depending on how it was + evaluated. + + Args: + quality_field (str): An ordinal annotation field. + + Returns: + list[str]: ``UNANNOTATED_BUCKET`` followed by the field's levels in ascending order. + """ + return [UNANNOTATED_BUCKET, *ordered_quality_levels(quality_field)] + + +def ordered_quality_levels(quality_field: str) -> tuple[str, ...]: + """Lists a field's levels from worst to best quality. + + Args: + quality_field (str): An ordinal annotation field. + + Returns: + tuple[str, ...]: Its declared levels in ascending order. + + Raises: + SelectionError: If the field has no declared ordinal scale, so its levels cannot be + ordered and no curve over them would mean anything. + """ + scale = ORDINAL_SCALES.get(quality_field) + if scale is None: + raise SelectionError( + f"{quality_field!r} has no declared ordinal scale, so its levels cannot be ordered worst " + f"to best. Ordinal fields: {sorted(ORDINAL_SCALES)}" + ) + return scale + + +def _predicate_factor( + cube: Cube, predicate: Predicate, policy: MissingPolicy +) -> tuple[np.ndarray, bool]: + """Computes one predicate's surviving fraction for every cube cell. + + Extracted so a predicate can be costed on its own, which is what attributing retention + to individual predicates needs, and not only as one term of a product. + + Args: + cube (Cube): The dataset's cube. + predicate (Predicate): The condition to apply. + policy (MissingPolicy): Policy for cells with no annotation. + + Returns: + tuple[np.ndarray, bool]: Per-cell factor in [0, 1], and whether this predicate was + answered exactly rather than interpolated inside a bin. + + Raises: + SelectionError: If the cube was not grouped on the predicate's field. + """ + table = cube.table + n_rows = table.num_rows + exact = True + + if predicate.is_numeric: + binning = cube.score_binnings.get(predicate.field) + if binning is None: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on native metric {predicate.field!r}; " + f"grouped metrics: {sorted(cube.score_binnings)}. Re-run with --exact to scan the sidecar." + ) + bins = table.column(f"native_{predicate.field}").to_numpy(zero_copy_only=False).astype(np.int64) + factor = np.empty(n_rows, dtype=np.float64) + for i, bin_index in enumerate(bins): + if bin_index < 0: + factor[i] = 1.0 if policy == MissingPolicy.KEEP else 0.0 + continue + low = binning.lower_bound_of(bin_index) + high = binning.upper_bound_of(bin_index) + if predicate.op == Op.GTE: + threshold = float(predicate.value) + if low >= threshold: + factor[i] = 1.0 + elif high <= threshold: + factor[i] = 0.0 + else: + factor[i] = _bin_fraction_above(binning, threshold, bin_index) + elif predicate.op == Op.LTE: + threshold = float(predicate.value) + if high <= threshold: + factor[i] = 1.0 + elif low >= threshold: + factor[i] = 0.0 + else: + factor[i] = 1.0 - _bin_fraction_above(binning, threshold, bin_index) + else: + lower, upper = float(predicate.values[0]), float(predicate.values[1]) + if low >= lower and high <= upper: + factor[i] = 1.0 + elif high <= lower or low >= upper: + factor[i] = 0.0 + else: + factor[i] = max( + 0.0, + _bin_fraction_above(binning, lower, bin_index) + - _bin_fraction_above(binning, upper, bin_index), + ) + if 0.0 < factor[i] < 1.0: + exact = False + return factor, exact + + if predicate.field not in cube.label_dimensions: + raise SelectionError( + f"cube for {cube.dataset!r} was not grouped on label {predicate.field!r}; " + f"grouped labels: {cube.label_dimensions}. Re-run with --exact to scan the sidecar." + ) + allowed = predicate.allowed_levels() + values = table.column(predicate.field).to_pylist() + keep_missing = policy == MissingPolicy.KEEP + factor = np.array( + [ + (1.0 if keep_missing else 0.0) if v is None or v == MISSING else (1.0 if v in allowed else 0.0) + for v in values + ], + dtype=np.float64, + ) + return factor, True + + +def _cube_weights( + cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy +) -> tuple[np.ndarray, bool, list[str]]: + """Computes each cube cell's surviving fraction under a dataset's predicates. + + Args: + cube (Cube): The dataset's cube. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + tuple[np.ndarray, bool, list[str]]: Per-cell weight in [0, 1], whether every + predicate was answered exactly, and the predicates that had to be interpolated. + + Raises: + SelectionError: If a predicate names a field the cube was not grouped on. + """ + weight = np.ones(cube.table.num_rows, dtype=np.float64) + result_exact = True + approximations: list[str] = [] + + for predicate in dataset.predicates: + factor, exact = _predicate_factor(cube, predicate, predicate.missing or missing_policy) + # Per predicate, not cumulative. The previous form tested the running flag, so once + # any predicate was interpolated every later one was reported as interpolated too. + if not exact: + result_exact = False + if predicate.describe() not in approximations: + approximations.append(predicate.describe()) + weight *= factor + + return weight, result_exact, approximations + + +@dataclass +class PredicateContribution: + """What one predicate does to a dataset, on its own and in company. + + Attributes: + description (str): The predicate, as written. + matched_tokens (int): Tokens satisfying this predicate alone, ignoring the others. + matched_documents (int): Documents satisfying it alone. + marginal_tokens (int): Extra tokens the dataset would keep if this predicate were + removed and the others left in place. This is the number that matters when + tuning: a predicate whose marginal is zero is being fully shadowed by its + neighbours and is only making the selection harder to read. + exact (bool): Whether the predicate was answered exactly rather than interpolated. + """ + + description: str + matched_tokens: int + matched_documents: int + marginal_tokens: int + exact: bool = True + + +@dataclass +class PredicateBreakdown: + """Per-predicate attribution for one dataset, plus how the predicates overlap. + + Attributes: + dataset (str): Dataset name. + total_tokens (int): Tokens before any predicate. + kept_tokens (int): Tokens surviving all of them. + contributions (list[PredicateContribution]): One per predicate, in config order. + overlap_tokens (list[list[int]]): Symmetric matrix of tokens satisfying both + predicates i and j; the diagonal is each predicate's own ``matched_tokens``. + Where either predicate was interpolated, an off-diagonal cell multiplies two + fractional survival rates and so assumes they are independent within a cube + cell, which is the same assumption the combined weight already makes. + """ + + dataset: str + total_tokens: int + kept_tokens: int + contributions: list[PredicateContribution] + overlap_tokens: list[list[int]] + + def describe(self) -> str: + """Renders the attribution as a small table. + + Returns: + str: One row per predicate, plus an overlap matrix when there are at least two. + """ + if not self.contributions: + return f" {self.dataset}: no predicates" + + def share(n: int) -> str: + return f"{n / self.total_tokens:>6.1%}" if self.total_tokens else " -" + + width = max(len(c.description) for c in self.contributions) + lines = [ + f" {self.dataset}: {self.kept_tokens:,} of {self.total_tokens:,} tokens kept", + f" {'predicate':<{width}} {'matches':>14} {'share':>7} {'marginal':>14}", + ] + for contribution in self.contributions: + marker = "" if contribution.exact else " ~" + lines.append( + f" {contribution.description:<{width}} {contribution.matched_tokens:>14,} " + f"{share(contribution.matched_tokens)} {contribution.marginal_tokens:>14,}{marker}" + ) + redundant = [c.description for c in self.contributions if c.marginal_tokens == 0] + if redundant and len(self.contributions) > 1: + lines.append(f" no effect given the others: {', '.join(redundant)}") + + if len(self.contributions) >= 2: + lines.append(f" overlap (tokens matching both), {len(self.contributions)} predicates:") + labels = [f"P{i + 1}" for i in range(len(self.contributions))] + lines.append(" " + " ".join(f"{label:>14}" for label in [""] + labels)) + for i, label in enumerate(labels): + cells = " ".join(f"{self.overlap_tokens[i][j]:>14,}" for j in range(len(labels))) + lines.append(f" {label:>14} {cells}") + for i, contribution in enumerate(self.contributions): + lines.append(f" P{i + 1} = {contribution.description}") + return "\n".join(lines) + + +def predicate_breakdown( + cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy +) -> PredicateBreakdown: + """Attributes a dataset's retention to its individual predicates. + + The per-dataset retention a preview reports says nothing about which condition caused + it. With two or three predicates per dataset, the interesting questions are which one + binds, which is redundant, and how much they overlap -- and all of them are answerable + from the cube in milliseconds, because the cell weights are already the whole + computation. + + Args: + cube (Cube): The dataset's cube. + dataset (DatasetSelection): The rule to attribute. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + PredicateBreakdown: Per-predicate matches, marginal effect, and pairwise overlap. + + Raises: + SelectionError: If a predicate names a field the cube was not grouped on. + """ + table = cube.table + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) + + factors: list[np.ndarray] = [] + exactness: list[bool] = [] + for predicate in dataset.predicates: + factor, exact = _predicate_factor(cube, predicate, predicate.missing or missing_policy) + factors.append(factor) + exactness.append(exact) + + combined = np.ones(table.num_rows, dtype=np.float64) + for factor in factors: + combined *= factor + kept = float((tokens * combined).sum()) + + contributions: list[PredicateContribution] = [] + for i, predicate in enumerate(dataset.predicates): + # What the dataset would keep with predicate i lifted, everything else in place. + without = np.ones(table.num_rows, dtype=np.float64) + for j, factor in enumerate(factors): + if j != i: + without *= factor + contributions.append( + PredicateContribution( + description=predicate.describe(), + matched_tokens=int(round(float((tokens * factors[i]).sum()))), + matched_documents=int(round(float((documents * factors[i]).sum()))), + marginal_tokens=int(round(float((tokens * without).sum()) - kept)), + exact=exactness[i], + ) + ) + + n = len(factors) + overlap: list[list[int]] = [] + for i in range(n): + row: list[int] = [] + for j in range(n): + if i == j: + # Not the product with itself: an interpolated predicate has fractional + # factors, and squaring them undercounts. On real data this read 8.19 M + # where the predicate matched 9.75 M. + row.append(contributions[i].matched_tokens) + else: + row.append(int(round(float((tokens * factors[i] * factors[j]).sum())))) + overlap.append(row) + + return PredicateBreakdown( + dataset=dataset.name, + total_tokens=int(round(float(tokens.sum()))), + kept_tokens=int(round(kept)), + contributions=contributions, + overlap_tokens=overlap, + ) + + +def evaluate_on_cube(cube: Cube, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: + """Evaluates a dataset's rule against its cube. + + Args: + cube (Cube): The dataset's cube. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + DatasetResult: Kept documents and tokens, flagged as exact or interpolated, plus the + solved curve when the dataset uses one. + + Raises: + SelectionError: If a predicate names a field the cube was not grouped on, or if a + curve cannot be solved for the dataset. + """ + weight, result_exact, approximations = _cube_weights(cube, dataset, missing_policy) + table = cube.table + documents = table.column("n_documents").to_numpy(zero_copy_only=False).astype(np.float64) + tokens = table.column("n_tokens").to_numpy(zero_copy_only=False).astype(np.float64) + + plan: Optional[UpsamplingPlan] = None + if dataset.upsampling is not None: + buckets = quality_buckets_from_cube(cube, dataset.upsampling.quality_field, weight) + try: + plan = solve_curve(buckets, dataset.upsampling) + except UpsamplingError as e: + raise SelectionError(f"dataset {dataset.name!r}: {e}") from e + + return DatasetResult( + name=dataset.name, + n_documents_total=int(documents.sum()), + n_documents_kept=plan.documents_kept if plan else int(round(float((documents * weight).sum()))), + tokens_total=int(tokens.sum()), + tokens_kept=int(round(float((tokens * weight).sum()))), + ratio=dataset.ratio, + exact=result_exact, + approximations=approximations, + plan=plan, + ) + + +def document_mask(table, dataset: DatasetSelection, missing_policy: MissingPolicy) -> np.ndarray: + """Evaluates a dataset's rule against per-document sidecar rows. + + Args: + table: A pyarrow table of sidecar rows. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + np.ndarray: Boolean mask, True where the document is kept. + + Raises: + SelectionError: If a predicate names a column the sidecar does not have. + """ + n_rows = table.num_rows + mask = np.ones(n_rows, dtype=bool) + names = set(table.schema.names) + + for predicate in dataset.predicates: + column = f"native_{predicate.field}" if predicate.is_numeric else predicate.field + if column not in names: + raise SelectionError( + f"sidecar has no column {column!r} for predicate on {predicate.field!r}; available: {sorted(names)}" + ) + policy = predicate.missing or missing_policy + if predicate.is_numeric: + values = table.column(column).to_numpy(zero_copy_only=False).astype(np.float64) + missing = np.isnan(values) + with np.errstate(invalid="ignore"): + if predicate.op == Op.GTE: + ok = values >= float(predicate.value) + elif predicate.op == Op.LTE: + ok = values <= float(predicate.value) + else: + ok = (values >= float(predicate.values[0])) & (values <= float(predicate.values[1])) + ok = np.where(missing, policy == MissingPolicy.KEEP, ok) + else: + allowed = predicate.allowed_levels() + keep_missing = policy == MissingPolicy.KEEP + ok = np.array( + [ + keep_missing if v is None or v == MISSING else (v in allowed) + for v in table.column(column).to_pylist() + ], + dtype=bool, + ) + mask &= ok + return mask + + +def evaluate_on_sidecar(sidecar_dir: Path, dataset: DatasetSelection, missing_policy: MissingPolicy) -> DatasetResult: + """Evaluates a dataset's rule exactly, by scanning its sidecar. + + Args: + sidecar_dir (Path): Directory of sidecar parts. + dataset (DatasetSelection): The rule to apply. + missing_policy (MissingPolicy): Policy for unannotated documents. + + Returns: + DatasetResult: Kept documents and tokens, always exact. + + Raises: + SelectionError: If the directory holds no sidecar parts. + """ + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + if not parts: + raise SelectionError(f"no sidecar parts found in {sidecar_dir}") + + spec = dataset.upsampling + # A curved dataset must be costed as a curve here too. Returning the flat ratio would + # report it at 1.0x -- the value the config validator forces when a curve is present -- + # so an --exact preview of a curved selection would silently understate its tokens, + # blend share and exposure. + bucket_labels = ordered_bucket_labels(spec.quality_field) if spec else [] + bucket_docs: dict[str, int] = dict.fromkeys(bucket_labels, 0) + bucket_tokens: dict[str, int] = dict.fromkeys(bucket_labels, 0) + + n_total = n_kept = tokens_total = tokens_kept = 0 + for part in parts: + parquet_file = pq.ParquetFile(part) + if spec and spec.quality_field not in parquet_file.schema_arrow.names: + raise SelectionError( + f"dataset {dataset.name!r}: sidecar has no column {spec.quality_field!r} to order " + f"quality by; join the annotations before costing a curve" + ) + for group_idx in range(parquet_file.metadata.num_row_groups): + table = parquet_file.read_row_group(group_idx) + tokens = table.column("est_tokens").to_numpy(zero_copy_only=False).astype(np.int64) + mask = document_mask(table, dataset, missing_policy) + n_total += table.num_rows + n_kept += int(mask.sum()) + tokens_total += int(tokens.sum()) + tokens_kept += int(tokens[mask].sum()) + + if spec: + levels = table.column(spec.quality_field).to_pylist() + for level, keep, token_count in zip(levels, mask, tokens): + if not keep: + continue + label = level if level in bucket_docs else UNANNOTATED_BUCKET + bucket_docs[label] += 1 + bucket_tokens[label] += int(token_count) + + plan: Optional[UpsamplingPlan] = None + if spec: + buckets = [ + QualityBucket( + label=label, + n_documents=bucket_docs[label], + n_tokens=bucket_tokens[label], + unannotated=label == UNANNOTATED_BUCKET, + ) + for label in bucket_labels + if bucket_tokens[label] > 0 + ] + try: + plan = solve_curve(buckets, spec) + except UpsamplingError as e: + raise SelectionError(f"dataset {dataset.name!r}: {e}") from e + n_kept = plan.documents_kept + + return DatasetResult( + name=dataset.name, + n_documents_total=n_total, + n_documents_kept=n_kept, + tokens_total=tokens_total, + tokens_kept=tokens_kept, + ratio=dataset.ratio, + exact=True, + plan=plan, + ) + + +def evaluate_blend( + config: SelectionConfig, + cubes: dict[str, Cube], + sidecar_dirs: Optional[dict[str, Path]] = None, + force_exact: bool = False, + allow_sidecar_fallback: bool = False, +) -> BlendResult: + """Evaluates a whole selection. + + Args: + config (SelectionConfig): The blend specification. + cubes (dict[str, Cube]): Cube per dataset name. + sidecar_dirs (Optional[dict[str, Path]]): Sidecar directory per dataset, used when + a cube cannot answer a predicate or when exactness is demanded. + force_exact (bool): Scan sidecars for every dataset instead of using cubes. + allow_sidecar_fallback (bool): Permit scanning a sidecar for datasets whose cube + cannot answer a predicate. Off by default, because that scan is thousands of + times more expensive than a cube lookup and a `preview` is meant to return in + seconds: a selection thresholding one ungrouped field turned a preview of this + blend into a read over 1.7 billion documents. + + Returns: + BlendResult: Per-dataset and total figures. + + Raises: + SelectionError: If a dataset can be evaluated neither from a cube nor from a + sidecar, or if a predicate needs a sidecar scan that was not permitted. Every + such dataset is reported together, so one run tells you everything to fix. + """ + results: list[DatasetResult] = [] + unanswerable: list[str] = [] + + for dataset in config.enabled_datasets(): + policy = config.policy_for(dataset) + sidecar_dir = (sidecar_dirs or {}).get(dataset.name) + + if force_exact: + if sidecar_dir is None: + raise SelectionError(f"exact evaluation of {dataset.name!r} needs a sidecar directory") + results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + continue + + cube = cubes.get(dataset.name) + if cube is None: + if sidecar_dir is None: + raise SelectionError(f"no cube and no sidecar for dataset {dataset.name!r}") + if not allow_sidecar_fallback: + unanswerable.append(f" {dataset.name}: no cube was built; run 'quality build-cube'") + continue + results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + continue + + try: + results.append(evaluate_on_cube(cube, dataset, policy)) + except SelectionError as e: + if sidecar_dir is None: + raise + if not allow_sidecar_fallback: + unanswerable.append( + f" {dataset.name}: {e} " + f"(answering it from the sidecar means reading {cube.n_documents:,} documents)" + ) + continue + results.append(evaluate_on_sidecar(sidecar_dir, dataset, policy)) + + if unanswerable: + raise SelectionError( + "these predicates cannot be answered from the cubes:\n" + + "\n".join(unanswerable) + + "\n\nEither drop or replace the offending predicate, rebuild the cube with that field as a " + "dimension (build-cube --label_dimension ...), or accept the cost with --allow-fallback. " + "A sidecar scan is exact but reads every document, so it takes minutes to hours rather than seconds." + ) + + return BlendResult(datasets=results, target_tokens=config.target_tokens) + + +@dataclass +class ExposureRow: + """How often one dataset, or one quality bucket, is actually seen. + + Attributes: + label (str): Dataset name, or ``dataset / bucket``. + factor (float): The repeat factor the selection asked for. + passes (float): Times the run traverses the whole blend. + cap (Optional[float]): Declared limit, if any. + """ + + label: str + factor: float + passes: float + cap: Optional[float] = None + + @property + def exposure(self) -> float: + """Times each document here is actually seen during the run. + + Returns: + float: The requested factor multiplied by the number of passes. + """ + return self.factor * self.passes + + @property + def exceeded(self) -> bool: + """Whether the declared cap is broken once wrapping is counted. + + Returns: + bool: True if a cap exists and the exposure is above it. + """ + return self.cap is not None and self.exposure > self.cap * (1.0 + 1e-9) + + +@dataclass +class ExposureReport: + """What the run really repeats, as opposed to what the ratios say. + + A ratio is relative to one pass over the blend. If the run consumes more tokens than the + blend yields, the loader comes round again and every factor is multiplied by the number + of passes -- so a bucket set to 7x under a curve becomes 14x on a second pass. That is + well past the point where repetition stops paying, and nothing in the pipeline noticed + it before: ``target_tokens`` was only ever printed. + + Attributes: + training_tokens (Optional[float]): Tokens the run consumes, if declared. + effective_tokens (float): Tokens one pass over the blend yields. + rows (list[ExposureRow]): One per dataset, or per bucket for curved datasets. + """ + + training_tokens: Optional[float] + effective_tokens: float + rows: list[ExposureRow] + + @property + def passes(self) -> float: + """Times the run traverses the blend. + + Returns: + float: Training tokens over effective tokens; 1.0 when nothing was declared. + """ + if not self.training_tokens or self.effective_tokens <= 0: + return 1.0 + return self.training_tokens / self.effective_tokens + + @property + def exceeded(self) -> list[ExposureRow]: + """Rows whose exposure breaks their declared cap. + + Returns: + list[ExposureRow]: Offending rows, worst first. + """ + return sorted( + (row for row in self.rows if row.exceeded), key=lambda r: r.exposure, reverse=True + ) + + def describe(self) -> str: + """Renders the exposure accounting. + + Returns: + str: A summary line, then the rows that repeat most, then any cap breaches. + """ + lines = [] + if self.training_tokens: + verb = "wraps" if self.passes > 1.0 else "uses" + lines.append( + f" run consumes {self.training_tokens:,.0f} tokens from {self.effective_tokens:,.0f} " + f"effective -- {verb} {self.passes:.2f} passes over the blend" + ) + else: + lines.append( + " no target_tokens declared, so repetition beyond one pass cannot be checked; " + "set it to the tokens the run will consume" + ) + top = sorted(self.rows, key=lambda r: r.exposure, reverse=True)[:6] + if top: + width = max(len(row.label) for row in top) + lines.append(f" {'':<{width}} {'factor':>8} {'exposure':>9} {'cap':>7}") + for row in top: + cap = f"{row.cap:g}" if row.cap is not None else "-" + flag = " OVER" if row.exceeded else "" + lines.append( + f" {row.label:<{width}} {row.factor:>8.2f} {row.exposure:>9.2f} {cap:>7}{flag}" + ) + for row in self.exceeded: + lines.append( + f" {row.label} is seen {row.exposure:.2f}x against a declared cap of {row.cap:g}x" + ) + return "\n".join(lines) + + +def exposure_report( + entries: list[tuple[str, float, Optional[float]]], + effective_tokens: float, + training_tokens: Optional[float], +) -> ExposureReport: + """Builds the exposure accounting for a blend. + + Args: + entries (list[tuple[str, float, Optional[float]]]): One ``(label, factor, cap)`` per + dataset, or per quality bucket where a curve splits one. + effective_tokens (float): Tokens one pass over the blend yields. + training_tokens (Optional[float]): Tokens the run consumes. + + Returns: + ExposureReport: The rows and the derived number of passes. + """ + report = ExposureReport( + training_tokens=training_tokens, effective_tokens=effective_tokens, rows=[] + ) + passes = report.passes + report.rows = [ + ExposureRow(label=label, factor=factor, passes=passes, cap=cap) for label, factor, cap in entries + ] + return report + + +def exposure_entries_from_blend( + result: "BlendResult", config: SelectionConfig +) -> list[tuple[str, float, Optional[float]]]: + """Collects exposure entries from an evaluated blend. + + Args: + result (BlendResult): The evaluated blend. + config (SelectionConfig): The selection, supplying caps. + + Returns: + list[tuple[str, float, Optional[float]]]: One entry per dataset, or per bucket for a + dataset with a curve. + """ + by_name = {dataset.name: dataset for dataset in config.datasets} + entries: list[tuple[str, float, Optional[float]]] = [] + for dataset_result in result.datasets: + selection = by_name.get(dataset_result.name) + spec = selection.upsampling if selection else None + if dataset_result.plan is not None: + for bucket in dataset_result.plan.buckets: + if bucket.factor > 0: + entries.append( + ( + f"{dataset_result.name} / {bucket.bucket.label}", + bucket.factor, + spec.max_factor if spec else None, + ) + ) + else: + entries.append((dataset_result.name, dataset_result.ratio, config.max_total_exposure)) + return entries + + +def format_blend_report( + result: BlendResult, + datasets_in_order: Optional[Iterable[str]] = None, + config: Optional[SelectionConfig] = None, +) -> str: + """Renders a blend result as a fixed-width table. + + Args: + result (BlendResult): The evaluated blend. + datasets_in_order (Optional[Iterable[str]]): Preferred row order by name. + config (Optional[SelectionConfig]): The selection. Supplies the repetition caps, so + the report can say how often documents are really seen; omitted, that block is + left out. + + Returns: + str: A table with per-dataset retention, ratio, effective tokens and share, + followed by the blend total and any accuracy caveats. + """ + + def humanise(n: float) -> str: + for unit, size in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("k", 1e3)): + if abs(n) >= size: + return f"{n / size:.2f}{unit}" + return f"{n:.0f}" + + rows = list(result.datasets) + if datasets_in_order is not None: + order = {name: i for i, name in enumerate(datasets_in_order)} + rows.sort(key=lambda d: order.get(d.name, len(order))) + + width = max([len(d.name) for d in rows] + [len("dataset")]) + header = ( + f"{'dataset':<{width}} {'docs kept':>11} {'row%':>6} {'tokens kept':>12} " + f"{'tok%':>6} {'ratio':>11} {'effective':>12} {'share':>6}" + ) + lines = [header, "-" * len(header)] + for d in rows: + marker = "" if d.exact else " ~" + lines.append( + f"{d.name:<{width}} {humanise(d.n_documents_kept):>11} {d.row_retention:>5.1%} " + f"{humanise(d.tokens_kept):>12} {d.token_retention:>5.1%} {d.ratio_label:>11} " + f"{humanise(d.effective_tokens):>12}{marker:<2} {result.share_of(d):>5.1%}" + ) + lines.append("-" * len(header)) + total = result.total_effective_tokens + lines.append( + f"{'TOTAL':<{width}} {'':>11} {'':>6} {'':>12} {'':>6} {'':>11} {humanise(total):>12} {1.0:>5.1%}" + ) + + exposure = exposure_report( + exposure_entries_from_blend(result, config) if config is not None else [], + effective_tokens=total, + training_tokens=result.target_tokens, + ) + if config is not None: + lines.append("\nrepetition, once wrapping is counted") + lines.append(exposure.describe()) + + curved = [d for d in rows if d.plan is not None] + if curved: + lines.append("\nquality-aware upsampling curves") + for d in curved: + lines.append(f"\n {d.name}") + lines.append(d.plan.describe()) + + if result.target_tokens: + gap = total - result.target_tokens + verb = "over" if gap >= 0 else "under" + lines.append( + f"\ntarget {humanise(result.target_tokens)} tokens -- {humanise(abs(gap))} {verb} " + f"({abs(gap) / result.target_tokens:.1%})" + ) + if not result.exact: + lines.append("\n~ interpolated: a threshold fell inside a cube bin rather than on an edge.") + for d in rows: + for approximation in d.approximations: + lines.append(f" {d.name}: {approximation}") + lines.append(" Re-run with --exact to scan the per-document sidecar instead.") + return "\n".join(lines) diff --git a/src/modalities/dataloader/preprocessing/quality/sidecar.py b/src/modalities/dataloader/preprocessing/quality/sidecar.py new file mode 100644 index 000000000..5a72382d4 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/sidecar.py @@ -0,0 +1,408 @@ +"""Builds the per-document table a blend's selection is computed from. + +This is the one pass that has to read the raw data. For every document it records +where the document lives, how many tokens it is expected to contribute, the key that +joins it to external annotations, and whatever quality signals its own record already +carries. Everything downstream -- annotation join, aggregation, previewing a selection, +writing a filtered index -- works from this table and never reads the JSONL again. + +The position columns are what make the later steps cheap: a selection is materialised +by writing out the ``(byte_offset, byte_len)`` pairs of the documents that survived, +which is exactly the on-disk format of a modalities index file. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any, Callable, Iterator, Optional + +import jq +import pyarrow as pa +import pyarrow.parquet as pq +from tqdm import tqdm + +from modalities.dataloader.create_index import IndexGenerator +from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, SourcePointerResolver +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration +from modalities.utils.logger_utils import get_logger + +# Columns every sidecar carries, whatever dataset it describes. Native metrics and +# annotation columns are added alongside these. +BASE_FIELDS: tuple[tuple[str, pa.DataType], ...] = ( + ("file_id", pa.uint32()), + ("line_no", pa.uint32()), + ("byte_offset", pa.uint64()), + ("byte_len", pa.uint32()), + ("text_bytes", pa.uint32()), + ("est_tokens", pa.uint32()), + ("join_key", pa.large_string()), +) + + +class SidecarWriteError(RuntimeError): + """Raised when a sidecar cannot be produced for a dataset.""" + + +# One path segment of a jq expression: a bare identifier, or a quoted key for names that +# are not valid identifiers (``."openlid-v3"``). +_PATH_SEGMENT = re.compile(r'\.(?:([A-Za-z_][A-Za-z0-9_]*)|"((?:[^"\\]|\\.)*)")') + + +def parse_simple_path(jq_pattern: str) -> Optional[list[str]]: + """Recognises a jq pattern that is nothing more than a chain of field lookups. + + Args: + jq_pattern (str): The pattern from a native-metric declaration. + + Returns: + Optional[list[str]]: The field names to walk, or None if the pattern uses + anything beyond plain field access -- filters, pipes, indexing, functions. + + Note: + This exists for speed, and the speed difference is not marginal. + ``jq.compile(...).input_value(record)`` converts the *whole* record into jq's + own representation on every call, so on a 21 KB document two such calls cost + more than twenty times the rest of building a sidecar row. Documents are read + by the billion here, so plain field access is used wherever the pattern allows + it and jq is kept only for patterns that genuinely need it. + """ + pattern = jq_pattern.strip() + if not pattern.startswith("."): + return None + keys: list[str] = [] + position = 0 + while position < len(pattern): + match = _PATH_SEGMENT.match(pattern, position) + if match is None: + return None + bare, quoted = match.group(1), match.group(2) + keys.append(bare if bare is not None else quoted.replace('\\"', '"').replace("\\\\", "\\")) + position = match.end() + return keys or None + + +def _lookup_path(record: Any, keys: list[str]) -> Any: + # Mirrors jq's behaviour for a field chain: a missing key, or a non-object where an + # object is needed, yields no value rather than an error. + current = record + for key in keys: + if not isinstance(current, dict): + return None + current = current.get(key) + if current is None: + return None + return current + + +def build_metric_extractor(jq_pattern: str) -> tuple[Callable[[dict[str, Any]], Any], bool]: + """Builds the fastest available extractor for a native-metric pattern. + + Args: + jq_pattern (str): The pattern from a native-metric declaration. + + Returns: + tuple[Callable[[dict[str, Any]], Any], bool]: A function pulling the value out + of a decoded record, and whether it took the plain-path route. The flag is + reported so a pattern that silently fell back to jq -- and therefore costs + twenty times more per document -- is visible rather than a mystery. + """ + keys = parse_simple_path(jq_pattern) + if keys is not None: + return (lambda record: _lookup_path(record, keys)), True + + program = jq.compile(jq_pattern) + + def extract_with_jq(record: dict[str, Any]) -> Any: + try: + return program.input_value(record).first() + except (ValueError, StopIteration): + return None + + return extract_with_jq, False + + +def _aggregate(values: Any, aggregation: Optional[str]) -> Optional[float]: + # Several corpora store a per-page array of scores rather than one document score. + # Without an aggregation the array cannot become a column, so the first element is + # used, which is the common case of a single-page document. + if values is None: + return None + if isinstance(values, (int, float)) and not isinstance(values, bool): + return float(values) + if isinstance(values, list): + numeric = [float(v) for v in values if isinstance(v, (int, float)) and not isinstance(v, bool)] + if not numeric: + return None + if aggregation in (None, "first"): + return numeric[0] + if aggregation == "min": + return min(numeric) + if aggregation == "max": + return max(numeric) + if aggregation == "mean": + return sum(numeric) / len(numeric) + raise ValueError(f"unknown aggregation {aggregation!r}") + return None + + +def ensure_index(jsonl_path: Path, index_path: Optional[Path] = None) -> Path: + """Returns the index for a JSONL file, creating it if it does not exist. + + Args: + jsonl_path (Path): The JSONL file. + index_path (Optional[Path]): Where the index lives. Defaults to the file's + own ``.idx`` sibling, matching ``LargeFileLinesReader``'s convention. + + Returns: + Path: Path to a usable index file. + """ + index_path = LargeFileLinesReader.default_index_path(jsonl_path, index_path) + if not index_path.is_file(): + get_logger(name="main").info(f"Creating missing index for {jsonl_path} ...") + index_path.parent.mkdir(parents=True, exist_ok=True) + IndexGenerator(jsonl_path).create_index(index_path) + return index_path + + +class SidecarBuilder: + """Produces one dataset's per-document table. + + The builder is deliberately per-file: each JSONL file yields an independent + parquet part, so a dataset of thousands of files can be built by as many parallel + tasks, and a failed shard can be rebuilt on its own. + """ + + def __init__( + self, + dataset: DatasetEntry, + calibration: TokenCalibration, + index_root: Optional[Path] = None, + row_group_size: int = 200_000, + ): + """ + Args: + dataset (DatasetEntry): The dataset being described. + calibration (TokenCalibration): Token estimator for this dataset. + index_root (Optional[Path]): Directory holding index files, if they are not + kept next to the JSONL. Source trees are often read-only. + row_group_size (int): Parquet row group size for the output. + """ + self._dataset = dataset + self._calibration = calibration + self._index_root = index_root + self._row_group_size = row_group_size + self._native_programs = [] + slow_patterns: list[str] = [] + for metric in dataset.native_metrics: + extractor, is_fast = build_metric_extractor(metric.jq_pattern) + self._native_programs.append((metric.name, extractor, metric.aggregation)) + if not is_fast: + slow_patterns.append(f"{metric.name}={metric.jq_pattern}") + if slow_patterns: + get_logger(name="main").warning( + f"{dataset.name}: {len(slow_patterns)} native metric(s) need jq and will dominate the pass " + f"({', '.join(slow_patterns)}). Rewrite as a plain field path if possible." + ) + + def _index_path_for(self, jsonl_path: Path) -> Path: + if self._index_root is None: + return LargeFileLinesReader.default_index_path(jsonl_path, None) + relative = jsonl_path.relative_to(self._dataset.jsonl_root) + return Path(self._index_root) / relative.with_suffix(".idx") + + def schema(self) -> pa.Schema: + """Builds the parquet schema for this dataset's sidecar. + + Returns: + pa.Schema: Base position/token columns plus one column per native metric. + """ + fields = [pa.field(name, dtype) for name, dtype in BASE_FIELDS] + fields += [pa.field(f"native_{name}", pa.float64()) for name, _, _ in self._native_programs] + return pa.schema(fields) + + def _rows_for_file(self, jsonl_path: Path, file_id: int) -> Iterator[dict[str, Any]]: + index_path = ensure_index(jsonl_path, self._index_path_for(jsonl_path)) + reader = LargeFileLinesReader(jsonl_path, index_path=index_path) + key_spec = self._dataset.key + text_field = self._dataset.text_field + try: + for line_no, (byte_offset, byte_len) in enumerate(reader.index): + try: + record = json.loads(reader[line_no]) + except (json.JSONDecodeError, UnicodeDecodeError): + # A single malformed line must not abort a multi-terabyte pass; it + # is dropped, and its absence shows up as a row-count mismatch + # against the index. + continue + + text = record.get(text_field) + text_bytes = len(text.encode("utf-8")) if isinstance(text, str) else 0 + + row: dict[str, Any] = { + "file_id": file_id, + "line_no": line_no, + "byte_offset": byte_offset, + "byte_len": byte_len, + "text_bytes": text_bytes, + "est_tokens": self._calibration.estimate(record, text_bytes), + "join_key": key_spec.derive(record) if key_spec is not None else None, + } + for name, extract, aggregation in self._native_programs: + row[f"native_{name}"] = _aggregate(extract(record), aggregation) + yield row + finally: + reader.close() + + def build_file(self, jsonl_path: Path, file_id: int, output_path: Path) -> int: + """Builds the sidecar part for one JSONL file. + + Args: + jsonl_path (Path): The JSONL file to describe. + file_id (int): Index of this file within the dataset's sorted file list. + Stored per row so a selection can be mapped back to its source file. + output_path (Path): Destination parquet path. Parents are created. + + Returns: + int: Number of documents written. + """ + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + schema = self.schema() + n_rows = 0 + batch: list[dict[str, Any]] = [] + writer = pq.ParquetWriter(output_path, schema, compression="zstd") + try: + for row in self._rows_for_file(jsonl_path, file_id): + batch.append(row) + if len(batch) >= self._row_group_size: + writer.write_table(pa.Table.from_pylist(batch, schema=schema)) + n_rows += len(batch) + batch = [] + if batch: + writer.write_table(pa.Table.from_pylist(batch, schema=schema)) + n_rows += len(batch) + finally: + writer.close() + return n_rows + + def build( + self, + output_dir: Path, + file_ids: Optional[list[int]] = None, + show_progress: bool = True, + ) -> dict[str, int]: + """Builds sidecar parts for the dataset's files. + + Args: + output_dir (Path): Directory receiving one parquet part per file. + file_ids (Optional[list[int]]): Restrict the build to these file ids, for + sharding the work across tasks. None builds every file. + show_progress (bool): Whether to show a progress bar. + + Returns: + dict[str, int]: Maps output part name to documents written. + + Raises: + SidecarWriteError: If the dataset matches no files, which usually means a + wrong root or glob rather than an empty corpus. + """ + files = self._dataset.iter_files() + if not files: + raise SidecarWriteError( + f"dataset {self._dataset.name!r}: no files matched {self._dataset.glob!r} " + f"under {self._dataset.jsonl_root}" + ) + selected = [(i, p) for i, p in enumerate(files) if file_ids is None or i in set(file_ids)] + + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + # Pin the file list the ids in these parts refer to. Every shard of a sharded + # build sees the same list and writes identical content, so a racing write is + # harmless; without this the ids are only meaningful while the source tree + # happens to be unchanged. + FileManifest.from_entry(self._dataset).write(output_dir) + written: dict[str, int] = {} + iterator = tqdm(selected, desc=f"sidecar {self._dataset.name}", disable=not show_progress) + for file_id, jsonl_path in iterator: + part_name = f"part-{file_id:06d}.parquet" + written[part_name] = self.build_file(jsonl_path, file_id, output_dir / part_name) + return written + + +def resolve_source_pointers( + sidecar_dir: Path, + dataset: DatasetEntry, + batch_size: int = 500_000, + only_parts: Optional[list[int]] = None, +) -> int: + """Rewrites pointer join keys into the annotation keys they stand for. + + A translated corpus stores a ``/`` pointer back to the document it was + translated from. The annotation belongs to that original, so the pointer has to be + exchanged for a hash of the original's text before the join can happen. + + Args: + sidecar_dir (Path): Directory of sidecar parts to rewrite in place. + dataset (DatasetEntry): The dataset, whose key spec supplies the source root. + batch_size (int): How many pointers to resolve per pass over the source files. + only_parts (Optional[list[int]]): Restrict to the parts of these file ids. Lets a + sharded build resolve only the parts it wrote, leaving the rest to the tasks + that own them. + + Returns: + int: Number of rows whose key was resolved. + + Raises: + ValueError: If the dataset's key is not a source pointer. + """ + if dataset.key is None or dataset.key.kind != KeyKind.SOURCE_POINTER: + raise ValueError(f"dataset {dataset.name!r} does not use a source-pointer key") + + resolver = SourcePointerResolver( + source_root=dataset.key.source_root, + text_field=dataset.key.text_field, + line_offset=dataset.key.source_line_offset, + ) + if only_parts is None: + parts = sorted(Path(sidecar_dir).glob("part-*.parquet")) + else: + candidates = (Path(sidecar_dir) / f"part-{file_id:06d}.parquet" for file_id in only_parts) + parts = [path for path in candidates if path.is_file()] + n_resolved = 0 + for part in tqdm(parts, desc=f"resolve pointers {dataset.name}"): + table = pq.read_table(part) + pointers = [p for p in table.column("join_key").to_pylist() if p is not None] + if not pointers: + continue + mapping: dict[str, str] = {} + unique_pointers = sorted(set(pointers)) + for start in range(0, len(unique_pointers), batch_size): + mapping.update(resolver.resolve(unique_pointers[start : start + batch_size])) + resolved = [mapping.get(p) if p is not None else None for p in table.column("join_key").to_pylist()] + n_part = sum(1 for r in resolved if r is not None) + if n_part == 0: + # Refuse before writing. The write-back replaces the pointer with the resolved + # key, so writing nulls destroys the only copy of the pointer and makes a retry + # impossible without rebuilding the whole sidecar -- which is what happened when + # the source corpus was moved to /data/annealing_unused: 225 M pointers were + # overwritten with nulls, logged as an INFO line, and the failure only surfaced + # as 0% join coverage two hours later. + raise SidecarWriteError( + f"dataset {dataset.name!r}: none of {len(unique_pointers):,} distinct pointers in " + f"{part.name} resolved against {dataset.key.source_root}. Check that the source " + f"corpus is there and that its files match the pointers " + f"(e.g. {unique_pointers[0]!r}). Refusing to write, because the write-back would " + f"replace the pointers with nulls and a retry would need a full rebuild." + ) + n_resolved += n_part + table = table.set_column( + table.schema.get_field_index("join_key"), + pa.field("join_key", pa.large_string()), + pa.array(resolved, type=pa.large_string()), + ) + pq.write_table(table, part, compression="zstd") + return n_resolved diff --git a/src/modalities/dataloader/preprocessing/quality/tokens.py b/src/modalities/dataloader/preprocessing/quality/tokens.py new file mode 100644 index 000000000..6527ec9ac --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/tokens.py @@ -0,0 +1,548 @@ +"""Per-document token estimates for datasets that have not been tokenized yet. + +Selecting documents before tokenizing means the token budget of a selection has to be +predicted rather than measured. Two properties matter: + +* The estimate must be **per document**, not per corpus. Quality correlates with + length, so a filter that keeps the better documents keeps longer ones too. Scaling a + corpus average by a row-retention rate therefore understates the surviving tokens, + sometimes badly. +* The estimate must be based on the **text**, not the stored line. Several corpora keep + several renderings of the same document in one record -- HPLT stores ``text``, ``xml`` + and ``md`` side by side -- so the JSON line can be three times the size of the text + that will actually be tokenized. + +Two estimators are supported, in order of preference: + +1. A token count the corpus already carries, rescaled to our tokenizer by a measured + factor. Most such fields were produced with a different tokenizer, so they are + proportional to our counts rather than equal to them. +2. The text byte length divided by a measured bytes-per-token ratio. + +Both factors are measured per dataset by tokenizing a sample, because they vary a lot +across languages and content types. +""" + +from __future__ import annotations + +import bisect +import json +import random +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Optional + +import yaml +from pydantic import BaseModel, Field + +if TYPE_CHECKING: # pragma: no cover + # Only `calibrate_dataset` needs a tokenizer. Importing it eagerly would make + # previewing a blend depend on transformers and sentencepiece, which the preview + # path has no use for -- estimates are applied from stored constants. + from modalities.tokenization.tokenizer_wrapper import TokenizerWrapper + + +# Document byte-length boundaries defining the calibration's size strata. Log-spaced, +# because the bytes-per-token ratio varies with document length and document length is +# heavy-tailed: a handful of very long documents dominate a corpus-wide sum ratio, so +# measuring one global constant makes every estimate hostage to whether the sample +# happened to catch them. Six strata are enough to flatten that without needing a large +# sample in each. +SIZE_STRATUM_BOUNDS: tuple[int, ...] = (1_024, 4_096, 16_384, 65_536, 262_144) + + +def size_stratum(text_bytes: int) -> int: + """Finds which size stratum a document belongs to. + + Args: + text_bytes (int): UTF-8 byte length of the document's text. + + Returns: + int: Index into the stratum list, in ``[0, len(SIZE_STRATUM_BOUNDS)]``. + """ + return bisect.bisect_right(SIZE_STRATUM_BOUNDS, text_bytes) + + +class TokenCalibration(BaseModel): + """Measured constants relating a dataset's records to our tokenizer's counts. + + Attributes: + dataset (str): Name of the dataset this calibration was measured on. + tokenizer (str): Identifier of the tokenizer used, recorded so a stale + calibration cannot be applied to a different tokenizer unnoticed. + bytes_per_token (float): Mean UTF-8 text bytes per token. + native_field (Optional[str]): Record field holding the corpus's own token + count, if it has one. + native_scale (Optional[float]): Multiplier turning that field into our + tokenizer's count. + native_coverage (float): Fraction of sampled documents that carried + ``native_field``. A field present in only some records is not usable as + the primary estimator. + sampled_documents (int): How many documents the calibration was measured on. + sampled_tokens (int): How many tokens those documents produced. + eod_tokens_per_document (int): Tokens the packer appends per document, added + to every estimate so the prediction matches what packing produces. + stratum_bytes_per_token (list[Optional[float]]): Ratio measured within each size + stratium of :data:`SIZE_STRATUM_BOUNDS`, or None for a stratum the sample did + not reach. Conditioning on length is what makes the estimate robust: a + document's own byte length is known exactly from the sidecar, so it can be + costed against documents of its own size class instead of a corpus-wide mean + dominated by the longest documents. + stratum_documents (list[int]): Sampled documents per stratum, so a ratio measured + from too few documents can be recognised. + """ + + dataset: str + tokenizer: str + bytes_per_token: float = Field(gt=0) + native_field: Optional[str] = None + native_scale: Optional[float] = Field(default=None, gt=0) + native_coverage: float = 0.0 + sampled_documents: int = 0 + sampled_tokens: int = 0 + eod_tokens_per_document: int = 1 + stratum_bytes_per_token: list[Optional[float]] = Field(default_factory=list) + stratum_documents: list[int] = Field(default_factory=list) + + # A native field present in fewer than this share of sampled documents is ignored, + # because falling back per record would mix two estimators with different biases. + MIN_NATIVE_COVERAGE: ClassVar[float] = 0.99 + + # A stratum measured from fewer documents than this falls back to the global ratio, + # so a stratum reached by one outlier does not become an estimator of its own. + MIN_STRATUM_DOCUMENTS: ClassVar[int] = 20 + + def uses_native_field(self) -> bool: + """Whether the corpus's own token count is the primary estimator. + + Returns: + bool: True if a native field was found on effectively every sampled + document and a scale factor was measured for it. + """ + return ( + self.native_field is not None + and self.native_scale is not None + and self.native_coverage >= self.MIN_NATIVE_COVERAGE + ) + + def estimate(self, record: dict[str, Any], text_bytes: int) -> int: + """Estimates the tokens one document will contribute. + + Args: + record (dict[str, Any]): The decoded JSONL record. + text_bytes (int): UTF-8 byte length of the document's text field. + + Returns: + int: Estimated token count, including the end-of-document token(s) the + packer appends. Never negative. + """ + if self.uses_native_field(): + native = record.get(self.native_field) + if isinstance(native, (int, float)): + return max(0, round(native * self.native_scale) + self.eod_tokens_per_document) + return max(0, round(text_bytes / self.ratio_for(text_bytes)) + self.eod_tokens_per_document) + + def ratio_for(self, text_bytes: int) -> float: + """Picks the bytes-per-token ratio to apply to a document of this length. + + Args: + text_bytes (int): UTF-8 byte length of the document's text. + + Returns: + float: The stratum's ratio if one was measured from enough documents, + otherwise the corpus-wide ratio. + """ + if not self.stratum_bytes_per_token: + return self.bytes_per_token + index = size_stratum(text_bytes) + if index >= len(self.stratum_bytes_per_token): + return self.bytes_per_token + ratio = self.stratum_bytes_per_token[index] + enough = index < len(self.stratum_documents) and self.stratum_documents[index] >= self.MIN_STRATUM_DOCUMENTS + return ratio if (ratio and enough) else self.bytes_per_token + + +class CalibrationSet(BaseModel): + """Calibrations for every dataset of a blend. + + Attributes: + tokenizer (str): Identifier of the tokenizer all entries were measured with. + calibrations (dict[str, TokenCalibration]): Keyed by dataset name. + """ + + tokenizer: str + calibrations: dict[str, TokenCalibration] = Field(default_factory=dict) + + @classmethod + def from_yaml(cls, path: Path) -> "CalibrationSet": + """Loads a calibration set from YAML. + + Args: + path (Path): Path to the calibration file. + + Returns: + CalibrationSet: The parsed calibrations. + """ + with Path(path).open() as f: + return cls.model_validate(yaml.safe_load(f)) + + def to_yaml(self, path: Path) -> None: + """Writes the calibration set to YAML. + + Args: + path (Path): Destination path. Parent directories are created. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w") as f: + yaml.safe_dump(self.model_dump(mode="json"), f, sort_keys=False) + + def get(self, dataset: str) -> TokenCalibration: + """Looks up one dataset's calibration. + + Args: + dataset (str): The dataset name. + + Returns: + TokenCalibration: The calibration for that dataset. + + Raises: + KeyError: If the dataset has not been calibrated. Estimating without a + calibration would silently invent a token budget, so this is an error + rather than a default. + """ + if dataset not in self.calibrations: + raise KeyError( + f"no token calibration for dataset {dataset!r}; " + f"run 'modalities data quality calibrate' first. Calibrated: {sorted(self.calibrations)}" + ) + return self.calibrations[dataset] + + +# Token-count fields observed across the corpora, most specific first. A field is only +# adopted if it is present on effectively every sampled document of a dataset. +NATIVE_TOKEN_FIELDS: tuple[str, ...] = ( + "token_count", + "num_tokens", + "total_tokens", + "len_cl100k_base", +) + + +def _probe_files(file_paths: list[Path], max_probe_files: int) -> list[Path]: + # An evenly spaced selection, so the sample spans the whole dataset without the read + # growing with the file count. Corpora here range from 4 files to 40 003; reading a + # fixed number of lines from every one of them would mean 26 TB for the largest. + if len(file_paths) <= max_probe_files: + return list(file_paths) + step = len(file_paths) / max_probe_files + return [file_paths[min(int(i * step), len(file_paths) - 1)] for i in range(max_probe_files)] + + +def _line_start(handle: Any, pos: int, window: int = 65_536) -> int: + """Finds the offset of the line containing a byte position. + + Args: + handle (Any): Binary file handle. + pos (int): Byte position somewhere inside the wanted line. + window (int): How many bytes to read backwards at a time looking for the previous + newline. Documents longer than this are rare, and the loop extends for them. + + Returns: + int: Offset of the first byte of the line containing ``pos``. + """ + if pos <= 0: + return 0 + cursor = pos + while cursor > 0: + back = max(0, cursor - window) + handle.seek(back) + chunk = handle.read(cursor - back) + index = chunk.rfind(b"\n") + if index != -1: + return back + index + 1 + cursor = back + return 0 + + +def _sample_one_span( + handle: Any, + start: int, + text_field: str, + max_lines: int = 64, +) -> Optional[dict[str, Any]]: + """Reads the document containing a byte offset. + + The *containing* document, not the one after it, and that distinction is the whole + point. Returning the following document makes selection uniform across documents, + because which document follows a position is independent of how long that position's + document is. Returning the containing one makes selection proportional to length, + which is what a bytes-per-token ratio needs: the ratio is a byte-weighted quantity, so + the documents that dominate it must be the ones most likely to be sampled. + + On the German FineWiki snapshot the nine documents of 256 KB and above hold 5.7 % of + all bytes and tokenize at 34.6 bytes per token against 3.6 for the small ones. Uniform + sampling found two of them in 2,000 draws and the estimate was 19 % out; length- + proportional sampling finds forty and it is within 0.5 %. + + Args: + handle (Any): Binary file handle to seek within. + start (int): Byte offset to sample at. + text_field (str): Field that must hold a string for the record to be usable. + max_lines (int): How many lines to try before abandoning this span, so a run of + records without the text field cannot turn one span into a full-file scan. + + Returns: + Optional[dict[str, Any]]: The decoded record, or None if the span yielded none. + """ + handle.seek(_line_start(handle, start)) + for _ in range(max_lines): + line = handle.readline() + if not line: + return None + try: + record = json.loads(line) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + if isinstance(record.get(text_field), str): + return record + return None + + +def _measure_slice(text: str, max_bytes: int, rnd: random.Random) -> str: + """Takes a bounded slice of a document to measure its ratio on. + + Sampling documents in proportion to length is what makes the estimate accurate, but it + also means the multi-megabyte documents get sampled, and tokenizing those dominates the + cost: the work scales with the mean of the squared length over the mean length, which + heavy tails make enormous. Calibrating five snapshot datasets ran past ten minutes. + + A document's own bytes-per-token ratio is far more uniform inside the document than it + is across documents, so a slice measures it well. The slice is taken from a random + position rather than the head, for the same reason the file sample is not a prefix: a + document's opening prose is not representative of a document that is mostly a table. + + Args: + text (str): The document text. + max_bytes (int): Largest slice to measure. + rnd (random.Random): Seeded source for the slice position. + + Returns: + str: The whole text if it is small enough, otherwise a slice of it. + """ + if len(text) <= max_bytes // 4: + # Even at 4 bytes per character this cannot exceed the cap, so skip the encode. + return text + encoded = text.encode("utf-8") + if len(encoded) <= max_bytes: + return text + start = rnd.randrange(len(encoded) - max_bytes + 1) + # A slice of encoded bytes can begin or end mid-character; drop the broken edges. + return encoded[start : start + max_bytes].decode("utf-8", errors="ignore") + + +def _sample_documents( + file_paths: Iterable[Path], + text_field: str, + sample_size: int, + seed: int, + max_probe_files: int, + max_lines_per_probe: int, +) -> list[dict[str, Any]]: + """Samples documents from across a dataset without reading it whole. + + Documents are taken at byte offsets spread evenly over each probe file, one per span, + rather than from the start of the file. Reading a prefix instead was a real defect: it + is deterministic, so it looked stable across seeds while being systematically wrong. + On the German FineWiki snapshot the first 2,000 documents gave 3.531 bytes per token + where the whole file gives 4.214 -- a 16 % error in every token estimate downstream -- + because a corpus ordered by article id, fetch time, or source is not homogeneous along + its length. + + Each span contributes the document containing its offset, so a document is sampled in + proportion to its length and a document spanning several spans is counted by each of + them. That multiplicity is deliberate -- it is the weighting that makes each stratum's + measured ratio byte-weighted, which is what predicting a token total from a byte total + requires. See :func:`_sample_one_span`. + + Args: + file_paths (Iterable[Path]): The dataset's JSONL files. + text_field (str): The field holding the document text. + sample_size (int): Target number of documents. + seed (int): Seed for placing the offset within each span, so a calibration is + reproducible while the offsets are not degenerate multiples of the span size. + max_probe_files (int): How many files to draw from, spread across the dataset. + max_lines_per_probe (int): Cap on lines read per span before abandoning it. + + Returns: + list[dict[str, Any]]: The sampled records. + """ + files = _probe_files(list(file_paths), max_probe_files) + if not files: + return [] + per_file = max(1, -(-sample_size // len(files))) + lines_per_span = max(1, min(64, max_lines_per_probe)) + rnd = random.Random(seed) + + collected: list[dict[str, Any]] = [] + for path in files: + try: + size = path.stat().st_size + if size == 0: + continue + with path.open("rb") as f: + for i in range(per_file): + lo = size * i // per_file + hi = size * (i + 1) // per_file + start = lo if hi <= lo + 1 else lo + rnd.randrange(hi - lo) + record = _sample_one_span(f, start, text_field, max_lines=lines_per_span) + if record is not None: + collected.append(record) + except OSError: + continue + + if len(collected) > sample_size: + collected = random.Random(seed).sample(collected, sample_size) + return collected + + +def calibrate_dataset( + dataset_name: str, + file_paths: list[Path], + tokenizer: "TokenizerWrapper", + tokenizer_name: str, + text_field: str = "text", + sample_size: int = 4000, + seed: int = 42, + max_probe_files: int = 32, + max_lines_per_probe: int = 100_000, + eod_tokens_per_document: int = 1, + max_measure_bytes: int = 65_536, +) -> TokenCalibration: + """Measures how a dataset's records relate to our tokenizer's token counts. + + Args: + dataset_name (str): Name recorded in the calibration. + file_paths (list[Path]): The dataset's JSONL files. + tokenizer (TokenizerWrapper): The tokenizer training will use. + tokenizer_name (str): Identifier recorded alongside the measurement. + text_field (str): The field holding the document text. + sample_size (int): How many documents to tokenize. The default leaves margin: the + stratified estimate held within 0.5 % across seeds at 2,000 documents on + FineWiki, but one 1,000-document sample was 9.7 % out, and the cost of a larger + sample is seconds per dataset. + seed (int): Seed for trimming the sample, so calibration is reproducible. + max_probe_files (int): How many files to draw the sample from, spread evenly + across the dataset. This bounds the read: the cost of calibrating is set by + the sample size, not by how many files the dataset happens to have. + max_lines_per_probe (int): Safety cap on lines scanned in one probe file, for a + file whose records mostly lack the text field. + eod_tokens_per_document (int): Tokens the packer appends per document. + max_measure_bytes (int): Largest slice of one document to tokenize. Bounds the cost + of calibrating, which would otherwise be set by the longest documents in the + corpus. See :func:`_measure_slice`. + + Returns: + TokenCalibration: The measured calibration. + + Raises: + ValueError: If no documents could be sampled, which means the files are empty, + unreadable, or the text field name is wrong. + """ + documents = _sample_documents( + file_paths=file_paths, + text_field=text_field, + sample_size=sample_size, + seed=seed, + max_probe_files=max_probe_files, + max_lines_per_probe=max_lines_per_probe, + ) + if not documents: + raise ValueError( + f"dataset {dataset_name!r}: no documents with a string {text_field!r} field were found in " + f"{len(file_paths)} file(s); check the registry's text_field and glob" + ) + + n_strata = len(SIZE_STRATUM_BOUNDS) + 1 + stratum_bytes = [0] * n_strata + stratum_tokens = [0] * n_strata + stratum_docs = [0] * n_strata + + # Documents are sampled in proportion to their length, so a plain sum ratio over the + # whole sample is weighted by the square of length and comes out dominated by the + # longest documents -- 62 % low on FineWiki. The corpus-wide ratio therefore uses + # inverse-probability weights, which is the unbiased estimator of sum(bytes) over + # sum(tokens) under this sampling scheme. Within a stratum the plain ratio is kept: it + # is byte-weighted, which is what costing a stratum's bytes calls for, and the length + # spread inside one stratum is small enough for the residual bias not to matter. + inverse_weighted_tokens = 0.0 + weighted_documents = 0 + + total_text_bytes = 0 + total_tokens = 0 + native_totals: dict[str, float] = {field: 0.0 for field in NATIVE_TOKEN_FIELDS} + native_counts: dict[str, int] = {field: 0 for field in NATIVE_TOKEN_FIELDS} + native_tokens: dict[str, float] = {field: 0.0 for field in NATIVE_TOKEN_FIELDS} + + measure_rnd = random.Random(seed) + for record in documents: + text = record[text_field] + measured = _measure_slice(text, max_measure_bytes, measure_rnd) + n_tokens = len(tokenizer.tokenize(measured)) + n_bytes = len(measured.encode("utf-8")) + # The stratum is the document's real size class, even when the ratio was measured + # on a slice of it: what is being estimated is the ratio for documents of that size. + full_bytes = len(text.encode("utf-8")) if measured is not text else n_bytes + total_text_bytes += n_bytes + total_tokens += n_tokens + stratum = size_stratum(full_bytes) + stratum_bytes[stratum] += n_bytes + stratum_tokens[stratum] += n_tokens + stratum_docs[stratum] += 1 + if n_bytes > 0: + inverse_weighted_tokens += n_tokens / n_bytes + weighted_documents += 1 + for field in NATIVE_TOKEN_FIELDS: + value = record.get(field) + if isinstance(value, (int, float)) and value > 0: + # The native scale relates a whole document's count to whole-document + # tokens, so a sliced measurement has to be extrapolated back up. + scaled_tokens = n_tokens * (full_bytes / n_bytes) if n_bytes else 0 + native_totals[field] += float(value) + native_counts[field] += 1 + native_tokens[field] += scaled_tokens + + if total_tokens == 0: + raise ValueError(f"dataset {dataset_name!r}: sampled documents produced zero tokens") + + global_bytes_per_token = ( + weighted_documents / inverse_weighted_tokens if inverse_weighted_tokens > 0 else total_text_bytes / total_tokens + ) + + native_field: Optional[str] = None + native_scale: Optional[float] = None + native_coverage = 0.0 + for field in NATIVE_TOKEN_FIELDS: + coverage = native_counts[field] / len(documents) + if coverage >= TokenCalibration.MIN_NATIVE_COVERAGE and native_totals[field] > 0: + native_field = field + # Scale is measured only over the documents that carry the field, so a + # partially present field cannot skew it. + native_scale = native_tokens[field] / native_totals[field] + native_coverage = coverage + break + + return TokenCalibration( + dataset=dataset_name, + tokenizer=tokenizer_name, + bytes_per_token=global_bytes_per_token, + native_field=native_field, + native_scale=native_scale, + native_coverage=native_coverage, + sampled_documents=len(documents), + sampled_tokens=total_tokens, + eod_tokens_per_document=eod_tokens_per_document, + stratum_bytes_per_token=[ + (stratum_bytes[i] / stratum_tokens[i]) if stratum_tokens[i] else None for i in range(n_strata) + ], + stratum_documents=stratum_docs, + ) diff --git a/src/modalities/dataloader/preprocessing/quality/upsampling.py b/src/modalities/dataloader/preprocessing/quality/upsampling.py new file mode 100644 index 000000000..7aa7e9884 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/upsampling.py @@ -0,0 +1,389 @@ +"""Turns one scalar up/downsample ratio per dataset into a quality-aware upsampling curve. + +A single ratio treats every surviving document alike: a dataset filtered to "educational +value at least basic" and set to 1.2 draws the barely-basic documents exactly as often as +the excellent ones. A curve instead makes the repeat factor rise with quality, so the token +budget is spent where the quality signal says it is worth spending. + +The idea and the functional form are from Dolma 3 / Olmo 3 (arXiv:2512.13961, §3.4.4 and +appendix A.2.4), which reports it beating flat quality filtering at every matched repetition +factor -- for instance 0.740 against 0.843-0.870 bits-per-byte on their Math suite. They +discard the bottom 40% of web text by quality and repeat the top 5% seven times. + +**The parameterisation.** Quality is placed on a [0, 1] axis, ordered worst to best, where a +bucket's width is its share of the dataset's tokens. The repeat factor is + + f(x) = 0 for x < a + f(x) = C * (x - a)**p for x >= a + +subject to three constraints, following the paper: the integral equals the target token +yield, no bucket averages more than ``max_factor``, and the curve is monotone. + +The paper's family carries an extra ``exp(lam * (x - a))`` factor. We fix ``lam = 0``, which +is a deliberate simplification: with two shape parameters the constraints define a curve of +feasible solutions rather than a point, so something further has to pick one. Dropping the +exponential makes the solution unique *and* every integral analytic, so no numerical +quadrature is involved. What is left is one degree of freedom, ``p``, and we spend it by +pushing the top bucket to exactly ``max_factor`` -- the steepest admissible curve, which is +also what the paper's own figure shows. + +**Why the cumulative form.** Writing ``q = p + 1`` and + + g(t) = ((t - a) / (1 - a))**q for t >= a, else 0 + +makes ``g`` the share of the token budget drawn from below quality ``t``, with ``g(1) = 1``. +The scale ``C`` cancels, so a bucket spanning ``[u, v]`` receives ``Z * (g(v) - g(u))`` +tokens and is repeated ``R * (g(v) - g(u)) / (v - u)`` times, where ``R = Z / X``. Every +quantity below is that expression. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +from pydantic import BaseModel, Field, model_validator + +# Ceiling on the exponent, and it earns its keep in the downsampling regime. Pushing the +# curve as steep as the cap allows is right when the cap binds -- it is what reproduces the +# published example, bottom 40% discarded and top 5% at 7x. But when the target is a small +# multiple of the pool, as it is whenever a blend draws far fewer tokens than it has, no +# exponent violates the cap and "steepest" is unbounded: the budget collapses onto the +# highest bucket and the curve degenerates into the hard top-k filtering it exists to beat. +# A modest ceiling keeps it a curve. Below it, nothing changes. +MAX_EXPONENT = 8.0 + +# Label for documents that carry no annotation and so cannot be placed on a quality axis. +UNANNOTATED_BUCKET = "" + + +class UpsamplingError(RuntimeError): + """Raised when no curve can satisfy the requested constraints.""" + + +class UpsamplingSpec(BaseModel): + """How to build a quality-aware curve for one dataset. + + Attributes: + quality_field (str): The field that orders documents from worst to best. Either an + ordinal annotation field, whose declared levels become the buckets, or a native + numeric metric, whose cube score bins become the buckets. + target_tokens (Optional[float]): Tokens to draw from this dataset in total. Exactly + one of this and ``target_ratio`` must be given. + target_ratio (Optional[float]): Tokens to draw as a multiple of the dataset's kept + tokens -- the average repeat factor over the whole dataset, discarded part + included. 1.0 means "as many tokens as the dataset has". + max_factor (float): No bucket may be repeated more than this on average. 7.0 is the + value Dolma 3 arrived at empirically; beyond a handful of repeats the returns + fall off sharply. + discard_below_percentile (float): Share of tokens, from the bottom of the quality + order, to drop entirely. This is on top of the dataset's predicates, which have + already been applied: the percentile is of what survived them. + exponent (Optional[float]): Fixes the curve's shape instead of solving for it. For + experiments; 1.0 is flat over the kept range, larger is steeper. + """ + + quality_field: str + target_tokens: Optional[float] = Field(default=None, gt=0) + target_ratio: Optional[float] = Field(default=None, gt=0) + max_factor: float = Field(default=7.0, gt=0) + discard_below_percentile: float = Field(default=0.0, ge=0.0, lt=100.0) + exponent: Optional[float] = Field(default=None, ge=1.0) + + @model_validator(mode="after") + def _check_one_target(self) -> "UpsamplingSpec": + if (self.target_tokens is None) == (self.target_ratio is None): + raise ValueError( + f"upsampling on {self.quality_field!r} needs exactly one of 'target_tokens' or " + f"'target_ratio'" + ) + return self + + @property + def discard_fraction(self) -> float: + """The discard threshold as a fraction rather than a percentage. + + Returns: + float: Value in [0, 1). + """ + return self.discard_below_percentile / 100.0 + + +@dataclass(frozen=True) +class QualityBucket: + """One quality level of a dataset, after its predicates have been applied. + + Attributes: + label (str): The level, or ``UNANNOTATED_BUCKET``. + n_documents (int): Documents in this bucket. + n_tokens (int): Estimated tokens in this bucket. + unannotated (bool): Whether this is the bucket of documents with no label, which + cannot be placed on the quality axis and so sit at the bottom of it. + """ + + label: str + n_documents: int + n_tokens: int + unannotated: bool = False + + +@dataclass(frozen=True) +class BucketPlan: + """What the curve prescribes for one bucket. + + Attributes: + bucket (QualityBucket): The bucket itself. + lower (float): Start of its interval on the quality axis. + upper (float): End of its interval on the quality axis. + factor (float): Repeat factor. 0.0 means discarded. + tokens_drawn (float): Tokens the blend takes from this bucket. + """ + + bucket: QualityBucket + lower: float + upper: float + factor: float + tokens_drawn: float + + +@dataclass(frozen=True) +class UpsamplingPlan: + """A solved curve together with its per-bucket consequences. + + Attributes: + field (str): The quality field the axis was built from. + exponent (float): The solved ``q = p + 1``. 1.0 is flat over the kept range, and + :data:`MAX_EXPONENT` means the cap never bound, so the shape came from the + ceiling rather than from the constraints. + discard_fraction (float): Share of tokens dropped from the bottom. + max_factor (float): The cap that was honoured. + target_ratio (float): Tokens drawn over tokens available. + buckets (list[BucketPlan]): Ordered worst to best quality. + saturated (bool): Whether the top bucket reached ``max_factor``. False means the + target was reachable without needing the steepest curve. + """ + + field: str + exponent: float + discard_fraction: float + max_factor: float + target_ratio: float + buckets: list[BucketPlan] + saturated: bool + + @property + def tokens_available(self) -> int: + """Tokens in the dataset after its predicates, before the curve. + + Returns: + int: Sum over all buckets. + """ + return sum(plan.bucket.n_tokens for plan in self.buckets) + + @property + def tokens_drawn(self) -> float: + """Tokens the blend takes from this dataset. + + Returns: + float: Sum over buckets of tokens drawn. + """ + return sum(plan.tokens_drawn for plan in self.buckets) + + @property + def documents_kept(self) -> int: + """Distinct documents that survive the curve's discard threshold. + + Returns: + int: Documents in buckets with a non-zero factor. Repetition does not + multiply this: the same documents are drawn more than once. + """ + return sum(plan.bucket.n_documents for plan in self.buckets if plan.factor > 0) + + def describe(self) -> str: + """Renders the curve as a small table. + + Returns: + str: One line per bucket, worst quality first. + """ + lines = [ + f" curve on {self.field}: exponent {self.exponent:.2f}, " + f"discard bottom {self.discard_fraction:.0%}, cap {self.max_factor:g}x" + + ("" if self.saturated else " (cap not reached)"), + f" {'bucket':<22} {'tokens':>14} {'share':>7} {'factor':>8} {'drawn':>14}", + ] + for plan in self.buckets: + share = plan.upper - plan.lower + factor = "discarded" if plan.factor == 0 else f"{plan.factor:.2f}x" + lines.append( + f" {plan.bucket.label:<22} {plan.bucket.n_tokens:>14,} {share:>6.1%} " + f"{factor:>8} {plan.tokens_drawn:>14,.0f}" + ) + return "\n".join(lines) + + +def _cumulative(exponent: float, discard: float, t: float) -> float: + """Share of the token budget drawn from quality below ``t``. + + Args: + exponent (float): ``q = p + 1``, at least 1. + discard (float): The discard threshold ``a``. + t (float): Point on the quality axis. + + Returns: + float: Value in [0, 1]. + """ + if t <= discard: + return 0.0 + if discard >= 1.0: + return 0.0 + return min(1.0, ((t - discard) / (1.0 - discard)) ** exponent) + + +def _edges(buckets: list[QualityBucket]) -> list[tuple[float, float]]: + """Places buckets on the [0, 1] quality axis, weighted by tokens. + + Args: + buckets (list[QualityBucket]): Ordered worst to best quality. + + Returns: + list[tuple[float, float]]: One (lower, upper) per bucket. + + Raises: + UpsamplingError: If the buckets hold no tokens, so no axis can be built. + """ + total = sum(b.n_tokens for b in buckets) + if total <= 0: + raise UpsamplingError("cannot build a quality axis: the surviving buckets hold no tokens") + edges: list[tuple[float, float]] = [] + cursor = 0.0 + for bucket in buckets: + width = bucket.n_tokens / total + edges.append((cursor, min(1.0, cursor + width))) + cursor += width + # Absorb rounding so the last bucket ends exactly at 1.0. + if edges: + edges[-1] = (edges[-1][0], 1.0) + return edges + + +def _top_factor(exponent: float, discard: float, target_ratio: float, edges: list[tuple[float, float]]) -> float: + """Repeat factor of the highest-quality bucket for a given exponent. + + Args: + exponent (float): ``q = p + 1``. + discard (float): The discard threshold. + target_ratio (float): ``Z / X``. + edges (list[tuple[float, float]]): Bucket intervals. + + Returns: + float: The top bucket's average repeat factor. + """ + lower, upper = edges[-1] + width = upper - lower + if width <= 0: + return float("inf") + drawn = _cumulative(exponent, discard, upper) - _cumulative(exponent, discard, lower) + return target_ratio * drawn / width + + +def solve_curve( + buckets: list[QualityBucket], + spec: UpsamplingSpec, +) -> UpsamplingPlan: + """Finds the steepest curve satisfying the spec, and its per-bucket factors. + + The exponent is chosen so the top bucket sits exactly at ``max_factor``. That is the + steepest admissible curve, which is the point of the exercise: a flatter one would spend + budget on weaker documents while leaving the cap unused. If even a flat curve over the + kept range exceeds the cap the request is infeasible, and if the cap cannot be reached + however steep the curve, the exponent is capped instead. + + Args: + buckets (list[QualityBucket]): Buckets ordered worst to best quality, holding the + documents that survived the dataset's predicates. + spec (UpsamplingSpec): The constraints. + + Returns: + UpsamplingPlan: The solved curve and what it draws from each bucket. + + Raises: + UpsamplingError: If no curve can meet the constraints, with the numbers needed to + see which constraint to relax. + """ + if not buckets: + raise UpsamplingError(f"no quality buckets for field {spec.quality_field!r}") + + edges = _edges(buckets) + available = sum(b.n_tokens for b in buckets) + discard = spec.discard_fraction + kept_share = 1.0 - discard + + if spec.target_ratio is not None: + target_ratio = spec.target_ratio + else: + target_ratio = spec.target_tokens / available + + # A flat curve over the kept range repeats everything target_ratio / kept_share times, + # which is the least any admissible curve can ask of its top bucket. + flat_factor = target_ratio / kept_share + if flat_factor > spec.max_factor + 1e-9: + raise UpsamplingError( + f"{spec.quality_field!r}: cannot draw {target_ratio:.2f}x this dataset's " + f"{available:,} tokens while discarding the bottom {discard:.0%} and repeating " + f"nothing more than {spec.max_factor:g}x -- even repeating every kept document " + f"equally needs {flat_factor:.2f}x. Discard less, raise max_factor, or lower the " + f"target." + ) + + if spec.exponent is not None: + exponent = spec.exponent + saturated = False + else: + # _top_factor rises monotonically with the exponent, from flat_factor towards + # target_ratio / top_width, so bisect for the cap. + low, high = 1.0, MAX_EXPONENT + if _top_factor(high, discard, target_ratio, edges) <= spec.max_factor: + # The cap does not bind anywhere in the admissible range, so there is no + # constraint left to pin the shape; take the steepest curve the ceiling allows. + exponent, saturated = high, False + else: + for _ in range(200): + mid = (low + high) / 2.0 + if _top_factor(mid, discard, target_ratio, edges) > spec.max_factor: + high = mid + else: + low = mid + exponent, saturated = (low + high) / 2.0, True + + plans: list[BucketPlan] = [] + for bucket, (lower, upper) in zip(buckets, edges): + drawn_share = _cumulative(exponent, discard, upper) - _cumulative(exponent, discard, lower) + tokens_drawn = target_ratio * drawn_share * available + factor = (tokens_drawn / bucket.n_tokens) if bucket.n_tokens > 0 else 0.0 + # A bucket entirely below the discard threshold draws nothing; float noise can leave + # a vanishing factor, which would materialise an index for no reason. + if factor < 1e-9: + factor, tokens_drawn = 0.0, 0.0 + plans.append( + BucketPlan( + bucket=bucket, lower=lower, upper=upper, factor=factor, tokens_drawn=tokens_drawn + ) + ) + + worst = max((p.factor for p in plans), default=0.0) + if worst > spec.max_factor * (1.0 + 1e-6): + raise UpsamplingError( + f"{spec.quality_field!r}: solved curve repeats a bucket {worst:.2f}x, above the " + f"{spec.max_factor:g}x cap. This happens when one bucket is much narrower than the " + f"others; widen the buckets or raise max_factor." + ) + + return UpsamplingPlan( + field=spec.quality_field, + exponent=exponent, + discard_fraction=discard, + max_factor=spec.max_factor, + target_ratio=target_ratio, + buckets=plans, + saturated=saturated, + ) diff --git a/src/modalities/dataloader/preprocessing/quality/verify.py b/src/modalities/dataloader/preprocessing/quality/verify.py new file mode 100644 index 000000000..bf4008c43 --- /dev/null +++ b/src/modalities/dataloader/preprocessing/quality/verify.py @@ -0,0 +1,217 @@ +"""Checks that a sidecar's byte offsets still describe the documents it claims. + +The manifest in ``file_manifest`` catches a source tree that moved, by comparing paths +and sizes. That is cheap and covers what we have actually seen go wrong, but it is +circumstantial: a file can be rewritten at identical size. This module does the direct +test instead -- seek to a recorded offset, read the recorded length, and check that the +bytes there are the document the sidecar says they are. + +Two details matter for the check to mean anything: + +* Rows at ``byte_offset == 0`` are skipped. The first document of any JSONL file parses + as JSON, so a row at offset 0 succeeds even when the file underneath is a completely + different one. Testing only those rows is how a broken sidecar looks healthy. +* The comparison is against ``text_bytes``, not the join key. It works for every dataset + including the two that have no identifier at all, and for the one whose key is a hash + resolved out of a third corpus. +""" + +from __future__ import annotations + +import json +import random +from dataclasses import dataclass, field +from pathlib import Path + +import pyarrow.parquet as pq + +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest, ManifestError, load_manifest +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry + + +@dataclass +class VerifyReport: + """Outcome of verifying one dataset's sidecar. + + Attributes: + dataset (str): Dataset name. + n_parts (int): Sidecar parts present. + n_sampled (int): Documents actually probed. + n_readable (int): Probes whose byte range yielded parseable JSON. + n_matching (int): Probes whose document had the recorded text length. + manifest_problems (list[str]): Drift reported by the manifest, if there is one. + has_manifest (bool): Whether a manifest was found. + notes (list[str]): Anything else worth saying. + """ + + dataset: str + n_parts: int = 0 + n_sampled: int = 0 + n_readable: int = 0 + n_matching: int = 0 + manifest_problems: list[str] = field(default_factory=list) + has_manifest: bool = False + notes: list[str] = field(default_factory=list) + + @property + def ok(self) -> bool: + """Whether the sidecar can be trusted. + + Returns: + bool: True if nothing drifted and every probe matched. + """ + return not self.manifest_problems and self.n_sampled > 0 and self.n_matching == self.n_sampled + + @property + def verdict(self) -> str: + """A short label for the outcome. + + Returns: + str: One of VALID, DRIFTED, BROKEN, PARTIAL, or EMPTY. + """ + if self.n_sampled == 0: + return "EMPTY" + if self.manifest_problems: + return "DRIFTED" + if self.n_matching == self.n_sampled: + return "VALID" + if self.n_readable == 0: + return "BROKEN" + return "PARTIAL" + + +def verify_sidecar( + sidecar_dir: Path, + entry: DatasetEntry, + n_parts: int = 8, + n_rows_per_part: int = 4, + seed: int = 0, + check_manifest: bool = True, +) -> VerifyReport: + """Probes a dataset's sidecar against its source files. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + entry (DatasetEntry): Registry entry supplying the source root. + n_parts (int): How many sidecar parts to sample. + n_rows_per_part (int): How many documents to probe per sampled part. + seed (int): Sampling seed, so a verdict is reproducible. + check_manifest (bool): Whether to compare the recorded file list against the + source tree. Turn off only when deliberately verifying a sidecar whose + manifest is expected to be stale, such as before adopting one. + + Returns: + VerifyReport: What was probed and what matched. + """ + sidecar_dir = Path(sidecar_dir) + report = VerifyReport(dataset=entry.name) + parts = sorted(sidecar_dir.glob("part-*.parquet")) + report.n_parts = len(parts) + if not parts: + report.notes.append(f"no sidecar parts in {sidecar_dir}") + return report + + manifest = load_manifest(sidecar_dir) + report.has_manifest = manifest is not None + if manifest is None: + report.notes.append("no file manifest; file ids resolved by re-globbing, which is what this checks") + if not entry.jsonl_root.exists(): + report.manifest_problems.append(f"source root {entry.jsonl_root} does not exist") + return report + source_files = entry.iter_files() + else: + if check_manifest: + report.manifest_problems = manifest.drift(entry) + source_files = manifest.resolve(entry) + + rnd = random.Random(seed) + for part in rnd.sample(parts, min(n_parts, len(parts))): + table = pq.read_table(part, columns=["file_id", "byte_offset", "byte_len", "text_bytes"]) + # Offset 0 parses in any JSONL file, so it is no evidence at all. + candidates = [i for i in range(table.num_rows) if table.column("byte_offset")[i].as_py() > 0] + if not candidates: + continue + for i in rnd.sample(candidates, min(n_rows_per_part, len(candidates))): + report.n_sampled += 1 + file_id = table.column("file_id")[i].as_py() + if file_id >= len(source_files): + continue + path = source_files[file_id] + try: + with open(path, "rb") as f: + f.seek(table.column("byte_offset")[i].as_py()) + record = json.loads(f.read(table.column("byte_len")[i].as_py())) + except (OSError, ValueError): + continue + report.n_readable += 1 + text = record.get(entry.text_field) + if isinstance(text, str) and len(text.encode("utf-8")) == table.column("text_bytes")[i].as_py(): + report.n_matching += 1 + + return report + + +def adopt_manifest(sidecar_dir: Path, entry: DatasetEntry, report: VerifyReport) -> Path: + """Records a file manifest for a sidecar built before manifests existed. + + Only sensible once the sidecar has been verified, which is why the report is a + required argument: stamping a manifest onto a stale sidecar would make a broken + sidecar look pinned and trustworthy. + + Args: + sidecar_dir (Path): The dataset's sidecar directory. + entry (DatasetEntry): Registry entry supplying the current file list. + report (VerifyReport): The verification that justifies adoption. + + Returns: + Path: The written manifest path. + + Raises: + ManifestError: If the report does not show the sidecar to be valid. + """ + if not report.ok: + raise ManifestError( + f"refusing to write a manifest for {entry.name!r}: verification says {report.verdict} " + f"({report.n_matching}/{report.n_sampled} probes matched). A manifest would make this " + f"sidecar look pinned when its offsets do not describe the current files. Rebuild it." + ) + return FileManifest.from_entry(entry).write(sidecar_dir) + + +def format_verify_report(reports: list[VerifyReport]) -> str: + """Renders verification results as a table. + + Args: + reports (list[VerifyReport]): One report per dataset. + + Returns: + str: A printable report. + """ + lines = [ + f"{'dataset':<18} {'parts':>7} {'probed':>7} {'read':>6} {'match':>6} {'manifest':>9} verdict", + "-" * 78, + ] + for r in reports: + lines.append( + f"{r.dataset:<18} {r.n_parts:>7} {r.n_sampled:>7} {r.n_readable:>6} {r.n_matching:>6} " + f"{('yes' if r.has_manifest else 'no'):>9} {r.verdict}" + ) + + broken = [r for r in reports if not r.ok] + if broken: + lines.append("") + lines.append(f"{len(broken)} of {len(reports)} dataset(s) cannot be trusted:") + for r in broken: + lines.append(f" {r.dataset} [{r.verdict}]") + for problem in r.manifest_problems[:5]: + lines.append(f" {problem}") + for note in r.notes: + lines.append(f" {note}") + if r.n_sampled and r.n_readable == 0: + lines.append(" every probed byte range was unreadable -- the source files changed") + lines.append("") + lines.append("Rebuild the sidecar, re-join and re-cube each of these before using them in a blend.") + else: + lines.append("") + lines.append(f"all {len(reports)} dataset(s) verified: byte offsets describe the current source files") + return "\n".join(lines) diff --git a/src/modalities/registry/components.py b/src/modalities/registry/components.py index 12a8cbcee..42e25687f 100644 --- a/src/modalities/registry/components.py +++ b/src/modalities/registry/components.py @@ -69,6 +69,7 @@ StepLRSchedulerConfig, TorchCheckpointLoadingConfig, WandBEvaluationResultSubscriberConfig, + WeightedCombinedDatasetConfig, WeightInitializedModelConfig, ) from modalities.dataloader.collate_fns.collator_fn_wrapper_for_loss_masking import ( @@ -316,6 +317,12 @@ class ComponentEntity: ), ComponentEntity("dataset", "dummy_dataset", DatasetFactory.get_dummy_dataset, DummyDatasetConfig), ComponentEntity("dataset", "combined", DatasetFactory.get_combined_dataset, CombinedDatasetConfig), + ComponentEntity( + "dataset", + "weighted_combined", + DatasetFactory.get_weighted_combined_dataset, + WeightedCombinedDatasetConfig, + ), # samplers ComponentEntity("sampler", "sequential_sampler", SequentialSampler, SequentialSamplerConfig), ComponentEntity("sampler", "distributed_sampler", DistributedSampler, DistributedSamplerConfig), diff --git a/tests/dataloader/preprocessing/quality/__init__.py b/tests/dataloader/preprocessing/quality/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/dataloader/preprocessing/quality/test_export.py b/tests/dataloader/preprocessing/quality/test_export.py new file mode 100644 index 000000000..9d63d5ae8 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_export.py @@ -0,0 +1,441 @@ +"""Exporting a selection as JSONL, with the sampling materialised in the bytes. + +The property under test throughout is that the output is the training set: what the packer +used to leave to `WeightedCombinedDataset` -- repeating a dataset 3x, or drawing 60% of it -- +now has to be visible as lines on disk, because a concatenation carries no weights. +""" + +import json +import subprocess +import sys +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +import yaml + +from modalities.dataloader.preprocessing.quality.annotation_join import bucket_annotations, join_annotations + +from modalities.dataloader.preprocessing.quality.export import ( + EXPORT_MANIFEST, + ExportError, + copies_for, + export_blend, + finalize_export, +) +from modalities.dataloader.preprocessing.quality.materialize import materialize_blend +from modalities.dataloader.preprocessing.quality.registry import ( + CorpusRegistry, + DatasetEntry, + KeyKind, + KeySpec, +) +from modalities.dataloader.preprocessing.quality.selection import DatasetSelection, SelectionConfig +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder +from modalities.dataloader.preprocessing.quality.tokens import calibrate_dataset +from modalities.dataloader.preprocessing.quality.upsampling import UpsamplingSpec + +LEVELS = ["none", "minimal", "basic", "moderate", "high"] +N_DOCS = 400 + + +class _Whitespace: + def tokenize(self, text): + return text.split() + + +@pytest.fixture +def corpus(tmp_path: Path) -> Path: + """Two shards, the last line of the second deliberately lacking a trailing newline.""" + directory = tmp_path / "corpus" + directory.mkdir() + for shard in range(2): + lines = [] + for i in range(N_DOCS // 2): + doc = shard * (N_DOCS // 2) + i + lines.append( + json.dumps( + { + "id": f"doc-{doc}", + "text": " ".join(["word"] * (5 + doc % 17)), + "educational_value": LEVELS[doc % len(LEVELS)], + } + ) + ) + body = "\n".join(lines) + # No trailing newline on the second shard: a real corpus contains both forms, and the + # index's byte_len excludes the terminator either way. + (directory / f"shard_{shard}.jsonl").write_text(body + ("\n" if shard == 0 else "")) + return directory + + +@pytest.fixture +def entry(corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + + +@pytest.fixture +def blend(tmp_path: Path, entry: DatasetEntry, corpus: Path): + """Builds a real sidecar, annotations and all, so curves can be exercised too.""" + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=entry.iter_files(), + tokenizer=_Whitespace(), + tokenizer_name="whitespace", + sample_size=100, + ) + sidecar_root = tmp_path / "sidecar" + SidecarBuilder(entry, calibration, index_root=tmp_path / "idx" / "toy").build( + sidecar_root / "toy", show_progress=False + ) + + # Quality labels reach the sidecar through the join, not from the source records, so a + # curve needs the annotation stages to have run. + rows = {"id": [], "educational_value": []} + for shard in sorted(corpus.glob("*.jsonl")): + for line in shard.read_text().splitlines(): + record = json.loads(line) + rows["id"].append(record["id"]) + rows["educational_value"].append(record["educational_value"]) + annotations = tmp_path / "annotations" + annotations.mkdir() + pq.write_table(pa.table(rows), annotations / "shard0.parquet") + bucket_annotations( + shard_paths=[annotations / "shard0.parquet"], + out_dir=tmp_path / "buckets", + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + join_annotations(sidecar_root / "toy", tmp_path / "buckets", "toy", "toy", show_progress=False) + + def build(selection: DatasetSelection, name: str = "mix"): + return materialize_blend( + config=SelectionConfig(datasets=[selection]), + registry=CorpusRegistry(datasets=[entry]), + sidecar_root=sidecar_root, + output_root=tmp_path / name, + show_progress=False, + ) + + return build + + +def _export(tmp_path: Path, entry: DatasetEntry, manifest_path: Path, out: str = "out", **kwargs): + registry_path = tmp_path / "registry.yaml" + registry_path.write_text( + yaml.safe_dump({"datasets": [json.loads(CorpusRegistry(datasets=[entry]).model_dump_json())["datasets"][0]]}) + ) + exports = export_blend( + manifest_path=manifest_path, + registry_path=registry_path, + output_root=tmp_path / out, + show_progress=False, + **kwargs, + ) + finalize_export(tmp_path / out) + return exports, tmp_path / out + + +def _lines(root: Path, dataset: str = "toy") -> list[str]: + """Every output line of a dataset, in shard order -- the concatenation.""" + out: list[str] = [] + for shard in sorted((root / dataset).rglob("*.jsonl")): + out.extend(shard.read_text().splitlines()) + return out + + +# --------------------------------------------------------------------------- the factor + + +@pytest.mark.parametrize("factor,expected", [(1.0, 1), (2.0, 2), (3.0, 3), (0.0, 0)]) +def test_a_whole_factor_writes_exactly_that_many_copies(factor, expected): + for offset in range(0, 5000, 137): + assert copies_for(factor, 42, "/data/x.jsonl", offset) == expected + + +def test_a_fractional_factor_splits_between_the_two_neighbouring_counts(): + counts = [copies_for(1.2, 42, "/data/x.jsonl", o) for o in range(0, 200_000, 10)] + assert set(counts) == {1, 2}, "1.2 must mean one copy or two, never three and never none" + assert sum(counts) / len(counts) == pytest.approx(1.2, abs=0.02) + + +def test_a_factor_below_one_drops_documents_rather_than_truncating_them(): + counts = [copies_for(0.6, 42, "/data/x.jsonl", o) for o in range(0, 200_000, 10)] + assert set(counts) == {0, 1} + assert sum(counts) / len(counts) == pytest.approx(0.6, abs=0.02) + + +def test_the_choice_is_reproducible_and_seed_dependent(): + a = [copies_for(1.5, 42, "/data/x.jsonl", o) for o in range(2000)] + assert a == [copies_for(1.5, 42, "/data/x.jsonl", o) for o in range(2000)] + assert a != [copies_for(1.5, 43, "/data/x.jsonl", o) for o in range(2000)] + + +def test_the_choice_is_stable_across_processes(): + # A per-process hash seed would redraw which documents were doubled on every run, so a + # resumed export would disagree with the shards it had already written. + script = ( + "from modalities.dataloader.preprocessing.quality.export import copies_for;" + "print([copies_for(1.5, 42, '/data/x.jsonl', o) for o in range(20)])" + ) + out = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + assert eval(out.stdout.strip()) == [copies_for(1.5, 42, "/data/x.jsonl", o) for o in range(20)] + + +def test_the_same_document_in_different_files_is_decided_separately(): + a = [copies_for(1.5, 42, "/data/a.jsonl", o) for o in range(500)] + b = [copies_for(1.5, 42, "/data/b.jsonl", o) for o in range(500)] + assert a != b + + +# --------------------------------------------------------------------------- the output + + +def test_every_document_is_written_once_at_ratio_one(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy")) + exports, root = _export(tmp_path, entry, manifest) + + lines = _lines(root) + assert len(lines) == N_DOCS + assert exports[0].n_lines == N_DOCS + assert [json.loads(line)["id"] for line in lines] == [f"doc-{i}" for i in range(N_DOCS)] + + +def test_output_lines_are_byte_identical_to_the_source(tmp_path: Path, entry: DatasetEntry, blend, corpus: Path): + # Including the last line of the shard that has no trailing newline, which the export + # must terminate itself without altering the record. + manifest = blend(DatasetSelection(name="toy")) + _, root = _export(tmp_path, entry, manifest) + + source = [] + for shard in sorted(corpus.glob("*.jsonl")): + source.extend(shard.read_text().splitlines()) + assert _lines(root) == source + + for shard in sorted((root / "toy").rglob("*.jsonl")): + assert shard.read_bytes().endswith(b"\n"), "every shard must be newline-terminated" + + +def test_an_upsampled_dataset_repeats_every_document(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=3.0)) + exports, root = _export(tmp_path, entry, manifest) + + lines = _lines(root) + assert len(lines) == 3 * N_DOCS + ids = [json.loads(line)["id"] for line in lines] + assert all(ids.count(f"doc-{i}") == 3 for i in range(N_DOCS)) + # Copies are adjacent, as chosen: a document's three copies sit together. + assert ids[:3] == ["doc-0"] * 3 + assert exports[0].n_documents == N_DOCS and exports[0].n_lines == 3 * N_DOCS + + +def test_a_downsampled_dataset_writes_fewer_whole_documents(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=0.5)) + _, root = _export(tmp_path, entry, manifest) + + lines = _lines(root) + assert 0.35 * N_DOCS < len(lines) < 0.65 * N_DOCS + for line in lines: + json.loads(line) # whole records, never a truncated one + assert len(set(json.loads(line)["id"] for line in lines)) == len(lines), "no duplicates below 1.0" + + +def test_a_fractional_upsample_writes_singles_and_doubles(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=1.5)) + _, root = _export(tmp_path, entry, manifest) + + ids = [json.loads(line)["id"] for line in _lines(root)] + counts = {i: ids.count(f"doc-{i}") for i in range(N_DOCS)} + assert set(counts.values()) == {1, 2} + assert sum(counts.values()) == pytest.approx(1.5 * N_DOCS, rel=0.1) + + +def test_a_curve_puts_every_bucket_in_one_directory(tmp_path: Path, entry: DatasetEntry, blend): + # materialize_dataset_buckets emits one row per quality level, each with its own factor. + # They describe one input dataset, so they must produce one output directory -- and a + # source file drawn from by several buckets must yield one shard fed by all of them. + manifest = blend( + DatasetSelection( + name="toy", upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0) + ) + ) + rows = yaml.safe_load(Path(manifest).read_text())["datasets"] + assert len(rows) > 1, "the fixture must actually produce several buckets" + + exports, root = _export(tmp_path, entry, manifest) + + assert [d.name for d in exports] == ["toy"] + assert sorted(p.name for p in (root / "toy").iterdir() if p.is_dir()) == [] + assert len(sorted((root / "toy").rglob("*.jsonl"))) == 2, "one shard per source file, not per bucket" + # Documents come out in source order rather than grouped by quality bucket, which is what + # writing bucket by bucket would have produced. The lowest bucket is discarded by the + # curve, so this checks the order of whatever survived, not of every document. + doc_numbers = [int(json.loads(line)["id"].removeprefix("doc-")) for line in _lines(root)] + assert doc_numbers == sorted(doc_numbers), "output must follow source order" + assert len(set(doc_numbers)) > 1, "the fixture must keep documents from more than one bucket" + + +def test_a_curves_buckets_each_get_their_own_factor(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend( + DatasetSelection( + name="toy", upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0) + ) + ) + rows = {r["name"]: r for r in yaml.safe_load(Path(manifest).read_text())["datasets"]} + _, root = _export(tmp_path, entry, manifest) + + ids = [json.loads(line)["id"] for line in _lines(root)] + # The highest-factor bucket must be repeated more than the lowest. + ranked = sorted(rows.values(), key=lambda r: r["ratio"]) + low, high = ranked[0], ranked[-1] + assert high["ratio"] > low["ratio"], "the fixture needs a spread of factors to be meaningful" + + def mean_copies(row): + levels = row["name"].rsplit("__", 1)[-1] + docs = [i for i in range(N_DOCS) if LEVELS[i % len(LEVELS)] == levels] + return sum(ids.count(f"doc-{i}") for i in docs) / len(docs) if docs else None + + assert mean_copies(high) > mean_copies(low) + + +# --------------------------------------------------------------------------- resume + + +def test_resume_leaves_a_complete_shard_alone(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=2.0)) + _, root = _export(tmp_path, entry, manifest) + shard = sorted((root / "toy").rglob("*.jsonl"))[0] + stamp = shard.stat().st_mtime_ns + + exports, _ = _export(tmp_path, entry, manifest) + + assert shard.stat().st_mtime_ns == stamp + assert all(s.skipped for e in exports for s in e.shards) + assert exports[0].n_lines == 2 * N_DOCS, "a skipped shard still reports what it holds" + + +def test_resume_rewrites_a_shard_that_never_finished(tmp_path: Path, entry: DatasetEntry, blend): + # Existence is not completion. A killed job leaves a plausible-looking partial file, and + # trusting it would silently shrink the training set. + manifest = blend(DatasetSelection(name="toy")) + _, root = _export(tmp_path, entry, manifest) + shard = sorted((root / "toy").rglob("*.jsonl"))[0] + shard.write_text(shard.read_text()[: len(shard.read_text()) // 2]) + + exports, _ = _export(tmp_path, entry, manifest) + + assert not exports[0].shards[0].skipped + assert len(_lines(root)) == N_DOCS + + +def test_no_resume_redoes_everything(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy")) + _export(tmp_path, entry, manifest) + exports, _ = _export(tmp_path, entry, manifest, resume=False) + assert not any(s.skipped for e in exports for s in e.shards) + + +def test_a_failed_shard_leaves_no_partial_file(tmp_path: Path, entry: DatasetEntry, blend, corpus: Path): + manifest = blend(DatasetSelection(name="toy")) + # Truncate a source file so its index names bytes the file cannot supply. + target = corpus / "shard_0.jsonl" + target.write_bytes(target.read_bytes()[:200]) + + with pytest.raises(ExportError, match="source has changed"): + _export(tmp_path, entry, manifest) + assert list((tmp_path / "out").rglob("*.partial")) == [] + + +# --------------------------------------------------------------------------- the manifest + + +def test_the_manifest_says_the_ratio_is_already_applied(tmp_path: Path, entry: DatasetEntry, blend): + # The footgun this closes: mix_manifest.yaml still says ratio 3.0, and carrying that into + # a weighted_combined config after the repetition is in the bytes trains it nine times. + manifest = blend(DatasetSelection(name="toy", ratio=3.0)) + _, root = _export(tmp_path, entry, manifest) + + exported = yaml.safe_load((root / EXPORT_MANIFEST).read_text()) + assert exported["repeat_factor_applied"] is True + assert exported["training_ratio"] == 1.0 + assert "do not apply" in exported["note"].lower() + + dataset = exported["datasets"][0] + assert dataset["ratio"] == 1.0 and dataset["repeat_factor_applied"] is True + assert dataset["factors_applied"] == {"toy": 3.0}, "what was applied is still recorded" + + +def test_the_manifest_line_count_matches_the_files(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=2.0)) + _, root = _export(tmp_path, entry, manifest) + + exported = yaml.safe_load((root / EXPORT_MANIFEST).read_text()) + assert exported["n_lines"] == len(_lines(root)) == 2 * N_DOCS + assert exported["n_bytes"] == sum(p.stat().st_size for p in (root / "toy").rglob("*.jsonl")) + + +def test_finalizing_without_any_export_is_an_error(tmp_path: Path): + (tmp_path / "empty").mkdir() + with pytest.raises(ExportError, match="nothing has been exported"): + finalize_export(tmp_path / "empty") + + +# --------------------------------------------------------------------------- sharding +# +# One array task per dataset is fine until one dataset holds far more files than the rest. +# On the real blend dolmino has 40,003 source files against a median near 500, which at a +# second per file is twelve hours against twenty minutes -- past the wall limit. + + +def test_shards_split_the_files_and_together_cover_everything(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=2.0)) + + for shard_id in range(2): + _export(tmp_path, entry, manifest, shard_id=shard_id, num_shards=2) + + # Two source files, one per task, and between them the whole dataset. + assert len(sorted((tmp_path / "out" / "toy").rglob("*.jsonl"))) == 2 + exported = yaml.safe_load((tmp_path / "out" / EXPORT_MANIFEST).read_text()) + assert exported["n_lines"] == 2 * N_DOCS + assert len(_lines(tmp_path / "out")) == 2 * N_DOCS + + +def test_a_sharded_export_matches_an_unsharded_one(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy", ratio=1.5)) + + _export(tmp_path, entry, manifest, out="whole") + for shard_id in range(3): + _export(tmp_path, entry, manifest, out="split", shard_id=shard_id, num_shards=3) + + assert _lines(tmp_path / "split") == _lines(tmp_path / "whole") + whole = yaml.safe_load((tmp_path / "whole" / EXPORT_MANIFEST).read_text()) + split = yaml.safe_load((tmp_path / "split" / EXPORT_MANIFEST).read_text()) + for key in ("n_lines", "n_documents", "n_bytes"): + assert split[key] == whole[key], key + assert split["datasets"][0]["n_shards"] == whole["datasets"][0]["n_shards"] + + +def test_each_task_records_its_own_counts(tmp_path: Path, entry: DatasetEntry, blend): + # A shared record would have concurrent tasks overwrite each other's counts, which is the + # race the per-dataset split already avoids at the blend level. + manifest = blend(DatasetSelection(name="toy")) + for shard_id in range(2): + _export(tmp_path, entry, manifest, shard_id=shard_id, num_shards=2) + + records = sorted((tmp_path / "out" / "toy").glob("_export*.yaml")) + assert len(records) == 2, "one record per task" + assert sum(yaml.safe_load(p.read_text())["n_lines"] for p in records) == N_DOCS + + +def test_a_shard_id_outside_the_range_is_refused(tmp_path: Path, entry: DatasetEntry, blend): + manifest = blend(DatasetSelection(name="toy")) + with pytest.raises(ValueError, match="not in"): + _export(tmp_path, entry, manifest, shard_id=3, num_shards=3) diff --git a/tests/dataloader/preprocessing/quality/test_file_manifest.py b/tests/dataloader/preprocessing/quality/test_file_manifest.py new file mode 100644 index 000000000..fe079aad9 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_file_manifest.py @@ -0,0 +1,189 @@ +"""Tests that a sidecar refuses to be used against a source tree that moved. + +These cover the failure that actually happened: a data transfer re-sharded four corpora +after their sidecars had been built. File ids are positions in a sorted file list, so the +same id then named a different file and every recorded byte offset pointed into the wrong +place. The only check that existed compared file *counts*, which a re-shard can leave +unchanged, so the pipeline carried on and would have produced a blend of garbage byte +ranges. +""" + +import json +import random +from pathlib import Path + +import pytest + +from modalities.dataloader.preprocessing.quality.file_manifest import FileManifest, ManifestError +from modalities.dataloader.preprocessing.quality.materialize import MaterializationError, materialize_dataset +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, KeySpec +from modalities.dataloader.preprocessing.quality.selection import DatasetSelection, MissingPolicy +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration +from modalities.dataloader.preprocessing.quality.verify import adopt_manifest, verify_sidecar + + +@pytest.fixture +def small_corpus(tmp_path: Path) -> Path: + """Three shards of documents, long enough that offsets are far from zero.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + rng = random.Random(5) + for shard in range(3): + with (corpus / f"shard_{shard}.jsonl").open("w") as f: + for i in range(60): + f.write( + json.dumps({"id": f"doc-{shard}-{i}", "text": " ".join(["word"] * rng.randint(20, 200))}) + "\n" + ) + return corpus + + +@pytest.fixture +def entry(small_corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=small_corpus, + glob="*.jsonl", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + + +@pytest.fixture +def sidecar(tmp_path: Path, entry: DatasetEntry) -> Path: + directory = tmp_path / "sidecar" + calibration = TokenCalibration(dataset="toy", tokenizer="t", bytes_per_token=4.0) + SidecarBuilder(entry, calibration, index_root=tmp_path / "idx").build(directory, show_progress=False) + return directory + + +def test_building_a_sidecar_records_the_file_list(sidecar: Path, entry: DatasetEntry): + manifest = FileManifest.read(sidecar) + assert manifest.dataset == "toy" + assert [f.path for f in manifest.files] == ["shard_0.jsonl", "shard_1.jsonl", "shard_2.jsonl"] + assert all(f.size > 0 for f in manifest.files) + assert not manifest.drift(entry) + + +def test_a_renamed_directory_is_detected(sidecar: Path, entry: DatasetEntry, tmp_path: Path): + # A tree that moved wholesale is fine: paths are relative, so only the file set has + # to agree. This is what lets a snapshot be taken without invalidating a sidecar. + moved = tmp_path / "corpus_moved" + entry.jsonl_root.rename(moved) + relocated = entry.model_copy(update={"jsonl_root": moved}) + assert not FileManifest.read(sidecar).drift(relocated) + + +def test_a_reshard_that_preserves_the_file_count_is_detected(sidecar: Path, entry: DatasetEntry): + # The real failure: file count unchanged, contents different. The old count check + # passed this and the byte offsets silently pointed past the end of the file. + files = sorted(entry.jsonl_root.glob("*.jsonl")) + assert len(files) == 3 + for path in files: + path.write_text(json.dumps({"id": "replaced", "text": "short"}) + "\n") + + problems = FileManifest.read(sidecar).drift(entry) + assert len(problems) == 3 + assert all("bytes, was" in p for p in problems) + with pytest.raises(ManifestError, match="source tree changed"): + FileManifest.read(sidecar).require_current(entry) + + +def test_a_removed_file_is_named(sidecar: Path, entry: DatasetEntry): + (entry.jsonl_root / "shard_1.jsonl").unlink() + problems = FileManifest.read(sidecar).drift(entry) + assert any("shard_1.jsonl is gone" in p for p in problems) + + +def test_an_added_file_does_not_shift_the_ids(sidecar: Path, entry: DatasetEntry): + # A file sorting *before* the existing ones is the dangerous case: re-globbing would + # renumber every id. Resolution goes through the recorded paths, so it cannot. + (entry.jsonl_root / "aaa_new.jsonl").write_text(json.dumps({"id": "new", "text": "hi"}) + "\n") + manifest = FileManifest.read(sidecar) + assert [p.name for p in manifest.resolve(entry)] == [ + "shard_0.jsonl", + "shard_1.jsonl", + "shard_2.jsonl", + ] + assert any("added to the source tree" in p for p in manifest.drift(entry)) + + +def test_materialize_refuses_a_drifted_source_tree(sidecar: Path, entry: DatasetEntry, tmp_path: Path): + (entry.jsonl_root / "shard_0.jsonl").write_text(json.dumps({"id": "x", "text": "tiny"}) + "\n") + with pytest.raises(MaterializationError, match="source tree changed"): + materialize_dataset( + sidecar_dir=sidecar, + dataset_entry=entry, + dataset_selection=DatasetSelection(name="toy", ratio=1.0), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "out", + show_progress=False, + ) + + +def test_materialize_explains_a_missing_manifest(sidecar: Path, entry: DatasetEntry, tmp_path: Path): + FileManifest.path_for(sidecar).unlink() + with pytest.raises(MaterializationError, match="verify-sidecar --adopt"): + materialize_dataset( + sidecar_dir=sidecar, + dataset_entry=entry, + dataset_selection=DatasetSelection(name="toy", ratio=1.0), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "out", + show_progress=False, + ) + + +def test_verify_passes_on_an_untouched_corpus(sidecar: Path, entry: DatasetEntry): + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.verdict == "VALID" + assert report.n_sampled > 0 + assert report.n_matching == report.n_sampled + + +def test_verify_reports_broken_when_the_bytes_moved(sidecar: Path, entry: DatasetEntry): + # Prepend a line to every file: same file names, similar sizes, every offset now + # points at the wrong document. Sizes catch it, and so do the byte probes. + for path in sorted(entry.jsonl_root.glob("*.jsonl")): + original = path.read_text() + path.write_text(json.dumps({"id": "inserted", "text": "x" * 500}) + "\n" + original) + + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.verdict == "DRIFTED" + assert not report.ok + + +def test_verify_ignores_offset_zero_rows(sidecar: Path, entry: DatasetEntry): + # The first document of any JSONL file parses, so a probe at offset 0 succeeds even + # against a completely different file. Sampling those is how a broken sidecar looked + # healthy during the incident. + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.n_sampled > 0 + manifest = FileManifest.read(sidecar) + for path in manifest.resolve(entry): + path.write_text(json.dumps({"id": "only", "text": "tiny"}) + "\n") + after = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8, check_manifest=False) + assert after.n_matching == 0 + + +def test_adopt_refuses_a_sidecar_that_does_not_verify(sidecar: Path, entry: DatasetEntry): + FileManifest.path_for(sidecar).unlink() + for path in sorted(entry.jsonl_root.glob("*.jsonl")): + path.write_text(json.dumps({"id": "x", "text": "tiny"}) + "\n") + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + with pytest.raises(ManifestError, match="refusing to write a manifest"): + adopt_manifest(sidecar, entry, report) + + +def test_adopt_stamps_a_manifest_onto_a_verified_sidecar(sidecar: Path, entry: DatasetEntry): + FileManifest.path_for(sidecar).unlink() + report = verify_sidecar(sidecar, entry, n_parts=3, n_rows_per_part=8) + assert report.ok and not report.has_manifest + adopt_manifest(sidecar, entry, report) + assert not FileManifest.read(sidecar).drift(entry) + + +def test_manifest_is_not_written_into_the_source_tree(sidecar: Path, entry: DatasetEntry): + # Source corpora are shared and read-only; the pipeline must leave them alone. + assert not list(entry.jsonl_root.glob("*.idx")) + assert not list(entry.jsonl_root.glob("_files.json")) + assert FileManifest.path_for(sidecar).exists() diff --git a/tests/dataloader/preprocessing/quality/test_production_regressions.py b/tests/dataloader/preprocessing/quality/test_production_regressions.py new file mode 100644 index 000000000..6de493400 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_production_regressions.py @@ -0,0 +1,178 @@ +"""Regressions for failures that surfaced only on the full production run. + +Each of these cost hours on a 20 TB blend and none of them had a test. They share a shape: +a stage fails, reports the failure somewhere quiet, and destroys or hides the evidence, so +the symptom appears much later and much further downstream. + + * pointer resolution wrote nulls over the pointers it could not resolve, making a retry + impossible without rebuilding the sidecar from source; + * the annotation scan held a fragment per bucket file, so a split with many small files + exceeded 160 GiB while a split with fifty times the data did not. + +A third belonged to the packing stage -- its skip-if-exists treated a truncated output as +finished. That stage is gone, but the lesson outlived it: `export.py` records a shard's line +and byte counts only after the file is in place, and treats a size mismatch as unfinished. +See `test_export.py::test_resume_rewrites_a_shard_that_never_finished`. +""" + +import json +from pathlib import Path + +import pyarrow.parquet as pq +import pytest + +from modalities.dataloader.preprocessing.quality import annotation_join +from modalities.dataloader.preprocessing.quality.annotation_join import ( + bucket_annotations, + join_annotations, +) +from modalities.dataloader.preprocessing.quality.registry import DatasetEntry, KeyKind, KeySpec +from modalities.dataloader.preprocessing.quality.sidecar import ( + SidecarBuilder, + SidecarWriteError, + resolve_source_pointers, +) +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration + + +EDUCATIONAL_LEVELS = ["none", "minimal", "basic", "moderate", "high"] + + +@pytest.fixture +def toy_corpus(tmp_path: Path) -> Path: + """Two shards of documents. Local rather than shared, so this file stands alone.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + for shard in range(2): + with (corpus / f"shard_{shard}.jsonl").open("w") as f: + for i in range(100): + f.write(json.dumps({"id": f"doc-{shard}-{i}", "text": "word " * (10 + i)}) + "\n") + return corpus + + +@pytest.fixture +def toy_entry(toy_corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=toy_corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + + +@pytest.fixture +def toy_annotations(tmp_path: Path, toy_corpus: Path) -> Path: + """Annotations for the first 150 of the 200 documents.""" + import pyarrow as pa + + rows = {"id": [], "educational_value": []} + for shard in range(2): + with (toy_corpus / f"shard_{shard}.jsonl").open() as f: + for line in f: + if len(rows["id"]) >= 150: + break + record = json.loads(line) + rows["id"].append(record["id"]) + rows["educational_value"].append(EDUCATIONAL_LEVELS[len(rows["id"]) % 5]) + out = tmp_path / "annotations" + out.mkdir() + pq.write_table(pa.table(rows), out / "shard0.parquet") + return out + + +def _pointer_corpus(tmp_path: Path, source_root: Path) -> DatasetEntry: + """A translated corpus whose ids point into another corpus, as KletterMix does.""" + corpus = tmp_path / "translated" + corpus.mkdir() + with (corpus / "part.jsonl").open("w") as f: + for i in range(20): + f.write(json.dumps({"id": f"part_0.jsonl/{i}", "text": f"uebersetzt {i}"}) + "\n") + return DatasetEntry( + name="translated", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="src", + key=KeySpec( + kind=KeyKind.SOURCE_POINTER, field="id", source_root=source_root, source_line_offset=0 + ), + ) + + +def test_pointer_resolution_refuses_to_overwrite_pointers_it_cannot_resolve(tmp_path: Path): + """The write-back replaces the pointer with the resolved key, so writing nulls destroys + the only copy. In production the source corpus had been moved and 225 M pointers were + overwritten, logged as INFO, and only noticed two hours later as 0% join coverage.""" + missing = tmp_path / "not_where_it_used_to_be" + missing.mkdir() + entry = _pointer_corpus(tmp_path, missing) + sidecar = tmp_path / "sidecar" + SidecarBuilder( + entry, TokenCalibration(dataset="translated", tokenizer="t", bytes_per_token=4.0), + index_root=tmp_path / "idx", + ).build(sidecar, show_progress=False) + + before = pq.read_table(sidecar / "part-000000.parquet").column("join_key").to_pylist() + assert all(v is not None for v in before), "the builder should have stored the raw pointers" + + with pytest.raises(SidecarWriteError, match="Refusing to write"): + resolve_source_pointers(sidecar, entry) + + after = pq.read_table(sidecar / "part-000000.parquet").column("join_key").to_pylist() + assert after == before, "the pointers must survive a failed resolution so a retry is possible" + + +def test_pointer_resolution_still_writes_when_it_resolves(tmp_path: Path): + source = tmp_path / "source" + source.mkdir() + with (source / "part_0.jsonl").open("w") as f: + for i in range(20): + f.write(json.dumps({"text": f"original {i}"}) + "\n") + entry = _pointer_corpus(tmp_path, source) + sidecar = tmp_path / "sidecar" + SidecarBuilder( + entry, TokenCalibration(dataset="translated", tokenizer="t", bytes_per_token=4.0), + index_root=tmp_path / "idx", + ).build(sidecar, show_progress=False) + + n = resolve_source_pointers(sidecar, entry) + assert n == 20 + keys = pq.read_table(sidecar / "part-000000.parquet").column("join_key").to_pylist() + # Pointers are replaced by content hashes of the source text. + assert all(k is not None and len(k) == 64 for k in keys) + + +def test_the_join_is_unchanged_by_how_many_bucket_files_it_scans_at_once( + tmp_path: Path, monkeypatch, toy_entry: DatasetEntry, toy_annotations: Path +): + """Chunking the fragment list is what stopped HPLT exceeding 160 GiB -- a split with + 65,536 files and 10.7 M documents died where one with 16,384 files and 504 M did not, so + the cost is per fragment. Chunking must not change what the join produces.""" + calibration = TokenCalibration(dataset="toy", tokenizer="t", bytes_per_token=4.0) + buckets = tmp_path / "buckets" + bucket_annotations( + shard_paths=sorted(toy_annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=16, + label_columns=["educational_value"], + show_progress=False, + ) + + results = {} + for group_size in (1, 3, 1024): + sidecar = tmp_path / f"sidecar_{group_size}" + SidecarBuilder(toy_entry, calibration, index_root=tmp_path / f"idx_{group_size}").build( + sidecar, show_progress=False + ) + monkeypatch.setattr(annotation_join, "SCAN_FILE_GROUP", group_size) + report = join_annotations(sidecar, buckets, "toy", "toy", show_progress=False) + labels = [] + for part in sorted(sidecar.glob("part-*.parquet")): + labels.extend(pq.read_table(part, columns=["educational_value"]).column(0).to_pylist()) + results[group_size] = (report.n_matched, report.n_documents, labels) + + reference = results[1024] + for group_size, outcome in results.items(): + assert outcome == reference, f"SCAN_FILE_GROUP={group_size} changed the join result" + assert reference[0] == 150, "the fixture annotates 150 of 200 documents" + diff --git a/tests/dataloader/preprocessing/quality/test_quality_pipeline.py b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py new file mode 100644 index 000000000..ae5b7628c --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_quality_pipeline.py @@ -0,0 +1,1484 @@ +import json +import pickle +import random +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from modalities.dataloader.large_file_lines_reader import LargeFileLinesReader +from modalities.dataloader.preprocessing.quality.annotation_join import bucket_annotations, join_annotations +from modalities.dataloader.preprocessing.quality.cube import build_cube +from modalities.dataloader.preprocessing.quality.materialize import ( + MaterializationError, + materialize_dataset, + materialize_dataset_buckets, +) +from modalities.dataloader.preprocessing.quality.registry import ( + CorpusRegistry, + DatasetEntry, + KeyKind, + KeySpec, + NativeMetric, + SourcePointerResolver, + sha256_text, + strip_urn_uuid, +) +from modalities.dataloader.preprocessing.quality.selection import ( + DatasetSelection, + MissingPolicy, + Op, + Predicate, + evaluate_on_cube, + evaluate_on_sidecar, + predicate_breakdown, +) +from modalities.dataloader.preprocessing.quality.upsampling import UNANNOTATED_BUCKET, UpsamplingSpec +from modalities.dataloader.preprocessing.quality.sidecar import SidecarBuilder +from modalities.dataloader.preprocessing.quality.tokens import TokenCalibration, calibrate_dataset + +EDUCATIONAL_LEVELS = ["none", "minimal", "basic", "moderate", "high"] + + +class _WhitespaceTokenizer: + """Stand-in tokenizer, so calibration can be tested without a model download.""" + + def tokenize(self, text: str) -> list[str]: + return text.split() + + +@pytest.fixture +def corpus(tmp_path: Path) -> Path: + """Two shards whose documents get longer as their quality rises.""" + corpus_dir = tmp_path / "corpus" + corpus_dir.mkdir() + rng = random.Random(3) + for shard in range(2): + with (corpus_dir / f"shard_{shard}.jsonl").open("w") as f: + for i in range(100): + quality = rng.random() + record = { + "id": f"doc-{shard}-{i}", + "text": " ".join(["word"] * int(10 + 200 * quality)), + "score": round(5 * quality, 4), + # Padding so the stored line is much longer than the text, which is + # what makes estimating from line length wrong. + "payload": "x" * 300, + } + f.write(json.dumps(record) + "\n") + return corpus_dir + + +@pytest.fixture +def annotations(tmp_path: Path, corpus: Path) -> Path: + """Annotations for the first 150 of the 200 documents.""" + rows = {"id": [], "educational_value": []} + random.Random(4) + for shard in range(2): + with (corpus / f"shard_{shard}.jsonl").open() as f: + for line in f: + record = json.loads(line) + if len(rows["id"]) >= 150: + break + rows["id"].append(record["id"]) + rows["educational_value"].append(EDUCATIONAL_LEVELS[min(4, int(record["score"] / 5 * 5))]) + annotation_dir = tmp_path / "annotations" + annotation_dir.mkdir() + pq.write_table(pa.table(rows), annotation_dir / "shard0.parquet") + return annotation_dir + + +@pytest.fixture +def dataset_entry(corpus: Path) -> DatasetEntry: + return DatasetEntry( + name="toy", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + native_metrics=[NativeMetric(name="score", jq_pattern=".score")], + ) + + +@pytest.fixture +def built_sidecar(tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path) -> Path: + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=100, + ) + sidecar_dir = tmp_path / "sidecar" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx").build(sidecar_dir, show_progress=False) + + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets", + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + join_annotations(sidecar_dir, tmp_path / "buckets", "toy", "toy", show_progress=False) + return sidecar_dir + + +def test_strip_urn_uuid_handles_both_stored_forms(): + assert strip_urn_uuid("") == "abc-123" + assert strip_urn_uuid("abc-123") == "abc-123" + + +def test_sha256_key_is_taken_over_the_exact_bytes(): + # Stripping or appending whitespace produces a key that matches nothing, so the + # digest must be over the unmodified text. + assert sha256_text("hello") != sha256_text("hello\n") + assert sha256_text("hello") != sha256_text(" hello ") + + +def test_calibration_prefers_a_native_count_when_every_record_has_one(tmp_path: Path): + corpus_dir = tmp_path / "native" + corpus_dir.mkdir() + with (corpus_dir / "a.jsonl").open("w") as f: + for i in range(50): + text = " ".join(["w"] * (i + 1)) + f.write(json.dumps({"text": text, "token_count": (i + 1) * 2}) + "\n") + + calibration = calibrate_dataset( + dataset_name="native", + file_paths=[corpus_dir / "a.jsonl"], + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + + assert calibration.uses_native_field() + assert calibration.native_field == "token_count" + # The corpus counts twice as many tokens as our tokenizer, so the scale halves them. + assert calibration.native_scale == pytest.approx(0.5, abs=1e-6) + + +def test_calibration_ignores_a_partially_present_native_count(tmp_path: Path): + corpus_dir = tmp_path / "partial" + corpus_dir.mkdir() + with (corpus_dir / "a.jsonl").open("w") as f: + for i in range(50): + record = {"text": " ".join(["w"] * (i + 1))} + if i % 2 == 0: + record["token_count"] = i + 1 + f.write(json.dumps(record) + "\n") + + calibration = calibrate_dataset( + dataset_name="partial", + file_paths=[corpus_dir / "a.jsonl"], + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + + assert not calibration.uses_native_field() + + +def test_estimate_falls_back_to_bytes_per_token(): + calibration = TokenCalibration( + dataset="toy", tokenizer="whitespace", bytes_per_token=4.0, eod_tokens_per_document=1 + ) + + assert calibration.estimate({}, text_bytes=400) == 101 + + +def test_sidecar_estimates_from_text_not_from_the_stored_line(built_sidecar: Path): + table = pq.read_table(sorted(built_sidecar.glob("part-*.parquet"))[0]) + + # Every record carries 300 bytes of padding, so a line-length estimate would be + # inflated by roughly that much on the shortest documents. + assert (table.column("byte_len").to_numpy() > table.column("text_bytes").to_numpy()).all() + shortest = min(table.column("text_bytes").to_pylist()) + assert shortest < 300 + + +def test_join_reports_partial_coverage(tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path): + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=100, + ) + sidecar_dir = tmp_path / "sidecar2" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx2").build(sidecar_dir, show_progress=False) + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets2", + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + + report = join_annotations(sidecar_dir, tmp_path / "buckets2", "toy", "toy", show_progress=False) + + assert report.n_documents == 200 + assert report.n_matched == 150 + assert report.coverage == pytest.approx(0.75) + + +@pytest.mark.parametrize("policy", [MissingPolicy.KEEP, MissingPolicy.DROP]) +def test_cube_agrees_with_the_per_document_scan(built_sidecar: Path, policy: MissingPolicy): + cube = build_cube(built_sidecar, "toy") + selection = DatasetSelection( + name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")] + ) + + from_cube = evaluate_on_cube(cube, selection, policy) + from_sidecar = evaluate_on_sidecar(built_sidecar, selection, policy) + + assert from_cube.n_documents_kept == from_sidecar.n_documents_kept + assert from_cube.tokens_kept == from_sidecar.tokens_kept + + +def test_materialized_index_is_loadable_and_selects_the_right_documents( + built_sidecar: Path, dataset_entry: DatasetEntry, tmp_path: Path +): + selection = DatasetSelection( + name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="moderate")] + ) + expected = evaluate_on_sidecar(built_sidecar, selection, MissingPolicy.DROP) + + result = materialize_dataset( + sidecar_dir=built_sidecar, + dataset_entry=dataset_entry, + dataset_selection=selection, + missing_policy=MissingPolicy.DROP, + output_dir=tmp_path / "indexes", + show_progress=False, + ) + + assert result.n_documents_kept == expected.n_documents_kept + + total_in_indexes = 0 + for source_path, index_path in result.index_files.items(): + entries = pickle.loads(Path(index_path).read_bytes()) + total_in_indexes += len(entries) + # The index must be readable by the same reader the packer uses, and every + # document it names must satisfy the predicate. + reader = LargeFileLinesReader(Path(source_path), index_path=Path(index_path)) + try: + assert len(reader) == len(entries) + for i in range(len(reader)): + record = json.loads(reader[i]) + assert record["id"].startswith("doc-") + finally: + reader.close() + assert total_in_indexes == expected.n_documents_kept + + +def test_materialized_index_entries_are_ordered_by_position( + built_sidecar: Path, dataset_entry: DatasetEntry, tmp_path: Path +): + result = materialize_dataset( + sidecar_dir=built_sidecar, + dataset_entry=dataset_entry, + dataset_selection=DatasetSelection(name="toy", predicates=[Predicate(field="score", op=Op.GTE, value=1.0)]), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "indexes2", + show_progress=False, + ) + + for index_path in result.index_files.values(): + entries = pickle.loads(Path(index_path).read_bytes()) + offsets = [offset for offset, _ in entries] + assert offsets == sorted(offsets), "a generated index lists documents in file order" + + +def test_source_pointer_resolution_uses_zero_indexed_lines(tmp_path: Path): + source_dir = tmp_path / "source" + source_dir.mkdir() + with (source_dir / "part_0.jsonl").open("w") as f: + for i in range(5): + f.write(json.dumps({"text": f"line {i}"}) + "\n") + + resolver = SourcePointerResolver(source_root=source_dir) + resolved = resolver.resolve(["part_0.jsonl/0", "part_0.jsonl/3"]) + + assert resolved["part_0.jsonl/0"] == sha256_text("line 0") + assert resolved["part_0.jsonl/3"] == sha256_text("line 3") + + +def test_source_pointer_resolution_skips_pointers_past_the_end(tmp_path: Path): + source_dir = tmp_path / "source2" + source_dir.mkdir() + (source_dir / "part_0.jsonl").write_text(json.dumps({"text": "only line"}) + "\n") + + resolved = SourcePointerResolver(source_root=source_dir).resolve(["part_0.jsonl/0", "part_0.jsonl/9"]) + + assert "part_0.jsonl/0" in resolved + assert "part_0.jsonl/9" not in resolved, "a pointer past the end must not map to a wrong key" + + +def test_registry_rejects_an_annotated_dataset_without_a_key(tmp_path: Path): + with pytest.raises(ValueError, match="no key spec"): + DatasetEntry(name="x", jsonl_root=tmp_path, annotation_split="some/split") + + +def test_registry_rejects_duplicate_dataset_names(tmp_path: Path): + with pytest.raises(ValueError, match="duplicate dataset name"): + CorpusRegistry( + datasets=[DatasetEntry(name="x", jsonl_root=tmp_path), DatasetEntry(name="x", jsonl_root=tmp_path)] + ) + + +# --------------------------------------------------------------- fast path extraction + + +@pytest.mark.parametrize( + "pattern,expected", + [ + (".score", ["score"]), + (".metadata.len_cl100k_base", ["metadata", "len_cl100k_base"]), + ('."openlid-v3".prob', ["openlid-v3", "prob"]), + ('.metadata.dclm_plus2."__label__1"', ["metadata", "dclm_plus2", "__label__1"]), + ], +) +def test_simple_paths_are_recognised(pattern, expected): + from modalities.dataloader.preprocessing.quality.sidecar import parse_simple_path + + assert parse_simple_path(pattern) == expected + + +@pytest.mark.parametrize("pattern", [".scores | max", ".a[0]", "select(.x)", ".a | length", "."]) +def test_non_path_patterns_fall_back_to_jq(pattern): + from modalities.dataloader.preprocessing.quality.sidecar import build_metric_extractor, parse_simple_path + + assert parse_simple_path(pattern) is None + _, is_fast = build_metric_extractor(pattern) + assert is_fast is False + + +def test_an_invalid_pattern_fails_when_the_builder_is_created(tmp_path: Path, corpus: Path): + # Better to fail here than to hand back an extractor that yields None for every + # document and look like a corpus with no metrics. + entry = DatasetEntry( + name="bad", + jsonl_root=corpus, + glob="*.jsonl", + native_metrics=[NativeMetric(name="broken", jq_pattern="score")], + ) + calibration = TokenCalibration(dataset="bad", tokenizer="w", bytes_per_token=4.0) + + with pytest.raises(ValueError, match="compile error"): + SidecarBuilder(entry, calibration, index_root=tmp_path / "bad_idx") + + +def test_fast_extractor_agrees_with_jq_on_real_shapes(): + import jq + + from modalities.dataloader.preprocessing.quality.sidecar import build_metric_extractor + + record = { + "score": 2.5, + "zero": 0, + "flag": False, + "openlid-v3": {"prob": [0.99, 0.01]}, + "metadata": {"dclm_plus2": {"__label__1": 0.94}, "nested": None}, + } + for pattern in ( + ".score", + ".zero", + ".flag", + '."openlid-v3".prob', + '.metadata.dclm_plus2."__label__1"', + ".missing", + ".metadata.nested.deeper", + ".score.deeper", + ): + fast, is_fast = build_metric_extractor(pattern) + assert is_fast, pattern + try: + expected = jq.compile(pattern).input_value(record).first() + except (ValueError, StopIteration): + expected = None + assert fast(record) == expected, f"{pattern}: {fast(record)!r} != {expected!r}" + + +def test_jq_fallback_still_extracts(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.sidecar import build_metric_extractor + + extract, is_fast = build_metric_extractor('."openlid-v3".prob | max') + assert not is_fast + assert extract({"openlid-v3": {"prob": [0.1, 0.9]}}) == 0.9 + + +def test_sidecar_uses_the_fast_path_and_still_records_metrics(tmp_path: Path, corpus: Path): + # A dataset whose metrics are all plain paths must produce the same values it would + # have produced through jq. + entry = DatasetEntry( + name="fast", + jsonl_root=corpus, + glob="*.jsonl", + native_metrics=[NativeMetric(name="score", jq_pattern=".score")], + ) + calibration = TokenCalibration(dataset="fast", tokenizer="w", bytes_per_token=4.0) + out = tmp_path / "fast_sidecar" + SidecarBuilder(entry, calibration, index_root=tmp_path / "fast_idx").build(out, show_progress=False) + + table = pq.read_table(sorted(out.glob("part-*.parquet"))[0]) + values = table.column("native_score").to_pylist() + assert len(values) == 100 + assert all(v is not None for v in values) + + +# --------------------------------------------------------------------- cube vectorising + + +def test_cube_is_independent_of_the_aggregation_batch_size(built_sidecar: Path): + # Batching row groups is a performance detail and must not change the result. + big = build_cube(built_sidecar, "toy", aggregate_batch_rows=10_000_000) + small = build_cube(built_sidecar, "toy", aggregate_batch_rows=1) + + assert big.n_documents == small.n_documents + assert big.n_tokens == small.n_tokens + assert big.table.num_rows == small.table.num_rows + + def as_set(cube): + cols = cube.dimensions + ["n_documents", "n_tokens"] + return {tuple(row[c] for c in cols) for row in cube.table.select(cols).to_pylist()} + + assert as_set(big) == as_set(small) + + +def test_cube_totals_match_the_sidecar(built_sidecar: Path): + cube = build_cube(built_sidecar, "toy") + total_docs = total_tokens = 0 + for part in sorted(built_sidecar.glob("part-*.parquet")): + table = pq.read_table(part, columns=["est_tokens"]) + total_docs += table.num_rows + total_tokens += sum(table.column("est_tokens").to_pylist()) + + assert cube.n_documents == total_docs + assert cube.n_tokens == total_tokens + assert sum(cube.table.column("n_documents").to_pylist()) == total_docs + assert sum(cube.table.column("n_tokens").to_pylist()) == total_tokens + + +# ------------------------------------------------------------------------- sharding + + +def test_plan_sidecar_work_partitions_completely_and_disjointly(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality.pipeline import plan_sidecar_work + + other = tmp_path / "other" + other.mkdir() + for i in range(5): + (other / f"f{i}.jsonl").write_text(json.dumps({"text": "x"}) + "\n") + registry = CorpusRegistry( + datasets=[ + DatasetEntry(name="a", jsonl_root=corpus, glob="*.jsonl"), + DatasetEntry(name="b", jsonl_root=other, glob="*.jsonl"), + ] + ) + + seen: list[tuple[str, int]] = [] + for shard_id in range(3): + for name, file_ids in plan_sidecar_work(registry, shard_id=shard_id, num_shards=3).items(): + seen.extend((name, file_id) for file_id in file_ids) + + expected = [("a", i) for i in range(2)] + [("b", i) for i in range(5)] + assert sorted(seen) == sorted(expected), "every file must be built exactly once across the array" + + +def test_plan_sidecar_work_rejects_an_out_of_range_shard(corpus: Path): + from modalities.dataloader.preprocessing.quality.pipeline import plan_sidecar_work + + registry = CorpusRegistry(datasets=[DatasetEntry(name="a", jsonl_root=corpus, glob="*.jsonl")]) + + with pytest.raises(ValueError, match="not in"): + plan_sidecar_work(registry, shard_id=3, num_shards=3) + + +def _build_sidecar_only(tmp_path: Path, dataset_entry: DatasetEntry, suffix: str) -> Path: + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=100, + ) + out = tmp_path / f"sidecar_{suffix}" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / f"idx_{suffix}").build(out, show_progress=False) + return out + + +def test_sharded_bucketing_gives_the_same_join_as_a_single_task( + tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path +): + # Spread the annotations over several files so there is something to shard. + split_dir = tmp_path / "annotations_split" + split_dir.mkdir() + table = pq.read_table(sorted(annotations.glob("*.parquet"))[0]) + for i in range(3): + pq.write_table(table.slice(i * 50, 50), split_dir / f"part{i}.parquet") + shards = sorted(split_dir.glob("*.parquet")) + + single_sidecar = _build_sidecar_only(tmp_path, dataset_entry, "single") + bucket_annotations( + shard_paths=shards, + out_dir=tmp_path / "buckets_single", + n_buckets=8, + label_columns=["educational_value"], + show_progress=False, + ) + single = join_annotations(single_sidecar, tmp_path / "buckets_single", "toy", "toy", show_progress=False) + + sharded_sidecar = _build_sidecar_only(tmp_path, dataset_entry, "sharded") + for shard_id in range(3): + bucket_annotations( + shard_paths=shards, + out_dir=tmp_path / "buckets_sharded", + n_buckets=8, + label_columns=["educational_value"], + shard_id=shard_id, + num_shards=3, + show_progress=False, + ) + sharded = join_annotations(sharded_sidecar, tmp_path / "buckets_sharded", "toy", "toy", show_progress=False) + + assert sharded.n_matched == single.n_matched + assert sharded.n_annotation_rows == single.n_annotation_rows + # The labels themselves, not just the counts. + single_labels = pq.read_table(sorted(single_sidecar.glob("part-*.parquet"))[0]).column("educational_value") + sharded_labels = pq.read_table(sorted(sharded_sidecar.glob("part-*.parquet"))[0]).column("educational_value") + assert single_labels.to_pylist() == sharded_labels.to_pylist() + + +def test_join_refuses_an_incomplete_bucketing_run(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError, read_bucket_metadata + + # One of three announced tasks ran, so two thirds of the annotations are absent. + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets_partial", + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=3, + show_progress=False, + ) + + with pytest.raises(AnnotationJoinError, match="incomplete"): + read_bucket_metadata(tmp_path / "buckets_partial") + + +def test_bucketing_rejects_an_out_of_range_shard(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError + + with pytest.raises(AnnotationJoinError, match="not in"): + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=tmp_path / "buckets_bad", + n_buckets=4, + shard_id=5, + num_shards=3, + show_progress=False, + ) + + +# ------------------------------------------------------- calibration read is bounded + + +def test_calibration_read_does_not_scale_with_file_count(tmp_path: Path): + # The whole point: a dataset with many files must not cost many times more to + # calibrate. Reading a fixed number of lines from every file made this stage read + # 30 TB over the real blend. + from modalities.dataloader.preprocessing.quality import tokens as tokens_module + + corpus_dir = tmp_path / "many_files" + corpus_dir.mkdir() + for f in range(200): + with (corpus_dir / f"shard_{f:04d}.jsonl").open("w") as fh: + for i in range(500): + fh.write(json.dumps({"text": " ".join(["w"] * 20)}) + "\n") + + opened: list[Path] = [] + original_open = Path.open + + def counting_open(self, *args, **kwargs): + opened.append(self) + return original_open(self, *args, **kwargs) + + monkey = pytest.MonkeyPatch() + monkey.setattr(Path, "open", counting_open) + try: + calibration = tokens_module.calibrate_dataset( + dataset_name="many", + file_paths=sorted(corpus_dir.glob("*.jsonl")), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=200, + max_probe_files=32, + ) + finally: + monkey.undo() + + assert calibration.sampled_documents == 200 + jsonl_opened = [p for p in opened if p.suffix == ".jsonl"] + assert len(jsonl_opened) <= 32, f"opened {len(jsonl_opened)} of 200 files; the read must be bounded" + + +def test_calibration_probes_files_across_the_whole_dataset(tmp_path: Path): + # Spread matters: a prefix would calibrate on whatever the corpus happens to be + # ordered by. Each file here has a distinct text length, so the sample reveals reach. + corpus_dir = tmp_path / "spread" + corpus_dir.mkdir() + for f in range(100): + with (corpus_dir / f"shard_{f:04d}.jsonl").open("w") as fh: + for _ in range(20): + fh.write(json.dumps({"text": " ".join(["w"] * (f + 1))}) + "\n") + + from modalities.dataloader.preprocessing.quality.tokens import _probe_files + + probed = _probe_files(sorted(corpus_dir.glob("*.jsonl")), max_probe_files=10) + + assert len(probed) == 10 + indices = [int(p.stem.split("_")[1]) for p in probed] + assert indices[0] < 10 and indices[-1] > 80, f"probe files clustered: {indices}" + + +def test_calibration_is_reproducible_for_a_given_seed(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality.tokens import calibrate_dataset as calibrate + + kwargs = dict( + dataset_name="toy", + file_paths=sorted(corpus.glob("*.jsonl")), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + first = calibrate(**kwargs, seed=7) + second = calibrate(**kwargs, seed=7) + + assert first.bytes_per_token == second.bytes_per_token + assert first.sampled_tokens == second.sampled_tokens + + +def test_calibration_is_written_after_each_dataset(tmp_path: Path, corpus: Path): + # Interrupting a long calibration must not throw away what it already measured. + from modalities.dataloader.preprocessing.quality import pipeline + from modalities.dataloader.preprocessing.quality.tokens import CalibrationSet + + other = tmp_path / "other_corpus" + other.mkdir() + (other / "a.jsonl").write_text(json.dumps({"text": "one two three"}) + "\n") + registry = CorpusRegistry( + datasets=[ + DatasetEntry(name="first", jsonl_root=corpus, glob="*.jsonl"), + DatasetEntry(name="second", jsonl_root=other, glob="*.jsonl"), + ] + ) + + work_dir = tmp_path / "work" + seen_after_first: list[list[str]] = [] + original = CalibrationSet.to_yaml + + def recording_to_yaml(self, path): + original(self, path) + seen_after_first.append(sorted(self.calibrations)) + + monkey = pytest.MonkeyPatch() + monkey.setattr(CalibrationSet, "to_yaml", recording_to_yaml) + try: + pipeline.calibrate_blend( + registry=registry, + work_dir=work_dir, + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=20, + ) + finally: + monkey.undo() + + assert seen_after_first == [["first"], ["first", "second"]], seen_after_first + + +# ------------------------------------------------------- bucket writer memory bound + + +def test_bucket_writer_memory_is_bounded_by_total_not_per_bucket(tmp_path: Path): + # The regression that OOM-killed all 64 tasks of a real run: the flush threshold was + # per bucket, so with many buckets no single bucket ever reached it and the whole + # input stayed in memory. The bound must be on the total held across buckets. + from modalities.dataloader.preprocessing.quality.annotation_join import _BucketWriter + + schema = pa.schema([pa.field("key", pa.large_string()), pa.field("label", pa.large_string())]) + writer = _BucketWriter(tmp_path / "buckets", schema, n_buckets=1024, max_buffered_rows=1000) + try: + # Spread 20,000 rows over 1024 buckets: ~20 per bucket, far below any sane + # per-bucket threshold, so a per-bucket rule would never flush. + for i in range(20_000): + writer.add(i % 1024, {"key": f"k{i}", "label": "x"}) + assert writer._buffered_rows < 1000 + 1, "total buffered rows exceeded the cap" + finally: + writer.close() + + written = sorted((tmp_path / "buckets").glob("*.parquet")) + assert written, "nothing was written" + total = sum(pq.ParquetFile(p).metadata.num_rows for p in written) + assert total == 20_000, f"rows lost or duplicated across flushes: {total}" + + +def test_bucket_writer_survives_repeated_flushes_of_the_same_bucket(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import _BucketWriter + + schema = pa.schema([pa.field("key", pa.large_string())]) + writer = _BucketWriter(tmp_path / "b", schema, n_buckets=2, max_buffered_rows=10) + try: + for i in range(100): + writer.add(0, {"key": f"k{i}"}) + finally: + writer.close() + + files = list((tmp_path / "b").glob("*.parquet")) + assert len(files) == 1, "one bucket must stay one file across flushes" + assert pq.ParquetFile(files[0]).metadata.num_rows == 100 + + +def test_bucketing_refuses_to_mix_runs_with_different_array_sizes(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError + + shards = sorted(annotations.glob("*.parquet")) + out = tmp_path / "mixed_buckets" + bucket_annotations( + shard_paths=shards, + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=4, + show_progress=False, + ) + + with pytest.raises(AnnotationJoinError, match="num_shards"): + bucket_annotations( + shard_paths=shards, + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=8, + show_progress=False, + ) + + +# --------------------------------------------- bucketing metadata write/read race + + +def test_a_truncated_metadata_file_does_not_stop_a_run(tmp_path: Path, annotations: Path): + # The regression: the guard read every _meta file with a bare json.loads, so a task + # that caught a sibling's file mid-write died. 12 of 64 tasks of a real run crashed + # this way. + out = tmp_path / "buckets_truncated" + out.mkdir() + (out / "_meta.0005.json").write_text("") + + n_rows, _ = bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=1, + show_progress=False, + ) + + assert n_rows == 150 + + +def test_read_bucket_metadata_skips_an_unreadable_file(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import AnnotationJoinError, read_bucket_metadata + + out = tmp_path / "buckets_partial_meta" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=2, + show_progress=False, + ) + # Shard 1's metadata exists but is corrupt: the run must read as incomplete rather + # than raising a decode error or, worse, joining without shard 1's annotations. + (out / "_meta.0001.json").write_text("{ truncated") + + with pytest.raises(AnnotationJoinError, match="incomplete"): + read_bucket_metadata(out) + + +def test_metadata_write_is_atomic_and_leaves_no_temp_file(tmp_path: Path, annotations: Path): + from modalities.dataloader.preprocessing.quality.annotation_join import read_bucket_metadata + + out = tmp_path / "buckets_atomic" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=1, + show_progress=False, + ) + + assert not list(out.glob("*.tmp")), "an atomic write must not leave its temporary file behind" + # A stray .tmp must be ignored rather than parsed as metadata. + (out / "_meta.0009.json.tmp").write_text("{ half written") + meta = read_bucket_metadata(out) + assert meta["n_rows"] == 150 + + +def test_guard_survives_metadata_being_rewritten_concurrently(tmp_path: Path, annotations: Path): + # Exercises the actual race: one thread rewriting metadata while bucketing tasks keep + # entering the directory and running the guard. + import threading + + from modalities.dataloader.preprocessing.quality.annotation_join import _write_metadata + + out = tmp_path / "buckets_concurrent" + out.mkdir() + payload = {"n_buckets": 4, "label_columns": ["educational_value"], "n_rows": 1, "shard_id": 7, "num_shards": 4} + + stop = threading.Event() + + def rewrite() -> None: + while not stop.is_set(): + _write_metadata(out / "_meta.0007.json", payload) + + writer = threading.Thread(target=rewrite, daemon=True) + writer.start() + try: + for _ in range(20): + # num_shards matches the payload, so the guard must pass rather than raise -- + # and must not blow up on a file being replaced underneath it. + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + shard_id=0, + num_shards=4, + show_progress=False, + ) + finally: + stop.set() + writer.join(timeout=5) + + +# ------------------------------------------------- join batching (read amplification) + + +def _sidecar_with_many_parts(tmp_path: Path, corpus: Path, suffix: str) -> tuple[Path, DatasetEntry]: + """A sidecar with one part per file, so batching has something to batch.""" + entry = DatasetEntry( + name="toy", + jsonl_root=corpus, + glob="*.jsonl", + annotation_split="toy", + key=KeySpec(kind=KeyKind.FIELD, field="id"), + ) + calibration = TokenCalibration(dataset="toy", tokenizer="w", bytes_per_token=4.0) + out = tmp_path / f"sidecar_{suffix}" + SidecarBuilder(entry, calibration, index_root=tmp_path / f"idx_{suffix}").build(out, show_progress=False) + return out, entry + + +def test_join_result_is_independent_of_the_batch_size(tmp_path: Path, corpus: Path, annotations: Path): + buckets = tmp_path / "b" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=8, + label_columns=["educational_value"], + show_progress=False, + ) + + results = {} + for label, batch in (("one_part_at_a_time", 1), ("all_at_once", 10_000_000)): + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, label) + report = join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=batch, show_progress=False) + labels = [] + for part in sorted(sidecar.glob("part-*.parquet")): + labels.extend(pq.read_table(part).column("educational_value").to_pylist()) + results[label] = (report.n_documents, report.n_matched, labels) + + small, large = results["one_part_at_a_time"], results["all_at_once"] + assert small[0] == large[0] == 200 + assert small[1] == large[1] == 150 + assert small[2] == large[2], "batching must not change which label each document gets" + + +def test_join_reads_each_bucket_once_per_batch_not_once_per_part( + tmp_path: Path, corpus: Path, annotations: Path, monkeypatch +): + # The regression this guards: reading the bucketed split once per sidecar part meant + # 454 TB of reads over the real blend, because every part's documents hash across + # every bucket. + buckets = tmp_path / "b2" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=8, + label_columns=["educational_value"], + show_progress=False, + ) + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "counted") + n_parts = len(list(sidecar.glob("part-*.parquet"))) + assert n_parts == 2 + + import pyarrow.parquet as pq_module + + reads: list[str] = [] + original = pq_module.read_table + + def counting_read_table(source, *args, **kwargs): + reads.append(str(source)) + return original(source, *args, **kwargs) + + monkeypatch.setattr(pq_module, "read_table", counting_read_table) + join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=10_000_000, show_progress=False) + + bucket_reads = [r for r in reads if "bucket-" in r] + # One batch covers both parts, so each populated bucket is read once -- not twice. + assert len(bucket_reads) == len(set(bucket_reads)), f"a bucket was read more than once: {bucket_reads}" + assert len(bucket_reads) <= 8 + + +def test_join_counts_a_duplicate_key_once_not_once_per_part(tmp_path: Path, corpus: Path): + # Duplicates used to be counted afresh on every part, inflating the figure by the part + # count. Only duplicates among keys the join actually wants should be reported. + rows = {"id": ["doc-0-0", "doc-0-0", "doc-0-1"], "educational_value": ["high", "basic", "high"]} + ann_dir = tmp_path / "dup_ann" + ann_dir.mkdir() + pq.write_table(pa.table(rows), ann_dir / "a.parquet") + + buckets = tmp_path / "dup_buckets" + bucket_annotations( + shard_paths=[ann_dir / "a.parquet"], + out_dir=buckets, + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "dup") + + report = join_annotations(sidecar, buckets, "toy", "toy", max_batch_keys=1, show_progress=False) + + assert report.n_duplicate_keys == 1, f"expected the one duplicate counted once, got {report.n_duplicate_keys}" + + +# ------------------------------------------------- resume, and cube consistency checks + + +def _buckets_with_label(tmp_path: Path, corpus: Path, name: str, level: str) -> Path: + """Buckets giving every document of `corpus` the same educational_value.""" + ids = [] + for shard in sorted(corpus.glob("*.jsonl")): + with shard.open() as f: + ids.extend(json.loads(line)["id"] for line in f) + ann = tmp_path / f"ann_{name}" + ann.mkdir() + pq.write_table(pa.table({"id": ids, "educational_value": [level] * len(ids)}), ann / "a.parquet") + out = tmp_path / f"buckets_{name}" + bucket_annotations( + shard_paths=[ann / "a.parquet"], + out_dir=out, + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + return out + + +def _labels_of(sidecar: Path) -> dict[str, list]: + return { + p.name: pq.read_table(p).column("educational_value").to_pylist() for p in sorted(sidecar.glob("part-*.parquet")) + } + + +def test_resume_skips_already_labelled_parts_and_keeps_their_values(tmp_path: Path, corpus: Path): + first = _buckets_with_label(tmp_path, corpus, "first", "high") + second = _buckets_with_label(tmp_path, corpus, "second", "none") + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "resume") + + join_annotations(sidecar, first, "toy", "toy", show_progress=False) + before = _labels_of(sidecar) + assert all(v == "high" for values in before.values() for v in values) + + # Drop the labels from one part, so a resumed run has exactly one part to do. + parts = sorted(sidecar.glob("part-*.parquet")) + stripped = parts[0] + table = pq.read_table(stripped) + pq.write_table(table.drop_columns(["educational_value"]), stripped) + + report = join_annotations(sidecar, second, "toy", "toy", resume=True, show_progress=False) + after = _labels_of(sidecar) + + assert all(v == "none" for v in after[stripped.name]), "the unlabelled part must be joined" + for other in parts[1:]: + assert after[other.name] == before[other.name], "an already-labelled part must be left alone" + # The report describes the whole sidecar, not just the part this run redid: a resumed + # join that counted only its own work reported 0% coverage, which reads as a failure. + assert report.n_parts_resumed == len(parts) - 1 + assert report.n_documents == sum(len(v) for v in after.values()) + assert report.n_matched == report.n_documents + + +def test_resume_off_redoes_every_part(tmp_path: Path, corpus: Path): + first = _buckets_with_label(tmp_path, corpus, "f2", "high") + second = _buckets_with_label(tmp_path, corpus, "s2", "none") + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "noresume") + + join_annotations(sidecar, first, "toy", "toy", show_progress=False) + join_annotations(sidecar, second, "toy", "toy", resume=False, show_progress=False) + + after = _labels_of(sidecar) + assert all(v == "none" for values in after.values() for v in values), "without resume the labels must be replaced" + + +def test_build_cube_rejects_a_partly_joined_sidecar(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality.cube import CubeError + + buckets = _buckets_with_label(tmp_path, corpus, "partial", "high") + sidecar, _ = _sidecar_with_many_parts(tmp_path, corpus, "partial") + join_annotations(sidecar, buckets, "toy", "toy", show_progress=False) + + # Simulate the interrupted join: one part never got its labels. + victim = sorted(sidecar.glob("part-*.parquet"))[0] + pq.write_table(pq.read_table(victim).drop_columns(["educational_value"]), victim) + + with pytest.raises(CubeError, match="sidecar parts are missing") as excinfo: + build_cube(sidecar, "toy") + message = str(excinfo.value) + assert "toy" in message and "educational_value" in message + assert "--resume" in message, "the error should say how to finish the join" + + +def test_build_cubes_builds_healthy_datasets_before_raising(tmp_path: Path, corpus: Path): + from modalities.dataloader.preprocessing.quality import pipeline + from modalities.dataloader.preprocessing.quality.cube import CubeError + + work = tmp_path / "work" + healthy_sidecar = pipeline.sidecar_dir(work, "healthy") + broken_sidecar = pipeline.sidecar_dir(work, "broken") + + buckets = _buckets_with_label(tmp_path, corpus, "cubes", "high") + for name, target in (("healthy", healthy_sidecar), ("broken", broken_sidecar)): + built, _ = _sidecar_with_many_parts(tmp_path, corpus, name) + target.parent.mkdir(parents=True, exist_ok=True) + target.mkdir(parents=True, exist_ok=True) + for p in sorted(built.glob("part-*.parquet")): + pq.write_table(pq.read_table(p), target / p.name) + join_annotations(target, buckets, name, "toy", show_progress=False) + victim = sorted(broken_sidecar.glob("part-*.parquet"))[0] + pq.write_table(pq.read_table(victim).drop_columns(["educational_value"]), victim) + + registry = CorpusRegistry( + datasets=[ + DatasetEntry(name="healthy", jsonl_root=corpus, glob="*.jsonl"), + DatasetEntry(name="broken", jsonl_root=corpus, glob="*.jsonl"), + ] + ) + + with pytest.raises(CubeError): + pipeline.build_cubes(registry, work) + + # The healthy dataset's cube must exist despite the other one failing. + assert pipeline.cube_path(work, "healthy").is_file() + assert not pipeline.cube_path(work, "broken").is_file() + + +def test_resumed_join_reports_the_sidecars_real_coverage( + tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path +): + """A resumed join must describe the sidecar, not just the work it happened to redo. + + The first version counted only the parts it re-joined, so a run that skipped + everything wrote a report saying 0 documents and 0% coverage. That reads as a failed + join, and on the real blend it overwrote genuine coverage figures with zeros. + """ + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + sidecar_dir = tmp_path / "sidecar" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx").build(sidecar_dir, show_progress=False) + buckets = tmp_path / "buckets" + bucket_annotations( + shard_paths=sorted(annotations.glob("*.parquet")), + out_dir=buckets, + n_buckets=4, + label_columns=["educational_value"], + show_progress=False, + ) + + first = join_annotations(sidecar_dir, buckets, "toy", "toy", show_progress=False) + assert first.n_documents == 200 + assert first.n_matched == 150 + assert first.n_parts_resumed == 0 + + resumed = join_annotations(sidecar_dir, buckets, "toy", "toy", resume=True, show_progress=False) + assert resumed.n_parts_resumed == 2 + assert resumed.n_documents == first.n_documents + assert resumed.n_matched == first.n_matched + assert resumed.coverage == first.coverage + + + +# --------------------------------------------------------------- quality-aware upsampling + + +def _curve_selection(target_ratio: float = 1.0, discard: float = 0.0) -> DatasetSelection: + return DatasetSelection( + name="toy", + upsampling=UpsamplingSpec( + quality_field="educational_value", + target_ratio=target_ratio, + max_factor=4.0, + discard_below_percentile=discard, + ), + ) + + +def test_a_curve_over_a_cube_hits_its_target_and_rises_with_quality(built_sidecar: Path, tmp_path: Path): + cube = build_cube(built_sidecar, "toy", label_dimensions=["educational_value"]) + result = evaluate_on_cube(cube, _curve_selection(target_ratio=1.5), MissingPolicy.KEEP) + + assert result.plan is not None + assert result.effective_tokens == pytest.approx(result.plan.tokens_available * 1.5, rel=1e-6) + factors = [b.factor for b in result.plan.buckets] + assert factors == sorted(factors) + # The unannotated fifty of two hundred documents cannot be ordered, so they form the + # bottom bucket rather than being silently dropped or silently kept at full weight. + assert result.plan.buckets[0].bucket.label == UNANNOTATED_BUCKET + + +def test_a_curve_and_a_flat_ratio_cannot_both_be_given(): + with pytest.raises(ValueError, match="already determines how much is drawn"): + DatasetSelection( + name="toy", + ratio=2.0, + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=1.0), + ) + + +def test_a_curve_needs_an_ordinal_quality_field(): + with pytest.raises(ValueError, match="needs an ordinal quality_field"): + DatasetSelection(name="toy", upsampling=UpsamplingSpec(quality_field="score", target_ratio=1.0)) + + +def test_materializing_a_curve_writes_one_index_tree_per_bucket( + built_sidecar: Path, dataset_entry: DatasetEntry, tmp_path: Path +): + selection = _curve_selection(target_ratio=1.0, discard=20.0) + results = materialize_dataset_buckets( + sidecar_dir=built_sidecar, + dataset_entry=dataset_entry, + dataset_selection=selection, + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "mix", + show_progress=False, + ) + + assert results, "a curve must produce at least one bucket" + # Each row is its own dataset for packing, but still points back at the registry entry. + assert {r.source_dataset for r in results} == {"toy"} + assert all(r.name.startswith("toy__") for r in results) + assert all(r.ratio > 0 for r in results) + # Factors rise with quality, and the whole point is that they differ. + assert len({round(r.ratio, 6) for r in results}) > 1 + + # No document may appear in two buckets: they are disjoint by construction, and an + # overlap would silently duplicate documents on top of the intended repetition. + seen: set[tuple[str, int, int]] = set() + for result in results: + for source, index_path in result.index_files.items(): + entries = pickle.loads(Path(index_path).read_bytes()) + for offset, length in entries: + key = (source, offset, length) + assert key not in seen, f"document {key} appears in more than one quality bucket" + seen.add(key) + + +def test_materializing_a_curve_refuses_a_sidecar_without_the_quality_column( + tmp_path: Path, dataset_entry: DatasetEntry, annotations: Path +): + """The curve orders by a joined label, so an unjoined sidecar cannot support one.""" + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=dataset_entry.iter_files(), + tokenizer=_WhitespaceTokenizer(), + tokenizer_name="whitespace", + sample_size=50, + ) + sidecar_dir = tmp_path / "unjoined" + SidecarBuilder(dataset_entry, calibration, index_root=tmp_path / "idx2").build( + sidecar_dir, show_progress=False + ) + with pytest.raises(MaterializationError, match="no column 'educational_value'"): + materialize_dataset_buckets( + sidecar_dir=sidecar_dir, + dataset_entry=dataset_entry, + dataset_selection=_curve_selection(), + missing_policy=MissingPolicy.KEEP, + output_dir=tmp_path / "mix2", + show_progress=False, + ) + + +# ------------------------------------------------------------ per-predicate attribution + + +def test_attribution_finds_the_predicate_that_does_nothing(built_sidecar: Path): + """A predicate matching everything the others already keep is the interesting case: it + changes no numbers while making the selection harder to read.""" + cube = build_cube(built_sidecar, "toy", label_dimensions=["educational_value"]) + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + # Every level is at least "none", so this one cannot remove anything. + Predicate(field="educational_value", op=Op.AT_LEAST, value="none"), + ], + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + + binding, redundant = breakdown.contributions + assert binding.marginal_tokens > 0 + assert redundant.marginal_tokens == 0 + assert "no effect given the others" in breakdown.describe() + + +def test_attribution_totals_agree_with_the_blend_evaluation(built_sidecar: Path): + cube = build_cube(built_sidecar, "toy", label_dimensions=["educational_value"]) + selection = DatasetSelection( + name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")] + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + result = evaluate_on_cube(cube, selection, MissingPolicy.KEEP) + + assert breakdown.kept_tokens == result.tokens_kept + assert breakdown.total_tokens == result.tokens_total + + +def test_the_overlap_diagonal_is_the_predicate_itself(built_sidecar: Path): + """Not the product with itself: an interpolated predicate has fractional factors, and + squaring them undercounts. On real data that read 8.19 M where it matched 9.75 M.""" + cube = build_cube(built_sidecar, "toy") + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="score", op=Op.GTE, value=2.5), + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + ], + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + for i, contribution in enumerate(breakdown.contributions): + assert breakdown.overlap_tokens[i][i] == contribution.matched_tokens + # And the matrix is symmetric. + assert breakdown.overlap_tokens[0][1] == breakdown.overlap_tokens[1][0] + + +def test_attribution_reports_which_predicate_was_interpolated(built_sidecar: Path): + cube = build_cube(built_sidecar, "toy") + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="score", op=Op.GTE, value=2.5), + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + ], + ) + breakdown = predicate_breakdown(cube, selection, MissingPolicy.KEEP) + numeric, ordinal = breakdown.contributions + # A threshold inside a quantile bin is interpolated; an ordinal level never is. + assert not numeric.exact + assert ordinal.exact + + +# --------------------------------------------------------------------------- reruns +# +# Both stages below write into a directory that a later stage globs for work. That makes +# leftovers dangerous rather than merely untidy: a stale index tree is packed, and a stale +# .pbin is trained on. These tests rerun each stage and check the directory afterwards +# describes only the current selection. + + +@pytest.fixture +def blend_inputs(tmp_path: Path, dataset_entry: DatasetEntry, built_sidecar: Path): + """A registry and a sidecar root laid out the way `materialize_blend` expects.""" + import shutil + + sidecar_root = tmp_path / "sidecar_root" + shutil.copytree(built_sidecar, sidecar_root / "toy") + return CorpusRegistry(datasets=[dataset_entry]), sidecar_root + + +def _blend_config(**overrides): + from modalities.dataloader.preprocessing.quality.selection import SelectionConfig + + settings = dict(datasets=[DatasetSelection(name="toy", ratio=2.0)]) + settings.update(overrides) + return SelectionConfig(**settings) + + +def _overexposed_config(): + """A selection whose repetition cap the run would blow through. + + Its predicate also selects a different, much smaller document set than + `_blend_config`. Without that the rejected run would rewrite byte-identical indexes + and corruption would be undetectable by comparing the directory. + """ + return _blend_config( + datasets=[ + DatasetSelection( + name="toy", + ratio=2.0, + predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="high")], + ) + ], + target_tokens=1e12, + max_total_exposure=1.0, + ) + + +def test_a_rejected_apply_leaves_the_previous_blend_exactly_as_it_was(tmp_path: Path, blend_inputs): + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + manifest_path = materialize_blend( + config=_blend_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + before = {p.relative_to(output_root): p.read_bytes() for p in sorted(output_root.rglob("*")) if p.is_file()} + assert before, "the first apply must have written something to compare against" + + # The exposure guard fires only after every index has been written, so this is the + # case where a destination-in-place apply would leave new indexes beside an old + # manifest: a directory that still looks complete but no longer agrees with itself. + with pytest.raises(MaterializationError, match="past its declared cap"): + materialize_blend( + config=_overexposed_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + + after = {p.relative_to(output_root): p.read_bytes() for p in sorted(output_root.rglob("*")) if p.is_file()} + assert after == before, "a rejected apply must not touch the blend that is already published" + assert manifest_path.exists() + leftovers = [p.name for p in tmp_path.iterdir() if p.name.startswith(".blend.")] + assert leftovers == [], f"staging directories must be cleaned up, found {leftovers}" + + +def test_a_rejected_first_apply_publishes_nothing_at_all(tmp_path: Path, blend_inputs): + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + with pytest.raises(MaterializationError): + materialize_blend( + config=_overexposed_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + + assert not output_root.exists(), "a failed apply must not leave a half-built blend behind" + assert [p.name for p in tmp_path.iterdir() if p.name.startswith(".blend.")] == [] + + +def test_a_successful_apply_replaces_the_previous_blend_rather_than_merging_into_it(tmp_path: Path, blend_inputs): + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + arguments = dict(registry=registry, sidecar_root=sidecar_root, output_root=output_root, show_progress=False) + + materialize_blend(config=_blend_config(), **arguments) + stale = output_root / "dropped-dataset" / "shard_0.idx" + stale.parent.mkdir(parents=True) + stale.write_bytes(b"from an earlier selection") + + materialize_blend(config=_blend_config(datasets=[DatasetSelection(name="toy", ratio=3.0)]), **arguments) + + assert not stale.exists(), "index trees the new selection does not name must not survive the rerun" + assert (output_root / "toy").is_dir() + + +def test_the_published_manifest_names_index_files_that_actually_exist(tmp_path: Path, blend_inputs): + # The indexes are written under a staging directory that is renamed away on + # publication. If the manifest kept the paths the writers physically used, every path + # in it would point inside a directory that no longer exists, and every packing config + # built from it would name a missing index. + import yaml + + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + manifest_path = materialize_blend( + config=_blend_config(), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + manifest = yaml.safe_load(manifest_path.read_text()) + + indexes = [Path(index) for dataset in manifest["datasets"] for index in dataset["index_files"].values()] + assert indexes, "the manifest must name at least one index" + missing = [str(index) for index in indexes if not index.exists()] + assert missing == [], f"the manifest names index files that are not on disk: {missing[:3]}" + for index in indexes: + assert output_root in index.parents, f"{index} is not under the published blend" + + +def test_a_curved_blend_also_publishes_usable_index_paths(tmp_path: Path, blend_inputs): + # Curves take a different code path, writing one index tree per quality bucket, so the + # rebasing has to hold there too. + import yaml + + from modalities.dataloader.preprocessing.quality.materialize import materialize_blend + + registry, sidecar_root = blend_inputs + output_root = tmp_path / "blend" + + manifest_path = materialize_blend( + config=_blend_config( + datasets=[ + DatasetSelection( + name="toy", + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=1.5), + ) + ] + ), + registry=registry, + sidecar_root=sidecar_root, + output_root=output_root, + show_progress=False, + ) + manifest = yaml.safe_load(manifest_path.read_text()) + + assert len(manifest["datasets"]) > 1, "a curve must split the dataset into per-bucket rows" + for dataset in manifest["datasets"]: + for index in dataset["index_files"].values(): + assert Path(index).exists(), f"{dataset['name']} names a missing index {index}" + diff --git a/tests/dataloader/preprocessing/quality/test_selection.py b/tests/dataloader/preprocessing/quality/test_selection.py new file mode 100644 index 000000000..25a33c8c6 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_selection.py @@ -0,0 +1,435 @@ +from pathlib import Path + +import pyarrow as pa +import pytest + +from modalities.dataloader.preprocessing.quality.cube import MISSING, Cube, ScoreBinning +from modalities.dataloader.preprocessing.quality.selection import ( + ORDINAL_SCALES, + DatasetSelection, + MissingPolicy, + Op, + Predicate, + SelectionConfig, + SelectionError, + document_mask, + evaluate_on_cube, + format_blend_report, +) + + +def test_at_least_uses_the_declared_order_not_the_alphabet(): + predicate = Predicate(field="information_density", op=Op.AT_LEAST, value="adequate") + + # `moderate` sorts before `adequate` alphabetically but ranks below it on the scale, + # so an alphabetical implementation would wrongly include it. + assert predicate.allowed_levels() == {"adequate", "dense"} + + +def test_at_most_includes_everything_up_to_the_level(): + predicate = Predicate(field="content_integrity", op=Op.AT_MOST, value="fragment") + + assert predicate.allowed_levels() == {"severely_degraded", "fragment"} + + +def test_ordinal_predicate_rejects_an_unknown_level(): + with pytest.raises(ValueError, match="is not a level of"): + Predicate(field="educational_value", op=Op.AT_LEAST, value="excellent") + + +def test_ordinal_predicate_rejects_a_field_without_a_scale(): + with pytest.raises(ValueError, match="no declared scale"): + Predicate(field="fw_edu", op=Op.AT_LEAST, value="high") + + +def test_every_declared_scale_has_unique_levels(): + for field, scale in ORDINAL_SCALES.items(): + assert len(set(scale)) == len(scale), f"{field} has duplicate levels" + + +@pytest.mark.parametrize("policy,expected", [(MissingPolicy.KEEP, True), (MissingPolicy.DROP, False)]) +def test_missing_value_follows_the_policy(policy, expected): + predicate = Predicate(field="educational_value", op=Op.AT_LEAST, value="basic") + + assert predicate.matches_value(None, policy) is expected + assert predicate.matches_value(MISSING, policy) is expected + + +def test_per_predicate_policy_overrides_the_selection_policy(): + predicate = Predicate(field="educational_value", op=Op.AT_LEAST, value="basic", missing=MissingPolicy.DROP) + + assert predicate.matches_value(None, MissingPolicy.KEEP) is False + + +def _cube_with_labels() -> Cube: + table = pa.table( + { + "educational_value": ["high", "basic", "none", MISSING], + "n_documents": [10, 20, 30, 40], + "n_tokens": [1000, 1500, 900, 2000], + } + ) + return Cube( + dataset="toy", + label_dimensions=["educational_value"], + score_binnings={}, + table=table, + n_documents=100, + n_tokens=5400, + ) + + +def test_cube_evaluation_sums_the_matching_cells(): + result = evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")]), + MissingPolicy.DROP, + ) + + assert result.n_documents_kept == 30 + assert result.tokens_kept == 2500 + assert result.exact + + +def test_cube_evaluation_keeps_unannotated_cells_when_told_to(): + result = evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")]), + MissingPolicy.KEEP, + ) + + assert result.n_documents_kept == 70 + assert result.tokens_kept == 4500 + + +def test_cube_evaluation_rejects_an_ungrouped_field(): + with pytest.raises(SelectionError, match="not grouped on"): + evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="content_quality", op=Op.AT_LEAST, value="good")]), + MissingPolicy.KEEP, + ) + + +def _cube_with_scores() -> Cube: + # Two bins with edges at 0, 2 and 4. + table = pa.table({"native_score": [0, 1], "n_documents": [100, 100], "n_tokens": [1000, 3000]}) + return Cube( + dataset="toy", + label_dimensions=[], + score_binnings={"score": ScoreBinning(column="native_score", edges=(0.0, 2.0, 4.0))}, + table=table, + n_documents=200, + n_tokens=4000, + ) + + +def test_numeric_threshold_on_a_bin_edge_is_exact(): + result = evaluate_on_cube( + _cube_with_scores(), + DatasetSelection(name="toy", predicates=[Predicate(field="score", op=Op.GTE, value=2.0)]), + MissingPolicy.KEEP, + ) + + assert result.exact + assert result.n_documents_kept == 100 + assert result.tokens_kept == 3000 + + +def test_numeric_threshold_inside_a_bin_is_reported_as_interpolated(): + result = evaluate_on_cube( + _cube_with_scores(), + DatasetSelection(name="toy", predicates=[Predicate(field="score", op=Op.GTE, value=3.0)]), + MissingPolicy.KEEP, + ) + + assert not result.exact, "a threshold splitting a bin cannot be answered exactly from a cube" + assert result.approximations == ["score gte 3.0"] + # Half of the upper bin, none of the lower one. + assert result.n_documents_kept == 50 + + +def test_document_mask_matches_cube_counts(): + table = pa.table( + { + "educational_value": ["high", "basic", "none", None], + "native_score": [3.0, 1.0, 0.5, float("nan")], + "est_tokens": [100, 200, 300, 400], + } + ) + selection = DatasetSelection( + name="toy", + predicates=[ + Predicate(field="educational_value", op=Op.AT_LEAST, value="basic"), + Predicate(field="score", op=Op.GTE, value=1.0), + ], + ) + + mask = document_mask(table, selection, MissingPolicy.DROP) + + assert list(mask) == [True, True, False, False] + + +def test_document_mask_rejects_a_missing_column(): + table = pa.table({"est_tokens": [1, 2]}) + + with pytest.raises(SelectionError, match="no column"): + document_mask( + table, + DatasetSelection(name="toy", predicates=[Predicate(field="nope", op=Op.GTE, value=1)]), + MissingPolicy.KEEP, + ) + + +def test_token_retention_can_exceed_row_retention(): + # Quality correlates with length, so a quality filter keeps a larger share of the + # tokens than of the documents. Scaling a corpus average by row retention would + # understate the surviving budget, which is why both are reported. + result = evaluate_on_cube( + _cube_with_labels(), + DatasetSelection(name="toy", predicates=[Predicate(field="educational_value", op=Op.AT_LEAST, value="basic")]), + MissingPolicy.DROP, + ) + + assert result.row_retention == pytest.approx(0.3) + assert result.token_retention > result.row_retention + + +def test_duplicate_dataset_in_a_selection_is_rejected(): + with pytest.raises(ValueError, match="appears twice"): + SelectionConfig(datasets=[DatasetSelection(name="toy"), DatasetSelection(name="toy")]) + + +def test_report_marks_interpolated_rows_and_shows_the_target_gap(): + from modalities.dataloader.preprocessing.quality.selection import BlendResult, DatasetResult + + result = BlendResult( + datasets=[ + DatasetResult("exact_one", 100, 50, 1000, 600, 2.0), + DatasetResult("fuzzy_one", 100, 50, 1000, 400, 1.0, exact=False, approximations=["score gte 3.0"]), + ], + target_tokens=2000, + ) + + report = format_blend_report(result) + + assert "~" in report + assert "score gte 3.0" in report + assert "under" in report + + +def _cube_without(field: str) -> Cube: + """A cube grouped on educational_value only, so any other field is unanswerable.""" + table = pa.table({"educational_value": ["high", "basic"], "n_documents": [10, 20], "n_tokens": [100, 200]}) + return Cube( + dataset="toy", + label_dimensions=["educational_value"], + score_binnings={}, + table=table, + n_documents=30, + n_tokens=300, + ) + + +def test_blend_refuses_a_silent_sidecar_scan_and_names_every_offender(tmp_path: Path): + # A predicate the cube cannot answer used to fall back to reading every document + # without saying so, which turned a "seconds" preview into a 1.7-billion-document + # scan. It must be reported instead. + from modalities.dataloader.preprocessing.quality.selection import evaluate_blend + + config = SelectionConfig( + datasets=[ + DatasetSelection( + name="toy", + predicates=[Predicate(field="commercial_bias", op=Op.AT_LEAST, value="minimal")], + ), + DatasetSelection( + name="other", + predicates=[Predicate(field="content_quality", op=Op.AT_LEAST, value="good")], + ), + ] + ) + cubes = {"toy": _cube_without("commercial_bias"), "other": _cube_without("content_quality")} + sidecars = {"toy": tmp_path, "other": tmp_path} + + with pytest.raises(SelectionError) as excinfo: + evaluate_blend(config, cubes, sidecar_dirs=sidecars) + + message = str(excinfo.value) + assert "commercial_bias" in message and "content_quality" in message, "both offenders must be reported at once" + assert "30" in message, "the message should say how many documents a fallback would read" + assert "--allow-fallback" in message + + +def test_blend_reports_a_dataset_with_no_cube_at_all(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_blend + + config = SelectionConfig(datasets=[DatasetSelection(name="toy")]) + + with pytest.raises(SelectionError, match="no cube was built"): + evaluate_blend(config, {}, sidecar_dirs={"toy": tmp_path}) + + +# --------------------------------------------------------------------------- curved selections +# +# A curved dataset has `ratio` pinned to 1.0 by the config validator, because its factors +# come from the solved curve rather than from a single number. That makes any code path +# which reports `ratio` instead of solving the curve look successful while costing the +# dataset at a flat 1.0x. The exact sidecar path did exactly this, so these tests check +# the two paths agree rather than only that each returns something. + + +def _curved(**overrides) -> DatasetSelection: + """A dataset upsampled along educational_value.""" + from modalities.dataloader.preprocessing.quality.upsampling import UpsamplingSpec + + return DatasetSelection( + name="toy", + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0, **overrides), + ) + + +_CURVE_ROWS = [ + # (educational_value, n_documents, tokens per document) + ("minimal", 40, 100), + ("basic", 30, 100), + ("moderate", 20, 100), + ("high", 10, 100), +] + + +def _curve_cube() -> Cube: + """A cube over `_CURVE_ROWS`, grouped on educational_value alone.""" + return Cube( + dataset="toy", + label_dimensions=["educational_value"], + score_binnings={}, + table=pa.table( + { + "educational_value": [level for level, _, _ in _CURVE_ROWS], + "n_documents": [n for _, n, _ in _CURVE_ROWS], + "n_tokens": [n * t for _, n, t in _CURVE_ROWS], + } + ), + n_documents=sum(n for _, n, _ in _CURVE_ROWS), + n_tokens=sum(n * t for _, n, t in _CURVE_ROWS), + ) + + +def _curve_sidecar(tmp_path: Path, extra_column: dict | None = None) -> Path: + """The same documents as `_curve_cube`, written out one row per document.""" + import pyarrow.parquet as pq + + levels: list[str] = [] + tokens: list[int] = [] + for level, n_documents, per_document in _CURVE_ROWS: + levels.extend([level] * n_documents) + tokens.extend([per_document] * n_documents) + + columns = { + "file_id": [0] * len(levels), + "line_no": list(range(len(levels))), + "byte_offset": [0] * len(levels), + "byte_len": [per for per in tokens], + "text_bytes": [per for per in tokens], + "est_tokens": tokens, + "educational_value": levels, + } + if extra_column: + columns.update({name: values * len(levels) for name, values in extra_column.items()}) + + sidecar_dir = tmp_path / "sidecar" + sidecar_dir.mkdir(parents=True, exist_ok=True) + # Several parts, because the exact path accumulates buckets across parts and row + # groups; a single-part fixture would not exercise that. + total = len(levels) + for part, start in enumerate(range(0, total, 37)): + chunk = {name: values[start : start + 37] for name, values in columns.items()} + pq.write_table(pa.table(chunk), sidecar_dir / f"part-{part:05d}.parquet", row_group_size=11) + return sidecar_dir + + +def test_the_exact_path_solves_the_curve_instead_of_reporting_a_flat_ratio(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + dataset = _curved() + assert dataset.ratio == 1.0, "the validator pins a curved dataset's flat ratio to 1.0" + + result = evaluate_on_sidecar(_curve_sidecar(tmp_path), dataset, MissingPolicy.KEEP) + + assert result.plan is not None, "an exact evaluation of a curved dataset must carry its solved curve" + assert result.effective_tokens > result.tokens_kept, "a 2.0x target must draw more than it holds" + assert result.ratio_label != "1.00", "reporting the pinned flat ratio hides the curve entirely" + + +def test_a_curve_costs_the_same_whether_it_is_evaluated_exactly_or_from_the_cube(tmp_path: Path): + # The regression that motivates this: the cube path solved the curve and the exact + # path did not, so `--exact` quietly reported a different, much smaller blend. + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + dataset = _curved() + + from_cube = evaluate_on_cube(_curve_cube(), dataset, MissingPolicy.KEEP) + from_sidecar = evaluate_on_sidecar(_curve_sidecar(tmp_path), dataset, MissingPolicy.KEEP) + + assert from_sidecar.tokens_total == from_cube.tokens_total + assert from_sidecar.n_documents_kept == from_cube.n_documents_kept + assert from_sidecar.effective_tokens == pytest.approx(from_cube.effective_tokens, rel=1e-9) + assert [b.factor for b in from_sidecar.plan.buckets] == pytest.approx( + [b.factor for b in from_cube.plan.buckets], rel=1e-9 + ) + + +def test_a_curve_that_discards_its_worst_bucket_drops_those_documents_exactly(tmp_path: Path): + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + # The 40 `minimal` documents hold 4,000 of the 10,000 tokens, so they are exactly the + # bottom 40% of the quality axis and a 40% discard should remove that bucket whole. + dataset = _curved(discard_below_percentile=40.0) + result = evaluate_on_sidecar(_curve_sidecar(tmp_path), dataset, MissingPolicy.KEEP) + + assert result.plan.discard_fraction > 0.0 + assert result.n_documents_kept == 60, "documents kept must reflect the curve's discard, not the predicates alone" + assert result.n_documents_kept < result.n_documents_total + + +def test_the_exact_path_names_the_missing_quality_column_rather_than_failing_obscurely(tmp_path: Path): + import pyarrow.parquet as pq + + from modalities.dataloader.preprocessing.quality.selection import evaluate_on_sidecar + + sidecar_dir = tmp_path / "unjoined" + sidecar_dir.mkdir() + pq.write_table(pa.table({"est_tokens": [10, 20]}), sidecar_dir / "part-00000.parquet") + + with pytest.raises(SelectionError, match="educational_value"): + evaluate_on_sidecar(sidecar_dir, _curved(), MissingPolicy.KEEP) + + +def test_the_fallback_path_keeps_the_curve_when_a_predicate_forces_a_sidecar_scan(tmp_path: Path): + # Deliberately unmocked. The previous version of this test stubbed out + # `evaluate_on_sidecar` and only covered a flat selection, which is precisely why a + # curve being flattened on that path went unnoticed. + from modalities.dataloader.preprocessing.quality.selection import evaluate_blend + from modalities.dataloader.preprocessing.quality.upsampling import UpsamplingSpec + + dataset = DatasetSelection( + name="toy", + # The cube below is not grouped on commercial_bias, so this predicate is + # unanswerable from it and the blend must fall back to the sidecar. + predicates=[Predicate(field="commercial_bias", op=Op.AT_LEAST, value="minimal")], + upsampling=UpsamplingSpec(quality_field="educational_value", target_ratio=2.0), + ) + sidecar_dir = _curve_sidecar(tmp_path, extra_column={"commercial_bias": ["minimal"]}) + + report = evaluate_blend( + SelectionConfig(datasets=[dataset]), + {"toy": _curve_cube()}, + sidecar_dirs={"toy": sidecar_dir}, + allow_sidecar_fallback=True, + ) + + result = report.datasets[0] + assert result.exact, "a sidecar scan is exact by construction" + assert result.plan is not None, "falling back to the sidecar must not discard the curve" + assert result.effective_tokens == pytest.approx(2.0 * result.tokens_kept, rel=0.05) diff --git a/tests/dataloader/preprocessing/quality/test_token_calibration.py b/tests/dataloader/preprocessing/quality/test_token_calibration.py new file mode 100644 index 000000000..ea3916dbb --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_token_calibration.py @@ -0,0 +1,192 @@ +"""Tests for the token estimator, which was quietly wrong in two ways at once. + +A corpus's bytes-per-token ratio is not one number. On the German FineWiki snapshot it runs +from 3.6 for documents under a kilobyte to 34.6 for the nine documents above 256 KB, and +those nine hold 5.7% of all bytes. Two consequences, both of which bit: + +* Measuring one global ratio makes every estimate hostage to whether the sample happened + to include those documents. Sampling uniformly by document count, it almost never does. +* Taking the sample from the *start* of each file is deterministic, so it looked perfectly + stable across seeds while being systematically wrong -- 16% on FineWiki, because a + corpus ordered by article id, fetch time or source is not homogeneous along its length. + +So documents are sampled at offsets spread across each file, the document *containing* +each offset is taken (which makes selection proportional to length), and the ratio is +measured per size stratum. These tests pin the properties that makes work, using a +synthetic corpus with the same shape: many small documents and a few enormous ones whose +tokens-per-byte differs sharply. +""" + +import json +from pathlib import Path + +import pytest + +from modalities.dataloader.preprocessing.quality.tokens import ( + SIZE_STRATUM_BOUNDS, + _sample_documents, + TokenCalibration, + calibrate_dataset, + size_stratum, +) + + +class _CharTokenizer: + """One token per character, so token counts are exactly predictable.""" + + def tokenize(self, text: str) -> list[str]: + return list(text) + + +class _WordTokenizer: + """One token per whitespace-separated word.""" + + def tokenize(self, text: str) -> list[str]: + return text.split() + + +@pytest.fixture +def skewed_corpus(tmp_path: Path) -> tuple[Path, int, int]: + """A corpus whose bytes live mostly in a few long documents that tokenize differently. + + Short documents are single characters repeated, so with the word tokenizer they are one + token each -- a high bytes-per-token ratio. Long documents are spaced words, so they + are many tokens -- a low ratio. The population ratio therefore depends heavily on the + long documents, which are rare by count and dominant by bytes: the exact shape that + defeated the original estimator. + """ + corpus = tmp_path / "corpus" + corpus.mkdir() + total_bytes = 0 + total_tokens = 0 + tokenizer = _WordTokenizer() + with (corpus / "shard.jsonl").open("w") as f: + for i in range(3000): + text = "x" * 200 + f.write(json.dumps({"id": f"s{i}", "text": text}) + "\n") + total_bytes += len(text.encode()) + total_tokens += len(tokenizer.tokenize(text)) + for i in range(12): + text = " ".join(["word"] * 120_000) + f.write(json.dumps({"id": f"l{i}", "text": text}) + "\n") + total_bytes += len(text.encode()) + total_tokens += len(tokenizer.tokenize(text)) + return corpus, total_bytes, total_tokens + + +def test_size_stratum_boundaries_are_inclusive_below(): + assert size_stratum(0) == 0 + assert size_stratum(SIZE_STRATUM_BOUNDS[0] - 1) == 0 + assert size_stratum(SIZE_STRATUM_BOUNDS[0]) == 1 + assert size_stratum(10**12) == len(SIZE_STRATUM_BOUNDS) + + +def test_the_sample_is_not_the_start_of_the_file(tmp_path: Path): + """The original sampler read a prefix, which is deterministic and so looked stable + across seeds while being systematically wrong on any corpus that varies along its + length -- 16% out on FineWiki.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + with (corpus / "shard.jsonl").open("w") as f: + for i in range(2000): + f.write(json.dumps({"id": f"first-{i}", "text": "a" * 400}) + "\n") + for i in range(2000): + f.write(json.dumps({"id": f"second-{i}", "text": "b" * 400}) + "\n") + + sample = _sample_documents( + file_paths=[corpus / "shard.jsonl"], + text_field="text", + sample_size=200, + seed=1, + max_probe_files=32, + max_lines_per_probe=100_000, + ) + halves = {record["id"].split("-")[0] for record in sample} + assert halves == {"first", "second"}, f"sample only reached {halves}; it is not spread over the file" + + +def test_stratified_estimate_beats_the_global_ratio_on_a_skewed_corpus( + skewed_corpus: tuple[Path, int, int], +): + corpus, total_bytes, total_tokens = skewed_corpus + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=[corpus / "shard.jsonl"], + tokenizer=_WordTokenizer(), + tokenizer_name="word", + sample_size=1000, + ) + + sizes = [] + with (corpus / "shard.jsonl").open() as f: + for line in f: + sizes.append(len(json.loads(line)["text"].encode())) + + stratified = sum(round(b / calibration.ratio_for(b)) for b in sizes) + global_only = sum(round(b / calibration.bytes_per_token) for b in sizes) + stratified_error = abs(stratified / total_tokens - 1) + global_error = abs(global_only / total_tokens - 1) + + assert stratified_error < 0.05, f"stratified estimate was {stratified_error:.1%} out" + assert stratified_error < global_error + + +def test_long_documents_are_actually_reached(skewed_corpus: tuple[Path, int, int]): + """Sampling by containing document, not the following one, is what finds them. + + Twelve documents out of 3,012 hold most of the bytes. Uniform-by-count sampling finds + them at a rate of 0.4%; length-proportional sampling finds them in proportion to their + share of the corpus, which is what the ratio needs. + """ + corpus, _, _ = skewed_corpus + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=[corpus / "shard.jsonl"], + tokenizer=_WordTokenizer(), + tokenizer_name="word", + sample_size=1000, + ) + long_stratum = size_stratum(len(" ".join(["word"] * 120_000).encode())) + assert calibration.stratum_documents[long_stratum] >= TokenCalibration.MIN_STRATUM_DOCUMENTS + + +def test_a_sparse_stratum_falls_back_to_the_global_ratio(): + calibration = TokenCalibration( + dataset="toy", + tokenizer="t", + bytes_per_token=4.0, + stratum_bytes_per_token=[3.0] + [99.0] * len(SIZE_STRATUM_BOUNDS), + stratum_documents=[500] + [1] * len(SIZE_STRATUM_BOUNDS), + ) + assert calibration.ratio_for(10) == 3.0 + # Measured from one document, so not trusted as an estimator of its own. + assert calibration.ratio_for(10**9) == 4.0 + + +def test_a_calibration_without_strata_still_estimates(): + """Calibrations written before strata existed must keep working.""" + calibration = TokenCalibration(dataset="toy", tokenizer="t", bytes_per_token=4.0) + assert calibration.ratio_for(10) == 4.0 + assert calibration.estimate({}, text_bytes=400) == 101 + + +def test_a_native_token_field_still_wins(tmp_path: Path): + """Stratification applies to the bytes-per-token path only; a corpus carrying its own + token count is estimated from that, per document, which needs no stratifying.""" + corpus = tmp_path / "corpus" + corpus.mkdir() + with (corpus / "shard.jsonl").open("w") as f: + for i in range(500): + text = " ".join(["word"] * (10 + i)) + f.write(json.dumps({"id": str(i), "text": text, "token_count": len(text.split())}) + "\n") + + calibration = calibrate_dataset( + dataset_name="toy", + file_paths=[corpus / "shard.jsonl"], + tokenizer=_WordTokenizer(), + tokenizer_name="word", + sample_size=200, + ) + assert calibration.uses_native_field() + assert calibration.native_field == "token_count" + assert calibration.estimate({"token_count": 100}, text_bytes=999_999) == 101 diff --git a/tests/dataloader/preprocessing/quality/test_upsampling.py b/tests/dataloader/preprocessing/quality/test_upsampling.py new file mode 100644 index 000000000..fd32e2fc4 --- /dev/null +++ b/tests/dataloader/preprocessing/quality/test_upsampling.py @@ -0,0 +1,191 @@ +"""Tests for quality-aware upsampling curves. + +The behaviour worth pinning is not "it produces numbers" but the four properties that make +the curve mean anything: the token target is hit, the repeat cap is honoured, the factors +rise with quality, and the bottom of the distribution is actually dropped. Everything here +is synthetic and runs in milliseconds -- no corpus, no cluster. +""" + +import pytest + +from modalities.dataloader.preprocessing.quality.selection import exposure_report +from modalities.dataloader.preprocessing.quality.upsampling import ( + UNANNOTATED_BUCKET, + QualityBucket, + UpsamplingError, + UpsamplingSpec, + solve_curve, +) + + +def vigintiles(tokens_each: int = 1_000_000) -> list[QualityBucket]: + """Twenty equal-token buckets, as Dolma 3 partitions web text.""" + return [ + QualityBucket(label=f"p{5 * i}-{5 * (i + 1)}", n_documents=1000, n_tokens=tokens_each) + for i in range(20) + ] + + +def test_the_token_target_is_hit(): + buckets = vigintiles() + available = sum(b.n_tokens for b in buckets) + plan = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_tokens=available * 1.5)) + assert plan.tokens_drawn == pytest.approx(available * 1.5, rel=1e-6) + + +def test_target_ratio_and_target_tokens_agree(): + buckets = vigintiles() + available = sum(b.n_tokens for b in buckets) + by_tokens = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_tokens=available * 2)) + by_ratio = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_ratio=2.0)) + assert by_tokens.exponent == pytest.approx(by_ratio.exponent) + assert by_tokens.tokens_drawn == pytest.approx(by_ratio.tokens_drawn) + + +def test_the_cap_is_honoured_and_reached(): + plan = solve_curve( + vigintiles(), + UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0, discard_below_percentile=40.0), + ) + factors = [b.factor for b in plan.buckets] + assert max(factors) == pytest.approx(7.0, rel=1e-4) + assert plan.saturated + + +def test_factors_rise_with_quality(): + plan = solve_curve(vigintiles(), UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0)) + factors = [b.factor for b in plan.buckets] + assert factors == sorted(factors), "a curve whose factors are not monotone is not quality-aware" + + +def test_the_discarded_share_is_actually_discarded(): + plan = solve_curve( + vigintiles(), + UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0, discard_below_percentile=40.0), + ) + # Eight of twenty equal-token buckets sit below the 40th percentile. + assert [b.factor for b in plan.buckets[:8]] == [0.0] * 8 + assert plan.documents_kept == sum(b.bucket.n_documents for b in plan.buckets[8:] if b.factor > 0) + + +def test_an_impossible_target_is_refused_with_the_numbers_to_fix_it(): + # Asking for 5x the data while discarding half of it and capping repeats at 2x cannot work: + # the kept half would have to be repeated 10x. + with pytest.raises(UpsamplingError, match="even repeating every kept document equally"): + solve_curve( + vigintiles(), + UpsamplingSpec( + quality_field="q", target_ratio=5.0, max_factor=2.0, discard_below_percentile=50.0 + ), + ) + + +def test_a_reachable_target_does_not_saturate_the_cap(): + # Drawing only a third of the data needs no repetition at all, so the cap stays unused. + plan = solve_curve( + vigintiles(), UpsamplingSpec(quality_field="q", target_ratio=0.33, max_factor=7.0) + ) + assert max(b.factor for b in plan.buckets) < 7.0 + + +def test_a_flat_exponent_reproduces_flat_filtering(): + plan = solve_curve( + vigintiles(), + UpsamplingSpec( + quality_field="q", target_ratio=0.6, discard_below_percentile=40.0, exponent=1.0 + ), + ) + kept = [b.factor for b in plan.buckets if b.factor > 0] + assert len(kept) == 12 + # Flat over the kept range: every surviving bucket repeated the same amount. + assert max(kept) == pytest.approx(min(kept), rel=1e-6) + assert max(kept) == pytest.approx(0.6 / 0.6, rel=1e-6) + + +def test_unequal_bucket_widths_are_weighted_by_tokens(): + # Ordinal levels are not equal-sized the way vigintiles are; a level holding most of the + # tokens must occupy most of the axis, or the curve prices it wrongly. + buckets = [ + QualityBucket(label="none", n_documents=10, n_tokens=100), + QualityBucket(label="basic", n_documents=900, n_tokens=9_000), + QualityBucket(label="excellent", n_documents=90, n_tokens=900), + ] + plan = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0)) + widths = [b.upper - b.lower for b in plan.buckets] + assert widths[1] == pytest.approx(0.9, abs=1e-9) + assert plan.tokens_drawn == pytest.approx(10_000, rel=1e-6) + + +def test_unannotated_documents_sit_at_the_bottom(): + buckets = [ + QualityBucket(label=UNANNOTATED_BUCKET, n_documents=50, n_tokens=500, unannotated=True), + QualityBucket(label="basic", n_documents=50, n_tokens=500), + QualityBucket(label="excellent", n_documents=50, n_tokens=500), + ] + plan = solve_curve(buckets, UpsamplingSpec(quality_field="q", target_ratio=1.0, max_factor=7.0)) + assert plan.buckets[0].bucket.label == UNANNOTATED_BUCKET + assert plan.buckets[0].factor <= plan.buckets[-1].factor + + +def test_empty_buckets_are_refused(): + with pytest.raises(UpsamplingError, match="no quality buckets"): + solve_curve([], UpsamplingSpec(quality_field="q", target_ratio=1.0)) + with pytest.raises(UpsamplingError, match="hold no tokens"): + solve_curve( + [QualityBucket(label="a", n_documents=0, n_tokens=0)], + UpsamplingSpec(quality_field="q", target_ratio=1.0), + ) + + +def test_a_spec_needs_exactly_one_target(): + with pytest.raises(ValueError, match="exactly one of"): + UpsamplingSpec(quality_field="q") + with pytest.raises(ValueError, match="exactly one of"): + UpsamplingSpec(quality_field="q", target_tokens=1, target_ratio=1.0) + + +# ------------------------------------------------------- exposure once wrapping is counted + + +def test_exposure_multiplies_the_factor_by_the_number_of_passes(): + report = exposure_report( + entries=[("a", 2.0, 2.5), ("b", 1.0, 2.5)], + effective_tokens=100.0, + training_tokens=200.0, + ) + assert report.passes == pytest.approx(2.0) + assert report.rows[0].exposure == pytest.approx(4.0) + # 2.0 asked for, 4.0 actually seen, against a 2.5 cap. + assert [row.label for row in report.exceeded] == ["a"] + + +def test_a_blend_that_covers_the_run_does_not_wrap(): + report = exposure_report( + entries=[("a", 2.0, 2.5)], effective_tokens=400.0, training_tokens=200.0 + ) + assert report.passes == pytest.approx(0.5) + assert report.rows[0].exposure == pytest.approx(1.0) + assert not report.exceeded + + +def test_without_a_declared_target_only_wrapping_is_unknown(): + """The cap still applies to the requested factor -- a ratio of 9 against a cap of 2 is a + violation at any number of passes. What is unknown without a target is only the + multiplier on top, and the report says so rather than implying one pass.""" + report = exposure_report(entries=[("a", 9.0, 2.0)], effective_tokens=100.0, training_tokens=None) + assert report.passes == 1.0 + assert report.rows[0].exposure == pytest.approx(9.0) + assert [row.label for row in report.exceeded] == ["a"] + assert "cannot be checked" in report.describe() + + +def test_a_factor_within_its_cap_at_one_pass_passes(): + report = exposure_report(entries=[("a", 1.8, 2.0)], effective_tokens=100.0, training_tokens=None) + assert not report.exceeded + + +def test_a_row_without_a_cap_is_never_exceeded(): + report = exposure_report( + entries=[("a", 50.0, None)], effective_tokens=1.0, training_tokens=100.0 + ) + assert not report.exceeded diff --git a/tests/dataloader/test_weighted_combined_dataset.py b/tests/dataloader/test_weighted_combined_dataset.py new file mode 100644 index 000000000..556992f75 --- /dev/null +++ b/tests/dataloader/test_weighted_combined_dataset.py @@ -0,0 +1,168 @@ +from collections import Counter +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from modalities.config.component_factory import ComponentFactory +from modalities.config.pydantic_if_types import PydanticDatasetIFType +from modalities.dataloader.dataset import WeightedCombinedDataset +from modalities.registry.components import COMPONENTS +from modalities.registry.registry import Registry + + +class _RangeDataset: + """Minimal dataset returning (tag, index), so samples can be traced to their source.""" + + def __init__(self, num_samples: int, tag: str): + self._num_samples = num_samples + self._tag = tag + + def __len__(self) -> int: + return self._num_samples + + def __getitem__(self, idx: int) -> tuple[str, int]: + return self._tag, idx + + +@pytest.fixture +def datasets() -> list[_RangeDataset]: + return [_RangeDataset(100, "a"), _RangeDataset(50, "b"), _RangeDataset(10, "c")] + + +def test_integer_repeat_factors_repeat_whole_datasets(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 2.0, 3.0]) + + assert len(dataset) == 100 + 100 + 30 + counts = Counter(tag for tag, _ in (dataset[i] for i in range(len(dataset)))) + assert counts == {"a": 100, "b": 100, "c": 30} + + +def test_fractional_repeat_factor_adds_a_partial_pass(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.0, 2.5, 0.0]) + + assert len(dataset) == 125 + drawn = [idx for _, idx in (dataset[i] for i in range(len(dataset)))] + counts = Counter(drawn) + # Two full passes plus half a pass: every document twice, half of them three times. + assert set(counts) == set(range(50)) + assert sorted(counts.values()) == [2] * 25 + [3] * 25 + + +def test_partial_pass_selects_distinct_documents(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.3, 0.0, 0.0]) + + drawn = [idx for _, idx in (dataset[i] for i in range(len(dataset)))] + assert len(drawn) == 30 + assert len(set(drawn)) == 30, "a partial pass must not draw the same document twice" + + +def test_downsampling_spreads_across_the_dataset(datasets): + # A prefix would over-sample whatever the corpus is ordered by, so the partial pass + # must reach into the whole index range rather than the front of it. + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.1, 0.0, 0.0]) + + drawn = sorted(idx for _, idx in (dataset[i] for i in range(len(dataset)))) + assert len(drawn) == 10 + assert max(drawn) > 50, f"partial pass stayed in the front of the dataset: {drawn}" + + +def test_repeat_factor_rounding_up_becomes_a_full_pass(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[0.999, 0.0, 0.0]) + + assert len(dataset) == 100 + drawn = [idx for _, idx in (dataset[i] for i in range(len(dataset)))] + assert sorted(drawn) == list(range(100)) + + +def test_zero_repeat_factor_excludes_a_dataset(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 0.0, 1.0]) + + assert len(dataset) == 110 + counts = Counter(tag for tag, _ in (dataset[i] for i in range(len(dataset)))) + assert "b" not in counts + + +def test_same_seed_gives_the_same_blend(datasets): + first = WeightedCombinedDataset(datasets, repeat_factors=[1.5, 0.4, 1.0], seed=7) + second = WeightedCombinedDataset(datasets, repeat_factors=[1.5, 0.4, 1.0], seed=7) + + assert [first[i] for i in range(len(first))] == [second[i] for i in range(len(second))] + + +def test_different_seed_changes_the_partial_pass(datasets): + first = WeightedCombinedDataset(datasets, repeat_factors=[0.5, 0.0, 0.0], seed=7) + second = WeightedCombinedDataset(datasets, repeat_factors=[0.5, 0.0, 0.0], seed=8) + + assert {idx for _, idx in (first[i] for i in range(len(first)))} != { + idx for _, idx in (second[i] for i in range(len(second))) + } + + +def test_out_of_bounds_index_raises(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 1.0, 1.0]) + + with pytest.raises(IndexError): + dataset[len(dataset)] + + +def test_negative_index_counts_from_the_end(datasets): + dataset = WeightedCombinedDataset(datasets, repeat_factors=[1.0, 1.0, 1.0]) + + assert dataset[-1] == dataset[len(dataset) - 1] + + +def test_mismatched_repeat_factors_are_rejected(datasets): + with pytest.raises(ValueError, match="repeat factors"): + WeightedCombinedDataset(datasets, repeat_factors=[1.0, 1.0]) + + +def test_negative_repeat_factor_is_rejected(datasets): + with pytest.raises(ValueError, match="non-negative"): + WeightedCombinedDataset(datasets, repeat_factors=[1.0, -1.0, 1.0]) + + +class _DatasetOnlyModel(BaseModel): + train_dataset: PydanticDatasetIFType + + +def test_weighted_combined_is_buildable_from_a_config(dummy_packed_data_path: Path): + # Guards the registry wiring: component key, variant key, config model and factory + # signature all have to line up for a config to resolve. + config_dict = { + "train_dataset": { + "component_key": "dataset", + "variant_key": "weighted_combined", + "config": { + "datasets": [ + { + "component_key": "dataset", + "variant_key": "packed_mem_map_dataset_continuous", + "config": { + "raw_data_path": str(dummy_packed_data_path), + "sequence_length": 4, + "sample_key": "input_ids", + "reuse_last_target": True, + }, + } + ], + "repeat_factors": [2.0], + "seed": 3, + }, + } + } + + component_factory = ComponentFactory(registry=Registry(COMPONENTS)) + components = component_factory.build_components(config_dict=config_dict, components_model_type=_DatasetOnlyModel) + + dataset = components.train_dataset + assert isinstance(dataset, WeightedCombinedDataset) + assert dataset.repeat_factors == [2.0] + assert len(dataset) == 2 * len(dataset.datasets[0]) + + +def test_negative_repeat_factor_is_rejected_by_the_config_model(): + from modalities.config.config import WeightedCombinedDatasetConfig + + with pytest.raises(ValueError): + WeightedCombinedDatasetConfig(datasets=[], repeat_factors=[-1.0])