| """Data pipeline for SplitBit LLM training. |
| |
| Streaming data loader — doesn't load full dataset into RAM. |
| Dynamic batching based on available memory. |
| Data augmentation: random cropping, context window sliding. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import os |
| import random |
| from typing import Iterator, List |
|
|
| import numpy as np |
|
|
| from ..model.tokenizer import BPETokenizer, BOS_ID, EOS_ID |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class DataPipeline: |
| """Streaming text data pipeline for training. |
| |
| Reads .txt files one at a time, encodes to tokens, and yields |
| batches of fixed-length sequences for training. |
| """ |
|
|
| def __init__( |
| self, |
| tokenizer: BPETokenizer, |
| seq_len: int = 256, |
| batch_size: int = 4, |
| data_dir: str | None = None, |
| ) -> None: |
| self.tokenizer = tokenizer |
| self.seq_len = seq_len |
| self.batch_size = batch_size |
| self.data_dir = data_dir |
|
|
| def load_texts(self, paths: List[str] | str) -> List[str]: |
| """Load text from files. Returns list of text strings.""" |
| if isinstance(paths, str): |
| paths = [paths] |
|
|
| texts = [] |
| for path in paths: |
| if os.path.isdir(path): |
| for fname in sorted(os.listdir(path)): |
| if fname.endswith((".txt", ".md", ".jsonl")): |
| fpath = os.path.join(path, fname) |
| with open(fpath, "r", encoding="utf-8", errors="replace") as f: |
| texts.append(f.read()) |
| elif os.path.isfile(path): |
| with open(path, "r", encoding="utf-8", errors="replace") as f: |
| texts.append(f.read()) |
| else: |
| logger.warning("Data path not found: %s", path) |
|
|
| total_chars = sum(len(t) for t in texts) |
| logger.info("Loaded %d texts, %d total chars", len(texts), total_chars) |
| return texts |
|
|
| def encode_texts(self, texts: List[str]) -> np.ndarray: |
| """Encode all texts into a single token array.""" |
| all_ids: List[int] = [] |
| for text in texts: |
| ids = self.tokenizer.encode(text, add_bos=True, add_eos=True) |
| all_ids.extend(ids) |
| logger.info("Encoded %d total tokens", len(all_ids)) |
| return np.array(all_ids, dtype=np.int64) |
|
|
| def create_batches(self, token_ids: np.ndarray, shuffle: bool = True) -> Iterator[tuple[np.ndarray, np.ndarray]]: |
| """Yield (input, target) batches for next-token prediction. |
| |
| input: [batch, seq_len] — tokens 0..seq_len-1 |
| target: [batch, seq_len] — tokens 1..seq_len (shifted by 1) |
| """ |
| n_tokens = len(token_ids) |
| n_seqs = n_tokens // (self.seq_len + 1) |
| if n_seqs == 0: |
| logger.warning("Not enough tokens (%d) for seq_len=%d", n_tokens, self.seq_len) |
| return |
|
|
| |
| usable = n_seqs * (self.seq_len + 1) |
| token_ids = token_ids[:usable] |
|
|
| |
| sequences = token_ids.reshape(n_seqs, self.seq_len + 1) |
|
|
| |
| if shuffle: |
| np.random.shuffle(sequences) |
|
|
| |
| for i in range(0, n_seqs, self.batch_size): |
| batch = sequences[i:i + self.batch_size] |
| if len(batch) < self.batch_size: |
| break |
| inputs = batch[:, :-1] |
| targets = batch[:, 1:] |
| yield inputs, targets |
|
|
| def stream_batches(self, paths: List[str] | str, shuffle: bool = True) -> Iterator[tuple[np.ndarray, np.ndarray]]: |
| """Stream batches from files without loading everything at once. |
| |
| Reads one file at a time, encodes, and yields batches. |
| """ |
| if isinstance(paths, str): |
| paths = [paths] |
|
|
| for path in paths: |
| texts = self.load_texts(path) |
| if not texts: |
| continue |
| token_ids = self.encode_texts(texts) |
| yield from self.create_batches(token_ids, shuffle=shuffle) |
|
|
