refactor: 从data 模块分离tokenizer

2026-04-04 16:12:58 +08:00
parent b531232a9b
commit bd9741dc5f
16 changed files with 108 additions and 92 deletions
@@ -0,0 +1,23 @@
+from astrai.tokenizer.tokenizer import (
+    BaseTokenizer,
+    BpeTokenizer,
+    BaseTrainer,
+    BpeTrainer,
+)
+from astrai.tokenizer.chat_template import (
+    HistoryType,
+    MessageType,
+    CHAT_TEMPLATES,
+    build_prompt,
+)
+
+__all__ = [
+    "BaseTokenizer",
+    "BpeTokenizer",
+    "BaseTrainer",
+    "BpeTrainer",
+    "HistoryType",
+    "MessageType",
+    "CHAT_TEMPLATES",
+    "build_prompt",
+]
@@ -0,0 +1,67 @@
+from typing import Dict, List, Optional, Tuple
+from jinja2 import Template
+
+HistoryType = List[Tuple[str, str]]
+MessageType = Dict[str, str]
+
+# Predefined chat templates using jinja2
+CHAT_TEMPLATES: Dict[str, str] = {
+    "chatml": """{%- if system_prompt -%}
+<｜im▁start｜>system
+{{ system_prompt }}<｜im▁end｜>
+{%- endif -%}
+{%- for message in messages -%}
+<｜im▁start｜>{{ message['role'] }}
+{{ message['content'] }}<｜im▁end｜>
+{%- endfor -%}
+<｜im▁start｜>assistant
+""",
+}
+
+
+def build_prompt(
+    query: str,
+    system_prompt: Optional[str] = None,
+    history: Optional[HistoryType] = None,
+    template: Optional[str] = None,
+) -> str:
+    """Build prompt using jinja2 template for query and history.
+
+    Args:
+        query (str): query string.
+        system_prompt (Optional[str]): system prompt string.
+        history (Optional[HistoryType]): history list of query and response.
+        template (Optional[str]): jinja2 template string. If None, uses default chatml template.
+
+    Returns:
+        str: prompt string formatted according to the template.
+
+    Example:
+        # Use default template
+        prompt = build_prompt(query="Hello", history=[...])
+
+        # Use custom template
+        custom_template = '''
+        {%- for msg in messages -%}
+        {{ msg['role'] }}: {{ msg['content'] }}
+        {%- endfor -%}
+        '''
+        prompt = build_prompt(query="Hello", template=custom_template)
+    """
+    # 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})
+
+    # Use provided template or default chatml template
+    template_str = template if template is not None else CHAT_TEMPLATES["chatml"]
+
+    # Render template
+    jinja_template = Template(template_str)
+    return jinja_template.render(
+        messages=messages,
+        system_prompt=system_prompt,
+    )
@@ -0,0 +1,212 @@
+from abc import ABC, abstractmethod
+from typing import List, Union
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer as BpeTrainerImpl
+
+
+class BaseTokenizer(ABC):
+    @abstractmethod
+    def _init_tokenizer(self):
+        pass
+
+    @abstractmethod
+    def save(self, path):
+        pass
+
+    @abstractmethod
+    def load(self, path):
+        pass
+
+    @abstractmethod
+    def encode(
+        self,
+        tokens: Union[str, List[str]],
+        out_ids: bool = True,
+        add_special_tokens: bool = False,
+    ) -> List:
+        pass
+
+    @abstractmethod
+    def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
+        pass
+
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def stop_ids(self) -> List[int]:
+        pass
+
+    @property
+    @abstractmethod
+    def bos_id(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def eos_id(self) -> int:
+        pass
+
+    @property
+    @abstractmethod
+    def pad_id(self) -> int:
+        pass
+
+
+class BaseTrainer(ABC):
+    def __init__(self, tokenizer: BaseTokenizer):
+        self.tokenizer = tokenizer
+
+    @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,
+    ):
+        self._control_tokens = control_tokens or [
+            "<｜begin▁of▁sentence｜>",
+            "<｜end▁of▁sentence｜>",
+            "<｜▁pad▁｜>",
+        ]
+        self._special_tokens = special_tokens or [
+            "<｜im▁start｜>",
+            "<｜im▁end｜>",
+        ]
+        self._tokenizer = None
+        self._init_tokenizer()
+        if path is not None:
+            self.load(path)
+
+    def _init_tokenizer(self):
+        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)
+
+    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
+        )