Blame for src/tokenizer.py

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

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

    @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("<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

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