format

understanding-search · Aug 6, 2023 · 16bc51e · 16bc51e
1 parent 144daa8
commit 16bc51e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/maze_transformer/evaluation/eval_model.py b/maze_transformer/evaluation/eval_model.py
@@ -208,7 +208,9 @@ def evaluate_model(
     }
 
     if dataset_tokens is None:
-        dataset_tokens = dataset.as_tokens(model.zanj_model_config.maze_tokenizer, join_tokens_individual_maze=False)
+        dataset_tokens = dataset.as_tokens(
+            model.zanj_model_config.maze_tokenizer, join_tokens_individual_maze=False
+        )
     else:
         assert len(dataset) == len(
             dataset_tokens

diff --git a/maze_transformer/evaluation/plot_attention.py b/maze_transformer/evaluation/plot_attention.py
@@ -15,8 +15,8 @@
 from circuitsvis.tokens import colored_tokens_multi
 from jaxtyping import Float
 from maze_dataset import CoordTup, MazeDataset, MazeDatasetConfig, SolvedMaze
-from maze_dataset.tokenization.token_utils import coord_str_to_tuple_noneable
 from maze_dataset.plotting import MazePlot
+from maze_dataset.tokenization.token_utils import coord_str_to_tuple_noneable
 
 # Utilities
 from muutils.json_serialize import SerializableDataclass, serializable_dataclass

diff --git a/maze_transformer/training/training.py b/maze_transformer/training/training.py
@@ -101,8 +101,7 @@ def train(
             evals_enabled = False
 
         val_dataset_tokens: list[list[str]] = val_dataset.as_tokens(
-            model.zanj_model_config.maze_tokenizer,
-            join_tokens_individual_maze=False
+            model.zanj_model_config.maze_tokenizer, join_tokens_individual_maze=False
         )
 
     # compute intervals