Support maze dataset tokenizers update

.gitignore

@@ -7,6 +7,7 @@ tests/_temp/**
 tests/**/_temp/**
 notebooks/data/**
 notebooks/plots/**
+# maze-dataset_test/**
 
 .coverage
 htmlcov/

makefile

@@ -79,6 +79,12 @@ clean:
 	python -Bc "import pathlib; [p.rmdir() for p in pathlib.Path('.').rglob('__pycache__')]"
 
 
+.PHONY: test_with_branch
+test_with_branch:
+	@echo "creating test environment"
+	cp pyproject.toml poetry.lock maze-dataset_test
+
+
 # listing targets, from stackoverflow
 # https://stackoverflow.com/questions/4219255/how-do-you-get-the-list-of-targets-in-a-makefile
 .PHONY: help

maze_transformer/evaluation/baseline_models.py

@@ -11,7 +11,7 @@
     LatticeMaze,
     SolvedMaze,
 )
-from maze_dataset.tokenization.token_utils import (
+from maze_dataset.token_utils import (
     get_origin_tokens,
     get_path_tokens,
     get_target_tokens,

maze_transformer/evaluation/eval_model.py

@@ -14,8 +14,9 @@
     MazeDatasetConfig,
     SolvedMaze,
 )
-from maze_dataset.tokenization import MazeTokenizer
-from maze_dataset.tokenization.token_utils import (
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
+from maze_dataset.token_utils import (
+    WhenMissing,
     get_context_tokens,
     get_path_tokens,
     remove_padding_from_token_str,
@@ -143,7 +144,7 @@ def predict_maze_paths(
             smart_max_new_tokens
         ), "if max_new_tokens is None, smart_max_new_tokens must be True"
 
-    maze_tokenizer: MazeTokenizer = model.tokenizer._maze_tokenizer
+    maze_tokenizer: MazeTokenizer | MazeTokenizerModular = model.config.maze_tokenizer
 
     contexts_lists: list[list[str]] = [
         get_context_tokens(tokens) for tokens in tokens_batch

maze_transformer/evaluation/eval_single_token_tasks.py

@@ -12,7 +12,7 @@
 # Our Code
 # dataset stuff
 from maze_dataset import MazeDataset
-from maze_dataset.tokenization import MazeTokenizer
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
 from muutils.json_serialize import SerializableDataclass, serializable_dataclass
 
 # TransformerLens imports
@@ -47,7 +47,7 @@ class TaskEvalResult(SerializableDataclass):
 
 def get_task_prompts_targets(
     dataset: MazeDataset,
-    maze_tokenizer: MazeTokenizer,
+    maze_tokenizer: MazeTokenizer | MazeTokenizerModular,
     tasks: dict[str, DLAProtocolFixed] = LOGIT_ATTRIB_TASKS,
 ) -> dict[str, TaskPrompt]:
     dataset_tokens: list[list[str]] = dataset.as_tokens(
@@ -63,7 +63,7 @@ def eval_model_task(
     task: TaskPrompt,
     do_cache: bool = False,
 ) -> TaskEvalResult:
-    maze_tokenizer: MazeTokenizer = model.tokenizer._maze_tokenizer
+    maze_tokenizer: MazeTokenizer | MazeTokenizerModular = model.config.maze_tokenizer
 
     prompts_joined: list[str] = [" ".join(prompt) for prompt in task.prompts]
 

maze_transformer/mechinterp/direct_logit_attribution.py

@@ -13,7 +13,7 @@
 
 # maze-datset stuff
 from maze_dataset import MazeDataset, MazeDatasetConfig
-from maze_dataset.tokenization import MazeTokenizer
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
 
 # TransformerLens imports
 from transformer_lens import ActivationCache
@@ -226,7 +226,7 @@ def create_report(
     # model and tokenizer
     if not isinstance(model, ZanjHookedTransformer):
         model = ZanjHookedTransformer.read(model)
-    tokenizer: MazeTokenizer = model.zanj_model_config.maze_tokenizer
+    tokenizer: MazeTokenizer | MazeTokenizerModular = model.zanj_model_config.maze_tokenizer
 
     # dataset cfg
     if dataset_cfg_source is None:

maze_transformer/mechinterp/plot_attention.py

@@ -18,8 +18,8 @@
 from maze_dataset.plotting import MazePlot
 from maze_dataset.plotting.plot_tokens import plot_colored_text
 from maze_dataset.plotting.print_tokens import color_tokens_cmap
-from maze_dataset.tokenization import MazeTokenizer
-from maze_dataset.tokenization.util import coord_str_to_tuple_noneable
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
+from maze_dataset.token_utils import coord_str_to_tuple_noneable
 
 # Utilities
 from muutils.json_serialize import SerializableDataclass, serializable_dataclass
@@ -377,7 +377,7 @@ def mazeplot_attention(
 def plot_attn_dist_correlation(
     tokens_context: list[list[str]],
     tokens_dist_to: list[str],  # either current or target token for each maze
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
     attention: Float[np.ndarray, "n_mazes n_tokens"],
     ax: plt.Axes | None = None,
     respect_topology: bool = False,  # manhattan distance if False
@@ -480,7 +480,7 @@ def plot_attention_final_token(
     prompts: list[list[str]],
     targets: list[str],
     mazes: list[SolvedMaze],
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
     n_mazes: int = 5,
     last_n_tokens: int = 20,
     # exponentiate_scores: bool = False,

maze_transformer/mechinterp/plot_logits.py

@@ -6,7 +6,7 @@
 from maze_dataset import CoordTup
 
 # Our Code
-from maze_dataset.tokenization import MazeTokenizer
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
 
 _DEFAULT_SUBPLOTS_KWARGS: dict = dict(
     figsize=(20, 20),
@@ -86,7 +86,7 @@ def plot_logit_histograms(
 
 def get_baseline_incorrect_group(
     prompts: list[list[str]],
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
     baseline: "RandomBaseline",
 ) -> Bool[torch.Tensor, "n_mazes d_vocab"]:
     """
@@ -116,7 +116,7 @@ def get_baseline_incorrect_group(
 def plot_logits(
     last_tok_logits: Float[torch.Tensor, "n_mazes d_vocab"],
     target_idxs: Int[torch.Tensor, "n_mazes"],
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
     n_bins: int = 50,
     mark_incorrect: bool = False,
     mark_correct: bool = True,

maze_transformer/mechinterp/residual_stream_structure.py

@@ -11,8 +11,8 @@
 
 # maze_dataset
 from maze_dataset.constants import _SPECIAL_TOKENS_ABBREVIATIONS
-from maze_dataset.tokenization import MazeTokenizer
-from maze_dataset.tokenization.util import strings_to_coords
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
+from maze_dataset.token_utils import strings_to_coords
 
 # scipy
 from scipy.spatial.distance import pdist, squareform
@@ -52,7 +52,7 @@ def coordinate_to_color(
 )
 
 
-def process_tokens_for_pca(tokenizer: MazeTokenizer) -> list[TokenPlottingInfo]:
+def process_tokens_for_pca(tokenizer: MazeTokenizer | MazeTokenizerModular) -> list[TokenPlottingInfo]:
     tokens_coords: list[str | tuple[int, int]] = strings_to_coords(
         tokenizer.token_arr, when_noncoord="include"
     )
@@ -227,7 +227,7 @@ def abs_dot_product(u, v):
 
 def compute_distances_and_correlation(
     embedding_matrix: Float[np.ndarray, "d_vocab d_model"],
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
     embedding_metric: str = "cosine",
     coordinate_metric: str = "euclidean",
     show: bool = True,
@@ -277,7 +277,7 @@ def compute_distances_and_correlation(
 
 def plot_distances_matrix(
     embedding_distances_matrix: Float[np.ndarray, "n_coord_tokens n_coord_tokens"],
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
     embedding_metric: str,
     show: bool = True,
     **kwargs,
@@ -313,7 +313,7 @@ def plot_distances_matrix(
 
 def compute_grid_distances(
     embedding_distances_matrix: Float[np.ndarray, "n_coord_tokens n_coord_tokens"],
-    tokenizer: MazeTokenizer,
+    tokenizer: MazeTokenizer | MazeTokenizerModular,
 ) -> Float[np.ndarray, "n n n n"]:
     n: int = tokenizer.max_grid_size
     grid_distances: Float[np.ndarray, "n n n n"] = np.full((n, n, n, n), np.nan)

maze_transformer/tokenizer.py

@@ -4,7 +4,7 @@
 import torch
 from maze_dataset import SPECIAL_TOKENS, LatticeMaze
 from maze_dataset.plotting import MazePlot
-from maze_dataset.tokenization import MazeTokenizer
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
 from muutils.tensor_utils import ATensor, NDArray
 from transformers import PreTrainedTokenizer
 from transformers.tokenization_utils import BatchEncoding
@@ -46,13 +46,13 @@ def apply_overrides(self) -> None:
     def __init__(
         self,
         seq_len_max: int,
-        maze_tokenizer: MazeTokenizer,
+        maze_tokenizer: MazeTokenizer | MazeTokenizerModular,
         **kwargs,
     ) -> None:
         """extension of PreTrainedTokenizer for mazes. takes maximum sequence length and maze_tokenizer. also, kwargs are passed to super `PreTrainedTokenizer`"""
         super().__init__(max_len=seq_len_max, **kwargs)
 
-        self._maze_tokenizer: MazeTokenizer = maze_tokenizer
+        self._maze_tokenizer: MazeTokenizer | MazeTokenizerModular = maze_tokenizer
         token_arr: list[str] = maze_tokenizer.token_arr
         self._token_arr: list[str] = token_arr
         self._seq_len_max: int = seq_len_max
@@ -81,8 +81,9 @@ def __init__(
 
         # We are having to do evil things here
         vocab: dict[str, int] = {token: i for i, token in enumerate(token_arr)}
-        vocab[self.unk_token] = len(vocab)
-        self.vocab: dict[str, int] = vocab
+        if self.unk_token not in vocab:  # maze-dataset ^1.0.0 includes <UNK> already
+            vocab[self.unk_token] = len(vocab)
+            self.vocab: dict[str, int] = vocab
 
         special_tokens = list(SPECIAL_TOKENS.values())
         normal_tokens = [x for x in token_arr if x not in special_tokens]

maze_transformer/training/config.py

@@ -10,7 +10,7 @@
 import torch
 from maze_dataset.dataset.configs import MAZE_DATASET_CONFIGS
 from maze_dataset.dataset.dataset import GPTDatasetConfig
-from maze_dataset.tokenization import MazeTokenizer, TokenizationMode
+from maze_dataset.tokenization import MazeTokenizer, TokenizationMode, MazeTokenizerModular
 from muutils.dictmagic import kwargs_to_nested_dict
 from muutils.json_serialize import (
     JSONitem,
@@ -370,17 +370,14 @@ def summary(self) -> dict:
 }
 
 
-def _load_maze_tokenizer(data: dict) -> MazeTokenizer:
+def _load_maze_tokenizer(data: dict) -> MazeTokenizerModular:
     """load the maze tokenizer, including vocab size from a legacy config"""
     if "maze_tokenizer" in data:
         # new style tokenizer
-        return load_item_recursive(data["maze_tokenizer"], path=tuple("maze_tokenizer"))
+        return MazeTokenizerModular.from_legacy(load_item_recursive(data["maze_tokenizer"], path=tuple("maze_tokenizer")))
     else:
         if "token_arr" in data["dataset_cfg"]:
-            output: MazeTokenizer = MazeTokenizer(
-                tokenization_mode=TokenizationMode.AOTP_UT_rasterized,
-                max_grid_size=None,
-            )
+            output: MazeTokenizerModular = MazeTokenizerModular()
         else:
             raise ValueError("Could not find vocab size in legacy config")
 
@@ -405,7 +402,7 @@ class ConfigHolder(SerializableDataclass):
     pretrainedtokenizer_kwargs: dict[str, JSONitem] | None = serializable_field(
         default=None
     )
-    maze_tokenizer: MazeTokenizer | None = serializable_field(
+    maze_tokenizer: MazeTokenizer | MazeTokenizerModular | None = serializable_field(
         default_factory=lambda: None,
         loading_fn=_load_maze_tokenizer,
     )
@@ -434,24 +431,22 @@ def n_heads(self) -> int:
         return self.model_cfg.n_heads
 
     def _set_tok_gridsize_from_dataset(self):
-        self.maze_tokenizer.max_grid_size = self.dataset_cfg.max_grid_n
-        self.maze_tokenizer.clear_cache()
+        if isinstance(self.maze_tokenizer, MazeTokenizer):
+            self.maze_tokenizer.max_grid_size = self.dataset_cfg.max_grid_n
+            self.maze_tokenizer.clear_cache()
 
     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,
-                )
+                self.maze_tokenizer = MazeTokenizerModular()
 
         # update the config of the maze tokenizer if there is no grid size
         # since we need the token array for the vocab size of the model
         if self.maze_tokenizer is not None:
-            if self.maze_tokenizer.max_grid_size is None:
+            if getattr(self.maze_tokenizer, "max_grid_size", None) is None:
                 self._set_tok_gridsize_from_dataset()
 
     def summary(self) -> str:

maze_transformer/training/train_save_files.py

@@ -6,7 +6,7 @@
 from maze_transformer.training.config import ConfigHolder
 
 
-@freeze
+# @freeze
 class TRAIN_SAVE_FILES:
     """namespace for filenames/formats for saving training data"""
 

maze_transformer/training/training.py

@@ -5,7 +5,7 @@
 import torch
 from jaxtyping import Float
 from maze_dataset import MazeDataset, SolvedMaze
-from maze_dataset.tokenization import MazeTokenizer
+from maze_dataset.tokenization import MazeTokenizer, MazeTokenizerModular
 from muutils.statcounter import StatCounter
 from torch.utils.data import DataLoader
 from transformer_lens.HookedTransformer import SingleLoss
@@ -19,7 +19,7 @@
 from maze_transformer.training.wandb_logger import WandbLogger
 
 
-def collate_batch(batch: list[SolvedMaze], maze_tokenizer: MazeTokenizer) -> list[str]:
+def collate_batch(batch: list[SolvedMaze], maze_tokenizer: MazeTokenizer | MazeTokenizerModular) -> list[str]:
     return [" ".join(maze.as_tokens(maze_tokenizer)) for maze in batch]
 
 

notebooks/appendix_figures.ipynb

@@ -23,7 +23,7 @@
     "# dataset stuff\n",
     "from maze_dataset import MazeDataset, MazeDatasetConfig, SolvedMaze, LatticeMaze, SPECIAL_TOKENS, LatticeMazeGenerators\n",
     "from maze_dataset.plotting import MazePlot, PathFormat\n",
-    "from maze_dataset.tokenization import MazeTokenizer, TokenizationMode\n",
+    "from maze_dataset.tokenization import MazeTokenizer, TokenizationMode, MazeTokenizerModular\n",
     "from maze_dataset.plotting.print_tokens import color_maze_tokens_AOTP\n",
     "\n",
     "# model stuff\n",
@@ -139,7 +139,7 @@
     "\tfig.savefig(plot_dir / \"rollouts.pdf\", bbox_inches=\"tight\")\n",
     "\tplt.show()\n",
     "\n",
-    "\ttokenizer: MazeTokenizer = model.zanj_model_config.maze_tokenizer\n",
+    "\ttokenizer: MazeTokenizer | MazeTokenizerModular = model.zanj_model_config.maze_tokenizer\n",
     "\ttask_prompts_targets: dict[str, TaskPrompt] = get_task_prompts_targets(\n",
     "\t\tdataset=dataset,\n",
     "\t\tmaze_tokenizer=tokenizer,\n",

notebooks/demo_dataset.ipynb

@@ -531,7 +531,7 @@
    "source": [
     "\n",
     "from maze_dataset.plotting import MazePlot\n",
-    "from maze_dataset.tokenization import MazeTokenizer, TokenizationMode\n",
+    "from maze_dataset.tokenization import MazeTokenizer, TokenizationMode, MazeTokenizerModular\n",
     "from maze_dataset.plotting.print_tokens import display_color_maze_tokens_AOTP, color_maze_tokens_AOTP\n",
     "\n",
     "maze: SolvedMaze = dataset[0]\n",
@@ -549,7 +549,7 @@
     "# as tokens\n",
     "\n",
     "# first, initialize a tokenizer -- more about this in the `notebooks/demo_tokenization.ipynb` notebook\n",
-    "tokenizer: MazeTokenizer = MazeTokenizer(tokenization_mode=TokenizationMode.AOTP_UT_rasterized, max_grid_size=100)\n",
+    "tokenizer: MazeTokenizerModular = MazeTokenizerModular()\n",
     "maze_tok = maze.as_tokens(maze_tokenizer=tokenizer)\n",
     "\n",
     "# you can view the tokens directly\n",

notebooks/direct_logit_attribution.ipynb

@@ -89,7 +89,7 @@
     "# Our Code\n",
     "# dataset stuff\n",
     "from maze_dataset import MazeDataset, MazeDatasetConfig, SolvedMaze, LatticeMaze, SPECIAL_TOKENS, LatticeMazeGenerators\n",
-    "from maze_dataset.tokenization import MazeTokenizer, TokenizationMode\n",
+    "from maze_dataset.tokenization import MazeTokenizer, TokenizationMode, MazeTokenizerModular\n",
     "from maze_dataset.plotting.print_tokens import color_maze_tokens_AOTP\n",
     "\n",
     "# model stuff\n",
@@ -287,7 +287,7 @@
     }
    ],
    "source": [
-    "TOKENIZER: MazeTokenizer = MODEL.zanj_model_config.maze_tokenizer\n",
+    "TOKENIZER: MazeTokenizer | MazeTokenizerModular = MODEL.zanj_model_config.maze_tokenizer\n",
     "DATASET_TOKENS: list[list[str]] = DATASET.as_tokens(TOKENIZER, join_tokens_individual_maze=False)\n",
     "\n",
     "# print some info\n",