refactor(tokenizer): 重构 BpeTokenizer 类并优化编码流程

modules/tokenizer.py

@@ -2,9 +2,8 @@ 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
-import concurrent.futures
+
 
 class BpeTokenizer:
     def __init__(self, path=None):
@@ -33,31 +32,29 @@ class BpeTokenizer:
         if path is not None:
             self._tokenizer = Tokenizer.from_file(path)
         
-    def __init_trainer(self, vocab_size, min_freq):
+    def _prepare_trainer(self, vocab_size: int, min_freq: int, reserved_token_size: int) -> tuple:
+        assert reserved_token_size > len(self._special_tokens)
+        reserved_tokens = [f"<|rsv{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 vocab_size > min_size
+        assert detail_vocab_size > min_size
         
         trainer = BpeTrainer(
-            vocab_size=vocab_size,
+            vocab_size=detail_vocab_size,
             min_frequency=min_freq,
-            limit_alphabet= vocab_size // 4,
+            limit_alphabet=detail_vocab_size // 2,
             max_token_length=18,
             special_tokens=self._control_tokens,
             show_progress=True,
             initial_alphabet=alphabet,
         )
-        return trainer
         
-    def _prepare_trainer_and_tokens(self, vocab_size: int, min_freq: int, reserved_token_size: int) -> tuple:
-        assert reserved_token_size > len(self._special_tokens)
-        reserved_tokens = [f"<|rsv{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))
-        trainer = self.__init_trainer(docab_size=detail_vocab_size, min_freq=min_freq)
         return trainer, detail_vocab_size, reserved_tokens
 
     def train(self, files, vocab_size, min_freq, reserved_token_size=100):
-        trainer, _, reserved_tokens = self._prepare_trainer_and_tokens(
+        trainer, _, reserved_tokens = self._prepare_trainer(
             vocab_size=vocab_size,
             min_freq=min_freq,
             reserved_token_size=reserved_token_size
@@ -66,7 +63,7 @@ class BpeTokenizer:
         self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
             
     def train_from_iterator(self, iterator, vocab_size, min_freq, reserved_token_size=100):
-        trainer, _, reserved_tokens = self._prepare_trainer_and_tokens(
+        trainer, _, reserved_tokens = self._prepare_trainer(
             vocab_size=vocab_size,
             min_freq=min_freq,
             reserved_token_size=reserved_token_size
@@ -80,20 +77,15 @@ class BpeTokenizer:
     def load(self, path):
         self._tokenizer = Tokenizer.from_file(path)
 
-    def encode(self, tokens: Union[str, List[str]], out_ids=True, num_threads=4) -> List:
+    def encode(self, tokens: Union[str, List[str]], out_ids: bool=True, add_special_tokens: bool=False) -> List:
         if isinstance(tokens, str):
-            encoded: Encoding = self._tokenizer.encode(tokens)
+            encoded: Encoding = self._tokenizer.encode(tokens, add_special_tokens=add_special_tokens)
             return encoded.ids if out_ids else encoded.tokens
-        else:
+        elif isinstance(tokens, list):
-            with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
+            encoded_list: List[Encoding] = self._tokenizer.encode_batch(tokens, add_special_tokens=add_special_tokens)
-                encodings: List[Encoding] = list(executor.map(self._tokenizer.encode, tokens))
+            return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
 
-            if out_ids:
+    def decode(self, tokens: List[int], skip_special_tokens: bool=True) -> str:
-                return [encoding.ids for encoding in encodings]
-            else:
-                return [encoding.tokens for encoding in encodings]
-
-    def decode(self, tokens: List[int], skip_special_tokens=True) -> str:
         return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
     
     def __len__(self) -> int: