Upload processor

processing_action_tokenizer.py

@@ -9,17 +9,6 @@ from tokenizers.trainers import BpeTrainer
 from transformers import PreTrainedTokenizerFast
 from transformers.processing_utils import ProcessorMixin
 
-ALPHABET = pre_tokenizers.ByteLevel.alphabet()
-REVERSE_ALPHABET = {c: i for i, c in enumerate(ALPHABET)}
-
-
-def _chr(i: int) -> str:
-    return ALPHABET[i]
-
-
-def _ord(c: str) -> int:
-    return REVERSE_ALPHABET[c]
-
 
 class UniversalActionProcessor(ProcessorMixin):
     attributes: ClassVar[list[str]] = ["bpe_tokenizer"]
@@ -67,7 +56,7 @@ class UniversalActionProcessor(ProcessorMixin):
         tokens = []
         for elem in dct_coeff:
             token_str = "".join(
-                map(_chr, np.maximum(elem.flatten() - self.min_token, 0).astype(int))
+                map(chr, np.maximum(elem.flatten() - self.min_token, 0).astype(int))
             )
             tokens.append(self.bpe_tokenizer(token_str)["input_ids"])
         return tokens
@@ -97,7 +86,7 @@ class UniversalActionProcessor(ProcessorMixin):
             try:
                 decoded_tokens = self.bpe_tokenizer.decode(token)
                 decoded_dct_coeff = (
-                    np.array(list(map(_ord, decoded_tokens))) + self.min_token
+                    np.array(list(map(ord, decoded_tokens))) + self.min_token
                 )
                 decoded_dct_coeff = decoded_dct_coeff.reshape(-1, self.action_dim)
                 assert decoded_dct_coeff.shape == (
@@ -159,21 +148,25 @@ class UniversalActionProcessor(ProcessorMixin):
                 tokens = dct_tokens.pop()
                 rounded_tokens = np.around(tokens * scale) - min_token
                 rounded_tokens = rounded_tokens.astype(int)
-                string = "".join(map(_chr, rounded_tokens))
+                string = "".join(map(chr, rounded_tokens))
                 yield string
 
         # Train BPE tokenizer
         tokenizer = Tokenizer(BPE())
-        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
+        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(
+            add_prefix_space=False, use_regex=False
+        )
         tokenizer.decoder = decoders.ByteLevel()
         tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)
 
+        alphabet = [chr(i) for i in range(256)]
         trainer = BpeTrainer(
             vocab_size=vocab_size,
             min_frequency=2,
             show_progress=True,
             special_tokens=[],
-            initial_alphabet=ALPHABET,
+            # initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+            initial_alphabet=alphabet,
             # max_token_length=256,
             max_token_length=10_000,
         )