Data ingest & normalization

Why this stage exists

Language models are trained on raw text. Before any neural network can touch it, two questions have to be answered:

1. Where does the text come from? v0 pulls from ~12 HuggingFace datasets spanning dialogue, reasoning traces, and encyclopedic content.
2. What shape does each sample have after this stage? Every sample is converted to one uniform schema so the rest of the pipeline does not care which dataset it came from.

Why a mix of datasets

A model trained on only one genre of text tends to talk like that genre — pure Wikipedia produces an encyclopedia, pure chat produces a chatbot that has no facts. Modern chat-LM recipes mix:

• Dialogue data — multi-turn conversations so the model learns turn-taking and response style.
• Reasoning traces — chain-of-thought style data from stronger teacher models; this is where the model learns to show its working.
• Encyclopedic text — Wikipedia-style dense factual prose keeps the model from losing general knowledge.

v0 ships a DATASETS registry in download_data.py covering all three roles. The full list with paper citations is in ../../techniques#datasets-v010-mix.

Download: download_data.py

The script iterates over the registry and, for each entry:

1. Calls datasets.load_dataset(path, split=..., streaming=...).
2. Saves the result to data/<name>/ via Dataset.save_to_disk(...) — HuggingFace's on-disk Arrow format, which is cheap to memory-map later.

Two implementation details carry weight:

• HF cache is redirected. HF_HOME and HF_DATASETS_CACHE are pinned to data/.hf_cache/ so downloads stay in the repo tree. Without this, HuggingFace would dump a multi-gigabyte cache into ~/.cache/huggingface/.
• Streaming vs. non-streaming. For datasets larger than ~1 million rows (Wikipedia is 6.4M articles), the registry sets streaming=False. The naive alternative — streaming=True + list.append(sample) + Dataset.from_list(rows) — accumulates every row into Python memory and OOMs. Non-streaming mode lets Arrow handle sharding via memory-mapped files.
• Wikipedia byte cap. The Wikipedia entry carries an extra max_bytes=32 * 1024**3 field. After HuggingFace materializes the memory-mapped Arrow snapshot, the download loop iterates ds.select_columns(["text"]) row-by-row to sum UTF-8 bytes; once the running total would exceed 32 GiB the scan stops and the final dataset is the ds.select(range(n)) prefix that fits. Iterating the text column via mmap is pointer arithmetic — no copy into Python — so the cap applies safely on machines with far less than 32 GiB of RAM. Today the full English dump is ~22 GiB of text and the cap is not hit; the field is a ceiling against future corpus growth.

See download_data.py:102 and download_data.py:173 for the streaming-flag plumbing.

Normalization: _norm_<name> functions

Every source dataset arrives in its own schema — WildChat has a conversation column, SODA has dialogue, Wikipedia has text. The notebook defines one _norm_<dataset> function per dataset, and a NORMALIZERS dict dispatches by dataset name.

The output shape is always a list of chat turns:

[
    {"role": "user", "content": "..."},
    {"role": "assistant", "content": "..."},
]

For non-dialogue sources (Wikipedia, reasoning dumps), the text is wrapped into a single- or two-turn form so downstream code reading {role, content} works unchanged.

The contract that cannot silently break

The tokenizer trains on whatever the normalizers emit, and the data loader reads by column name. If you rename a column in download_data.py without updating its normalizer in the notebook, or vice versa, training continues without raising — it just trains on garbage. This is the single most common source of silent bugs at this stage.

Reading the code

• download_data.py — the dataset registry and save loop.
• full_llm_pipeline.ipynb — search for _norm_ to see every normalizer and the NORMALIZERS dispatch dict.

Pre-flight size estimate (v0.1.1)

With SAMPLE_LIMIT=-1 the notebook pulls full corpora, which can take hours and tens of GB of disk before you see any feedback. v0.1.1 adds a pre-flight block to the dataset-loading cell so the cost is known before the ingest loop starts.

For each registry entry, the block does one of three things:

• If the dataset is already on disk under data/<name>/ (a prior python download_data.py run), it walks the directory and reports the on-disk size — no network hit.
• If the entry uses a parquet= pointer at HuggingFace's auto- converted parquet branch, it marks the size as unknown. There is no cheap metadata probe for a raw parquet URL; the size is only discovered at stream time.
• Otherwise it calls datasets.load_dataset_builder(path, name=...).info. That's a metadata-only call — no data is downloaded — and it returns download_size (bytes transferred) and dataset_size (bytes extracted into the HF cache). Gated datasets (LMSYS-Chat, Arena) fail here without HF_TOKEN and are flagged (gated -- needs HF_TOKEN) in the output.

Then it rolls up known totals across probed sources:

• Network download — sum of download_size. This lands in the HF cache directory (which v0 pins to data/.hf_cache/ via HF_HOME / HF_DATASETS_CACHE).
• On-disk footprint — sum of dataset_size.
• Estimated RAM for train_data — 2× to 3× the on-disk footprint. That multiplier accounts for Python's per-object overhead: the normalized list holds dicts of strs, and short-string overhead is substantial relative to the payload. The Wikipedia normalizer truncates each article to 4 KB of text (see _norm_wikipedia), so this is an upper bound for that source.

Finally, when psutil and shutil.disk_usage are available, it compares the projections against the host's vm.available and disk.free for DATA_DIR, emitting explicit WARNING lines if either projection exceeds what the host can provide. Parquet sources are counted as unknown, so the printed totals are a lower bound, not a guarantee.

The practical outcome: on a host with, say, 32 GB of RAM and 500 GB of free disk, you will see at session start whether the run is safe to let go overnight or needs a SAMPLE_LIMIT cap before it starts.

Tokenised-dataset cache (v0.1.3)

Once the normalizer stage produces train_data (raw conversations) and the upsample stage produces train_texts (weighted-expanded flat strings rendered in the GPT-4o ChatML template — see Tokenizer for the template details), the next stage is tokenising both into tensors for the three training regimes:

• PretrainDataset flattens train_texts into one long token stream and chunks it into fixed-length windows of size MAX_LEN.
• SFTDataset tokenises every conversation in train_data into a padded (ids, mask) pair, with the loss mask set only on assistant turns.
• DPODataset tokenises each conversation twice (chosen + rejected-with-shuffled-content) and produces a 4-tuple of padded tensors.

On 1M-conversation corpora each of these stages takes hours. A mid-stage crash used to wipe all that work. v0.1.3 adds a resumable on-disk cache under data/cache/ (inherits the data/ gitignore entry, so nothing ever leaks into a commit):

File	Contents	Resumable granularity
data/cache/tokenizer.json	Trained BPE	whole tokenizer (load or retrain)
data/cache/pretrain_tokens.bin	Flat int32 token stream (np.memmap)	per conversation
data/cache/sft_ids.bin + sft_mask.bin	int32 ids + int8 mask, one row per sample	per conversation
data/cache/dpo_ids.bin + dpo_mask.bin	int32 ids + int8 mask, one row per pair (concat chosen + rejected)	per conversation

Each *.bin file is accompanied by a *.meta.json that records the cache key + last committed byte offset; those are the source of truth for "what on disk is actually valid."

Storage format changed in v0.1.4. v0.1.3 used torch.save pickles of Python lists of tensors, which required holding the whole corpus in RAM before saving. v0.1.4 streams raw int32 / int8 bytes to the .bin files in small write batches (2 K docs / 500 convs / 500 convs respectively) and memory-maps them at training time via np.memmap, so peak construction memory is O(write batch) regardless of corpus size.

Cache-save cadence

Each resumable cache is re-saved every 30 minutes of wall time during its tokenisation loop. A crash therefore costs at most that much work, not the whole run. Saves are atomic (torch.save(..., path.tmp) followed by os.replace(tmp, path)), so an interrupted save leaves the previous valid cache in place rather than corrupting anything.

Cache-key invalidation

Each cache file embeds a key dict at save time. The next run rebuilds that key from the current inputs and reuses the cache only on exact match. The fields are:

• tok_sha: sha256 of tokenizer.to_str(). Retraining the tokenizer produces a different SHA and invalidates every downstream cache — so retokenising costs as much as it should and never less.
• max_len, pad_id, eos_id, limit: the numeric knobs that change the tensor shape or masking.
• texts_fp / raw_fp: a cheap content fingerprint ({n, total_chars|head, head_sha16, tail_sha16}). Catches the case where n_texts is preserved but the content differs (e.g. a different SAMPLE_LIMIT or a changed normalizer).

Fast path

When the final cache file's key matches the current run's key, the __init__ of the dataset class skips every loop and returns directly from the cached tensor list in seconds. This is the normal case on a clean restart.

Forcing a full rebuild

Delete data/cache/ (or the relevant *.bin + *.meta.json pair) to force a clean rebuild. Do this after any intentional data-pipeline change that the fingerprint might miss (a normalizer edit that preserves length, for example). Old *.pt files from v0.1.3 are simply ignored by v0.1.4 code — you can rm data/cache/*.pt to reclaim disk.

Next stage: Tokenizer — turning the normalized text into integer IDs.