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