tokenization_stellarai.py

# coding=utf-8
"""
StellarAI Tokenizer - Hugging Face compatible
Wraps SimpleTokenizer to match PreTrainedTokenizer interface.
"""
import os
import json
from typing import List, Optional, Dict, Tuple, Union, Any

from transformers import PreTrainedTokenizer
from transformers.tokenization_utils import AddedToken

# Import SimpleTokenizer from sibling stellarai package
import sys
_CURDIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_CURDIR, ".."))
from stellarai.tokenizer import SimpleTokenizer as _StellarTokenizer
sys.path.pop(0)

VOCAB_FILES_NAMES = {
    "vocab_file": "backend_tokenizer.json",
}

PRETRAINED_VOCAB_FILES_MAP = {}
PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
    "stellarai-tiny": 1024,
}
PRETRAINED_INIT_CONFIGURATION = {}


def _find_backend_vocab(search_path: Optional[str]) -> Optional[str]:
    """Locate the SimpleTokenizer format JSON file (backend_tokenizer.json)."""
    candidates = []
    if search_path is not None:
        if os.path.isfile(search_path):
            return search_path
        if os.path.isdir(search_path):
            candidates.append(os.path.join(search_path, "backend_tokenizer.json"))
            candidates.append(os.path.join(search_path, "tokenizer.json"))
    # default: same directory as this file
    candidates.append(os.path.join(_CURDIR, "backend_tokenizer.json"))
    candidates.append(os.path.join(_CURDIR, "..", "outputs", "stellar_pt", "tokenizer.json"))
    for c in candidates:
        if c and os.path.isfile(c):
            # Make sure it's the SimpleTokenizer format (has "token_to_id")
            try:
                with open(c, "r", encoding="utf-8") as f:
                    head = f.read(256)
                    if '"token_to_id"' in head or "'token_to_id'" in head:
                        return c
            except Exception:
                continue
    return None


class StellarAITokenizer(PreTrainedTokenizer):
    """
    Hugging Face compatible tokenizer for StellarAI.
    Wraps the original SimpleTokenizer for 100% training-consistent tokenization.
    """
    vocab_files_names = VOCAB_FILES_NAMES
    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file=None,
        unk_token="[UNK]",
        bos_token="[BOS]",
        eos_token="[EOS]",
        sep_token="[SEP]",
        pad_token="[PAD]",
        cls_token="[CLS]",
        mask_token="[MASK]",
        additional_special_tokens=None,
        model_max_length=1024,
        do_lower_case=False,
        **kwargs,
    ):
        if additional_special_tokens is None:
            additional_special_tokens = ["[IMG]", "[BOI]", "[EOI]"]
        # Wrap mask_token to AddedToken for HF compatibility
        mask_token = AddedToken(mask_token, lstrip=False, rstrip=False) if isinstance(mask_token, str) else mask_token

        # === IMPORTANT: create backend BEFORE super().__init__() ===
        resolved = _find_backend_vocab(vocab_file)
        self._tok = _StellarTokenizer(vocab_size=32000)
        if resolved is not None:
            try:
                self._tok.load(resolved)
            except Exception:
                # Fall back to base charset without pre-trained merges
                pass
        self.do_lower_case = do_lower_case

        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            sep_token=sep_token,
            pad_token=pad_token,
            cls_token=cls_token,
            mask_token=mask_token,
            additional_special_tokens=additional_special_tokens,
            model_max_length=model_max_length,
            do_lower_case=do_lower_case,
            **kwargs,
        )

    @property
    def vocab_size(self) -> int:
        return len(self._tok.token_to_id)

    def get_vocab(self) -> Dict[str, int]:
        return dict(self._tok.token_to_id)

    def _tokenize(self, text: str, **kwargs) -> List[str]:
        """Tokenize a string into BPE token strings (used by encode/decode pipeline)."""
        if self.do_lower_case:
            text = text.lower()
        ids = self._tok.encode(text, add_bos=False, add_eos=False)
        return [self._tok.id_to_token.get(i, self.unk_token) for i in ids]

    def _convert_token_to_id(self, token: str) -> int:
        return self._tok.token_to_id.get(token, self._tok.token_to_id.get(self.unk_token, 3))

    def _convert_id_to_token(self, index: int) -> str:
        return self._tok.id_to_token.get(index, self.unk_token)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        ids = [self._convert_token_to_id(t) for t in tokens]
        return self._tok.decode(ids, skip_special=False)

    # --- direct encode / decode overrides ---
    def _encode_plus(
        self,
        text,
        text_pair=None,
        add_special_tokens=True,
        padding_strategy="do_not_pad",
        truncation_strategy="longest_first",
        max_length=None,
        stride=0,
        is_split_into_words=False,
        pad_to_multiple_of=None,
        return_tensors=None,
        return_token_type_ids=None,
        return_attention_mask=None,
        return_overflowing_tokens=False,
        return_special_tokens_mask=False,
        return_offsets_mapping=False,
        return_length=False,
        verbose=True,
        **kwargs,
    ):
        if is_split_into_words:
            text = "".join(text) if isinstance(text, list) else text
        if self.do_lower_case:
            text = text.lower()
            if text_pair is not None:
                text_pair = text_pair.lower() if not isinstance(text_pair, list) else "".join(text_pair)

        ids = list(self._tok.encode(text, add_bos=False, add_eos=False))
        if text_pair is not None:
            pair_ids = list(self._tok.encode(text_pair, add_bos=False, add_eos=False))
        else:
            pair_ids = None

        if add_special_tokens:
            bos_id = self._tok.SPECIAL_TOKENS.get("[BOS]", 1)
            eos_id = self._tok.SPECIAL_TOKENS.get("[EOS]", 2)
            sep_id = self._tok.SPECIAL_TOKENS.get("[SEP]", 6)
            if pair_ids is None:
                ids = [bos_id] + ids + [eos_id]
            else:
                ids = [bos_id] + ids + [sep_id] + pair_ids + [eos_id]

        # Truncation
        if max_length is not None and len(ids) > max_length:
            if truncation_strategy == "longest_first":
                ids = ids[:max_length]

        input_ids = ids
        attention_mask = [1] * len(ids)

        # Padding
        if padding_strategy != "do_not_pad" and max_length is not None and len(ids) < max_length:
            pad_id = self._tok.SPECIAL_TOKENS.get("[PAD]", 0)
            pad_len = max_length - len(ids)
            input_ids = input_ids + [pad_id] * pad_len
            attention_mask = attention_mask + [0] * pad_len

        encoding = {"input_ids": input_ids, "attention_mask": attention_mask}
        if return_token_type_ids:
            tti = [0] * len(input_ids)
            if pair_ids is not None and add_special_tokens:
                sep_id = self._tok.SPECIAL_TOKENS.get("[SEP]", 6)
                sep_idx = None
                for i, _id in enumerate(input_ids):
                    if _id == sep_id and sep_idx is None:
                        sep_idx = i
                if sep_idx is not None:
                    for j in range(sep_idx + 1, len(tti)):
                        tti[j] = 1
            encoding["token_type_ids"] = tti
        if return_length:
            encoding["length"] = len(input_ids)

        if return_tensors is not None:
            import torch
            for k, v in list(encoding.items()):
                if isinstance(v, list) and all(isinstance(x, int) for x in v):
                    encoding[k] = torch.tensor([v], dtype=torch.long)
                elif isinstance(v, int):
                    encoding[k] = torch.tensor([v], dtype=torch.long)

        return encoding

    def decode(
        self,
        token_ids: Union[int, List[int], Any],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = None,
        **kwargs,
    ) -> str:
        if hasattr(token_ids, "tolist"):
            token_ids = token_ids.tolist()
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        if (
            isinstance(token_ids, list)
            and len(token_ids) == 1
            and isinstance(token_ids[0], list)
        ):
            token_ids = token_ids[0]
        if not isinstance(token_ids, list):
            token_ids = list(token_ids)
        int_ids = [int(x) for x in token_ids]
        return self._tok.decode(int_ids, skip_special=skip_special_tokens)

    def batch_decode(self, sequences, **kwargs):
        return [self.decode(seq, **kwargs) for seq in sequences]

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str, ...]:
        if not os.path.isdir(save_directory):
            raise ValueError(f"Vocabulary path ({save_directory}) should be a directory")
        fname = (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
        out_path = os.path.join(save_directory, fname)
        self._tok.save(out_path)
        return (out_path,)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
        # Ensure vocab_file points to a resolved backend tokenizer JSON (SimpleTokenizer format)
        if "vocab_file" not in kwargs or kwargs["vocab_file"] is None:
            search = pretrained_model_name_or_path
            if isinstance(search, str) and os.path.isdir(search):
                candidate = os.path.join(search, "backend_tokenizer.json")
                if os.path.isfile(candidate):
                    kwargs["vocab_file"] = candidate
                else:
                    # fallback to outputs dir for local development
                    alt = os.path.join(_CURDIR, "backend_tokenizer.json")
                    if os.path.isfile(alt):
                        kwargs["vocab_file"] = alt
        return super().from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)