xtrain/crates/xtrain-train/src/data.rs

//! Data pipeline: load the GPT-2 BPE (reusing xserv's from-scratch tokenizer),
//! tokenize a text corpus into one flat token stream, and sample fixed-length
//! `(input, target)` windows for next-token prediction. Host-only (no GPU).
//!
//! For the scaling runs the corpus is large (full TinyStories ≈ 2 GB / ~470 M
//! tokens), and the from-scratch BPE is slow, so [`Corpus::load_cached`]
//! tokenizes ONCE and caches the id stream to a `<corpus>.u16.bin` next to the
//! text (GPT-2 vocab = 50257 < 65536, so u16 is exact). Subsequent runs mmap-read
//! the cache instead of re-tokenizing.

use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use xserv_tokenizer::Tokenizer;

/// A tokenized corpus: one flat stream of token ids, plus the vocab size.
pub struct Corpus {
    pub tokens: Vec<i32>,
    pub vocab_size: usize,
}

impl Corpus {
    /// Load `tokenizer.json` (GPT-2 BPE) and tokenize the UTF-8 text at
    /// `corpus_path` into a single id stream. TinyStories separates stories with
    /// `<|endoftext|>`; the GPT-2 tokenizer emits that as a single special token,
    /// so document boundaries are preserved in the stream.
    pub fn load(tokenizer_path: &Path, corpus_path: &Path) -> Self {
        let tok = Tokenizer::from_file(tokenizer_path);
        let text = std::fs::read_to_string(corpus_path)
            .unwrap_or_else(|e| panic!("failed to read corpus {}: {e}", corpus_path.display()));
        // The range-fetched corpus may start/end mid-story; drop a leading partial
        // line and a trailing partial story so we only train on whole sentences.
        let text = trim_to_whole_stories(&text);
        let ids: Vec<i32> = tok.encode(text).into_iter().map(|t| t as i32).collect();
        Self {
            tokens: ids,
            vocab_size: tok.vocab_size(),
        }
    }

    /// Like [`load`](Self::load) but caches the tokenized id stream to
    /// `<corpus_path>.u16.bin`. On the first run it tokenizes the (large) corpus
    /// and writes the cache; on later runs it reads the cache directly, skipping
    /// the slow BPE. The cache is just a flat little-endian `[u16]` (no header) —
    /// it is keyed only by path, so delete it if the corpus or tokenizer changes.
    pub fn load_cached(tokenizer_path: &Path, corpus_path: &Path) -> Self {
        let cache = cache_path(corpus_path);
        let vocab_size = Tokenizer::from_file(tokenizer_path).vocab_size();
        if cache.exists() {
            let tokens = read_u16_cache(&cache);
            println!(
                "corpus: read {} cached tokens from {}",
                tokens.len(),
                cache.display()
            );
            return Self { tokens, vocab_size };
        }
        let me = Self::load(tokenizer_path, corpus_path);
        write_u16_cache(&cache, &me.tokens);
        println!(
            "corpus: tokenized {} tokens → cached to {}",
            me.tokens.len(),
            cache.display()
        );
        me
    }

    /// Split off the last `n` tokens as a held-out validation corpus, leaving the
    /// rest as the train corpus. Returns `(train, valid)`. Used for periodic val
    /// loss during training without leaking the eval window into training.
    pub fn split_tail(self, n: usize) -> (Self, Self) {
        let n = n.min(self.tokens.len() / 10); // never hand off more than 10%
        let cut = self.tokens.len() - n;
        let valid = self.tokens[cut..].to_vec();
        let mut train = self.tokens;
        train.truncate(cut);
        (
            Self {
                tokens: train,
                vocab_size: self.vocab_size,
            },
            Self {
                tokens: valid,
                vocab_size: self.vocab_size,
            },
        )
    }

    /// Total number of tokens.
    pub fn len(&self) -> usize {
        self.tokens.len()
    }

    pub fn is_empty(&self) -> bool {
        self.tokens.is_empty()
    }

    /// Sample one `(input, target)` pair of length `seq` for next-token
    /// prediction: a window `[s, s+seq+1)` → input `[s, s+seq)`, target shifted
    /// by one. `rng_state` is advanced (a tiny LCG, so sampling is reproducible
    /// from a seed without pulling in an RNG crate).
    pub fn sample(&self, seq: usize, rng_state: &mut u64) -> (Vec<i32>, Vec<i32>) {
        assert!(self.tokens.len() > seq + 1, "corpus shorter than a window");
        let max_start = self.tokens.len() - seq - 1;
        let start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
        let input = self.tokens[start..start + seq].to_vec();
        let target = self.tokens[start + 1..start + seq + 1].to_vec();
        (input, target)
    }
}

/// Drop a leading partial line (before the first newline) and everything after
/// the last `<|endoftext|>` marker, so a byte-range download still yields only
/// complete stories. Falls back to the raw text if no marker is present.
fn trim_to_whole_stories(text: &str) -> &str {
    let start = text.find('\n').map(|i| i + 1).unwrap_or(0);
    let body = &text[start..];
    match body.rfind("<|endoftext|>") {
        Some(end) => &body[..end + "<|endoftext|>".len()],
        None => body,
    }
}

/// `<corpus_path>.u16.bin` — the token-id cache beside the corpus text.
fn cache_path(corpus_path: &Path) -> PathBuf {
    let mut s = corpus_path.as_os_str().to_os_string();
    s.push(".u16.bin");
    PathBuf::from(s)
}

/// Read a flat little-endian `[u16]` cache into an `i32` id stream.
fn read_u16_cache(path: &Path) -> Vec<i32> {
    let mut r = BufReader::new(
        std::fs::File::open(path).unwrap_or_else(|e| panic!("open cache {}: {e}", path.display())),
    );
    let mut buf = Vec::new();
    r.read_to_end(&mut buf).expect("read cache");
    assert!(buf.len() % 2 == 0, "corrupt u16 cache (odd byte count)");
    buf.chunks_exact(2)
        .map(|b| u16::from_le_bytes([b[0], b[1]]) as i32)
        .collect()
}

/// Write an id stream as a flat little-endian `[u16]` cache. Ids must fit in u16
/// (GPT-2 vocab = 50257 < 65536); asserts otherwise.
fn write_u16_cache(path: &Path, tokens: &[i32]) {
    let mut w = BufWriter::new(
        std::fs::File::create(path)
            .unwrap_or_else(|e| panic!("create cache {}: {e}", path.display())),
    );
    for &t in tokens {
        assert!((0..=u16::MAX as i32).contains(&t), "token id {t} > u16");
        w.write_all(&(t as u16).to_le_bytes()).expect("write cache");
    }
    w.flush().expect("flush cache");
}

/// Tiny LCG (same constants as the model tests' deterministic fill) so dataset
/// sampling is reproducible from a single u64 seed.
fn next_rand(state: &mut u64) -> u64 {
    *state = state
        .wrapping_mul(6364136223846793005)
        .wrapping_add(1442695040888963407);
    *state >> 16
}