refactor: 修改模型架构

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(

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 文档请参阅对应的源文件。

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",

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

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:

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]:

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

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｜>",

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]:

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",

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:

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",
+]

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]

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

pipeline/tokenizer.py

@@ -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▁｜>")

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:

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