Update app.py

app.py

@@ -167,17 +167,70 @@ def __(ChessBoard, dropdown_fen, dropdown_moves):
 
 
 @app.cell
-def __(focus_square):
+def __(torch):
+    def rollout(x, skip_connection=True, parse="min"):
+        attns = []
+        for k, v in x.items():
+            v = v[0, :, ::-1, :]
+            v = torch.tensor(v.copy())
+            if parse == "min": 
+                item = torch.min(v, dim=0).values
+            elif parse == "max": 
+                item = torch.max(v, dim=0).values
+            elif parse == "mean":
+                item = torch.mean(v, dim=0)
+        attns.append(item + int(skip_connection) * torch.eye(item.shape[0]))
+        roll = torch.prod(torch.stack(attns), dim=0)
+        return roll.numpy()
+
+    return (rollout,)
+
+
+@app.cell
+def __(mo):
+    METHODS = ["Attention visualization", "Attention rollout (MIN)", "Attention rollout (MEAN)", "Attention rollout (MAX)"]
+    dropdown_method = mo.ui.dropdown(options=METHODS, value=METHODS[0], label="Select XAI method")
+    dropdown_method
+    return METHODS, dropdown_method
+
+
+@app.cell
+def __(mo):
+    switch = mo.ui.switch(value=False, label="To use skip connection in rollout")
+    switch
+    return (switch,)
+
+
+@app.cell
+def __(switch):
+    switch.value
+    return
+
+
+@app.cell
+def __(dropdown_method, focus_square, rollout, switch, torch):
     import chess 
     from global_data import global_data
 
+    def parse_activations(act, layer_number=None):
+        if dropdown_method.value == "Attention visualization":
+            layer_key = [k for k in act.keys() if "0" in k][0].replace("0", f"{layer_number}")
+            a = act[layer_key][0, :, ::-1 , :]
+        elif dropdown_method.value == "Attention rollout (MIN)":
+            a = rollout(act, skip_connection=switch.value, parse="min")
+            a = torch.stack([a for _ in range(32)])
+        elif dropdown_method.value == "Attention rollout (MAX)":
+            a = rollout(act, skip_connection=switch.value, parse="max")
+            a = torch.stack([a for _ in range(32)])
+        elif dropdown_method.value == "Attention rollout (MEAN)":
+            a = rollout(act, skip_connection=switch.value, parse="mean")
+            a = torch.stack([a for _ in range(32)])
+        return a
     focus_square_ind = 8 * (int(focus_square.value[1]) - 1) + ord(focus_square.value[0]) - ord("a")
 
-    def set_plotting_parameters(act, layer_number, fen):
-        layer_key = [k for k in act.keys() if "0" in k][0].replace("0", f"{layer_number}")
-        print(act.keys())
+    def set_plotting_parameters(act, fen):
         global_data.model = 'test'
-        global_data.activations = act[layer_key][0, :, ::-1 , :]
+        global_data.activations = act
         print(global_data.activations.shape)
         global_data.subplot_rows = 8
         global_data.subplot_cols = 4
@@ -191,7 +244,13 @@ def __(focus_square):
         global_data.visualization_mode_is_64x64 = False
         global_data.colorscale_mode = "mode1"
         global_data.show_colorscale = False
-    return chess, focus_square_ind, global_data, set_plotting_parameters
+    return (
+        chess,
+        focus_square_ind,
+        global_data,
+        parse_activations,
+        set_plotting_parameters,
+    )
 
 
 @app.cell
@@ -200,6 +259,7 @@ def __(
     dropdown_layer,
     get_activations_from_model,
     get_models,
+    parse_activations,
     set_plotting_parameters,
 ):
     # FEN = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
@@ -210,12 +270,13 @@ def __(
     # PATTERN = "smolgen_weights"
     MODEL = get_models()[-1]
     ACTIVATIONS = get_activations_from_model(MODEL, PATTERN, FEN)
-    set_plotting_parameters(ACTIVATIONS, int(dropdown_layer.value), FEN)
+    a = parse_activations(ACTIVATIONS, layer_number=int(dropdown_layer.value))
+    set_plotting_parameters(a, FEN)
     from activation_heatmap import heatmap_figure
     fig = heatmap_figure()
     fig.update_layout(height=1500, width=1200)
     fig
-    return ACTIVATIONS, MODEL, PATTERN, fig, heatmap_figure
+    return ACTIVATIONS, MODEL, PATTERN, a, fig, heatmap_figure
 
 
 @app.cell