Zanj integration: datasets & training

maze_transformer/evaluation/eval_model.py

@@ -50,6 +50,7 @@ def find_config(folder: Path) -> Path | tuple[Path, Path] | None:
 def load_model_with_configs(
     model_path: Path,
     verbose: bool = False,
+    fold_ln: bool = True,
 ) -> tuple[HookedTransformer, ConfigHolder]:
     """
     Load a model and associated config files from a path.
@@ -89,7 +90,7 @@ def load_model_with_configs(
     # will complain about the fact that we deleted layernorm from the state_dict
     # NOTE temporary fix until https://github.com/neelnanda-io/TransformerLens/issues/219 is resolved
 
-    model.process_weights_(fold_ln=True)
+    model.process_weights_(fold_ln=fold_ln)
     model.setup()  # Re-attach layernorm hooks by calling setup
     model.eval()
 
@@ -152,8 +153,9 @@ def evaluate_model(
     for batch in chunks(dataset.mazes_tokens, batch_size):
         # TODO: This won't be needed after #124, then we can call mazes_objs instead
         # https://github.com/orgs/AISC-understanding-search/projects/1/views/1?pane=issue&itemId=23879308
-        solved_mazes = [SolvedMaze.from_tokens(tokens, dataset.cfg) for tokens in batch]
-        mazes, solutions = zip(*solved_mazes)
+        solved_mazes: SolvedMaze = [
+            SolvedMaze.from_tokens(tokens, dataset.cfg) for tokens in batch
+        ]
 
         predictions = predict_maze_paths(
             tokens_batch=batch,
@@ -166,12 +168,12 @@ def evaluate_model(
         for name, func in eval_functions.items():
             score_counters[name].update(
                 func(
-                    maze=maze,
-                    solution=np.array(solution),
+                    maze=sm.maze,
+                    solution=np.array(sm.solution),
                     prediction=np.array(prediction),
                     model=model,
                 )
-                for maze, solution, prediction in zip(mazes, solutions, predictions)
+                for sm, prediction in zip(solved_mazes, predictions)
             )
 
     return score_counters

maze_transformer/evaluation/path_evals.py

@@ -1,27 +1,25 @@
-from typing import Iterable, Optional, Protocol, TypeAlias
+from typing import Iterable, Optional, Protocol
 
 import numpy as np
-from jaxtyping import Int
 
-from maze_transformer.generation.constants import Coord, CoordTup
+from maze_transformer.generation.constants import Coord, CoordArray, CoordTup
 from maze_transformer.generation.lattice_maze import LatticeMaze
 from maze_transformer.utils.utils import register_method
 
 # pylint: disable=unused-argument
-MazePath: TypeAlias = Int[np.ndarray, "node x_y_pos"]
 
 
 class PathEvalFunction(Protocol):
     def __call__(
         self,
         maze: Optional[LatticeMaze] = None,
-        solution: Optional[MazePath] = None,
-        prediction: Optional[MazePath] = None,
+        solution: Optional[CoordArray] = None,
+        prediction: Optional[CoordArray] = None,
     ) -> float:
         ...
 
 
-def path_as_segments_iter(path: MazePath) -> Iterable[tuple]:
+def path_as_segments_iter(path: CoordArray) -> Iterable[tuple]:
     """
     Iterate over the segments of a path (ie each consecutive pair).
     """
@@ -44,7 +42,7 @@ class PathEvals:
 
     @register_method(evals)
     @staticmethod
-    def node_overlap(solution: MazePath, prediction: MazePath, **_) -> float:
+    def node_overlap(solution: CoordArray, prediction: CoordArray, **_) -> float:
         """number of shared nodes (any order) / total number of (unique) nodes in solution"""
 
         solution_set = {tuple(coord) for coord in solution}
@@ -54,7 +52,7 @@ def node_overlap(solution: MazePath, prediction: MazePath, **_) -> float:
 
     @register_method(evals)
     @staticmethod
-    def num_connections_adjacent_lattice(prediction: MazePath, **_) -> float:
+    def num_connections_adjacent_lattice(prediction: CoordArray, **_) -> float:
         """number of the connections in prediction which actually connect nodes that are adjacent on the lattice, ignoring if they are adjacent on the maze"""
         n_adj: float = 0.0
         for step_start, step_end in path_as_segments_iter(prediction):
@@ -65,14 +63,16 @@ def num_connections_adjacent_lattice(prediction: MazePath, **_) -> float:
 
     @register_method(evals)
     @staticmethod
-    def fraction_connections_adjacent_lattice(prediction: MazePath, **_) -> float:
+    def fraction_connections_adjacent_lattice(prediction: CoordArray, **_) -> float:
         """fraction of the connections in prediction which actually connect nodes that are adjacent on the lattice, ignoring if they are adjacent on the maze"""
 
         return PathEvals.num_connections_adjacent_lattice(prediction) / len(prediction)
 
     @register_method(evals)
     @staticmethod
-    def num_connections_adjacent(maze: LatticeMaze, prediction: MazePath, **_) -> float:
+    def num_connections_adjacent(
+        maze: LatticeMaze, prediction: CoordArray, **_
+    ) -> float:
         """number of connections in prediction which are are valid paths on the maze"""
         n_connected: float = 0.0
         for step_start, step_end in path_as_segments_iter(prediction):
@@ -84,7 +84,7 @@ def num_connections_adjacent(maze: LatticeMaze, prediction: MazePath, **_) -> fl
     @register_method(evals)
     @staticmethod
     def fraction_connections_adjacent(
-        maze: LatticeMaze, prediction: MazePath, **_
+        maze: LatticeMaze, prediction: CoordArray, **_
     ) -> float:
         """fraction of connections in prediction which are are valid paths on the maze"""
 
@@ -95,20 +95,22 @@ def fraction_connections_adjacent(
 
     @register_method(evals)
     @staticmethod
-    def exact_path_predicted(solution: MazePath, prediction: MazePath, **_) -> float:
+    def exact_path_predicted(
+        solution: CoordArray, prediction: CoordArray, **_
+    ) -> float:
         """Was the maze successfully solved?"""
         return float(np.array_equal(solution, prediction))
 
     @register_method(evals)
     @staticmethod
-    def solution_length(solution: MazePath, **_) -> float:
+    def solution_length(solution: CoordArray, **_) -> float:
         return float(len(solution))
 
     @register_method(evals)
     @staticmethod
     def streak_length_until_incorrect(
-        solution: MazePath,
-        prediction: MazePath,
+        solution: CoordArray,
+        prediction: CoordArray,
         **_,
     ) -> float:
         """How many moves until the predicted path deviates from the solution"""

maze_transformer/evaluation/plot_maze.py

@@ -6,13 +6,19 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
-from jaxtyping import Float
+from jaxtyping import Bool, Float
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import ListedColormap, Normalize
 from muutils.tensor_utils import NDArray
 
 from maze_transformer.generation.constants import Coord, CoordArray, CoordList
-from maze_transformer.generation.lattice_maze import Coord, CoordArray, LatticeMaze
+from maze_transformer.generation.lattice_maze import (
+    Coord,
+    CoordArray,
+    LatticeMaze,
+    SolvedMaze,
+    TargetedLatticeMaze,
+)
 
 MAX_NODE_VALUE_EPSILON: float = 1e-10
 
@@ -113,16 +119,16 @@ class MazePlot:
         "slategrey",
     ]
 
-    def __init__(self, maze: LatticeMaze) -> None:
+    def __init__(self, maze: LatticeMaze, unit_length: int = 14) -> None:
         """
         UNIT_LENGTH: Set ratio between node size and wall thickness in image.
         Wall thickness is fixed to 1px
         A "unit" consists of a single node and the right and lower connection/wall.
         Example: ul = 14 yields 13:1 ratio between node size and wall thickness
         """
-        self.unit_length: int = 14
+        self.unit_length: int = unit_length
         self.maze: LatticeMaze = maze
-        self.true_path: StyledPath = None
+        self.true_path: StyledPath | None = None
         self.predicted_paths: list[StyledPath] = []
         self.node_values: Float[np.ndarray, "grid_n grid_n"] = None
         self.custom_node_value_flag: bool = False
@@ -131,6 +137,23 @@ def __init__(self, maze: LatticeMaze) -> None:
         self.target_token_coord: Coord = None
         self.preceding_tokens_coords: CoordArray = None
 
+        if isinstance(maze, TargetedLatticeMaze):
+            self.add_true_path(SolvedMaze.from_targeted_lattice_maze(maze).solution)
+
+        if isinstance(maze, SolvedMaze):
+            self.add_true_path(maze.solution)
+
+    @property
+    def solved_maze(self) -> SolvedMaze:
+        if self.true_path is None:
+            raise ValueError(
+                "Cannot return SolvedMaze object without true path. Add true path with add_true_path method."
+            )
+        return SolvedMaze.from_lattice_maze(
+            lattice_maze=self.maze,
+            solution=self.true_path.path,
+        )
+
     def add_true_path(
         self,
         path: CoordList | CoordArray | StyledPath,
@@ -227,12 +250,6 @@ def plot(self, dpi: int = 100, title: str = "") -> MazePlot:
         self.ax.set_ylabel("row")
         self.fig.suptitle(title)
 
-    def show(self, dpi: int = 100, title: str = "") -> None:
-        """Plot the maze and paths and show the plot. DONT USE THIS IN NOTEBOOKS WHICH NEED TO BE TESTED IN CI!!!"""
-        self.plot(dpi=dpi, title=title)
-        plt.show()
-        return self
-
     def _rowcol_to_coord(self, point: Coord) -> NDArray:
         """Transform Point from MazeTransformer (row, column) notation to matplotlib default (x, y) notation where x is the horizontal axis."""
         point = np.array([point[1], point[0]])
@@ -294,7 +311,10 @@ def _plot_maze(self) -> None:
 
             self.ax.imshow(img, cmap=cmap, vmin=-1, vmax=1)
 
-    def _lattice_maze_to_img(self) -> NDArray["row col", bool]:
+    def _lattice_maze_to_img(
+        self,
+        connection_val_scale: float = 0.93,
+    ) -> Bool[np.ndarray, "row col"]:
         """
         Build an image to visualise the maze.
         Each "unit" consists of a node and the right and lower adjacent wall/connection. Its area is ul * ul.
@@ -318,7 +338,7 @@ def _lattice_maze_to_img(self) -> NDArray["row col", bool]:
         # Set node and connection values
         if self.node_values is None:
             scaled_node_values = np.ones(self.maze.grid_shape)
-            connection_values = scaled_node_values * 0.93
+            connection_values = scaled_node_values * connection_val_scale
         else:
             # Normalizing node colors to match color_map running in (-1, 1) (defined in ._plot_maze()).
             scaled_node_values = self.node_values / self.max_node_value
@@ -403,30 +423,14 @@ def _plot_path(self, path_format: PathFormat) -> None:
             ms=10,
         )
 
-    def as_ascii(self, start=None, end=None):
-        """
-        Returns an ASCII visualization of the maze.
-        Courtesy of ChatGPT
-        """
-        wall_char = "#"
-        path_char = " "
-        self.unit_length = 2
-
-        # Determine the size of the maze
-        maze = self._lattice_maze_to_img()
-        n_rows, n_cols = maze.shape
-        maze_str = ""
-
-        # Iterate through each element of the maze and print the appropriate symbol
-        for i in range(n_rows):
-            for j in range(n_cols):
-                if start is not None and start[0] == i - 1 and start[1] == j - 1:
-                    maze_str += "S"
-                elif end is not None and end[0] == i - 1 and end[1] == j - 1:
-                    maze_str += "E"
-                elif maze[i, j] == -1:
-                    maze_str += wall_char
-                else:
-                    maze_str += path_char
-            maze_str += "\n"  # Start a new line after each row
-        return maze_str
+    def to_ascii(
+        self,
+        show_endpoints: bool = True,
+        show_solution: bool = True,
+    ) -> str:
+        if self.true_path:
+            return self.solved_maze.as_ascii(
+                show_endpoints=show_endpoints, show_solution=show_solution
+            )
+        else:
+            return self.maze.as_ascii(show_endpoints=show_endpoints)

maze_transformer/evaluation/util.py

@@ -0,0 +1,44 @@
+import torch
+from muutils.zanj.torchutil import ConfigMismatchException, assert_model_cfg_equality
+
+from maze_transformer.training.config import ZanjHookedTransformer
+
+
+def assert_model_output_equality(
+    model_a: ZanjHookedTransformer, model_b: ZanjHookedTransformer
+):
+    try:
+        assert_model_cfg_equality(model_a, model_b)
+    except ConfigMismatchException as e:
+        if e.diff == {
+            "model_cfg": {"are_weights_processed": {"self": False, "other": True}}
+        } or e.diff == {
+            "model_cfg": {
+                "are_layernorms_folded": {"self": False, "other": True},
+                "are_weights_processed": {"self": False, "other": True},
+            }
+        }:
+            pass
+        else:
+            raise e
+
+    # Random input tokens
+    dataset_cfg = model_a.zanj_model_config.dataset_cfg
+    input_sequence = torch.randint(
+        low=0,
+        high=len(dataset_cfg.token_arr),
+        size=(1, min(dataset_cfg.seq_len_max, 10)),
+    )
+
+    # (copied from `test_eval_model.py`)
+    # Check for equality in argsort (absolute values won't be equal due to centering the unembedding weight matrix)
+    assert torch.all(
+        model_a(input_sequence.clone()).argsort()
+        == model_b(input_sequence.clone()).argsort()
+    )
+    # apply normalization (e.g. softmax) and check with atol v-small
+    # (roughly 1E-7 for float error on logexp I think)
+    output_a = torch.nn.functional.softmax(model_a(input_sequence.clone()), dim=-1)
+    output_b = torch.nn.functional.softmax(model_b(input_sequence.clone()), dim=-1)
+
+    assert torch.allclose(output_a, output_b, atol=1e-7)

maze_transformer/generation/constants.py

@@ -1,9 +1,9 @@
 import numpy as np
-from muutils.tensor_utils import NDArray
+from jaxtyping import Int8
 
-Coord = NDArray["x y", np.int8]
+Coord = Int8[np.ndarray, "x y"]
 CoordTup = tuple[int, int]
-CoordArray = NDArray["coords", np.int8]
+CoordArray = Int8[np.ndarray, "coord x y"]
 CoordList = list[CoordTup]
 
 SPECIAL_TOKENS: dict[str, str] = dict(
@@ -20,7 +20,7 @@
     padding="<PADDING>",
 )
 
-DIRECTIONS_MAP: NDArray["direction axes", int] = np.array(
+DIRECTIONS_MAP: Int8[np.ndarray, "direction axes"] = np.array(
     [
         [0, 1],  # down
         [0, -1],  # up
@@ -30,7 +30,7 @@
 )
 
 
-NEIGHBORS_MASK: NDArray["coord point", int] = np.array(
+NEIGHBORS_MASK: Int8[np.ndarray, "coord point"] = np.array(
     [
         [0, 1],  # down
         [0, -1],  # up