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 labels: Option<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,
            labels: None,
            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,
                labels: None,
                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
    }

    /// Load assistant-only SFT data from a two-column TSV:
    ///
    /// ```text
    /// user<TAB>assistant
    /// ```
    ///
    /// Literal `\n` and `\t` escapes are decoded. Each row is formatted as
    /// `User: ...\nAssistant:` + answer + `<|endoftext|>`. Labels are `-100`
    /// for prompt tokens and the token id itself for answer/EOS tokens, so the
    /// cross-entropy op ignores prompt rows while still training the assistant
    /// answer and stop token.
    pub fn load_sft_tsv_cached(tokenizer_path: &Path, corpus_path: &Path) -> Self {
        let token_cache = cache_path(corpus_path);
        let label_cache = label_cache_path(corpus_path);
        let vocab_size = Tokenizer::from_file(tokenizer_path).vocab_size();
        if token_cache.exists() && label_cache.exists() {
            let tokens = read_u16_cache(&token_cache);
            let labels = read_i32_cache(&label_cache);
            assert_eq!(
                tokens.len(),
                labels.len(),
                "SFT cache token/label length mismatch"
            );
            println!(
                "corpus: read {} cached SFT tokens from {} (+ labels {})",
                tokens.len(),
                token_cache.display(),
                label_cache.display()
            );
            return Self {
                tokens,
                labels: Some(labels),
                vocab_size,
            };
        }

        let tok = Tokenizer::from_file(tokenizer_path);
        let text = std::fs::read_to_string(corpus_path)
            .unwrap_or_else(|e| panic!("failed to read SFT corpus {}: {e}", corpus_path.display()));
        let mut tokens = Vec::new();
        let mut labels = Vec::new();
        for (lineno, line) in text.lines().enumerate() {
            if line.trim().is_empty() {
                continue;
            }
            let (user, assistant) = line
                .split_once('\t')
                .unwrap_or_else(|| panic!("SFT TSV line {} missing tab", lineno + 1));
            let user = decode_tsv_escapes(user);
            let assistant = decode_tsv_escapes(assistant);
            let prompt = format!("User: {user}\nAssistant:");
            let answer = format!(" {assistant}\n<|endoftext|>");
            let prompt_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
            let answer_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
            let (row_tokens, row_labels) = sft_row(&prompt_ids, &answer_ids);
            tokens.extend(row_tokens);
            labels.extend(row_labels);
        }
        assert_eq!(tokens.len(), labels.len(), "SFT tokens/labels mismatch");
        write_u16_cache(&token_cache, &tokens);
        write_i32_cache(&label_cache, &labels);
        println!(
            "corpus: tokenized {} SFT tokens → cached to {} (+ labels {})",
            tokens.len(),
            token_cache.display(),
            label_cache.display()
        );
        Self {
            tokens,
            labels: Some(labels),
            vocab_size: tok.vocab_size(),
        }
    }

    /// 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_tokens = self.tokens[cut..].to_vec();
        let valid_labels = self.labels.as_ref().map(|labels| labels[cut..].to_vec());
        let mut train = self.tokens;
        train.truncate(cut);
        let train_labels = self.labels.map(|mut labels| {
            labels.truncate(cut);
            labels
        });
        (
            Self {
                tokens: train,
                labels: train_labels,
                vocab_size: self.vocab_size,
            },
            Self {
                tokens: valid_tokens,
                labels: valid_labels,
                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 mut start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
        if let Some(labels) = &self.labels {
            for _ in 0..16 {
                if labels[start + 1..start + seq + 1].iter().any(|&t| t >= 0) {
                    break;
                }
                start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
            }
        }
        let input = self.tokens[start..start + seq].to_vec();
        let target = self.target_window(start, seq);
        (input, target)
    }

    /// Deterministic target labels for an input window starting at `start`.
    pub fn target_window(&self, start: usize, seq: usize) -> Vec<i32> {
        match &self.labels {
            Some(labels) => labels[start + 1..start + seq + 1].to_vec(),
            None => self.tokens[start + 1..start + seq + 1].to_vec(),
        }
    }
}

/// 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)
}

fn label_cache_path(corpus_path: &Path) -> PathBuf {
    let mut s = corpus_path.as_os_str().to_os_string();
    s.push(".labels.i32.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()
}

fn read_i32_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() % 4 == 0, "corrupt i32 cache (odd byte count)");
    buf.chunks_exact(4)
        .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
        .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");
}

fn write_i32_cache(path: &Path, labels: &[i32]) {
    let mut w = BufWriter::new(
        std::fs::File::create(path)
            .unwrap_or_else(|e| panic!("create cache {}: {e}", path.display())),
    );
    for &t in labels {
        w.write_all(&t.to_le_bytes()).expect("write cache");
    }
    w.flush().expect("flush cache");
}

fn decode_tsv_escapes(s: &str) -> String {
    s.replace("\\n", "\n").replace("\\t", "\t")
}

/// Build one SFT example's `(tokens, labels)` from already-tokenized prompt/answer
/// ids: prompt tokens are masked to the ignore-index (`-100`, which `cross_entropy`
/// skips) so only the answer + EOS tokens contribute to the loss. Pure (no tokenizer
/// / no CUDA) so the assistant-only masking is unit-testable directly.
fn sft_row(prompt_ids: &[i32], answer_ids: &[i32]) -> (Vec<i32>, Vec<i32>) {
    let mut tokens = Vec::with_capacity(prompt_ids.len() + answer_ids.len());
    tokens.extend_from_slice(prompt_ids);
    tokens.extend_from_slice(answer_ids);
    let mut labels = Vec::with_capacity(prompt_ids.len() + answer_ids.len());
    labels.extend(std::iter::repeat(-100).take(prompt_ids.len()));
    labels.extend_from_slice(answer_ids);
    (tokens, labels)
}

/// 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
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sft_row_masks_prompt_supervises_answer() {
        let prompt = [5, 6, 7];
        let answer = [8, 9]; // includes the EOS token in real use
        let (tokens, labels) = sft_row(&prompt, &answer);
        // Tokens are prompt then answer, in order.
        assert_eq!(tokens, vec![5, 6, 7, 8, 9]);
        // Prompt positions are ignore-index (-100); answer positions are supervised.
        assert_eq!(labels, vec![-100, -100, -100, 8, 9]);
        assert_eq!(tokens.len(), labels.len());
    }

    #[test]
    fn sft_row_handles_empty_answer() {
        let (tokens, labels) = sft_row(&[1, 2], &[]);
        assert_eq!(tokens, vec![1, 2]);
        assert_eq!(labels, vec![-100, -100]);
    }
}