feat: 实现模型动态注册机制
This commit is contained in:
@@ -1,22 +1,23 @@
|
||||
from astrai.tokenize.tokenizer import (
|
||||
BaseTokenizer,
|
||||
TextTokenizer,
|
||||
BpeTokenizer,
|
||||
BaseTrainer,
|
||||
BpeTrainer,
|
||||
)
|
||||
from astrai.tokenize.trainer import BpeTrainer
|
||||
from astrai.tokenize.chat_template import (
|
||||
ChatTemplate,
|
||||
HistoryType,
|
||||
MessageType,
|
||||
build_prompt,
|
||||
)
|
||||
|
||||
# Alias for compatibility
|
||||
AutoTokenizer = TextTokenizer
|
||||
|
||||
__all__ = [
|
||||
"BaseTokenizer",
|
||||
"TextTokenizer",
|
||||
"AutoTokenizer",
|
||||
"BpeTokenizer",
|
||||
"BaseTrainer",
|
||||
"BpeTrainer",
|
||||
"ChatTemplate",
|
||||
"HistoryType",
|
||||
"MessageType",
|
||||
"CHAT_TEMPLATES",
|
||||
"build_prompt",
|
||||
]
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
from jinja2 import Template
|
||||
from dataclasses import dataclass
|
||||
from astrai.factory import Registry
|
||||
|
||||
HistoryType = List[Tuple[str, str]]
|
||||
MessageType = Dict[str, str]
|
||||
@@ -75,213 +74,3 @@ class ChatTemplate:
|
||||
|
||||
jinja_template = Template(self.template_str)
|
||||
return jinja_template.render(**variables)
|
||||
|
||||
|
||||
# Global registry instance
|
||||
_default_registry = Registry()
|
||||
|
||||
# Default template name
|
||||
_default_template_name = "chatml"
|
||||
|
||||
|
||||
# Convenience functions
|
||||
def register_chat_template(
|
||||
name: str,
|
||||
template_str: str,
|
||||
description: str = "",
|
||||
default_variables: Optional[Dict[str, Any]] = None,
|
||||
special_tokens: Optional[Dict[str, str]] = None,
|
||||
) -> ChatTemplate:
|
||||
"""Register a chat template in the global registry."""
|
||||
template = ChatTemplate(
|
||||
name=name,
|
||||
template_str=template_str,
|
||||
description=description,
|
||||
default_variables=default_variables,
|
||||
special_tokens=special_tokens,
|
||||
)
|
||||
_default_registry.register(name, template, category=None, priority=0)
|
||||
return template
|
||||
|
||||
|
||||
def set_default_chat_template(name: str) -> None:
|
||||
"""Set the default chat template name globally."""
|
||||
global _default_template_name
|
||||
if not _default_registry.contains(name):
|
||||
raise KeyError(
|
||||
f"Chat template '{name}' not found. Available: {list(_default_registry.list_names())}"
|
||||
)
|
||||
_default_template_name = name
|
||||
|
||||
|
||||
def get_default_chat_template_name() -> str:
|
||||
"""Get the current default chat template name."""
|
||||
return _default_template_name
|
||||
|
||||
|
||||
def get_chat_template(name: str) -> ChatTemplate:
|
||||
"""Get a chat template from the global registry."""
|
||||
return _default_registry.get(name)
|
||||
|
||||
|
||||
def list_chat_templates() -> List[str]:
|
||||
"""List all registered chat template names."""
|
||||
return _default_registry.list_names()
|
||||
|
||||
|
||||
def chat_template_exists(name: str) -> bool:
|
||||
"""Check if a chat template exists."""
|
||||
return _default_registry.contains(name)
|
||||
|
||||
|
||||
def build_prompt(
|
||||
query: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
history: Optional[HistoryType] = None,
|
||||
template: Optional[str] = None,
|
||||
template_name: Optional[str] = None,
|
||||
**extra_variables: Any,
|
||||
) -> str:
|
||||
"""Build prompt using a registered chat template or a custom template string.
|
||||
|
||||
This function maintains backward compatibility with the previous API.
|
||||
|
||||
Args:
|
||||
query: The current user query.
|
||||
system_prompt: Optional system prompt.
|
||||
history: Optional list of (user_msg, assistant_msg) pairs.
|
||||
template: If provided, uses this exact Jinja2 template string (overrides template_name).
|
||||
template_name: Name of a registered template to use (ignored if `template` is given).
|
||||
If None, uses the globally set default template (see `set_default_chat_template`).
|
||||
**extra_variables: Additional variables to pass to the template.
|
||||
|
||||
Returns:
|
||||
Rendered prompt string.
|
||||
|
||||
Raises:
|
||||
KeyError: If `template_name` is not registered.
|
||||
"""
|
||||
# Convert history to message format
|
||||
messages: List[MessageType] = []
|
||||
if history:
|
||||
for user_msg, assistant_msg in history:
|
||||
messages.append({"role": "user", "content": user_msg})
|
||||
messages.append({"role": "assistant", "content": assistant_msg})
|
||||
messages.append({"role": "user", "content": query})
|
||||
|
||||
if template is not None:
|
||||
# Use the provided template string directly
|
||||
jinja_template = Template(template)
|
||||
variables = {"messages": messages, **extra_variables}
|
||||
if system_prompt is not None:
|
||||
variables["system_prompt"] = system_prompt
|
||||
return jinja_template.render(**variables)
|
||||
else:
|
||||
# Determine which template name to use
|
||||
if template_name is None:
|
||||
template_name = _default_template_name
|
||||
# Use a registered template
|
||||
chat_template = get_chat_template(template_name)
|
||||
return chat_template.render(
|
||||
messages=messages,
|
||||
system_prompt=system_prompt,
|
||||
**extra_variables,
|
||||
)
|
||||
|
||||
|
||||
# Predefined templates
|
||||
# ChatML template (original)
|
||||
register_chat_template(
|
||||
name="chatml",
|
||||
template_str=(
|
||||
"{%- if system_prompt -%}\n"
|
||||
"{{ bos_token }}system\n"
|
||||
"{{ system_prompt }}{{ eos_token }}\n"
|
||||
"{%- endif -%}\n"
|
||||
"{%- for message in messages -%}\n"
|
||||
"{{ bos_token }}{{ message['role'] }}\n"
|
||||
"{{ message['content'] }}{{ eos_token }}\n"
|
||||
"{%- endfor -%}\n"
|
||||
"{{ bos_token }}assistant\n"
|
||||
),
|
||||
description="ChatML format with configurable special tokens.",
|
||||
special_tokens={"bos_token": "<｜im▁start｜>", "eos_token": "<｜im▁end｜>"},
|
||||
)
|
||||
|
||||
# Simplified template without special tokens (plain text)
|
||||
register_chat_template(
|
||||
name="plain",
|
||||
template_str=(
|
||||
"{%- if system_prompt -%}\n"
|
||||
"System: {{ system_prompt }}\n"
|
||||
"{%- endif -%}\n"
|
||||
"{%- for message in messages -%}\n"
|
||||
"{{ message['role']|capitalize }}: {{ message['content'] }}\n"
|
||||
"{%- endfor -%}\n"
|
||||
"Assistant:"
|
||||
),
|
||||
description="Plain text format with role labels.",
|
||||
)
|
||||
|
||||
# Alpaca-style template
|
||||
register_chat_template(
|
||||
name="alpaca",
|
||||
template_str=(
|
||||
"{%- if system_prompt -%}\n"
|
||||
"### Instruction:\n"
|
||||
"{{ system_prompt }}\n"
|
||||
"{%- endif -%}\n"
|
||||
"### Input:\n"
|
||||
"{{ messages[-1]['content'] }}\n"
|
||||
"### Response:"
|
||||
),
|
||||
description="Alpaca instruction‑response format (single‑turn).",
|
||||
default_variables={},
|
||||
)
|
||||
|
||||
# OpenAI chat format (approximation)
|
||||
register_chat_template(
|
||||
name="openai",
|
||||
template_str=(
|
||||
"{%- if system_prompt -%}\n"
|
||||
"{{ bos_token }}system\n"
|
||||
"{{ system_prompt }}{{ eos_token }}\n"
|
||||
"{%- endif -%}\n"
|
||||
"{%- for message in messages -%}\n"
|
||||
"{{ bos_token }}{{ message['role'] }}\n"
|
||||
"{{ message['content'] }}{{ eos_token }}\n"
|
||||
"{%- endfor -%}\n"
|
||||
"{{ bos_token }}assistant\n"
|
||||
),
|
||||
description="OpenAI‑compatible chat format with configurable special tokens.",
|
||||
special_tokens={"bos_token": "<｜im▁start｜>", "eos_token": "<｜im▁end｜>"},
|
||||
)
|
||||
|
||||
# Llama‑2 style with [INST] tags
|
||||
register_chat_template(
|
||||
name="llama2",
|
||||
template_str=(
|
||||
"{%- if system_prompt -%}\n"
|
||||
"<<SYS>>\n"
|
||||
"{{ system_prompt }}\n"
|
||||
"<</SYS>>\n"
|
||||
"{%- endif -%}\n"
|
||||
"[INST] {{ messages[-1]['content'] }} [/INST]"
|
||||
),
|
||||
description="Llama‑2 style with [INST] tags (single‑turn).",
|
||||
default_variables={},
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ChatTemplate",
|
||||
"register_chat_template",
|
||||
"get_chat_template",
|
||||
"list_chat_templates",
|
||||
"chat_template_exists",
|
||||
"build_prompt",
|
||||
"set_default_chat_template",
|
||||
"get_default_chat_template_name",
|
||||
"HistoryType",
|
||||
"MessageType",
|
||||
]
|
||||
|
||||
+218
-166
@@ -1,97 +1,254 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Union
|
||||
"""
|
||||
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, decoders, normalizers, pre_tokenizers, processors
|
||||
from tokenizers.models import BPE
|
||||
from tokenizers.trainers import BpeTrainer as BpeTrainerImpl
|
||||
from astrai.tokenize.chat_template import ChatTemplate
|
||||
|
||||
|
||||
class BaseTokenizer(ABC):
|
||||
@abstractmethod
|
||||
def _init_tokenizer(self):
|
||||
pass
|
||||
class TextTokenizer:
|
||||
"""Base tokenizer class with automatic loading support"""
|
||||
|
||||
@abstractmethod
|
||||
def save(self, path):
|
||||
pass
|
||||
TOKENIZER_CLASSES = {} # Registry for auto-loading
|
||||
|
||||
@abstractmethod
|
||||
def load(self, path):
|
||||
pass
|
||||
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], **kwargs) -> "TextTokenizer":
|
||||
"""Load tokenizer from pretrained directory."""
|
||||
instance = cls(path)
|
||||
return instance
|
||||
|
||||
def save_pretrained(self, tokenizer, save_path: str):
|
||||
"""
|
||||
Save tokenizer to pretrained directory.
|
||||
|
||||
Args:
|
||||
tokenizer: Tokenizer instance to save
|
||||
save_path: Path to save the tokenizer
|
||||
"""
|
||||
save_path = Path(save_path)
|
||||
save_path.mkdir(parents=True, exist_ok=True)
|
||||
self._tokenizer.save(tokenizer, save_path)
|
||||
|
||||
@classmethod
|
||||
def register_tokenizer(cls, name: str, tokenizer_class: type):
|
||||
"""
|
||||
Register a new tokenizer class.
|
||||
|
||||
Args:
|
||||
name: Name to register the tokenizer class under
|
||||
tokenizer_class: The tokenizer class to register
|
||||
"""
|
||||
cls.TOKENIZER_CLASSES[name] = tokenizer_class
|
||||
|
||||
@abstractmethod
|
||||
def encode(
|
||||
self,
|
||||
tokens: Union[str, List[str]],
|
||||
out_ids: bool = True,
|
||||
add_special_tokens: bool = False,
|
||||
is_pretokenized: bool = False,
|
||||
add_special_tokens: bool = True,
|
||||
) -> List:
|
||||
pass
|
||||
"""Encode text to tokens or token IDs."""
|
||||
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
|
||||
else:
|
||||
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
|
||||
]
|
||||
|
||||
@abstractmethod
|
||||
def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
|
||||
pass
|
||||
"""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)
|
||||
|
||||
@abstractmethod
|
||||
def __len__(self) -> int:
|
||||
pass
|
||||
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
|
||||
"""
|
||||
# Handle stop_ids
|
||||
if key == "stop_ids":
|
||||
return [
|
||||
self._special_token_map.get(val)
|
||||
for val in self._special_token_map.values()
|
||||
]
|
||||
|
||||
# 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
|
||||
@abstractmethod
|
||||
def stop_ids(self) -> List[int]:
|
||||
pass
|
||||
def vocab_size(self) -> int:
|
||||
return len(self)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def bos_id(self) -> int:
|
||||
pass
|
||||
def pad_id(self) -> Optional[int]:
|
||||
"""Return the pad token ID if available."""
|
||||
pad_token = self._special_token_map.get("pad")
|
||||
if pad_token is None or self._tokenizer is None:
|
||||
return None
|
||||
return self._tokenizer.token_to_id(pad_token)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def eos_id(self) -> int:
|
||||
pass
|
||||
def set_chat_template(self, template: Union[str, ChatTemplate]):
|
||||
"""
|
||||
Set the chat template for the tokenizer.
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def pad_id(self) -> int:
|
||||
pass
|
||||
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: List[Dict[str, str]],
|
||||
system_prompt: Optional[str] = None,
|
||||
tokenize: bool = True,
|
||||
**kwargs,
|
||||
) -> Union[str, List[int]]:
|
||||
"""
|
||||
Apply the chat template to messages and optionally tokenize the result.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
system_prompt: Optional system prompt string.
|
||||
tokenize: Whether to return token IDs (True) or raw string (False).
|
||||
**kwargs: Additional variables to pass to the template.
|
||||
|
||||
Returns:
|
||||
Either the rendered string or list of token IDs.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If chat template is not set.
|
||||
"""
|
||||
if self._chat_template is None:
|
||||
raise RuntimeError(
|
||||
"Chat template not set. Use set_chat_template() to set a template first."
|
||||
)
|
||||
|
||||
# Render the template
|
||||
rendered = self._chat_template.render(
|
||||
messages=messages,
|
||||
system_prompt=system_prompt,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if tokenize:
|
||||
return self.encode(rendered)
|
||||
|
||||
return rendered
|
||||
|
||||
|
||||
class BaseTrainer(ABC):
|
||||
def __init__(self, tokenizer: BaseTokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
class BpeTokenizer(TextTokenizer):
|
||||
"""BPE tokenizer implementation."""
|
||||
|
||||
@abstractmethod
|
||||
def train(self, files, vocab_size, min_freq, **kwargs):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def train_from_iterator(self, iterator, vocab_size, min_freq, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
class BpeTokenizer(BaseTokenizer):
|
||||
def __init__(
|
||||
self,
|
||||
control_tokens: List[str] = None,
|
||||
special_tokens: List[str] = None,
|
||||
path=None,
|
||||
special_token_map: Dict[str, str] = None,
|
||||
path: Optional[str] = None,
|
||||
chat_template: Optional[str] = None,
|
||||
):
|
||||
self._control_tokens = control_tokens or [
|
||||
"<｜begin▁of▁sentence｜>",
|
||||
"<｜end▁of▁sentence｜>",
|
||||
"<｜▁pad▁｜>",
|
||||
]
|
||||
self._special_tokens = special_tokens or [
|
||||
"<｜im▁start｜>",
|
||||
"<｜im▁end｜>",
|
||||
]
|
||||
special_token_map = special_token_map or {
|
||||
"bos": "<｜begin▁of▁sentence｜>",
|
||||
"eos": "<｜end▁of▁sentence｜>",
|
||||
"pad": "<｜▁pad▁｜>",
|
||||
"im_start": "<｜im▁start｜>",
|
||||
"im_end": "<｜im▁end｜>",
|
||||
}
|
||||
self._tokenizer = None
|
||||
self._init_tokenizer()
|
||||
if path is not None:
|
||||
self.load(path)
|
||||
super().__init__(
|
||||
path, special_token_map=special_token_map, chat_template=chat_template
|
||||
)
|
||||
|
||||
def _init_tokenizer(self):
|
||||
"""Initialize a new BPE tokenizer with default settings."""
|
||||
model = BPE()
|
||||
self._tokenizer = Tokenizer(model)
|
||||
self._tokenizer.normalizer = normalizers.Sequence(
|
||||
@@ -105,108 +262,3 @@ class BpeTokenizer(BaseTokenizer):
|
||||
)
|
||||
self._tokenizer.decoder = decoders.ByteLevel()
|
||||
self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
|
||||
|
||||
def save(self, path):
|
||||
self._tokenizer.save(path)
|
||||
|
||||
def load(self, path):
|
||||
self._tokenizer = Tokenizer.from_file(path)
|
||||
|
||||
def encode(
|
||||
self,
|
||||
tokens: Union[str, List[str]],
|
||||
out_ids: bool = True,
|
||||
add_special_tokens: bool = False,
|
||||
) -> List:
|
||||
if isinstance(tokens, str):
|
||||
encoded = self._tokenizer.encode(
|
||||
tokens, add_special_tokens=add_special_tokens
|
||||
)
|
||||
return encoded.ids if out_ids else encoded.tokens
|
||||
else:
|
||||
encoded_list = self._tokenizer.encode_batch(
|
||||
tokens, 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:
|
||||
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._tokenizer.get_vocab_size()
|
||||
|
||||
@property
|
||||
def stop_ids(self) -> List[int]:
|
||||
stop_token = self._control_tokens + self._special_tokens
|
||||
return [self._tokenizer.token_to_id(tok) for tok in stop_token]
|
||||
|
||||
@property
|
||||
def bos_id(self) -> int:
|
||||
return self._tokenizer.token_to_id(self._control_tokens[0])
|
||||
|
||||
@property
|
||||
def eos_id(self) -> int:
|
||||
return self._tokenizer.token_to_id(self._control_tokens[1])
|
||||
|
||||
@property
|
||||
def pad_id(self) -> int:
|
||||
return self._tokenizer.token_to_id(self._control_tokens[2])
|
||||
|
||||
|
||||
class BpeTrainer(BaseTrainer):
|
||||
def __init__(self, tokenizer: BaseTokenizer):
|
||||
super().__init__(tokenizer)
|
||||
|
||||
def _prepare_trainer(
|
||||
self,
|
||||
vocab_size: int,
|
||||
min_freq: int,
|
||||
reserved_token_size: int,
|
||||
max_token_length=18,
|
||||
):
|
||||
assert reserved_token_size > len(self.tokenizer._special_tokens)
|
||||
reserved_tokens = [
|
||||
f"<|reserve{i:02d}|>"
|
||||
for i in range(reserved_token_size - len(self.tokenizer._special_tokens))
|
||||
]
|
||||
detail_vocab_size = vocab_size - (
|
||||
len(reserved_tokens) + len(self.tokenizer._special_tokens)
|
||||
)
|
||||
alphabet = pre_tokenizers.ByteLevel.alphabet()
|
||||
min_size = len(alphabet) + len(self.tokenizer._control_tokens)
|
||||
assert detail_vocab_size > min_size
|
||||
|
||||
trainer = BpeTrainerImpl(
|
||||
vocab_size=detail_vocab_size,
|
||||
min_frequency=min_freq,
|
||||
limit_alphabet=detail_vocab_size // 6,
|
||||
max_token_length=max_token_length,
|
||||
special_tokens=self.tokenizer._control_tokens,
|
||||
initial_alphabet=alphabet,
|
||||
show_progress=True,
|
||||
)
|
||||
return trainer, reserved_tokens
|
||||
|
||||
def train(self, files, vocab_size, min_freq, reserved_token_size=100, **kwargs):
|
||||
trainer, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size, min_freq, reserved_token_size, **kwargs
|
||||
)
|
||||
self.tokenizer._tokenizer.train(files=files, trainer=trainer)
|
||||
self.tokenizer._tokenizer.add_special_tokens(
|
||||
self.tokenizer._special_tokens + reserved_tokens
|
||||
)
|
||||
|
||||
def train_from_iterator(
|
||||
self, iterator, vocab_size, min_freq, reserved_token_size=100, **kwargs
|
||||
):
|
||||
trainer, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size, min_freq, reserved_token_size, **kwargs
|
||||
)
|
||||
self.tokenizer._tokenizer.train_from_iterator(
|
||||
iterator=iterator, trainer=trainer
|
||||
)
|
||||
self.tokenizer._tokenizer.add_special_tokens(
|
||||
self.tokenizer._special_tokens + reserved_tokens
|
||||
)
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
"""
|
||||
BPE Tokenizer Trainer module.
|
||||
|
||||
Provides training functionality for BPE tokenizers.
|
||||
"""
|
||||
|
||||
from typing import List, Union
|
||||
|
||||
from tokenizers import pre_tokenizers
|
||||
from tokenizers.trainers import BpeTrainer as BpeTrainerImpl
|
||||
|
||||
|
||||
class BpeTrainer:
|
||||
"""BPE tokenizer trainer."""
|
||||
|
||||
def __init__(self, tokenizer):
|
||||
"""Initialize trainer with a tokenizer instance.
|
||||
|
||||
Args:
|
||||
tokenizer: A BpeTokenizer instance
|
||||
"""
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def _prepare_trainer(
|
||||
self,
|
||||
vocab_size: int,
|
||||
min_freq: int,
|
||||
reserved_token_size: int,
|
||||
max_token_length: int = 18,
|
||||
):
|
||||
"""Prepare the BPE trainer with proper configuration."""
|
||||
assert reserved_token_size > len(self.tokenizer._special_tokens)
|
||||
reserved_tokens = [
|
||||
f"<|reserve{i:02d}|>"
|
||||
for i in range(reserved_token_size - len(self.tokenizer._special_tokens))
|
||||
]
|
||||
detail_vocab_size = vocab_size - (
|
||||
len(reserved_tokens) + len(self.tokenizer._special_tokens)
|
||||
)
|
||||
alphabet = pre_tokenizers.ByteLevel.alphabet()
|
||||
min_size = len(alphabet) + len(self.tokenizer._control_tokens)
|
||||
assert detail_vocab_size > min_size
|
||||
|
||||
trainer = BpeTrainerImpl(
|
||||
vocab_size=detail_vocab_size,
|
||||
min_frequency=min_freq,
|
||||
limit_alphabet=detail_vocab_size // 6,
|
||||
max_token_length=max_token_length,
|
||||
special_tokens=self.tokenizer._control_tokens,
|
||||
initial_alphabet=alphabet,
|
||||
show_progress=True,
|
||||
)
|
||||
return trainer, reserved_tokens
|
||||
|
||||
def train(
|
||||
self,
|
||||
files: Union[str, List[str]],
|
||||
vocab_size: int,
|
||||
min_freq: int,
|
||||
reserved_token_size: int = 100,
|
||||
**kwargs,
|
||||
):
|
||||
"""Train tokenizer from files.
|
||||
|
||||
Args:
|
||||
files: Path or list of paths to training files
|
||||
vocab_size: Target vocabulary size
|
||||
min_freq: Minimum frequency for tokens
|
||||
reserved_token_size: Number of reserved tokens
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
trainer, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size, min_freq, reserved_token_size, **kwargs
|
||||
)
|
||||
self.tokenizer._tokenizer.train(files=files, trainer=trainer)
|
||||
self.tokenizer._tokenizer.add_special_tokens(
|
||||
self.tokenizer._special_tokens + reserved_tokens
|
||||
)
|
||||
|
||||
def train_from_iterator(
|
||||
self,
|
||||
iterator,
|
||||
vocab_size: int,
|
||||
min_freq: int,
|
||||
reserved_token_size: int = 100,
|
||||
**kwargs,
|
||||
):
|
||||
"""Train tokenizer from iterator.
|
||||
|
||||
Args:
|
||||
iterator: Iterator yielding training strings
|
||||
vocab_size: Target vocabulary size
|
||||
min_freq: Minimum frequency for tokens
|
||||
reserved_token_size: Number of reserved tokens
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
trainer, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size, min_freq, reserved_token_size, **kwargs
|
||||
)
|
||||
self.tokenizer._tokenizer.train_from_iterator(
|
||||
iterator=iterator, trainer=trainer
|
||||
)
|
||||
self.tokenizer._tokenizer.add_special_tokens(
|
||||
self.tokenizer._special_tokens + reserved_tokens
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["BpeTrainer"]
|
||||
Reference in New Issue
Block a user