""" CRAYON Tokenizer Wrapper Wraps CrayonVocab into a full-featured tokenizer compatible with FORGE training. Exploits CRAYON's 24M+ tok/sec CPU throughput via parallel multi-process pipeline. """ from __future__ import annotations import os import gc import math import time import logging import threading import multiprocessing as mp from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed from typing import List, Union, Optional, Dict, Any, Iterator, Tuple from pathlib import Path import queue # Fix Windows console UTF-8 printing issues for external packages (e.g., Crayon print statements with emoji) import sys import io if sys.platform.startswith("win"): try: sys.stdout.reconfigure(encoding='utf-8') sys.stderr.reconfigure(encoding='utf-8') except AttributeError: sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace') sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace') import torch import numpy as np from crayon import CrayonVocab logger = logging.getLogger(__name__) # ───────────────────────────────────────────────────────────────────────────── # THREAD-LOCAL TOKENIZER POOL # Each worker thread gets its own CrayonVocab instance to avoid any GIL # contention on the underlying C++ DAT engine. # ───────────────────────────────────────────────────────────────────────────── _thread_local = threading.local() _init_lock = threading.Lock() def _get_thread_local_tokenizer(profile: str = "standard") -> CrayonVocab: """Return a thread-local CrayonVocab instance, creating one if needed.""" if not hasattr(_thread_local, "tokenizer"): with _init_lock: # CPU is 20× faster than CUDA for CRAYON (23.84M vs 1.22M tok/s) _thread_local.tokenizer = CrayonVocab(device="cpu") _thread_local.tokenizer.load_profile(profile) logger.debug(f"Thread {threading.current_thread().name}: " f"CrayonVocab loaded with profile='{profile}'") return _thread_local.tokenizer # ───────────────────────────────────────────────────────────────────────────── # SPECIAL TOKEN REGISTRY # We manage special tokens separately since CRAYON's vocabulary may not have # them natively — we append them to the vocab ID space. # ───────────────────────────────────────────────────────────────────────────── SPECIAL_TOKENS = { "<|pad|>": 0, "<|bos|>": 1, "<|eos|>": 2, "<|sep|>": 3, "<|sys|>": 4, "<|/sys|>": 5, "<|usr|>": 6, "<|/usr|>": 7, "<|asst|>": 8, "<|/asst|>": 9, "<|tool|>": 10, "<|tool_resp|>": 11, "<|unk|>": 12, } # For the standard profile: actual vocab starts at ID 13 # Effective vocab_size = 206,373 (crayon) + 13 (specials) = 206,386 → pad to 206,464 SPECIAL_TOKEN_OFFSET = len(SPECIAL_TOKENS) # 13 CRAYON_STANDARD_VOCAB = 206_373 CRAYON_LITE_VOCAB = 50_000 # Reverse mapping for decode ID_TO_SPECIAL = {v: k for k, v in SPECIAL_TOKENS.items()} class ForgeTokenizer: """ Production-grade FORGE tokenizer built on top of CRAYON. Features: - Thread-safe, zero-copy multi-threaded tokenization - Batched encoding with automatic chunking for throughput - Full special token support with chat template - Streaming decode - Vocabulary size adapter between CRAYON profile and model vocab - Offline preprocessing at 24M+ tokens/sec on CPU Usage: tokenizer = ForgeTokenizer(profile="standard") ids = tokenizer.encode("Hello world") text = tokenizer.decode(ids) # Chat format ids = tokenizer.encode_chat([ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "What is 2+2?"}, ]) """ def __init__( self, profile: str = "standard", device: str = "cpu", n_workers: int = None, max_length: Optional[int] = None, add_bos: bool = True, add_eos: bool = True, ): self.profile = profile self.device_str = device # "cpu" recommended for CRAYON throughput self.n_workers = n_workers or max(1, mp.cpu_count() - 2) self.max_length = max_length self.add_bos = add_bos self.add_eos = add_eos # Primary tokenizer (used for single-threaded ops) self._tokenizer = CrayonVocab(device=device) self._tokenizer.load_profile(profile) # Special tokens self.special_tokens = SPECIAL_TOKENS.copy() self.id_to_special = ID_TO_SPECIAL.copy() self.special_token_offset = SPECIAL_TOKEN_OFFSET # Derive vocab size if profile == "standard": self._crayon_vocab_size = CRAYON_STANDARD_VOCAB else: self._crayon_vocab_size = CRAYON_LITE_VOCAB # Effective vocab size = specials + crayon tokens # Padded to nearest multiple of 128 for efficient matmuls raw_vocab = self._crayon_vocab_size + SPECIAL_TOKEN_OFFSET self.vocab_size = math.ceil(raw_vocab / 128) * 128 # Lock for thread safety on single-instance operations self._lock = threading.Lock() # Thread pool for parallel tokenization self._thread_pool = ThreadPoolExecutor( max_workers=self.n_workers, thread_name_prefix="forge_tok" ) logger.info( f"ForgeTokenizer initialized: profile={profile}, " f"vocab_size={self.vocab_size} (crayon={self._crayon_vocab_size} " f"+ specials={SPECIAL_TOKEN_OFFSET} + padding={self.vocab_size - raw_vocab}), " f"n_workers={self.n_workers}, throughput ~24M tok/s (CPU)" ) # ───────────────────────────────────────────────────────────────────────── # TOKEN ID CONVERSION HELPERS # CRAYON returns IDs in [0, crayon_vocab_size). # We shift them by SPECIAL_TOKEN_OFFSET to [SPECIAL_TOKEN_OFFSET, vocab_size). # Special tokens occupy IDs [0, SPECIAL_TOKEN_OFFSET). # ───────────────────────────────────────────────────────────────────────── def _shift_crayon_ids(self, ids: List[int]) -> List[int]: """Shift CRAYON token IDs to make room for special tokens at the front.""" return [i + self.special_token_offset for i in ids] def _unshift_crayon_ids(self, ids: List[int]) -> List[int]: """Inverse of _shift_crayon_ids for decode path.""" out = [] for i in ids: if i < self.special_token_offset: out.append(None) # special token — handled separately else: out.append(i - self.special_token_offset) return out # ───────────────────────────────────────────────────────────────────────── # CORE ENCODE / DECODE # ───────────────────────────────────────────────────────────────────────── def encode( self, text: str, add_bos: Optional[bool] = None, add_eos: Optional[bool] = None, truncate: bool = True, ) -> List[int]: """ Encode a single string to token IDs. Thread-safe via CRAYON's internal lock or thread-local tokenizers. """ add_bos_ = self.add_bos if add_bos is None else add_bos add_eos_ = self.add_eos if add_eos is None else add_eos # CRAYON tokenization with self._lock: raw_ids = self._tokenizer.tokenize(text) # Handle both list and tensor returns from CRAYON if isinstance(raw_ids, torch.Tensor): raw_ids = raw_ids.tolist() ids = self._shift_crayon_ids(raw_ids) # Prepend/append special tokens if add_bos_: ids = [self.special_tokens["<|bos|>"]] + ids if add_eos_: ids = ids + [self.special_tokens["<|eos|>"]] # Truncate if needed if truncate and self.max_length is not None: ids = ids[:self.max_length] return ids def encode_batch( self, texts: List[str], add_bos: Optional[bool] = None, add_eos: Optional[bool] = None, pad: bool = True, pad_to_multiple_of: int = 8, return_tensors: str = "pt", ) -> Dict[str, torch.Tensor]: """ Encode a batch of texts in parallel using thread pool. Leverages CRAYON's 24M tok/s CPU throughput across all threads. Returns dict with 'input_ids' and 'attention_mask'. """ add_bos_ = self.add_bos if add_bos is None else add_bos add_eos_ = self.add_eos if add_eos is None else add_eos # ── Parallel tokenization ────────────────────────────────────────── def _tokenize_one(text: str) -> List[int]: tok = _get_thread_local_tokenizer(self.profile) raw = tok.tokenize(text) if isinstance(raw, torch.Tensor): raw = raw.tolist() ids = self._shift_crayon_ids(raw) if add_bos_: ids = [self.special_tokens["<|bos|>"]] + ids if add_eos_: ids = ids + [self.special_tokens["<|eos|>"]] if self.max_length is not None: ids = ids[:self.max_length] return ids if self.n_workers <= 1: # Run sequentially on the main thread using the primary tokenizer instance # to completely avoid any C++ multi-threading deadlocks. all_ids = [] for t in texts: raw = self._tokenizer.tokenize(t) if isinstance(raw, torch.Tensor): raw = raw.tolist() ids = self._shift_crayon_ids(raw) if add_bos_: ids = [self.special_tokens["<|bos|>"]] + ids if add_eos_: ids = ids + [self.special_tokens["<|eos|>"]] if self.max_length is not None: ids = ids[:self.max_length] all_ids.append(ids) else: # Submit all texts to thread pool futures = [self._thread_pool.submit(_tokenize_one, t) for t in texts] all_ids = [f.result() for f in futures] if not pad or return_tensors is None: return {"input_ids": all_ids} # ── Padding ──────────────────────────────────────────────────────── max_len = max(len(ids) for ids in all_ids) # Pad to multiple for efficiency if pad_to_multiple_of > 1: max_len = math.ceil(max_len / pad_to_multiple_of) * pad_to_multiple_of pad_id = self.special_tokens["<|pad|>"] padded_ids = [] attention_masks = [] for ids in all_ids: n_pad = max_len - len(ids) padded_ids.append(ids + [pad_id] * n_pad) attention_masks.append([1] * len(ids) + [0] * n_pad) if return_tensors == "pt": return { "input_ids": torch.tensor(padded_ids, dtype=torch.long), "attention_mask": torch.tensor(attention_masks, dtype=torch.long), } return {"input_ids": padded_ids, "attention_mask": attention_masks} def decode( self, ids: Union[List[int], torch.Tensor], skip_special_tokens: bool = True, ) -> str: """Decode token IDs back to text.""" if isinstance(ids, torch.Tensor): ids = ids.tolist() # Separate special tokens from CRAYON tokens crayon_ids = [] special_strs = [] for tok_id in ids: if tok_id < self.special_token_offset: if not skip_special_tokens and tok_id in self.id_to_special: special_strs.append(self.id_to_special[tok_id]) else: crayon_ids.append(tok_id - self.special_token_offset) # CRAYON decode if crayon_ids: with self._lock: text = self._tokenizer.decode(crayon_ids) else: text = "" # Re-insert special tokens if not skipping if special_strs and not skip_special_tokens: text = " ".join(special_strs) + " " + text return text def decode_batch( self, batch_ids: Union[List[List[int]], torch.Tensor], skip_special_tokens: bool = True, ) -> List[str]: """Decode a batch in parallel.""" if isinstance(batch_ids, torch.Tensor): batch_ids = batch_ids.tolist() if self.n_workers <= 1: return [self.decode(ids, skip_special_tokens) for ids in batch_ids] futures = [ self._thread_pool.submit(self.decode, ids, skip_special_tokens) for ids in batch_ids ] return [f.result() for f in futures] # ───────────────────────────────────────────────────────────────────────── # CHAT TEMPLATE # ───────────────────────────────────────────────────────────────────────── def encode_chat( self, messages: List[Dict[str, str]], add_generation_prompt: bool = False, return_as_tensor: bool = True, ) -> Union[List[int], torch.Tensor]: """ Encode a conversation using FORGE chat template. Template: <|bos|><|sys|>{system}<|/sys|> <|usr|>{user}<|/usr|> <|asst|>{assistant}<|/asst|> ... Args: messages: List of {"role": "system|user|assistant", "content": str} add_generation_prompt: If True, append <|asst|> at end (inference) """ token_ids = [self.special_tokens["<|bos|>"]] role_to_open = { "system": "<|sys|>", "user": "<|usr|>", "assistant": "<|asst|>", } role_to_close = { "system": "<|/sys|>", "user": "<|/usr|>", "assistant": "<|/asst|>", } for msg in messages: role = msg["role"] content = msg["content"] # Open tag open_tag = role_to_open.get(role, "<|usr|>") close_tag = role_to_close.get(role, "<|/usr|>") token_ids.append(self.special_tokens.get(open_tag, self.special_tokens["<|unk|>"])) # Encode content (no bos/eos, no truncation) content_ids = self.encode(content, add_bos=False, add_eos=False, truncate=False) token_ids.extend(content_ids) token_ids.append(self.special_tokens.get(close_tag, self.special_tokens["<|unk|>"])) if add_generation_prompt: token_ids.append(self.special_tokens["<|asst|>"]) if self.max_length is not None: token_ids = token_ids[:self.max_length] if return_as_tensor: return torch.tensor(token_ids, dtype=torch.long) return token_ids def encode_chat_with_labels( self, messages: List[Dict[str, str]], loss_on_prompt: bool = False, ) -> Dict[str, torch.Tensor]: """ Encode a conversation and produce labels for SFT. Labels are -100 (ignored) for non-assistant turns unless loss_on_prompt=True. Returns dict with 'input_ids' and 'labels'. """ input_ids = [] labels = [] input_ids.append(self.special_tokens["<|bos|>"]) labels.append(-100 if not loss_on_prompt else self.special_tokens["<|bos|>"]) role_to_open = { "system": "<|sys|>", "user": "<|usr|>", "assistant": "<|asst|>", } role_to_close = { "system": "<|/sys|>", "user": "<|/usr|>", "assistant": "<|/asst|>", } for msg in messages: role = msg["role"] content = msg["content"] is_assistant = (role == "assistant") open_tag_id = self.special_tokens.get(role_to_open.get(role, "<|usr|>"), self.special_tokens["<|unk|>"]) close_tag_id = self.special_tokens.get(role_to_close.get(role, "<|/usr|>"), self.special_tokens["<|unk|>"]) content_ids = self.encode(content, add_bos=False, add_eos=False, truncate=False) turn_ids = [open_tag_id] + content_ids + [close_tag_id] input_ids.extend(turn_ids) if is_assistant or loss_on_prompt: # Compute loss on assistant content + close tag, NOT on open tag turn_labels = [-100] + content_ids + [close_tag_id] labels.extend(turn_labels) else: labels.extend([-100] * len(turn_ids)) if self.max_length is not None: input_ids = input_ids[:self.max_length] labels = labels[:self.max_length] return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), } # ───────────────────────────────────────────────────────────────────────── # HIGH-THROUGHPUT OFFLINE PREPROCESSING # ───────────────────────────────────────────────────────────────────────── def tokenize_file_streaming( self, input_path: str, output_path: str, chunk_size: int = 100_000, # documents per chunk seq_len: int = 2048, append_eos: bool = True, n_processes: int = None, ) -> int: """ High-throughput offline tokenization of a large text file. Architecture: - Producer process reads text lines - N worker processes each run their own CrayonVocab instance - Consumer aggregates and packs sequences Throughput: ~5-10GB/min depending on hardware. Returns: total number of tokens written. """ n_processes = n_processes or max(1, mp.cpu_count() - 2) total_tokens = 0 logger.info(f"Offline tokenization: {input_path} → {output_path}, " f"n_processes={n_processes}, seq_len={seq_len}") # We use process-based parallelism to bypass GIL fully with ProcessPoolExecutor(max_workers=n_processes) as pool: with open(input_path, "r", encoding="utf-8", errors="replace") as fin: chunk_buffer = [] futures = [] for line in fin: chunk_buffer.append(line.strip()) if len(chunk_buffer) >= chunk_size: futures.append( pool.submit( _tokenize_chunk_worker, chunk_buffer.copy(), self.profile, self.special_token_offset, SPECIAL_TOKENS["<|eos|>"] if append_eos else None, ) ) chunk_buffer.clear() # Submit remaining if chunk_buffer: futures.append( pool.submit( _tokenize_chunk_worker, chunk_buffer, self.profile, self.special_token_offset, SPECIAL_TOKENS["<|eos|>"] if append_eos else None, ) ) # Collect and pack token_stream = [] packed_sequences = [] for future in as_completed(futures): chunk_tokens = future.result() # List[int] token_stream.extend(chunk_tokens) # Pack into fixed-length sequences while len(token_stream) >= seq_len: packed_sequences.append(token_stream[:seq_len]) token_stream = token_stream[seq_len:] total_tokens += seq_len # Flush to disk in chunks if len(packed_sequences) >= 10_000: _flush_to_disk(packed_sequences, output_path, total_tokens) packed_sequences.clear() # Flush remaining if packed_sequences: _flush_to_disk(packed_sequences, output_path, total_tokens) logger.info(f"Tokenization complete: {total_tokens:,} tokens written to {output_path}") return total_tokens def __repr__(self) -> str: return (f"ForgeTokenizer(profile={self.profile!r}, " f"vocab_size={self.vocab_size:,}, " f"n_workers={self.n_workers})") def save_pretrained(self, path: str): """Save tokenizer config for later loading.""" import json os.makedirs(path, exist_ok=True) config = { "profile": self.profile, "device": self.device_str, "vocab_size": self.vocab_size, "special_tokens": self.special_tokens, "special_token_offset": self.special_token_offset, "add_bos": self.add_bos, "add_eos": self.add_eos, "max_length": self.max_length, } with open(os.path.join(path, "tokenizer_config.json"), "w") as f: json.dump(config, f, indent=2) logger.info(f"Tokenizer config saved to {path}") @classmethod def from_pretrained(cls, path: str) -> "ForgeTokenizer": import json with open(os.path.join(path, "tokenizer_config.json")) as f: config = json.load(f) return cls( profile=config["profile"], device=config.get("device", "cpu"), max_length=config.get("max_length"), add_bos=config.get("add_bos", True), add_eos=config.get("add_eos", True), ) def __del__(self): """Graceful cleanup of thread pool.""" try: self._thread_pool.shutdown(wait=False) except Exception: pass # ───────────────────────────────────────────────────────────────────────────── # WORKER FUNCTIONS (must be module-level for pickling with ProcessPoolExecutor) # ───────────────────────────────────────────────────────────────────────────── def _tokenize_chunk_worker( texts: List[str], profile: str, offset: int, eos_id: Optional[int], ) -> List[int]: """ Worker function for parallel tokenization. Runs in a separate process with its own CrayonVocab instance. """ tok = CrayonVocab(device="cpu") tok.load_profile(profile) all_ids = [] for text in texts: if not text: continue raw = tok.tokenize(text) if isinstance(raw, torch.Tensor): raw = raw.tolist() ids = [i + offset for i in raw] if eos_id is not None: ids.append(eos_id) all_ids.extend(ids) return all_ids def _flush_to_disk(sequences: List[List[int]], output_path: str, offset: int): """Append packed sequences to a memory-mapped numpy file.""" arr = np.array(sequences, dtype=np.uint32) mode = "ab" # append binary with open(output_path + ".npy", mode) as f: np.save(f, arr) # Also write offset index for random access with open(output_path + ".idx", "ab") as f: np.save(f, np.array([offset], dtype=np.int64))