From 324fe089c8f870fb0622b23dd7c5c2c633cd76be Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 6 Aug 2025 13:49:08 +0800
Subject: [PATCH] =?UTF-8?q?refactor(tokenizer):=20=E9=87=8D=E6=9E=84=20Bpe?=
 =?UTF-8?q?Tokenizer=20=E7=B1=BB=E5=B9=B6=E4=BC=98=E5=8C=96=E7=BC=96?=
 =?UTF-8?q?=E7=A0=81=E6=B5=81=E7=A8=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 modules/tokenizer.py | 48 ++++++++++++++++++--------------------------
 1 file changed, 20 insertions(+), 28 deletions(-)

diff --git a/modules/tokenizer.py b/modules/tokenizer.py
index ba52a31..107f001 100644
--- a/modules/tokenizer.py
+++ b/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):
@@ -32,32 +31,30 @@ 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
+        min_size = len(alphabet) + len(self._control_tokens)
+        assert detail_vocab_size > min_size
         
         trainer = BpeTrainer(
-            vocab_size=vocab_size,
-            min_frequency=min_freq, 
-            limit_alphabet= vocab_size // 4,
+            vocab_size=detail_vocab_size,
+            min_frequency=min_freq,
+            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:
-            with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
-                encodings: List[Encoding] = list(executor.map(self._tokenizer.encode, tokens))
-            
-            if out_ids:
-                return [encoding.ids for encoding in encodings]
-            else:
-                return [encoding.tokens for encoding in encodings]
+        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=True) -> str:
+    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: