fix broken padding setting in HuggingMazeTokenizer

maze_transformer/tokenizer.py

@@ -29,11 +29,23 @@ class HuggingMazeTokenizer(PreTrainedTokenizer):
     ]
 
     # Overwrite class attributes
-    padding_side = "left"
-    truncation_side = "left"  #! strange choice, but it's what we did in pad_sequence
+    # as of https://github.com/neelnanda-io/TransformerLens/pull/344 this gets overwritten to "right" on `HookedTransformer.__init__()`
+    # so, we have to do this overwriting in a weird way
+    _true_padding_side = "left"
+    _true_truncation_side = (
+        "left"  #! strange choice, but it's what we did in pad_sequence
+    )
 
     name_or_path = "hugging_maze_tokenizer"
 
+    def apply_overrides(self) -> None:
+        """Overwrite class attributes to deal with padding direction issues
+
+        see https://github.com/neelnanda-io/TransformerLens/pull/344
+        """
+        self.padding_side = self._true_padding_side
+        self.truncation_side = self._true_truncation_side
+
     def __init__(
         self,
         seq_len_max: int,
@@ -50,6 +62,7 @@ def __init__(
         self._vocab_size: int = maze_tokenizer.vocab_size
         self.vocab_size = self._vocab_size
         self._tokenizer_map = maze_tokenizer.tokenizer_map
+        self.apply_overrides()
 
         assert isinstance(
             seq_len_max, int

maze_transformer/training/config.py

@@ -408,6 +408,12 @@ class ConfigHolder(SerializableDataclass):
         loading_fn=_load_maze_tokenizer,
     )
 
+    _tokenizer: PreTrainedTokenizer | None = serializable_field(
+        default=None,
+        serialization_fn=lambda x: str(x),
+        loading_fn=lambda data: None,
+    )
+
     # shortcut properties
     @property
     def d_model(self) -> int:
@@ -433,6 +439,8 @@ def __post_init__(self):
         # fallback to default maze tokenizer if no kwargs are provided
         if self.pretrainedtokenizer_kwargs is None:
             if self.maze_tokenizer is None:
+                # TODO: is this the right default? maybe set it to AOTP_UT_rasterized
+                # since thats what legacy models are likely to be?
                 self.maze_tokenizer = MazeTokenizer(
                     tokenization_mode=TokenizationMode.AOTP_UT_uniform,
                     max_grid_size=None,
@@ -460,18 +468,21 @@ def summary(self) -> str:
     def seed(self) -> int:
         return self.dataset_cfg.seed
 
-    @cached_property
+    @property
     def tokenizer(self) -> PreTrainedTokenizer:
         """get a tokenizer via a pretrainedtokenizer_kwargs, or a hugging maze tokenizer"""
-        if self.pretrainedtokenizer_kwargs is not None:
-            return PreTrainedTokenizer(**self.pretrainedtokenizer_kwargs)
-        elif self.maze_tokenizer is not None:
-            return HuggingMazeTokenizer(
-                seq_len_max=self.dataset_cfg.seq_len_max,
-                maze_tokenizer=self.maze_tokenizer,
-            )
+        if self._tokenizer is None:
+            if self.pretrainedtokenizer_kwargs is not None:
+                return PreTrainedTokenizer(**self.pretrainedtokenizer_kwargs)
+            elif self.maze_tokenizer is not None:
+                return HuggingMazeTokenizer(
+                    seq_len_max=self.dataset_cfg.seq_len_max,
+                    maze_tokenizer=self.maze_tokenizer,
+                )
+            else:
+                raise ValueError("no tokenizer specified")
         else:
-            raise ValueError("no tokenizer specified")
+            return self._tokenizer
 
     @cached_property
     def hooked_transformer_cfg(self) -> HookedTransformerConfig:
@@ -585,6 +596,19 @@ def __init__(self, cfg_holder: ConfigHolder) -> None:
             tokenizer=cfg_holder.tokenizer,
         )
 
+        # update the tokenizer attributes (evil)
+        # see `apply_overrides()` code for info
+        if isinstance(self.zanj_model_config.tokenizer, HuggingMazeTokenizer):
+            self.zanj_model_config.tokenizer.apply_overrides()
+            self.set_tokenizer(
+                self.zanj_model_config.tokenizer,
+                default_padding_side=self.zanj_model_config.tokenizer.padding_side,
+            )
+        else:
+            warnings.warn(
+                "tokenizer is not a HuggingMazeTokenizer, so we can't apply overrides. this might break padding and your whole model"
+            )
+
     @property
     def config(self) -> ConfigHolder:
         return self.zanj_model_config