ViperEkura
fc47319240
- Extract _resolve_base_type and _validate_component as module-level helpers - Replace ForwardRef._evaluate private API with eval in module namespace - Remove broad except Exception in __init_subclass__, _component_base always set - Replace direct _entries mutation in strategy.py with register() call form - Remove dead TOKENIZER_CLASSES registry from AutoTokenizer - Extract ModelFactory(BaseFactory[nn.Module]) as pure factory - AutoModel now inherits only nn.Module, no factory state - Move @AutoModel.register to @ModelFactory.register in transformer.py and encoder.py
290 lines
9.9 KiB
Python
290 lines
9.9 KiB
Python
"""
|
|
Tokenizer module with implementation and auto-loading support.
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
from tokenizers import Tokenizer
|
|
|
|
from astrai.tokenize.chat_template import ChatTemplate
|
|
|
|
Message = Dict[str, str]
|
|
"""Single chat message with ``role`` and ``content`` keys."""
|
|
|
|
Messages = List[Message]
|
|
"""Single conversation — a list of messages."""
|
|
|
|
|
|
class AutoTokenizer:
|
|
"""Base tokenizer class with automatic loading support"""
|
|
|
|
def __init__(
|
|
self,
|
|
path: Optional[Union[str, Path]] = None,
|
|
special_token_map: Optional[Dict[str, str]] = None,
|
|
chat_template: Optional[str] = None,
|
|
):
|
|
self._tokenizer: Tokenizer = None
|
|
self._chat_template: Optional[ChatTemplate] = None
|
|
self._special_token_map: Optional[Dict] = special_token_map or {}
|
|
|
|
if chat_template:
|
|
self.set_chat_template(chat_template)
|
|
|
|
if path:
|
|
self.load(path)
|
|
|
|
def load(self, path: Union[str, Path]):
|
|
"""Load tokenizer from directory."""
|
|
path = Path(path)
|
|
tokenizer_file = path / "tokenizer.json"
|
|
config_file = path / "tokenizer_config.json"
|
|
self._tokenizer = Tokenizer.from_file(str(tokenizer_file))
|
|
|
|
if config_file.exists():
|
|
with open(config_file, "r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
|
|
if "special_tokens" in config:
|
|
self._special_token_map.update(config["special_tokens"])
|
|
|
|
# Load chat template from config
|
|
if "chat_template" in config:
|
|
self.set_chat_template(config["chat_template"])
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, path: Union[str, Path]) -> "AutoTokenizer":
|
|
"""Load tokenizer from pretrained directory.
|
|
|
|
Raises:
|
|
FileNotFoundError: If tokenizer.json is missing.
|
|
RuntimeError: If tokenizer failed to initialize.
|
|
"""
|
|
path = Path(path)
|
|
tokenizer_file = path / "tokenizer.json"
|
|
if not tokenizer_file.exists():
|
|
raise FileNotFoundError(
|
|
f"Tokenizer file not found: {tokenizer_file}. "
|
|
"A valid tokenizer.json is required."
|
|
)
|
|
instance = cls(path)
|
|
if instance._tokenizer is None:
|
|
raise RuntimeError(
|
|
f"Failed to load tokenizer from {path}. "
|
|
"The tokenizer.json may be corrupted or incompatible."
|
|
)
|
|
return instance
|
|
|
|
def save_pretrained(self, save_path: str):
|
|
"""
|
|
Save tokenizer to pretrained directory.
|
|
|
|
Args:
|
|
save_path: Path to save the tokenizer
|
|
"""
|
|
|
|
if self._tokenizer is None:
|
|
raise RuntimeError(
|
|
"Tokenizer not initialized. Load or create a tokenizer first."
|
|
)
|
|
|
|
save_path = Path(save_path)
|
|
save_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Save tokenizer
|
|
self._tokenizer.save(str(save_path / "tokenizer.json"))
|
|
|
|
# Save tokenizer config
|
|
config = {}
|
|
if self._special_token_map is not None:
|
|
config["special_tokens"] = self._special_token_map
|
|
if self._chat_template is not None:
|
|
config["chat_template"] = self._chat_template.template_str
|
|
|
|
with open(save_path / "tokenizer_config.json", "w", encoding="utf-8") as f:
|
|
json.dump(config, f, ensure_ascii=False, indent=2)
|
|
|
|
def encode(
|
|
self,
|
|
tokens: Union[str, List[str]],
|
|
out_ids: bool = True,
|
|
is_pretokenized: bool = False,
|
|
add_special_tokens: bool = True,
|
|
) -> List:
|
|
"""Encode text to token IDs.
|
|
|
|
Accepts both single strings and batches:
|
|
|
|
- ``encode("hello")`` → ``[123, 456]``
|
|
- ``encode(["hello", "world"])`` → ``[[123, 456], [789]]``
|
|
|
|
Batches are tokenised in parallel via the Rust backend's
|
|
``encode_batch`` (uses all available CPU cores).
|
|
"""
|
|
if self._tokenizer is None:
|
|
raise RuntimeError(
|
|
"Tokenizer not initialized. Load or create a tokenizer first."
|
|
)
|
|
|
|
if isinstance(tokens, str):
|
|
encoded = self._tokenizer.encode(
|
|
tokens,
|
|
is_pretokenized=is_pretokenized,
|
|
add_special_tokens=add_special_tokens,
|
|
)
|
|
return encoded.ids if out_ids else encoded.tokens
|
|
|
|
encoded_list = self._tokenizer.encode_batch(
|
|
tokens,
|
|
is_pretokenized=is_pretokenized,
|
|
add_special_tokens=add_special_tokens,
|
|
)
|
|
return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
|
|
|
|
def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
|
|
"""Decode token IDs to text."""
|
|
if self._tokenizer is None:
|
|
raise RuntimeError(
|
|
"Tokenizer not initialized. Load or create a tokenizer first."
|
|
)
|
|
|
|
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
|
|
|
|
def __len__(self) -> int:
|
|
if self._tokenizer is None:
|
|
return 0
|
|
return self._tokenizer.get_vocab_size()
|
|
|
|
def __getattr__(self, key: str):
|
|
"""
|
|
Dynamically intercept special token attribute access.
|
|
Supports three forms:
|
|
- tokenizer.bos_token → returns string
|
|
- tokenizer.bos_token_id → returns corresponding integer ID
|
|
- tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
|
|
|
|
Internal/private attrs are not intercepted: during unpickle
|
|
``__dict__`` is empty, so probing ``self._special_token_map``
|
|
would recurse infinitely.
|
|
"""
|
|
if key.startswith("_"):
|
|
raise AttributeError(key)
|
|
|
|
# Handle stop_ids - return IDs for all special tokens
|
|
if key == "stop_ids":
|
|
stop_ids = []
|
|
|
|
if self._tokenizer is None:
|
|
return stop_ids
|
|
|
|
for val in self._special_token_map.values():
|
|
token_id = self._tokenizer.token_to_id(val)
|
|
if token_id is not None:
|
|
stop_ids.append(token_id)
|
|
|
|
return stop_ids
|
|
|
|
# Handle _id suffix (e.g., bos_token_id -> bos_token)
|
|
if key.endswith("_id"):
|
|
base_attr = key[:-3] # Remove "_id"
|
|
token_str = self._special_token_map.get(base_attr)
|
|
if token_str is None:
|
|
return None
|
|
if self._tokenizer is None:
|
|
raise RuntimeError("Tokenizer not loaded, cannot convert token to id.")
|
|
return self._tokenizer.token_to_id(token_str)
|
|
|
|
# Handle regular string attributes
|
|
if key in self._special_token_map:
|
|
return self._special_token_map.get(key)
|
|
|
|
# Other attributes trigger default AttributeError
|
|
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{key}'")
|
|
|
|
@property
|
|
def vocab_size(self) -> int:
|
|
return len(self)
|
|
|
|
def set_chat_template(self, template: Union[str, ChatTemplate]):
|
|
"""
|
|
Set the chat template for the tokenizer.
|
|
|
|
Args:
|
|
template: Either a template name (str) registered in the global registry,
|
|
or a ChatTemplate instance, or a Jinja2 template string.
|
|
|
|
Raises:
|
|
KeyError: If template name is not registered.
|
|
"""
|
|
if isinstance(template, str):
|
|
self._chat_template = ChatTemplate.from_string(template)
|
|
elif isinstance(template, ChatTemplate):
|
|
self._chat_template = template
|
|
else:
|
|
raise ValueError("Invalid template type, must be str or ChatTemplate.")
|
|
|
|
def apply_chat_template(
|
|
self,
|
|
messages: Union[Messages, List[Messages]],
|
|
system_prompt: Optional[str] = None,
|
|
tokenize: bool = True,
|
|
add_generation_prompt: bool = True,
|
|
**kwargs,
|
|
) -> Union[str, List[int], List[str], List[List[int]]]:
|
|
"""Apply the chat template and optionally tokenize.
|
|
|
|
Accepts both single conversations and batches:
|
|
|
|
- ``apply_chat_template([msg1, msg2])`` → ``"..."`` or ``[ids]``
|
|
- ``apply_chat_template([[msg1, msg2], [msg3]])`` → ``["..", ".."]``
|
|
or ``[[ids], [ids]]``
|
|
|
|
Batches render each conversation list and tokenise all at once via
|
|
:meth:`encode` (``List[str]`` → Rust parallel ``encode_batch``).
|
|
|
|
Args:
|
|
messages: Single conversation (``Messages``) or batch of
|
|
conversations (``BatchMessages``).
|
|
system_prompt: Optional system prompt prepended (single mode only).
|
|
tokenize: Whether to return token IDs (True) or raw string (False).
|
|
add_generation_prompt: Whether to add the generation prompt.
|
|
**kwargs: Additional template variables.
|
|
|
|
Returns:
|
|
Single mode: ``str`` or ``List[int]``.
|
|
Batch mode: ``List[str]`` or ``List[List[int]]``.
|
|
"""
|
|
if self._chat_template is None:
|
|
raise RuntimeError(
|
|
"Chat template not set. Use set_chat_template() to set a template first."
|
|
)
|
|
|
|
is_batch = bool(messages) and isinstance(messages[0], list)
|
|
|
|
if is_batch:
|
|
rendered = [
|
|
self._chat_template.render(
|
|
messages=msgs,
|
|
add_generation_prompt=add_generation_prompt,
|
|
**kwargs,
|
|
)
|
|
for msgs in messages
|
|
]
|
|
if tokenize:
|
|
return self.encode(rendered) # List[str] → batch encode
|
|
return rendered
|
|
|
|
# Single conversation
|
|
if system_prompt:
|
|
messages = [{"role": "system", "content": system_prompt}] + list(messages)
|
|
rendered = self._chat_template.render(
|
|
messages=messages,
|
|
add_generation_prompt=add_generation_prompt,
|
|
**kwargs,
|
|
)
|
|
if tokenize:
|
|
return self.encode(rendered)
|
|
return rendered
|