File size: 4,096 Bytes
32112fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Data pipeline for Singularity 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
# Truncate to full sequences
usable = n_seqs * (self.seq_len + 1)
token_ids = token_ids[:usable]
# Reshape into sequences
sequences = token_ids.reshape(n_seqs, self.seq_len + 1)
# Shuffle
if shuffle:
np.random.shuffle(sequences)
# Yield batches
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] # [batch, seq_len]
targets = batch[:, 1:] # [batch, seq_len]
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)
|