Spaces:
Sleeping
Sleeping
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| Canonical Tokenizer for Q-TensorFormer (`src/tokenizer.py`). | |
| Provides: | |
| - Deterministic hybrid subword + character vocabulary construction | |
| - Explicit special token handling: <pad>=0, <bos>=1, <eos>=2, <unk>=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 = "<pad>" | |
| BOS_TOKEN = "<bos>" | |
| EOS_TOKEN = "<eos>" | |
| UNK_TOKEN = "<unk>" | |
| 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"<extra_id_{idx}>" | |
| 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("<extra_id_") and tok.endswith(">")): | |
| 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)) | |
| 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"<extra_id_{idx}>" | |
| 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) | |
| 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("<extra_id_") and tok.endswith(">"): | |
| 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 | |
| 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), | |
| } | |
| 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) | |