fixed concu & best move

app.py

@@ -3,7 +3,6 @@ Gradio demo for lczerolens.
 """
 
 import gradio as gr
-import subprocess
 
 from demo.interfaces import (
     board,

demo/interfaces/activations.py

@@ -129,21 +129,26 @@ with gr.Blocks() as interface:
         initial_load,
         inputs=[model_name, game_pgn, board_fen, layer_index, channel_index],
         outputs=[model, board, activations, image_board, colorbar],
+        concurrency_limit=1, 
+        concurrency_id="trace_queue"
     )
     game_pgn.submit(
         on_board_change,
         inputs=[model, game_pgn, board_fen, layer_index, channel_index],
         outputs=[board, activations, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     board_fen.submit(
         on_board_change,
         inputs=[model, game_pgn, board_fen, layer_index, channel_index],
         outputs=[board, activations, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     model_name.change(
         on_model_change,
         inputs=[model_name, board, layer_index, channel_index],
         outputs=[model, activations, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     layer_index.change(
         render_activations,

demo/interfaces/gradients.py

@@ -10,17 +10,20 @@ import os
 
 from lczerolens import LczeroBoard, LczeroModel, Lens
 
-from .. import constants
+from demo import constants
 
 def get_model(model_name: str):
     return LczeroModel.from_onnx_path(os.path.join(constants.ONNX_MODEL_DIRECTORY, model_name))
 
-def get_gradients(model: LczeroModel, board: LczeroBoard, wdl_target: str):
+def get_gradients(model: LczeroModel, board: LczeroBoard, target: str):
     lens = Lens.from_name("gradient")
-    wdl_index = {"win": 0, "draw": 1, "loss": 2}[wdl_target]
 
     def init_target(model):
-        return getattr(model, "output/wdl").output[:, wdl_index]
+        if target == "best_move":
+            return getattr(model, "output/policy").output.max(dim=1).values
+        else:
+            wdl_index = {"win": 0, "draw": 1, "loss": 2}[target]
+            return getattr(model, "output/wdl").output[:, wdl_index]
     results = lens.analyse(model, board, init_target=init_target)
 
     return results["input_grad"]
@@ -57,27 +60,27 @@ def render_gradients(board: LczeroBoard, gradients, average_over_planes:bool, be
     )
     return  f"{constants.FIGURE_DIRECTORY}/gradients_board.svg", f"{constants.FIGURE_DIRECTORY}/gradients_colorbar.svg"
 
-def initial_load(model_name: str, board_fen: str, game_pgn: str, wdl_target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+def initial_load(model_name: str, board_fen: str, game_pgn: str, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
     model = get_model(model_name)
     board = get_board(game_pgn, board_fen)
-    gradients = get_gradients(model, board, wdl_target)
+    gradients = get_gradients(model, board, target)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return model, board, gradients, *plots
 
-def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, wdl_target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
     board = get_board(game_pgn, board_fen)
-    gradients = get_gradients(model, board, wdl_target)
+    gradients = get_gradients(model, board, target)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return board, gradients, *plots
 
-def on_model_change(model_name: str, board: LczeroBoard, wdl_target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+def on_model_change(model_name: str, board: LczeroBoard, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
     model = get_model(model_name)
-    gradients = get_gradients(model, board, wdl_target)
+    gradients = get_gradients(model, board, target)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return model, gradients, *plots
 
-def on_wdl_target_change(model: LczeroModel, board: LczeroBoard, wdl_target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
-    gradients = get_gradients(model, board, wdl_target)
+def on_target_change(model: LczeroModel, board: LczeroBoard, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+    gradients = get_gradients(model, board, target)
     plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
     return gradients, *plots
 
@@ -104,8 +107,8 @@ with gr.Blocks() as interface:
                     label="Model",
                     choices=constants.ONNX_MODEL_NAMES,
                 )
-                wdl_target = gr.Radio(
-                    ["win", "draw", "loss"], label="WDL target",
+                target = gr.Radio(
+                    ["win", "draw", "loss", "best_move"], label="Target",
                     value="win",
                 )
             with gr.Group():
@@ -143,28 +146,33 @@ with gr.Blocks() as interface:
 
     interface.load(
         initial_load,
-        inputs=[model_name, game_pgn, board_fen, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model_name, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[model, board, gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     game_pgn.submit(
         on_board_change,
-        inputs=[model, game_pgn, board_fen, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[board, gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     board_fen.submit(
         on_board_change,
-        inputs=[model, game_pgn, board_fen, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[board, gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     model_name.change(
         on_model_change,
-        inputs=[model_name, board, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        inputs=[model_name, board, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[model, gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
-    wdl_target.change(
-        on_wdl_target_change,
-        inputs=[model, board, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+    target.change(
+        on_target_change,
+        inputs=[model, board, target, average_over_planes, begin_average_index, end_average_index, plane_index],
         outputs=[gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
     )
     for render_arg in [average_over_planes, begin_average_index, end_average_index, plane_index]:
         render_arg.change(