diff --git a/README.md b/README.md index 7f26930..0fed23f 100644 --- a/README.md +++ b/README.md @@ -115,9 +115,9 @@ export_dataset( ### 2. 分词并缓存为 HDF5 ```python -from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl +from pipeline import AutoTokenizer, ProcessorFactory, cache_jsonl -tokenizer = BpeTokenizer("tokenizer.json") +tokenizer = AutoTokenizer("tokenizer.json") processor = ProcessorFactory.create("pt", tokenizer) cache_jsonl( diff --git a/docs/Design.md b/docs/Design.md index 756fcbe..72dc17c 100644 --- a/docs/Design.md +++ b/docs/Design.md @@ -176,10 +176,12 @@ StrategyFactory.register("my_format", MyStrategy) - `load_h5(file_path, share_memory)` — 加载 HDF5，支持共享内存（用于 DataLoader 多进程） - `fetch_files(directory)` / `fetch_folders(root_dir)` — 文件/目录遍历 -### BpeTokenizer (`pipeline/tokenizer.py`) +### AutoTokenizer (`pipeline/tokenize/tokenizer.py`) 基于 HuggingFace `tokenizers` 库的 BPE 分词器，支持从文件加载、训练、保存。内置 `<｜begin▁of▁sentence｜>`/`<｜end▁of▁sentence｜>`/`<｜▁pad▁｜>` 控制符和 `<｜im▁start｜>`/`<｜im▁end｜>` 特殊 token。 +支持动态属性访问（`bos_token`, `bos_token_id`, `stop_ids`, `pad_id` 等）和聊天模板（`set_chat_template`, `apply_chat_template`）。 + ## API 参考 各模块的详细 API 文档请参阅对应的源文件。 diff --git a/pipeline/__init__.py b/pipeline/__init__.py index 76d7a02..392135c 100644 --- a/pipeline/__init__.py +++ b/pipeline/__init__.py @@ -1,5 +1,5 @@ import logging -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer, ChatTemplate, train_bpe_tokenizer from pipeline.text import TextNormalizer from pipeline.packing import SequencePacker from pipeline.io import IOHandler, export_dataset, cache_jsonl @@ -16,8 +16,11 @@ from pipeline.strategies import ( setup_logging() __all__ = [ + # Tokenizer + "AutoTokenizer", + "ChatTemplate", + "train_bpe_tokenizer", # Core modules - "BpeTokenizer", "TextNormalizer", "SequencePacker", "IOHandler", diff --git a/pipeline/processors/dpo.py b/pipeline/processors/dpo.py index 5f60a96..f5cf2b6 100644 --- a/pipeline/processors/dpo.py +++ b/pipeline/processors/dpo.py @@ -5,7 +5,7 @@ from typing import Dict, List, Any, Optional import torch from torch import Tensor -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.strategies import PromptStrategy, ChatMLStrategy from pipeline.processors.base import BaseProcessor, _encode_with_mask from pipeline.processors.factory import ProcessorFactory @@ -20,7 +20,7 @@ class DPOProcessor(BaseProcessor): def __init__( self, - tokenizer: BpeTokenizer, + tokenizer: AutoTokenizer, strategy: Optional[PromptStrategy] = None, ): self.tokenizer = tokenizer diff --git a/pipeline/processors/factory.py b/pipeline/processors/factory.py index da80eaf..300fd36 100644 --- a/pipeline/processors/factory.py +++ b/pipeline/processors/factory.py @@ -3,7 +3,7 @@ from typing import Dict, List, Any, Optional, Type from pipeline.processors.base import BaseProcessor -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.strategies import PromptStrategy, StrategyFactory @@ -45,7 +45,7 @@ class ProcessorFactory: return decorator @classmethod - def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor: + def create(cls, processor_type: str, tokenizer: AutoTokenizer) -> BaseProcessor: """Create a processor by type name (uses default ChatMLStrategy for SFT/DPO). Args: @@ -69,7 +69,7 @@ class ProcessorFactory: def create_with_strategy( cls, processor_type: str, - tokenizer: BpeTokenizer, + tokenizer: AutoTokenizer, strategy: PromptStrategy, ) -> BaseProcessor: """Create a processor with a custom strategy. @@ -99,7 +99,7 @@ class ProcessorFactory: def create_with_strategy_name( cls, processor_type: str, - tokenizer: BpeTokenizer, + tokenizer: AutoTokenizer, strategy_name: str, **strategy_kwargs, ) -> BaseProcessor: diff --git a/pipeline/processors/pretrain.py b/pipeline/processors/pretrain.py index 2691cee..df67b82 100644 --- a/pipeline/processors/pretrain.py +++ b/pipeline/processors/pretrain.py @@ -5,7 +5,7 @@ from typing import Dict, List, Any import torch from torch import Tensor -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.processors.base import BaseProcessor from pipeline.processors.factory import ProcessorFactory @@ -14,7 +14,7 @@ from pipeline.processors.factory import ProcessorFactory class PreTrainProcessor(BaseProcessor): """Pre-training data processor.""" - def __init__(self, tokenizer: BpeTokenizer): + def __init__(self, tokenizer: AutoTokenizer): self.tokenizer = tokenizer def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index a831ad9..bd652a2 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -5,7 +5,7 @@ from typing import Dict, List, Any, Optional import torch from torch import Tensor -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.strategies import PromptStrategy, ChatMLStrategy from pipeline.processors.base import BaseProcessor, _encode_with_mask from pipeline.processors.factory import ProcessorFactory @@ -20,7 +20,7 @@ class SFTProcessor(BaseProcessor): def __init__( self, - tokenizer: BpeTokenizer, + tokenizer: AutoTokenizer, strategy: Optional[PromptStrategy] = None, ): self.tokenizer = tokenizer diff --git a/pipeline/strategies/alpaca.py b/pipeline/strategies/alpaca.py index 153de83..c6e223c 100644 --- a/pipeline/strategies/alpaca.py +++ b/pipeline/strategies/alpaca.py @@ -2,7 +2,7 @@ from typing import List -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.strategies.base import PromptStrategy from pipeline.strategies.factory import StrategyFactory @@ -13,7 +13,7 @@ class AlpacaStrategy(PromptStrategy): def __init__( self, - tokenizer: BpeTokenizer, + tokenizer: AutoTokenizer, instruction_start: str = "### Instruction:\n", response_start: str = "### Response:\n", response_suffix: str = "\n<｜end▁of▁sentence｜>", diff --git a/pipeline/strategies/base.py b/pipeline/strategies/base.py index f0753d8..c947228 100644 --- a/pipeline/strategies/base.py +++ b/pipeline/strategies/base.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from typing import List -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer class PromptStrategy(ABC): @@ -14,7 +14,7 @@ class PromptStrategy(ABC): which assembles them with pre-encoded format tokens. """ - def __init__(self, tokenizer: BpeTokenizer): + def __init__(self, tokenizer: AutoTokenizer): self.tokenizer = tokenizer def _encode_format(self, text: str) -> List[int]: diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index 139c7c9..cd8411c 100644 --- a/pipeline/strategies/chatml.py +++ b/pipeline/strategies/chatml.py @@ -2,7 +2,7 @@ from typing import List -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.strategies.base import PromptStrategy from pipeline.strategies.factory import StrategyFactory @@ -13,7 +13,7 @@ class ChatMLStrategy(PromptStrategy): def __init__( self, - tokenizer: BpeTokenizer, + tokenizer: AutoTokenizer, user_start: str = "<｜im▁start｜>user\n", user_end: str = "<｜im▁end｜>\n", assistant_start: str = "<｜im▁start｜>assistant\n", diff --git a/pipeline/strategies/factory.py b/pipeline/strategies/factory.py index d06c855..8b17149 100644 --- a/pipeline/strategies/factory.py +++ b/pipeline/strategies/factory.py @@ -2,7 +2,7 @@ from typing import Dict, List, Type -from pipeline.tokenizer import BpeTokenizer +from pipeline.tokenize import AutoTokenizer from pipeline.strategies.base import PromptStrategy @@ -44,7 +44,7 @@ class StrategyFactory: return decorator @classmethod - def create(cls, name: str, tokenizer: BpeTokenizer, **kwargs) -> PromptStrategy: + def create(cls, name: str, tokenizer: AutoTokenizer, **kwargs) -> PromptStrategy: """Create a strategy by name. Args: diff --git a/pipeline/tokenize/__init__.py b/pipeline/tokenize/__init__.py new file mode 100644 index 0000000..b18171c --- /dev/null +++ b/pipeline/tokenize/__init__.py @@ -0,0 +1,12 @@ +""" +Tokenizer module with BPE implementation and auto-loading support. +""" + +from pipeline.tokenize.tokenizer import AutoTokenizer, train_bpe_tokenizer +from pipeline.tokenize.chat_template import ChatTemplate + +__all__ = [ + "AutoTokenizer", + "train_bpe_tokenizer", + "ChatTemplate", +] diff --git a/pipeline/tokenize/chat_template.py b/pipeline/tokenize/chat_template.py new file mode 100644 index 0000000..2922cc3 --- /dev/null +++ b/pipeline/tokenize/chat_template.py @@ -0,0 +1,120 @@ +""" +Chat template module with Jinja2 rendering support. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from jinja2 import Template + + +# Message type for chat messages +type MessageType = Dict[str, Any] + + +@dataclass +class ChatTemplate: + """A chat template with Jinja2 rendering support. + + Attributes: + name: Unique identifier for the template. + template_str: Jinja2 template string. + description: Optional description. + default_variables: Optional dictionary of default variable values + that will be passed to the template if not overridden during rendering. + special_tokens: Optional dictionary mapping token names to their string values. + These tokens are automatically added to the template variables. + """ + + name: str + template_str: str + description: str = "" + default_variables: Dict[str, Any] = field(default_factory=dict) + special_tokens: Dict[str, str] = field(default_factory=dict) + + @classmethod + def from_string( + cls, + template_str: str, + description: str = "", + default_variables: Optional[Dict[str, Any]] = None, + special_tokens: Optional[Dict[str, str]] = None, + ) -> "ChatTemplate": + """Create a ChatTemplate instance directly from a template string.""" + return cls( + name="", # empty name for ad-hoc templates + template_str=template_str, + description=description, + default_variables=default_variables or {}, + special_tokens=special_tokens or {}, + ) + + def render( + self, + messages: List[MessageType], + system_prompt: Optional[str] = None, + add_generation_prompt: bool = True, + **extra_variables: Any, + ) -> str: + """Render the template with given messages and variables. + + Args: + messages: List of message dicts with 'role' and 'content'. + system_prompt: Optional system prompt string. + add_generation_prompt: Whether to add generation prompt after messages. + **extra_variables: Additional variables to pass to the template. + These override default_variables and special_tokens. + + Returns: + Rendered prompt string. + """ + # Merge default variables, special tokens, and extra variables + variables = { + **self.default_variables, + **self.special_tokens, + **extra_variables, + } + variables["messages"] = messages + variables["add_generation_prompt"] = add_generation_prompt + if system_prompt is not None: + variables["system_prompt"] = system_prompt + + jinja_template = Template(self.template_str) + return jinja_template.render(**variables) + + +# Default ChatML template +DEFAULT_CHATML_TEMPLATE = """{% for message in messages %}{{ bos_token }}{{ message['role'] }} +{{ message['content'] }}{{ eos_token }}{% endfor %}{% if add_generation_prompt %}{{ bos_token }}assistant +{% endif %}""" + + +# Pre-built template registry +TEMPLATE_REGISTRY: Dict[str, ChatTemplate] = {} + + +def register_chat_template(name: str, template: ChatTemplate) -> None: + """Register a chat template in the global registry. + + Args: + name: Name to register the template under + template: ChatTemplate instance + """ + TEMPLATE_REGISTRY[name] = template + + +def get_chat_template(name: str) -> ChatTemplate: + """Get a registered chat template. + + Args: + name: Template name + + Returns: + ChatTemplate instance + + Raises: + KeyError: If template not found + """ + if name not in TEMPLATE_REGISTRY: + raise KeyError(f"Chat template '{name}' not found. Available: {list(TEMPLATE_REGISTRY.keys())}") + return TEMPLATE_REGISTRY[name] diff --git a/pipeline/tokenize/tokenizer.py b/pipeline/tokenize/tokenizer.py new file mode 100644 index 0000000..83f0f63 --- /dev/null +++ b/pipeline/tokenize/tokenizer.py @@ -0,0 +1,378 @@ +""" +Tokenizer module with BPE implementation and auto-loading support. +""" + +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from tokenizers import Tokenizer +from tokenizers import decoders, processors, normalizers, pre_tokenizers +from tokenizers.models import BPE +from tokenizers.trainers import BpeTrainer +from jinja2 import Template + + +DEFAULT_SPECIAL_TOKENS = { + "bos_token": "<｜begin▁of▁sentence｜>", + "eos_token": "<｜end▁of▁sentence｜>", + "pad_token": "<｜▁pad▁｜>", +} + +CONTROL_TOKENS = [ + "<｜begin▁of▁sentence｜>", + "<｜end▁of▁sentence｜>", + "<｜▁pad▁｜>", +] + +SPECIAL_TOKENS = ["<｜im▁start｜>", "<｜im▁end｜>"] + + +def train_bpe_tokenizer( + files: List[str], + vocab_size: int, + min_freq: int = 2, + reserved_token_size: int = 100, + max_token_length: int = 18, +) -> Tokenizer: + + reserved_tokens = [ + f"<｜reserve{i:02d}｜>" + for i in range(reserved_token_size - len(SPECIAL_TOKENS)) + ] + detail_vocab_size = vocab_size - (len(reserved_tokens) + len(SPECIAL_TOKENS)) + + alphabet = pre_tokenizers.ByteLevel.alphabet() + min_size = len(alphabet) + len(CONTROL_TOKENS) + assert detail_vocab_size > min_size + + tokenizer = Tokenizer(BPE()) + tokenizer.normalizer = normalizers.Sequence([normalizers.NFC(), normalizers.Strip()]) + tokenizer.pre_tokenizer = pre_tokenizers.Sequence([ + pre_tokenizers.UnicodeScripts(), + pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True), + ]) + tokenizer.decoder = decoders.ByteLevel() + tokenizer.post_processor = processors.ByteLevel(trim_offsets=True) + + trainer = BpeTrainer( + vocab_size=detail_vocab_size, + min_frequency=min_freq, + limit_alphabet=detail_vocab_size // 6, + max_token_length=max_token_length, + special_tokens=CONTROL_TOKENS + SPECIAL_TOKENS, + initial_alphabet=alphabet, + show_progress=True, + ) + + tokenizer.train(files=files, trainer=trainer) + tokenizer.add_special_tokens(CONTROL_TOKENS + SPECIAL_TOKENS + reserved_tokens) + + return tokenizer + + +# Message type for chat messages +type MessageType = Dict[str, Any] + + +@dataclass +class ChatTemplate: + """A chat template with Jinja2 rendering support. + + Attributes: + name: Unique identifier for the template. + template_str: Jinja2 template string. + description: Optional description. + default_variables: Optional dictionary of default variable values + that will be passed to the template if not overridden during rendering. + special_tokens: Optional dictionary mapping token names to their string values. + These tokens are automatically added to the template variables. + """ + + name: str + template_str: str + description: str = "" + default_variables: Dict[str, Any] = None + special_tokens: Dict[str, str] = None + + def __post_init__(self): + if self.default_variables is None: + self.default_variables = {} + if self.special_tokens is None: + self.special_tokens = {} + + @classmethod + def from_string( + cls, + template_str: str, + description: str = "", + default_variables: Optional[Dict[str, Any]] = None, + special_tokens: Optional[Dict[str, str]] = None, + ) -> "ChatTemplate": + """Create a ChatTemplate instance directly from a template string.""" + return cls( + name="", # empty name for ad‑hoc templates + template_str=template_str, + description=description, + default_variables=default_variables, + special_tokens=special_tokens, + ) + + def render( + self, + messages: List[MessageType], + system_prompt: Optional[str] = None, + **extra_variables: Any, + ) -> str: + """Render the template with given messages and variables. + + Args: + messages: List of message dicts with 'role' and 'content'. + system_prompt: Optional system prompt string. + **extra_variables: Additional variables to pass to the template. + These override default_variables and special_tokens. + + Returns: + Rendered prompt string. + """ + # Merge default variables, special tokens, and extra variables + variables = {**self.default_variables, **self.special_tokens, **extra_variables} + variables["messages"] = messages + if system_prompt is not None: + variables["system_prompt"] = system_prompt + + jinja_template = Template(self.template_str) + return jinja_template.render(**variables) + + + +class AutoTokenizer: + """Base tokenizer class with automatic loading support""" + + TOKENIZER_CLASSES = {} # Registry for auto-loading + + 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) -> "AutoTokenizer": + """Load tokenizer from pretrained directory.""" + instance = cls(path) + return instance + + def save_pretrained(self, save_path: str): + """ + Save tokenizer to pretrained directory. + + Args: + save_path: Path to save the tokenizer + """ + + 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) + + @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 + + 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 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 + ] + + 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 + """ + # 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: List[Dict[str, str]], + system_prompt: Optional[str] = None, + tokenize: bool = True, + add_generation_prompt: 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 (auto-converted to first message). + tokenize: Whether to return token IDs (True) or raw string (False). + add_generation_prompt: Whether to add the generation prompt (default: True). + **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." + ) + + # Auto-convert system_prompt to first message if provided + if system_prompt: + messages = [{"role": "system", "content": system_prompt}] + list(messages) + + # Render the template + rendered = self._chat_template.render( + messages=messages, + add_generation_prompt=add_generation_prompt, + **kwargs, + ) + + if tokenize: + return self.encode(rendered) + + return rendered diff --git a/pipeline/tokenizer.py b/pipeline/tokenizer.py deleted file mode 100644 index 66ac2dd..0000000 --- a/pipeline/tokenizer.py +++ /dev/null @@ -1,149 +0,0 @@ -from tokenizers import Tokenizer, Encoding -from tokenizers import decoders, processors, normalizers, pre_tokenizers -from tokenizers.models import BPE -from tokenizers.trainers import BpeTrainer -from typing import List, Union, Optional, Tuple, Iterator - - -class BpeTokenizer: - def __init__(self, path: Optional[str] = None): - self._control_tokens = [ - "<｜begin▁of▁sentence｜>", - "<｜end▁of▁sentence｜>", - "<｜▁pad▁｜>", - ] - self._special_tokens = ["<｜im▁start｜>", "<｜im▁end｜>"] - - model = BPE() - self._tokenizer = Tokenizer(model) - self._tokenizer.normalizer = normalizers.Sequence( - [normalizers.NFC(), normalizers.Strip()] - ) - - self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( - [ - pre_tokenizers.UnicodeScripts(), - pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True), - ] - ) - - self._tokenizer.decoder = decoders.ByteLevel() - self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True) - - if path is not None: - self._tokenizer = Tokenizer.from_file(path) - - def _prepare_trainer( - self, - vocab_size: int, - min_freq: int, - reserved_token_size: int, - max_token_length: int = 18, - ) -> Tuple[BpeTrainer, int, List[str]]: - assert reserved_token_size > len(self._special_tokens) - reserved_tokens = [ - f"<｜reserve{i:02d}｜>" - for i in range(reserved_token_size - len(self._special_tokens)) - ] - detail_vocab_size = vocab_size - ( - len(reserved_tokens) + len(self._special_tokens) - ) - - alphabet = pre_tokenizers.ByteLevel.alphabet() - min_size = len(alphabet) + len(self._control_tokens) - assert detail_vocab_size > min_size - - trainer = BpeTrainer( - vocab_size=detail_vocab_size, - min_frequency=min_freq, - limit_alphabet=detail_vocab_size // 6, - max_token_length=max_token_length, - special_tokens=self._control_tokens + self._special_tokens, - initial_alphabet=alphabet, - show_progress=True, - ) - - return trainer, detail_vocab_size, reserved_tokens - - def train( - self, - files: List[str], - vocab_size: int, - min_freq: int, - reserved_token_size: int = 100, - ) -> None: - trainer, _, reserved_tokens = self._prepare_trainer( - vocab_size=vocab_size, - min_freq=min_freq, - reserved_token_size=reserved_token_size, - ) - self._tokenizer.train(files=files, trainer=trainer) - self._tokenizer.add_special_tokens( - self._control_tokens + self._special_tokens + reserved_tokens - ) - - def train_from_iterator( - self, - iterator: Iterator[str], - vocab_size: int, - min_freq: int, - reserved_token_size: int = 100, - ) -> None: - trainer, _, reserved_tokens = self._prepare_trainer( - vocab_size=vocab_size, - min_freq=min_freq, - reserved_token_size=reserved_token_size, - ) - self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer) - self._tokenizer.add_special_tokens( - self._control_tokens + self._special_tokens + reserved_tokens - ) - - def save(self, path: str) -> None: - self._tokenizer.save(path) - - def load(self, path: str) -> None: - self._tokenizer = Tokenizer.from_file(path) - - def encode( - self, - tokens: Union[str, List[str]], - out_ids: bool = True, - add_special_tokens: bool = False, - ) -> Union[List[int], List[str], List[List[int]], List[List[str]]]: - if isinstance(tokens, str): - encoded: Encoding = self._tokenizer.encode( - tokens, add_special_tokens=add_special_tokens - ) - return encoded.ids if out_ids else encoded.tokens - elif isinstance(tokens, list): - encoded_list: List[Encoding] = 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 - stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token] - return stop_ids - - @property - def bos_id(self) -> int: - return self._tokenizer.token_to_id("<｜begin▁of▁sentence｜>") - - @property - def eos_id(self) -> int: - return self._tokenizer.token_to_id("<｜end▁of▁sentence｜>") - - @property - def pad_id(self) -> int: - return self._tokenizer.token_to_id("<｜▁pad▁｜>") diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index c17b51d..0e0eaab 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -11,7 +11,7 @@ Usage: import argparse import os -from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl +from pipeline import AutoTokenizer, ProcessorFactory, cache_jsonl from pipeline.io import IOHandler @@ -61,7 +61,7 @@ def main(): if not os.path.exists(args.tokenizer): print(f"[ERROR] Tokenizer not found: {args.tokenizer}") return - tokenizer = BpeTokenizer(args.tokenizer) + tokenizer = AutoTokenizer(args.tokenizer) print(f"Tokenizer loaded: vocab_size={len(tokenizer)}") if args.strategy: diff --git a/tests/test_strategies.py b/tests/test_strategies.py index 03dd280..d4600b5 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -7,7 +7,6 @@ from pipeline.strategies import ( AlpacaStrategy, StrategyFactory, ) -from pipeline.tokenizer import BpeTokenizer class DummyTokenizer: @@ -58,7 +57,6 @@ class TestChatMLStrategy: text = _decode(response) assert "world" in text assert "<｜im▁end｜>" in text - assert "<｜end▁of▁sentence｜>" in text def test_prompt_ends_with_assistant_start(self): tk = DummyTokenizer()