training-data-packer

Packaging annotated datasets into final training data. Purpose is to have a repeatable process and well-packaged data to remove any data management in the training step.

The packer consists of the following tools:

• oellm-package-data - Take source files and apply decontamination, PII-masking, and sampling. Each file in the source directory gets a correspondent file with the data processed. The tool is idempotent, if it fails, then run it again, and it will take of where it left. This is the main tool. The tool contain three sub tools: lint, sample, release (default), and merge.
• oellm-collect-metrics - Collect and summarize metrics from a collection directory. This does also happens automatically when run oellm-package-data in mode merge.
• oellm-propella-structure - Structure Propella data based on source data structure. For each record in the source files, if its ID exists in the Propella data, it is written to the output. This arranges Propella data in the same oßrder as source data.
• oellm-propella-merge - If Propella data is to big to be in memory oellm-propella-structure can run on individual Propella-parquet files. Then use this tool to merge the results from all Propella-parquet files.

All tools read a file metadata.yaml containing metadata about the structure and processing of the data.

TL;DR; I just want to run!

Here are steps to run packaging and reduce the number of files. Instructions are intended for Lumi and using latest Singularity container.

1. Check out https://github.com/OpenEuroLLM/training-data-packer
2. In tools/slurm there are scripts to schedule jobs with SLURM. Run them from the root of the checkout, since job ids and logs are written to the logs directory in the current directory.
3. Start the pipeline with package, merge, and tokenization: ./tools/slurm/pipeline-auto-plan.sh /scratch/project_465002530/training/collection/flag/finepdfs-edu-1.0.0 Change path to the dataset you want to process. The script lints the metadata first and stops if it finds errors. The array size and time of each job are calculated from the files in the collection.
4. Check the status of the scheduled pipelines with ./tools/slurm/pipeline-status.sh. It prints a Markdown table with one row per collection and one column per step: 🏃🏻 running, ✅ completed, ❌ failed, 🕘 pending or waiting for a dependency, ❔ unknown to SLURM. Completed steps also show their run time.
5. Check logs and metrics for each step.
6. When the pipeline is done, move logs from logs directory within the directory where you checked out the code into /scratch/project_465002530/training/collection/flag/log/release

The two first steps are a one time operation or when you update the packager.

To set the array sizes yourself, use ./tools/slurm/pipeline.sh <collection-dir> <package-array> <merge-array> <tokenize-array> instead, e.g. ./tools/slurm/pipeline.sh /scratch/project_465002530/training/collection/flag/finepdfs-edu-1.0.0 0-10 0-38 0-38.

Dataset Directory Structure

A dataset directory managed by the packager contains input sources, annotations, configuration, and generated outputs. The directories relate to each other as follows:

dataset-directory/
├── metadata.yaml        # Dataset configuration, schema mappings, and pipeline rules
├── source/              # Original raw data files (e.g., .jsonl.gz), arranged in hierarchical subdirectories
├── open-privacy-filter/ # PII annotation files mirroring source/ paths (documents PII spans to mask). Output from open-privacy-filter
├── nemo-curator/        # Contamination annotation files mirroring source/ paths (documents benchmark-contaminated records to remove). Output from Nemo-Curator
├── propella-4b/         # Optional Propella quality score parquet files used for quality filtering
├── release-raw/         # Output of oellm-package-data, mode release: per-file processed data with PII masked, decontaminated, and sampled
├── release/             # Output of oellm-package-data, mode merge: consolidated and merged files ready for tokenization
├── megatron-lm/         # Output of oellm-package-data, mode tokenize: one subdirectory per tokenizer in the megatron-lm section
└── sample/              # Output of oellm-package-data, mode sample: files sampled on quality using propella-4b and other quality metrics

Directory Roles and Data Flow

1. Input & Annotations:

   • source/: The primary input data. Files are typically organized by domain, language, or partition.
   • open-privacy-filter/, propella-4b & nemo-curator/: Annotation companion directories. Every annotation file mirrors the relative path and filename of the corresponding file in source/.
   • metadata.yaml: Coordinates the pipeline by specifying dataset name, source URLs, annotation file formats, sampling ratios, and shard parameters.

2. Packaging (oellm-package-data):

   • Reads source/, applies masking annotations from pii/ and exclusions from contamination/, and writes the packaged files to release-raw/.
   • Output files in release-raw/ preserve the directory hierarchy of source/ and include hidden metric sidecars (.*.metrics.json).

3. Consolidation (oellm-package-data -m merge):

   • Reads files from release-raw/ and merges smaller chunks into target-sized files in release/, maintaining path semantics while reducing file handle overhead for tokenization and pretraining.
   • Compile metrics

This is the beginning of the README. Read the rest at the original. Read the full README ↗