""" Canonical Tokenizer for Q-TensorFormer (`src/tokenizer.py`). Provides: - Deterministic hybrid subword + character vocabulary construction - Explicit special token handling: =0, =1, =2, =3 - Lossless text -> token IDs -> text round-trip for printable ASCII - Batch encoding with padding, truncation, and attention_mask generation - Versioned JSON artifact persistence (`tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json`) - Strict vocabulary-size alignment check against `ModelConfig` / `QTensorFormer` """ import hashlib import json import re from collections import Counter from pathlib import Path from typing import Dict, List, Optional, Sequence, Union import torch class QTensorFormerTokenizer: """ Deterministic, self-contained tokenizer for Q-TensorFormer language modeling. """ FORMAT_VERSION = "1.0.0" PAD_TOKEN = "" BOS_TOKEN = "" EOS_TOKEN = "" UNK_TOKEN = "" PAD_TOKEN_ID = 0 BOS_TOKEN_ID = 1 EOS_TOKEN_ID = 2 UNK_TOKEN_ID = 3 def __init__( self, vocab_size: int = 256, token_to_id: Optional[Dict[str, int]] = None, model_max_length: int = 128, ): if vocab_size < 104: raise ValueError( f"vocab_size must be >= 104 to accommodate special tokens and printable ASCII, got {vocab_size}." ) self.model_max_length = int(model_max_length) self.pad_token = self.PAD_TOKEN self.bos_token = self.BOS_TOKEN self.eos_token = self.EOS_TOKEN self.unk_token = self.UNK_TOKEN self.pad_token_id = self.PAD_TOKEN_ID self.bos_token_id = self.BOS_TOKEN_ID self.eos_token_id = self.EOS_TOKEN_ID self.unk_token_id = self.UNK_TOKEN_ID if token_to_id is not None: self.token_to_id = dict(token_to_id) self.id_to_token = {int(v): k for k, v in self.token_to_id.items()} self._vocab_size = len(self.token_to_id) else: self.token_to_id = {} self.id_to_token = {} self._vocab_size = int(vocab_size) self._init_default_vocab(self._vocab_size) self._rebuild_lookup_structures() def _base_tokens(self) -> List[str]: specials = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN] # Deterministic printable ASCII + standard whitespace ascii_chars = ["\t", "\n", "\r"] + [chr(i) for i in range(32, 127)] return specials + ascii_chars def _init_default_vocab(self, target_vocab_size: int) -> None: base = self._base_tokens() # Deterministic common English subwords/bigrams/words to fill up to target_vocab_size common_subwords = [ " the", " and", " of", " to", " in", " is", " for", " that", " with", " on", " as", " are", " by", " be", " at", " or", " from", " an", " this", " which", " model", " tensor", " quantum", " attention", " rank", " cache", " token", " layer", " network", " memory", " latency", " energy", " compute", " state", " value", " adaptive", " inference", " training", " matrix", " vector", "ing", "tion", "ed", "er", "es", "al", "ly", "en", "an", "in", "re", "on", "st", "nd", "or", "te", "ti", "ar", "to", "nt", "is", "of", "it", "ha", "se", "et", "ou", "le", "sa", "ve", "he", "th", "in", "er", "an", "re", " Q-TensorFormer", " Transformer", " Tensor-Train", " KV-cache", " FFN", " core", " cores", " head", " heads", " block", " blocks", " budget", " allocation", " controller", " entropy", " perplexity", " loss", " gradient", " optimization", " compression", " deployment", " hardware", " edge", " precision", " quantization", " routing", " policy", " benchmark", " evaluation", " validation", " accuracy", " sequence", " context", " position", " embedding", " projection", " linear", " activation", " distribution", " probability", " uncertainty", " information", " dynamic", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", ] tokens: List[str] = [] seen = set() for tok in base + common_subwords: if tok not in seen and len(tokens) < target_vocab_size: seen.add(tok) tokens.append(tok) idx = 0 while len(tokens) < target_vocab_size: extra_tok = f"" if extra_tok not in seen: seen.add(extra_tok) tokens.append(extra_tok) idx += 1 self.token_to_id = {tok: i for i, tok in enumerate(tokens)} self.id_to_token = {i: tok for i, tok in enumerate(tokens)} self._vocab_size = len(self.token_to_id) def _rebuild_lookup_structures(self) -> None: self._special_set = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN} self._multi_char_tokens_by_first: Dict[str, List[str]] = {} for tok in self.token_to_id: if tok in self._special_set or (tok.startswith("")): continue if len(tok) > 1: first_ch = tok[0] self._multi_char_tokens_by_first.setdefault(first_ch, []).append(tok) for first_ch, tok_list in self._multi_char_tokens_by_first.items(): # Sort longest first for greedy longest-prefix match, breaking ties lexicographically tok_list.sort(key=lambda s: (-len(s), s)) @classmethod def from_corpus( cls, texts: Union[str, Sequence[str]], vocab_size: int = 256, model_max_length: int = 128, ) -> "QTensorFormerTokenizer": """ Deterministically construct a vocabulary of exact size `vocab_size` from `texts`. """ if isinstance(texts, str): corpus_list = [texts] else: corpus_list = list(texts) full_text = "\n".join(t for t in corpus_list if isinstance(t, str)) if not full_text.strip(): raise ValueError("Cannot build tokenizer vocabulary from an empty corpus.") instance = cls(vocab_size=vocab_size, model_max_length=model_max_length) base = instance._base_tokens() seen = set(base) tokens = list(base) # Extract frequent space-prefixed words and character n-grams deterministically word_Pieces = re.findall(r" ?[A-Za-z0-9_-]{2,16}", full_text) counts = Counter(word_Pieces) sorted_candidates = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0])) for cand, _ in sorted_candidates: if len(tokens) >= vocab_size: break if cand not in seen and all(ch in seen for ch in cand): seen.add(cand) tokens.append(cand) idx = 0 while len(tokens) < vocab_size: extra_tok = f"" if extra_tok not in seen: seen.add(extra_tok) tokens.append(extra_tok) idx += 1 token_to_id = {tok: i for i, tok in enumerate(tokens)} return cls(vocab_size=vocab_size, token_to_id=token_to_id, model_max_length=model_max_length) @property def vocab_size(self) -> int: return self._vocab_size def __len__(self) -> int: return self._vocab_size def validate_model_alignment(self, model_or_config) -> bool: """ Verify that the tokenizer vocabulary size matches the model config and output head. """ cfg_vocab = getattr(model_or_config, "vocab_size", None) if cfg_vocab is None and hasattr(model_or_config, "config"): cfg_vocab = getattr(model_or_config.config, "vocab_size", None) if cfg_vocab is not None and int(cfg_vocab) != self.vocab_size: raise ValueError( f"VocabularyMismatchError: Tokenizer vocab_size ({self.vocab_size}) != " f"Model config vocab_size ({cfg_vocab})." ) if hasattr(model_or_config, "lm_head") and hasattr(model_or_config.lm_head, "out_features"): head_out = int(model_or_config.lm_head.out_features) if head_out != self.vocab_size: raise ValueError( f"VocabularyMismatchError: Tokenizer vocab_size ({self.vocab_size}) != " f"Model lm_head.out_features ({head_out})." ) return True def encode( self, text: str, add_bos: bool = False, add_eos: bool = False, max_length: Optional[int] = None, truncation: bool = False, ) -> List[int]: """ Encode a single string into a list of integer token IDs. """ if not isinstance(text, str): raise TypeError(f"Tokenizer.encode expected a string, got {type(text).__name__}.") if len(text) == 0 and not add_bos and not add_eos: raise ValueError("Cannot encode an empty string when add_bos=False and add_eos=False.") ids: List[int] = [] if add_bos: ids.append(self.bos_token_id) i = 0 n = len(text) while i < n: ch = text[i] matched = False candidates = self._multi_char_tokens_by_first.get(ch) if candidates: for cand in candidates: if text.startswith(cand, i): ids.append(self.token_to_id[cand]) i += len(cand) matched = True break if not matched: ids.append(self.token_to_id.get(ch, self.unk_token_id)) i += 1 if add_eos: ids.append(self.eos_token_id) if truncation and max_length is not None: if max_length <= 0: raise ValueError(f"max_length must be > 0 when truncation=True, got {max_length}.") if len(ids) > max_length: if add_eos and max_length >= 2: ids = ids[: max_length - 1] + [self.eos_token_id] else: ids = ids[:max_length] return ids def decode( self, token_ids: Union[Sequence[int], torch.Tensor], skip_special_tokens: bool = True, ) -> str: """ Decode a 1D sequence of token IDs back into a string. """ if isinstance(token_ids, torch.Tensor): if token_ids.dim() > 1: raise ValueError( f"Tokenizer.decode expects a 1D sequence of token IDs, got tensor of shape {tuple(token_ids.shape)}. " "Use batch_decode for 2D batches." ) id_list = token_ids.detach().cpu().tolist() elif isinstance(token_ids, (list, tuple)): id_list = list(token_ids) else: raise TypeError(f"Unsupported token_ids type: {type(token_ids).__name__}") pieces: List[str] = [] for raw_id in id_list: tid = int(raw_id) if tid < 0 or tid >= self._vocab_size: raise ValueError( f"Out-of-vocabulary token ID {tid} encountered during decode (valid range: [0, {self._vocab_size - 1}])." ) tok = self.id_to_token.get(tid, self.UNK_TOKEN) if skip_special_tokens: if tok in (self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN): continue if tok == self.UNK_TOKEN: pieces.append("?") continue if tok.startswith(""): continue pieces.append(tok) return "".join(pieces) def batch_decode( self, sequences: Union[Sequence[Sequence[int]], torch.Tensor], skip_special_tokens: bool = True, ) -> List[str]: if isinstance(sequences, torch.Tensor): if sequences.dim() == 1: return [self.decode(sequences, skip_special_tokens=skip_special_tokens)] seq_list = sequences.detach().cpu().tolist() else: seq_list = list(sequences) return [self.decode(seq, skip_special_tokens=skip_special_tokens) for seq in seq_list] def batch_encode( self, texts: Sequence[str], add_bos: bool = True, add_eos: bool = False, padding: Union[bool, str] = True, truncation: bool = True, max_length: Optional[int] = None, padding_side: str = "right", return_tensors: Optional[str] = "pt", ) -> Dict[str, Union[torch.Tensor, List[List[int]]]]: """ Encode a batch of strings with optional padding, truncation, and attention_mask. """ if not isinstance(texts, (list, tuple)) or len(texts) == 0: raise ValueError("batch_encode requires a non-empty list or tuple of strings.") if padding_side not in ("right", "left"): raise ValueError(f"padding_side must be 'right' or 'left', got {padding_side!r}.") effective_max = max_length if max_length is not None else (self.model_max_length if truncation else None) encoded_batch: List[List[int]] = [] for idx, t in enumerate(texts): if not isinstance(t, str): raise TypeError(f"Batch item at index {idx} is not a string (got {type(t).__name__}).") if len(t) == 0 and not add_bos and not add_eos: raise ValueError(f"Batch item at index {idx} is an empty string with add_bos=False and add_eos=False.") ids = self.encode( t, add_bos=add_bos, add_eos=add_eos, max_length=effective_max, truncation=truncation, ) encoded_batch.append(ids) if padding is True or padding == "longest": target_len = max(len(seq) for seq in encoded_batch) elif padding == "max_length": if effective_max is None: raise ValueError("max_length must be specified when padding='max_length'.") target_len = effective_max elif padding is False or padding == "do_not_pad": target_len = None else: raise ValueError(f"Unsupported padding mode: {padding!r}") input_ids_list: List[List[int]] = [] attention_mask_list: List[List[int]] = [] for seq in encoded_batch: if target_len is not None and len(seq) < target_len: pad_len = target_len - len(seq) if padding_side == "right": padded_seq = seq + [self.pad_token_id] * pad_len mask_seq = [1] * len(seq) + [0] * pad_len else: padded_seq = [self.pad_token_id] * pad_len + seq mask_seq = [0] * pad_len + [1] * len(seq) else: padded_seq = list(seq) mask_seq = [1] * len(seq) input_ids_list.append(padded_seq) attention_mask_list.append(mask_seq) if return_tensors == "pt": if target_len is None and len({len(s) for s in input_ids_list}) > 1: raise ValueError("Cannot return PyTorch tensors for ragged sequences without padding=True.") return { "input_ids": torch.tensor(input_ids_list, dtype=torch.long), "attention_mask": torch.tensor(attention_mask_list, dtype=torch.long), } return { "input_ids": input_ids_list, "attention_mask": attention_mask_list, } def __call__( self, text: Union[str, Sequence[str]], add_bos: bool = True, add_eos: bool = False, padding: Union[bool, str] = True, truncation: bool = True, max_length: Optional[int] = None, return_tensors: Optional[str] = "pt", ) -> Dict[str, Union[torch.Tensor, List[List[int]]]]: if isinstance(text, str): batch = [text] else: batch = list(text) return self.batch_encode( batch, add_bos=add_bos, add_eos=add_eos, padding=padding, truncation=truncation, max_length=max_length, return_tensors=return_tensors, ) def fingerprint(self) -> str: """Compute deterministic SHA-256 fingerprint of vocabulary mapping.""" ordered = sorted(self.token_to_id.items(), key=lambda kv: kv[1]) payload = json.dumps(ordered, ensure_ascii=True, separators=(",", ":")) return hashlib.sha256(payload.encode("utf-8")).hexdigest() def to_dict(self) -> Dict: return { "format": "QTensorFormerTokenizer", "version": self.FORMAT_VERSION, "vocab_size": self._vocab_size, "model_max_length": self.model_max_length, "special_tokens": { "pad_token": self.PAD_TOKEN, "pad_token_id": self.PAD_TOKEN_ID, "bos_token": self.BOS_TOKEN, "bos_token_id": self.BOS_TOKEN_ID, "eos_token": self.EOS_TOKEN, "eos_token_id": self.EOS_TOKEN_ID, "unk_token": self.UNK_TOKEN, "unk_token_id": self.UNK_TOKEN_ID, }, "vocab_sha256": self.fingerprint(), "vocab": self.token_to_id, } def save(self, filepath: Union[str, Path]) -> Path: path = Path(filepath) path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(self.to_dict(), f, indent=2, ensure_ascii=True) return path @classmethod def load(cls, filepath: Union[str, Path]) -> "QTensorFormerTokenizer": path = Path(filepath) if not path.exists(): raise FileNotFoundError(f"Tokenizer artifact not found at: {path}") if path.is_dir(): return cls.from_pretrained(path) with open(path, "r", encoding="utf-8") as f: data = json.load(f) if "vocab" not in data or not isinstance(data["vocab"], dict): raise ValueError(f"Invalid tokenizer JSON format in {path}: missing 'vocab' dictionary.") vocab = {str(k): int(v) for k, v in data["vocab"].items()} tok = cls( vocab_size=int(data.get("vocab_size", len(vocab))), token_to_id=vocab, model_max_length=int(data.get("model_max_length", 128)), ) expected_hash = data.get("vocab_sha256") if expected_hash and tok.fingerprint() != expected_hash: raise ValueError( f"Tokenizer checksum mismatch in {path}: expected {expected_hash}, got {tok.fingerprint()}" ) return tok def save_pretrained(self, save_directory: Union[str, Path]) -> Dict[str, str]: save_dir = Path(save_directory) save_dir.mkdir(parents=True, exist_ok=True) tokenizer_json = save_dir / "tokenizer.json" tokenizer_config_json = save_dir / "tokenizer_config.json" special_tokens_json = save_dir / "special_tokens_map.json" self.save(tokenizer_json) with open(tokenizer_config_json, "w", encoding="utf-8") as f: json.dump( { "tokenizer_class": "QTensorFormerTokenizer", "version": self.FORMAT_VERSION, "vocab_size": self._vocab_size, "model_max_length": self.model_max_length, "pad_token": self.PAD_TOKEN, "bos_token": self.BOS_TOKEN, "eos_token": self.EOS_TOKEN, "unk_token": self.UNK_TOKEN, "vocab_sha256": self.fingerprint(), }, f, indent=2, ) with open(special_tokens_json, "w", encoding="utf-8") as f: json.dump( { "pad_token": self.PAD_TOKEN, "bos_token": self.BOS_TOKEN, "eos_token": self.EOS_TOKEN, "unk_token": self.UNK_TOKEN, }, f, indent=2, ) return { "tokenizer_file": str(tokenizer_json), "tokenizer_config_file": str(tokenizer_config_json), "special_tokens_map_file": str(special_tokens_json), } @classmethod def from_pretrained(cls, pretrained_path: Union[str, Path]) -> "QTensorFormerTokenizer": path = Path(pretrained_path) if path.is_file(): return cls.load(path) cand = path / "tokenizer.json" if not cand.exists(): raise FileNotFoundError( f"No 'tokenizer.json' found in directory '{path}'. " "Save a tokenizer via tokenizer.save_pretrained(dir) first." ) return cls.load(cand)