working interface

src/helpers/generator.py

@@ -20,6 +20,7 @@ class OutputGenerator:
         self.wrapper = wrapper
         self.lens = ActivationLens(module_exp=module_exp)
 
+    @torch.no_grad
     def generate(
         self,
         root_fen: Optional[str] = None,
@@ -35,17 +36,16 @@ class OutputGenerator:
             input_encoding = InputEncoding.INPUT_CLASSICAL_112_PLANE_REPEATED
         else:
             raise ValueError
-        iter_boards = iter([[root_board, traj_board]])
-        act_dict, (model_output,) = self.lens.analyse_batched_boards(
+        iter_boards = iter([([root_board, traj_board],)])
+        result_iter = self.lens.analyse_batched_boards(
             iter_boards,
             self.wrapper,
-            {
-                "return_output": True,
-                "wrapper_kwargs": {
-                    "input_encoding": input_encoding,
-                }
+            return_output=True,
+            wrapper_kwargs={
+                "input_encoding": input_encoding,
             }
         )
+        act_dict, (model_output,) = next(result_iter)
         if len(act_dict) == 0:
             raise ValueError("No module matced the given expression.")
         elif len(act_dict) > 1:

src/interfaces/feature_interface.py

@@ -5,6 +5,7 @@ Gradio interface for plotting policy.
 import chess
 import gradio as gr
 import uuid
+import torch
 
 from lczerolens.encodings import encode_move
 
@@ -19,28 +20,57 @@ def compute_features_fn(
     traj_fen, 
     feature_index
 ):
-    model_output, _, sae_output = global_variables.generator.generate(
+    model_output, pixel_acts, sae_output = global_variables.generator.generate(
         root_fen=root_fen,
         traj_fen=traj_fen
     )
-    features = sae_output["f"]
+    features = sae_output["features"]
+    x_hat = sae_output["x_hat"]
     first_output = render_feature_index(
         features,
         model_output,
         file_id,
-        feature_index,
-        traj_fen,
+        traj_fen, 
+        feature_index
     )
-    game_info = f"WDL: {model_output.get('wdl')}"
-    return *first_output, game_info
+
+    half_a_dim = constants.ACTIVATION_DIM // 2
+    half_f_dim = constants.DICTIONARY_SIZE // 2
+    pixel_f_avg = features.mean(dim=0)
+    pixel_f_active = (features > 0).float().mean(dim=0)
+    pixel_p_avg = features.mean(dim=1)
+    pixel_p_active = (features > 0).float().mean(dim=1)
+
+    board = chess.Board(traj_fen)
+    if board.turn:
+        most_avg_pixels = pixel_p_avg.topk(5).indices.tolist()
+        most_active_pixels = pixel_p_active.topk(5).indices.tolist()
+    else:
+        most_avg_pixels = pixel_p_avg.view(8,8).flip(0).view(64).topk(5).indices.tolist()
+        most_active_pixels = pixel_p_active.view(8,8).flip(0).view(64).topk(5).indices.tolist()
+
+    info = f"Root WDL: {model_output['wdl'][0]}\n"
+    info += f"Traj WDL: {model_output['wdl'][1]}\n"
+    info += f"MSE loss: {torch.nn.functional.mse_loss(x_hat, pixel_acts, reduction='none').sum(dim=1).mean()}\n"
+    info += f"MSE loss (root): {torch.nn.functional.mse_loss(x_hat[:,:half_a_dim], pixel_acts[:,:half_a_dim], reduction='none').sum(dim=1).mean()}\n"
+    info += f"MSE loss (traj): {torch.nn.functional.mse_loss(x_hat[:,half_a_dim:], pixel_acts[:,half_a_dim:], reduction='none').sum(dim=1).mean()}\n"
+    info += f"L0 loss: {(features>0).sum(dim=1).float().mean()}\n"
+    info += f"L0 loss (c): {(features[:,:half_f_dim]>0).sum(dim=1).float().mean()}\n"
+    info += f"L0 loss (d): {(features[:,half_f_dim:]>0).sum(dim=1).float().mean()}\n"
+    info += f"Most active features (avg): {pixel_f_avg.topk(5).indices.tolist()}\n"
+    info += f"Most active features (active): {pixel_f_active.topk(5).indices.tolist()}\n"
+    info += f"Most active pixels (avg): {[chess.SQUARE_NAMES[p] for p in most_avg_pixels]}\n"
+    info += f"Most active pixels (active): {[chess.SQUARE_NAMES[p] for p in most_active_pixels]}"
+    
+    return *first_output, info
 
 
 def render_feature_index(
     features,
     model_output,
     file_id,
-    feature_index,
-    traj_fen,
+    traj_fen, 
+    feature_index
 ):
     if file_id is None:
         file_id = str(uuid.uuid4())
@@ -98,14 +128,14 @@ with gr.Blocks() as interface:
                     feature_index = gr.Slider(
                         label="Feature index",
                         minimum=0,
-                        maximum=constants.N_FEATURES,
+                        maximum=constants.DICTIONARY_SIZE-1,
                         step=1,
                         value=0,
                     )
                     
             with gr.Group():
                 with gr.Row():    
-                    game_info = gr.Textbox(label="Game info", lines=1, max_lines=1, value="")
+                    info = gr.Textbox(label="Info", lines=1, max_lines=20, value="")
                 with gr.Row():
                     colorbar = gr.Plot(label="Colorbar")
         with gr.Column():
@@ -114,8 +144,14 @@ with gr.Blocks() as interface:
     features = gr.State(None)
     model_output = gr.State(None)
     file_id = gr.State(None)
+    
     compute_features.click(
         compute_features_fn,
         inputs=[features, model_output, file_id, root_fen, traj_fen, feature_index],
-        outputs=[features, model_output, file_id, board_image, colorbar, game_info],
-    )
+        outputs=[features, model_output, file_id, board_image, colorbar, info],
+    )
+    feature_index.change(
+        render_feature_index,
+        inputs=[features, model_output, file_id, traj_fen, feature_index],
+        outputs=[features, model_output, file_id, board_image, colorbar],
+    )