Upload processor

processing_action_tokenizer.py

@@ -1,12 +1,14 @@
 import logging
+from sre_parse import Tokenizer
 from typing import ClassVar
 
 import numpy as np
 from scipy.fft import dct, idct
-from tokenizers import ByteLevelBPETokenizer
+from tokenizers import ByteLevelBPETokenizer, AddedToken, Tokenizer, decoders, pre_tokenizers, processors, trainers
 from tokenizers.trainers import BpeTrainer
 from transformers.processing_utils import ProcessorMixin
 from transformers import PreTrainedTokenizerFast
+from tokenizers.models import BPE
 
 
 class UniversalActionProcessor(ProcessorMixin):
@@ -131,22 +133,27 @@ class UniversalActionProcessor(ProcessorMixin):
                 yield string
 
         # Train BPE tokenizer
-        bpe = ByteLevelBPETokenizer()
-
+        # bpe = ByteLevelBPETokenizer()
+        tokenizer = Tokenizer(BPE())
+        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
+        tokenizer.decoder = decoders.ByteLevel()
+        tokenizer.post_processor = processors.ByteLevel()
+        
         # Set up the entire range of possible tokens as the initial alphabet
-        alphabet = [chr(i) for i in range(max_token - min_token + 1)]
+        # alphabet = [chr(i) for i in range(max_token - min_token + 1)]
         trainer = BpeTrainer(
             vocab_size=vocab_size,
             min_frequency=2,
             show_progress=True,
             special_tokens=[],
-            initial_alphabet=alphabet,
+            initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
             max_token_length=256,
         )
+        tokenizer.train_from_iterator(_token_iter(), trainer=trainer, length=len(dct_tokens))
 
         # Train the inner tokenizer (don't use ByteLevelBPETokenizer.train_from_iterator()
         # because it doesn't support custom alphabets)
-        bpe._tokenizer.train_from_iterator(_token_iter(), trainer=trainer)
+        # bpe._tokenizer.train_from_iterator(_token_iter(), trainer=trainer, length=len(dct_tokens))
         # trainer.train_from_iterator(_token_iter(), trainer=trainer)
         # bpe.train_from_iterator(
         #     _token_iter(), 
@@ -157,7 +164,7 @@ class UniversalActionProcessor(ProcessorMixin):
         # )
 
         return cls(
-            PreTrainedTokenizerFast(tokenizer_object=bpe, clean_up_tokenization_spaces=False),
+            PreTrainedTokenizerFast(tokenizer_object=tokenizer, clean_up_tokenization_spaces=False),
             scale=scale,
             vocab_size=vocab_size,
             min_token=min_token,