Upload processor

processing_action_tokenizer.py

@@ -9,6 +9,17 @@ 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"]
@@ -56,7 +67,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
@@ -86,7 +97,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 == (
@@ -148,27 +159,9 @@ 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
-        # bpe = ByteLevelBPETokenizer()
-
-        # # Set up the entire range of possible tokens as the initial alphabet
-        # 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,
-        #     max_token_length=10000,
-        # )
-
-        # # 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)
-
         # Train BPE tokenizer
         tokenizer = Tokenizer(BPE())
         tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(
@@ -181,15 +174,12 @@ class UniversalActionProcessor(ProcessorMixin):
             add_prefix_space=False, trim_offsets=False, use_regex=False
         )
 
-        # Set up the entire range of possible tokens as the initial alphabet
-        alphabet = [chr(i) for i in range(256)]
         trainer = BpeTrainer(
             vocab_size=vocab_size,
             min_frequency=2,
             show_progress=True,
             special_tokens=[],
-            # initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
-            initial_alphabet=alphabet,
+            initial_alphabet=ALPHABET,
             # max_token_length=256,
             max_token_length=10_000,
         )

processor_config.json

@@ -7,5 +7,5 @@
   "processor_class": "UniversalActionProcessor",
   "scale": 10.0,
   "time_horizon": 50,
-  "vocab_size": 2048
+  "vocab_size": 1024
 }