src/taoTrain/data/tokenizer.py

"""SentencePiece tokenizer wrapper and special-token metadata helpers."""

import json
from pathlib import Path
from typing import Optional


BUILTIN_SPECIAL_TOKENS = ("<UNK>", "<BOS>", "<EOS>", "<PAD>")


def get_special_token_metadata_path(tokenizer_path: str | Path) -> Path:
    """Return the sidecar metadata path for a tokenizer artifact."""
    tokenizer_path = Path(tokenizer_path)
    return tokenizer_path.with_suffix(".special_tokens.json")


def load_special_token_metadata(tokenizer_path: str | Path) -> dict[str, int]:
    """Load resolved special-token IDs saved alongside a tokenizer."""
    metadata_path = get_special_token_metadata_path(tokenizer_path)
    if not metadata_path.exists():
        return {}

    with metadata_path.open("r", encoding="utf-8") as handle:
        data = json.load(handle)

    special_tokens = data.get("special_tokens", {})
    return {str(token): int(token_id) for token, token_id in special_tokens.items()}


def load_special_token_metadata_payload(tokenizer_path: str | Path) -> dict[str, object]:
    """Load the raw special-token metadata payload, if available."""
    metadata_path = get_special_token_metadata_path(tokenizer_path)
    if not metadata_path.exists():
        return {}

    with metadata_path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def load_tokenizer_runtime_info(
    tokenizer_path: str,
    tokenizer_type: Optional[str] = None,
) -> dict[str, int]:
    """Load basic runtime tokenizer information needed for model/data validation."""
    if tokenizer_type is None:
        tokenizer_type = "sentencepiece" if tokenizer_path.endswith(".model") else "huggingface"

    if tokenizer_type == "sentencepiece":
        try:
            import sentencepiece as spm
        except ImportError:
            raise ImportError("SentencePiece not installed. Install with: pip install sentencepiece")

        processor = spm.SentencePieceProcessor()
        processor.Load(tokenizer_path)
        return {
            "vocab_size": int(processor.vocab_size()),
            "unk_id": int(processor.unk_id()),
            "bos_id": int(processor.bos_id()),
            "eos_id": int(processor.eos_id()),
            "pad_id": int(processor.pad_id()),
        }

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    vocab_size = getattr(tokenizer, "vocab_size", None)
    if vocab_size is None:
        vocab_size = len(tokenizer)
    return {
        "vocab_size": int(vocab_size),
        "unk_id": int(getattr(tokenizer, "unk_token_id", -1)),
        "bos_id": int(getattr(tokenizer, "bos_token_id", -1)),
        "eos_id": int(getattr(tokenizer, "eos_token_id", -1)),
        "pad_id": int(getattr(tokenizer, "pad_token_id", -1)),
    }


def resolve_special_token_id(tokenizer, token: str) -> Optional[int]:
    """Compatibility helper that delegates to tokenizer-native special-token lookup."""
    if hasattr(tokenizer, "get_special_token_id"):
        token_id = tokenizer.get_special_token_id(token)
        if token_id is not None:
            return int(token_id)

    if hasattr(tokenizer, "convert_tokens_to_ids"):
        token_id = tokenizer.convert_tokens_to_ids(token)
        unk_token_id = getattr(tokenizer, "unk_token_id", None)
        if token_id is not None and token_id != unk_token_id:
            return int(token_id)

    return None


def require_special_token_id(tokenizer, token: str) -> int:
    """Resolve a special token ID or raise a clear error."""
    if hasattr(tokenizer, "require_special_token_id"):
        return int(tokenizer.require_special_token_id(token))

    token_id = resolve_special_token_id(tokenizer, token)
    if token_id is None:
        raise ValueError(
            f"Special token `{token}` is not registered in the tokenizer. "
            "Retrain or regenerate the tokenizer with that token registered as a special token."
        )
    return int(token_id)


class SentencePieceTokenizerWrapper:
    """Wrapper to make SentencePiece tokenizer compatible with a HuggingFace-like interface."""

    def __init__(self, sp_processor, special_token_ids: Optional[dict[str, int]] = None):
        """Initialize wrapper."""
        self.sp = sp_processor
        self.vocab_size = self.sp.vocab_size()
        self.special_token_ids = self._build_special_token_ids(special_token_ids or {})
        self.pad_token_id = self.special_token_ids["<PAD>"]
        self.eos_token_id = self.special_token_ids["<EOS>"]
        self.bos_token_id = self.special_token_ids["<BOS>"]
        self.unk_token_id = self.special_token_ids["<UNK>"]

    def _build_special_token_ids(self, special_token_ids: dict[str, int]) -> dict[str, int]:
        """Merge processor-native built-ins with any saved custom special-token IDs."""
        resolved_ids = {
            "<UNK>": int(self.sp.unk_id()),
            "<BOS>": int(self.sp.bos_id()),
            "<EOS>": int(self.sp.eos_id()),
            "<PAD>": int(self.sp.pad_id()),
        }
        for token, token_id in special_token_ids.items():
            resolved_ids[str(token)] = int(token_id)

        invalid_builtins = [token for token in BUILTIN_SPECIAL_TOKENS if resolved_ids.get(token, -1) < 0]
        if invalid_builtins:
            raise ValueError(
                "Tokenizer is missing required built-in special token IDs: "
                f"{invalid_builtins}. Retrain or regenerate the tokenizer with built-in PAD/BOS/EOS/UNK tokens enabled."
            )
        return resolved_ids

    def get_special_token_id(self, token: str) -> Optional[int]:
        """Resolve a special token ID from the wrapper registry or exact SentencePiece pieces."""
        token_id = self.special_token_ids.get(token)
        if token_id is not None:
            return int(token_id)

        if hasattr(self.sp, "piece_to_id") and hasattr(self.sp, "id_to_piece"):
            piece_id = self.sp.piece_to_id(token)
            if piece_id is not None and piece_id >= 0 and self.sp.id_to_piece(piece_id) == token:
                return int(piece_id)

        return None

    def require_special_token_id(self, token: str) -> int:
        """Require that a special token resolve to a single known token ID."""
        token_id = self.get_special_token_id(token)
        if token_id is None:
            raise ValueError(
                f"Special token `{token}` is not registered in the tokenizer wrapper. "
                "Retrain or regenerate the tokenizer with that token registered as a special token."
            )
        return int(token_id)

    def __call__(self, text, **kwargs):
        """Tokenize text."""
        is_single = isinstance(text, str)
        texts = [text] if is_single else text

        max_length = kwargs.get("max_length")
        padding = kwargs.get("padding")
        truncation = kwargs.get("truncation", False)
        return_attention_mask = kwargs.get("return_attention_mask", True)

        all_input_ids = []
        for t in texts:
            if isinstance(t, str):
                special_token_id = self.get_special_token_id(t)
                if special_token_id is not None:
                    tokens = [special_token_id]
                else:
                    tokens = self.sp.encode(t, out_type=int)
            else:
                tokens = self.sp.encode(t, out_type=int)

            if truncation and max_length and len(tokens) > max_length:
                tokens = tokens[:max_length]

            all_input_ids.append(tokens)

        if padding or max_length:
            target_length = max_length or (max(len(ids) for ids in all_input_ids) if all_input_ids else 0)
            padded_input_ids = []
            padded_attention_masks = []

            for ids in all_input_ids:
                pad_length = target_length - len(ids)
                if pad_length > 0:
                    padded_ids = ids + [self.pad_token_id] * pad_length
                else:
                    padded_ids = ids[:target_length]

                padded_input_ids.append(padded_ids)
                attention_mask = [1] * min(len(ids), target_length) + [0] * max(0, target_length - len(ids))
                padded_attention_masks.append(attention_mask)

            result = {
                "input_ids": padded_input_ids if not is_single else padded_input_ids[0],
            }
            if return_attention_mask:
                result["attention_mask"] = (
                    padded_attention_masks if not is_single else padded_attention_masks[0]
                )
            return result

        result = {
            "input_ids": all_input_ids[0] if is_single else all_input_ids,
        }
        if return_attention_mask:
            attention_masks = [[1] * len(ids) for ids in all_input_ids]
            result["attention_mask"] = attention_masks[0] if is_single else attention_masks
        return result

    def encode(self, text, return_tensors=None, **kwargs):
        """Encode text to token IDs."""
        result = self(text, **kwargs)
        input_ids = result["input_ids"]

        if return_tensors == "pt":
            import torch

            if isinstance(input_ids[0], list):
                input_ids = input_ids[0]
            return torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)

        return input_ids

    def encode_plus(self, text, **kwargs):
        """Encode text with additional information."""
        return self(text, **kwargs)

    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
        """Decode token IDs to text."""
        if hasattr(token_ids, "tolist"):
            token_ids = token_ids.tolist()

        if isinstance(token_ids, (list, tuple)) and token_ids and isinstance(token_ids[0], (list, tuple)):
            token_ids = token_ids[0]

        if not isinstance(token_ids, list):
            token_ids = [int(t) for t in token_ids]

        if skip_special_tokens:
            special_token_ids = {int(token_id) for token_id in self.special_token_ids.values()}
            token_ids = [int(token_id) for token_id in token_ids if int(token_id) not in special_token_ids]

        return self.sp.decode(token_ids)