""" FORGE-3B Streaming Tokenizer. A thin wrapper around ForgeTokenizer that streams a large corpus file and writes tokenized output to disk as memory-mapped numpy arrays. Designed for preprocessing pipelines where the full corpus does not fit in RAM. Processes documents lazily in chunks, writing each completed shard atomically. Usage: python -m tokenizer.streaming_tokenizer \\ --input /data/raw/fineweb_edu \\ --output /data/tokenized/fineweb_edu \\ --profile standard \\ --seq_len 2048 \\ --shard_size 100000 """ from __future__ import annotations import os import json import time import logging import argparse from pathlib import Path from typing import Iterator, List import numpy as np logger = logging.getLogger(__name__) _BOS = 1 _EOS = 2 _PAD = 0 _DTYPE = np.uint32 # ───────────────────────────────────────────────────────────────────────────── # Streaming Tokenizer # ───────────────────────────────────────────────────────────────────────────── class StreamingTokenizer: """ Tokenizes a corpus of JSONL files in a streaming fashion, packing tokens into fixed-length sequences and writing them to numbered .npy shards. Key properties: - O(shard_size * seq_len * 4 bytes) peak RAM usage — stream-safe - Atomic shard writes (temp file → rename) - Resumable: detects and skips existing shards - Thread-safe per instance (no shared state) """ def __init__( self, tokenizer, output_dir: str, seq_len: int = 2048, shard_size: int = 100_000, split: str = "train", ): self.tok = tokenizer self.seq_len = seq_len self.shard_size = shard_size self.split = split self._out_dir = Path(output_dir) self._out_dir.mkdir(parents=True, exist_ok=True) # Rolling packed-token buffer self._buffer: List[int] = [] # Shard accumulation buffer self._shard_seqs: List[np.ndarray] = [] self._shard_idx = self._detect_next_shard() # Stats self._n_docs = 0 self._n_tokens = 0 self._n_seqs = 0 self._t0 = time.perf_counter() # ── Internal helpers ────────────────────────────────────────────────────── def _detect_next_shard(self) -> int: """Skip past shards that already exist on disk (resumable run).""" existing = sorted(self._out_dir.glob(f"{self.split}_shard_*.npy")) if not existing: return 0 idx = int(existing[-1].stem.split("_")[-1]) logger.info(f"StreamingTokenizer: resuming after shard {idx}") return idx + 1 def _shard_path(self, idx: int) -> Path: return self._out_dir / f"{self.split}_shard_{idx:04d}.npy" def _flush_shard(self): """Write the shard buffer to disk atomically via temp→rename.""" if not self._shard_seqs: return arr = np.stack(self._shard_seqs, axis=0) # (N, seq_len) uint32 path = self._shard_path(self._shard_idx) tmp = path.with_suffix(".tmp.npy") np.save(str(tmp), arr) tmp.rename(path) logger.info( f" Shard {self._shard_idx:04d}: {len(arr)} seqs → {path.name} " f"({arr.nbytes / 1e6:.1f} MB)" ) self._shard_idx += 1 self._shard_seqs = [] def _drain_buffer(self): """Carve complete sequences from the rolling buffer.""" while len(self._buffer) >= self.seq_len: seq = np.array(self._buffer[: self.seq_len], dtype=_DTYPE) self._buffer = self._buffer[self.seq_len :] self._shard_seqs.append(seq) self._n_seqs += 1 if len(self._shard_seqs) >= self.shard_size: self._flush_shard() # ── Public API ──────────────────────────────────────────────────────────── def add_document(self, text: str): """Tokenize one text document and pack its tokens into the buffer.""" try: token_ids = self.tok.encode( text, add_bos=False, add_eos=False, truncate=False, ) except Exception as exc: logger.warning(f"Tokenization failed: {exc} — skipping document") return if not token_ids: return # BOS + tokens + EOS self._buffer.extend([_BOS] + list(token_ids) + [_EOS]) self._n_tokens += len(token_ids) + 2 self._n_docs += 1 self._drain_buffer() def add_token_array(self, token_ids: np.ndarray): """Add a pre-tokenized uint32 array directly (bypasses the tokenizer).""" self._buffer.extend([_BOS] + token_ids.tolist() + [_EOS]) self._n_tokens += len(token_ids) + 2 self._n_docs += 1 self._drain_buffer() def finalize(self) -> dict: """ Flush all remaining sequences (partial last shard) and write metadata. Must be called exactly once after all documents have been added. Returns a summary dict. """ # Drain any remaining full sequences if len(self._buffer) >= self.seq_len: self._drain_buffer() # Write partial shard if there is anything left if self._shard_seqs: self._flush_shard() elapsed = time.perf_counter() - self._t0 speed = self._n_tokens / max(elapsed, 1e-6) meta = { "n_documents": self._n_docs, "n_tokens": self._n_tokens, "n_sequences": self._n_seqs, "seq_len": self.seq_len, "shard_size": self.shard_size, "n_shards": self._shard_idx, "elapsed_s": round(elapsed, 1), "tokens_per_sec": int(speed), } meta_path = self._out_dir / f"{self.split}_streaming_meta.json" with open(str(meta_path), "w") as f: json.dump(meta, f, indent=2) logger.info( f"StreamingTokenizer finalized: " f"{self._n_seqs:,} seqs | " f"{self._n_tokens / 1e9:.2f}B tokens | " f"{speed / 1e6:.1f}M tok/s | " f"{elapsed:.1f}s" ) return meta # ───────────────────────────────────────────────────────────────────────────── # Corpus Iterator # ───────────────────────────────────────────────────────────────────────────── def _iter_texts(input_path: Path, text_key: str = "text") -> Iterator[str]: """Yield text strings from all JSONL / plain-text files under input_path.""" files: list = [] for ext in (".jsonl", ".json", ".txt"): files.extend(sorted(input_path.rglob(f"*{ext}"))) if not files: raise FileNotFoundError(f"No JSONL/txt files found under {input_path}") logger.info(f"_iter_texts: found {len(files)} files") for path in files: try: if path.suffix == ".txt": text = path.read_text(encoding="utf-8", errors="replace").strip() if text: yield text else: with open(str(path), encoding="utf-8", errors="replace") as fh: for line in fh: line = line.strip() if not line: continue try: obj = json.loads(line) text = obj.get(text_key, "") if text: yield text except json.JSONDecodeError: continue except OSError as exc: logger.warning(f"Cannot read {path}: {exc}") # ───────────────────────────────────────────────────────────────────────────── # CLI Entry Point # ───────────────────────────────────────────────────────────────────────────── if __name__ == "__main__": logging.basicConfig( level=logging.INFO, format="[%(asctime)s][%(levelname)s] %(message)s", datefmt="%H:%M:%S", ) parser = argparse.ArgumentParser( description="FORGE-3B Streaming Tokenizer", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument("--input", required=True, help="Input dir with JSONL/txt files") parser.add_argument("--output", required=True, help="Output dir for .npy shards") parser.add_argument("--profile", default="standard", choices=["standard", "lite"]) parser.add_argument("--seq_len", type=int, default=2048) parser.add_argument("--shard_size", type=int, default=100_000) parser.add_argument("--split", default="train") parser.add_argument("--text_key", default="text") parser.add_argument("--n_workers", type=int, default=min(8, os.cpu_count() or 1)) args = parser.parse_args() from tokenizer.crayon_wrapper import ForgeTokenizer tok = ForgeTokenizer( profile=args.profile, device="cpu", n_workers=args.n_workers, max_length=args.seq_len * 8, ) streamer = StreamingTokenizer( tokenizer=tok, output_dir=args.output, seq_len=args.seq_len, shard_size=args.shard_size, split=args.split, ) input_path = Path(args.input) n = 0 for text in _iter_texts(input_path, text_key=args.text_key): streamer.add_document(text) n += 1 if n % 100_000 == 0: logger.info(f"Processed {n:,} documents...") meta = streamer.finalize() print(json.dumps(meta, indent=2))