tokenizer/streaming_tokenizer.py · Phase-Technologies/FORGE-500M-Chat at main

FORGE-500M-Chat / tokenizer /streaming_tokenizer.py

Add tokenizer component: streaming_tokenizer.py

6b836f4 verified 4 days ago

10.9 kB

"""
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))

	"""
	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))