old files

app/attention_interface.py

@@ -0,0 +1,292 @@
+"""
+Gradio interface for plotting attention.
+"""
+
+import copy
+
+import chess
+import gradio as gr
+
+from demo import constants, utils, visualisation
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info(leela=False)
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def compute_cache(
+    board_fen,
+    action_seq,
+    model_name,
+    attention_layer,
+    attention_head,
+    square,
+    state_board_index,
+    state_boards,
+    state_cache,
+):
+    if model_name == "":
+        gr.Warning("No model selected.")
+        return None, None, None, state_boards, state_cache
+
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        board = chess.Board()
+        gr.Warning("Invalid FEN, using starting position.")
+    state_boards = [board.copy()]
+    if action_seq:
+        try:
+            if action_seq.startswith("1."):
+                for action in action_seq.split():
+                    if action.endswith("."):
+                        continue
+                    board.push_san(action)
+                    state_boards.append(board.copy())
+            else:
+                for action in action_seq.split():
+                    board.push_uci(action)
+                    state_boards.append(board.copy())
+        except ValueError:
+            gr.Warning(f"Invalid action {action} stopping before it.")
+    try:
+        wrapper, lens = utils.get_wrapper_lens_from_state(
+            model_name,
+            "activation",
+            lens_name="attention",
+            module_exp=r"encoder\d+/mha/QK/softmax",
+        )
+    except ValueError:
+        gr.Warning("Could not load model.")
+        return None, None, None, state_boards, state_cache
+    state_cache = []
+    for board in state_boards:
+        attention_cache = copy.deepcopy(lens.analyse_board(board, wrapper))
+        state_cache.append(attention_cache)
+    return (
+        *make_plot(
+            attention_layer,
+            attention_head,
+            square,
+            state_board_index,
+            state_boards,
+            state_cache,
+        ),
+        state_boards,
+        state_cache,
+    )
+
+
+def make_plot(
+    attention_layer,
+    attention_head,
+    square,
+    state_board_index,
+    state_boards,
+    state_cache,
+):
+    if state_cache == []:
+        gr.Warning("No cache available.")
+        return None, None, None
+
+    board = state_boards[state_board_index]
+    num_attention_layers = len(state_cache[state_board_index])
+    if attention_layer > num_attention_layers:
+        gr.Warning(
+            f"Attention layer {attention_layer} does not exist, " f"using layer {num_attention_layers} instead."
+        )
+        attention_layer = num_attention_layers
+
+    key = f"encoder{attention_layer-1}/mha/QK/softmax"
+    try:
+        attention_tensor = state_cache[state_board_index][key]
+    except KeyError:
+        gr.Warning(f"Combination {key} does not exist.")
+        return None, None, None
+    if attention_head > attention_tensor.shape[1]:
+        gr.Warning(
+            f"Attention head {attention_head} does not exist, " f"using head {attention_tensor.shape[1]+1} instead."
+        )
+        attention_head = attention_tensor.shape[1]
+    try:
+        square_index = chess.SQUARE_NAMES.index(square)
+    except ValueError:
+        gr.Warning(f"Invalid square {square}, using a1 instead.")
+        square_index = 0
+        square = "a1"
+    if board.turn == chess.BLACK:
+        square_index = chess.square_mirror(square_index)
+
+    heatmap = attention_tensor[0, attention_head - 1, square_index]
+    if board.turn == chess.BLACK:
+        heatmap = heatmap.view(8, 8).flip(0).view(64)
+    svg_board, fig = visualisation.render_heatmap(board, heatmap, square=square)
+    with open(f"{constants.FIGURE_DIRECTORY}/attention.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/attention.svg", board.fen(), fig
+
+
+def previous_board(
+    attention_layer,
+    attention_head,
+    square,
+    state_board_index,
+    state_boards,
+    state_cache,
+):
+    state_board_index -= 1
+    if state_board_index < 0:
+        gr.Warning("Already at first board.")
+        state_board_index = 0
+    return (
+        *make_plot(
+            attention_layer,
+            attention_head,
+            square,
+            state_board_index,
+            state_boards,
+            state_cache,
+        ),
+        state_board_index,
+    )
+
+
+def next_board(
+    attention_layer,
+    attention_head,
+    square,
+    state_board_index,
+    state_boards,
+    state_cache,
+):
+    state_board_index += 1
+    if state_board_index >= len(state_boards):
+        gr.Warning("Already at last board.")
+        state_board_index = len(state_boards) - 1
+    return (
+        *make_plot(
+            attention_layer,
+            attention_head,
+            square,
+            state_board_index,
+            state_boards,
+            state_cache,
+        ),
+        state_board_index,
+    )
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board starting FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            action_seq = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                max_lines=1,
+                value=("e2e3 b8c6 d2d4 e7e5 g1f3 d8e7 " "d4d5 e5e4 f3d4 c6e5 f2f4 e5g6"),
+            )
+            compute_cache_button = gr.Button("Compute cache")
+
+            with gr.Group():
+                with gr.Row():
+                    attention_layer = gr.Slider(
+                        label="Attention layer",
+                        minimum=1,
+                        maximum=24,
+                        step=1,
+                        value=1,
+                    )
+                    attention_head = gr.Slider(
+                        label="Attention head",
+                        minimum=1,
+                        maximum=24,
+                        step=1,
+                        value=1,
+                    )
+                with gr.Row():
+                    square = gr.Textbox(
+                        label="Square",
+                        lines=1,
+                        max_lines=1,
+                        value="a1",
+                        scale=1,
+                    )
+                with gr.Row():
+                    previous_board_button = gr.Button("Previous board")
+                    next_board_button = gr.Button("Next board")
+            current_board_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+            )
+            colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            image = gr.Image(label="Board")
+
+    state_board_index = gr.State(0)
+    state_boards = gr.State([])
+    state_cache = gr.State([])
+    base_inputs = [
+        attention_layer,
+        attention_head,
+        square,
+        state_board_index,
+        state_boards,
+        state_cache,
+    ]
+    outputs = [image, current_board_fen, colorbar]
+
+    compute_cache_button.click(
+        compute_cache,
+        inputs=[board_fen, action_seq, model_name] + base_inputs,
+        outputs=outputs + [state_boards, state_cache],
+    )
+
+    previous_board_button.click(
+        previous_board,
+        inputs=base_inputs,
+        outputs=outputs + [state_board_index],
+    )
+    next_board_button.click(next_board, inputs=base_inputs, outputs=outputs + [state_board_index])
+
+    attention_layer.change(make_plot, inputs=base_inputs, outputs=outputs)
+    attention_head.change(make_plot, inputs=base_inputs, outputs=outputs)
+    square.submit(make_plot, inputs=base_inputs, outputs=outputs)

app/backend_interface.py

@@ -0,0 +1,207 @@
+"""
+Gradio interface for visualizing the policy of a model.
+"""
+
+import chess
+import chess.svg
+import gradio as gr
+import torch
+from lczero.backends import Backend, GameState, Weights
+
+from demo import constants, utils, visualisation
+from lczerolens import move_encodings
+from lczerolens.model import lczero as lczero_utils
+from lczerolens.xai import PolicyLens
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info(onnx=False)
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def make_policy_plot(
+    board_fen,
+    action_seq,
+    view,
+    model_name,
+    depth,
+    use_softmax,
+    aggregate_topk,
+    render_bestk,
+    only_legal,
+):
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return (
+            None,
+            None,
+            "",
+        )
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        board = chess.Board()
+        gr.Warning("Invalid FEN, using starting position.")
+    if action_seq:
+        try:
+            for action in action_seq.split():
+                board.push_uci(action)
+        except ValueError:
+            gr.Warning("Invalid action sequence, using starting position.")
+            board = chess.Board()
+    lczero_weights = Weights(f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}")
+    lczero_backend = Backend(lczero_weights)
+    uci_moves = [move.uci() for move in board.move_stack]
+    lczero_game = GameState(moves=uci_moves)
+    policy, value = lczero_utils.prediction_from_backend(
+        lczero_backend,
+        lczero_game,
+        softmax=use_softmax,
+        only_legal=only_legal,
+        illegal_value=0,
+    )
+    pickup_agg, dropoff_agg = PolicyLens.aggregate_policy(policy, int(aggregate_topk))
+
+    if view == "from":
+        if board.turn == chess.WHITE:
+            heatmap = pickup_agg
+        else:
+            heatmap = pickup_agg.view(8, 8).flip(0).view(64)
+    else:
+        if board.turn == chess.WHITE:
+            heatmap = dropoff_agg
+        else:
+            heatmap = dropoff_agg.view(8, 8).flip(0).view(64)
+    us_them = (board.turn, not board.turn)
+    if only_legal:
+        legal_moves = [move_encodings.encode_move(move, us_them) for move in board.legal_moves]
+        filtered_policy = torch.zeros(1858)
+        filtered_policy[legal_moves] = policy[legal_moves]
+        if (filtered_policy < 0).any():
+            gr.Warning("Some legal moves have negative policy.")
+        topk_moves = torch.topk(filtered_policy, render_bestk)
+    else:
+        topk_moves = torch.topk(policy, render_bestk)
+    arrows = []
+    for move_index in topk_moves.indices:
+        move = move_encodings.decode_move(move_index, us_them)
+        arrows.append((move.from_square, move.to_square))
+    svg_board, fig = visualisation.render_heatmap(board, heatmap, arrows=arrows)
+    with open(f"{constants.FIGURE_DIRECTORY}/policy.svg", "w") as f:
+        f.write(svg_board)
+    raw_policy, _ = lczero_utils.prediction_from_backend(
+        lczero_backend,
+        lczero_game,
+        softmax=False,
+        only_legal=False,
+        illegal_value=0,
+    )
+    fig_dist = visualisation.render_policy_distribution(
+        raw_policy,
+        [move_encodings.encode_move(move, us_them) for move in board.legal_moves],
+    )
+    return (
+        f"{constants.FIGURE_DIRECTORY}/policy.svg",
+        fig,
+        (f"Value: {value:.2f}"),
+        fig_dist,
+    )
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            action_seq = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                max_lines=1,
+                value=("e2e3 b8c6 d2d4 e7e5 g1f3 d8e7 " "d4d5 e5e4 f3d4 c6e5 f2f4 e5g6"),
+            )
+            with gr.Group():
+                with gr.Row():
+                    depth = gr.Radio(label="Depth", choices=[0], value=0)
+                    use_softmax = gr.Checkbox(label="Use softmax", value=True)
+                with gr.Row():
+                    aggregate_topk = gr.Slider(
+                        label="Aggregate top k",
+                        minimum=1,
+                        maximum=1858,
+                        step=1,
+                        value=1858,
+                        scale=3,
+                    )
+                    view = gr.Radio(
+                        label="View",
+                        choices=["from", "to"],
+                        value="from",
+                        scale=1,
+                    )
+                with gr.Row():
+                    render_bestk = gr.Slider(
+                        label="Render best k",
+                        minimum=1,
+                        maximum=5,
+                        step=1,
+                        value=5,
+                        scale=3,
+                    )
+                    only_legal = gr.Checkbox(label="Only legal", value=True, scale=1)
+
+            policy_button = gr.Button("Plot policy")
+            colorbar = gr.Plot(label="Colorbar")
+            game_info = gr.Textbox(label="Game info", lines=1, max_lines=1, value="")
+        with gr.Column():
+            image = gr.Image(label="Board")
+            density_plot = gr.Plot(label="Density")
+
+    policy_inputs = [
+        board_fen,
+        action_seq,
+        view,
+        model_name,
+        depth,
+        use_softmax,
+        aggregate_topk,
+        render_bestk,
+        only_legal,
+    ]
+    policy_outputs = [image, colorbar, game_info, density_plot]
+    policy_button.click(make_policy_plot, inputs=policy_inputs, outputs=policy_outputs)

app/board_interface.py

@@ -0,0 +1,80 @@
+"""
+Gradio interface for plotting a board.
+"""
+
+import chess
+import gradio as gr
+
+from demo import constants
+
+
+def make_board_plot(board_fen, arrows, square):
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        board = chess.Board()
+        gr.Warning("Invalid FEN, using starting position.")
+    try:
+        if arrows:
+            arrows_list = arrows.split(" ")
+            chess_arrows = []
+            for arrow in arrows_list:
+                from_square, to_square = arrow[:2], arrow[2:]
+                chess_arrows.append(
+                    (
+                        chess.parse_square(from_square),
+                        chess.parse_square(to_square),
+                    )
+                )
+        else:
+            chess_arrows = []
+    except ValueError:
+        chess_arrows = []
+        gr.Warning("Invalid arrows, using none.")
+
+    color_dict = {chess.parse_square(square): "#FF0000"} if square else {}
+    svg_board = chess.svg.board(
+        board,
+        size=350,
+        arrows=chess_arrows,
+        fill=color_dict,
+    )
+    with open(f"{constants.FIGURE_DIRECTORY}/board.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/board.svg"
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board starting FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            arrows = gr.Textbox(
+                label="Arrows",
+                lines=1,
+                max_lines=1,
+                value="",
+                placeholder="e2e4 e7e5",
+            )
+            square = gr.Textbox(
+                label="Square",
+                lines=1,
+                max_lines=1,
+                value="",
+                placeholder="e4",
+            )
+        with gr.Column():
+            image = gr.Image(label="Board", interactive=False)
+
+    inputs = [
+        board_fen,
+        arrows,
+        square,
+    ]
+    board_fen.submit(make_board_plot, inputs=inputs, outputs=image)
+    arrows.submit(make_board_plot, inputs=inputs, outputs=image)
+    interface.load(make_board_plot, inputs=inputs, outputs=image)

app/constants.py

@@ -0,0 +1,7 @@
+"""
+Constants for the demo.
+"""
+
+MODEL_DIRECTORY = "demo/onnx_models"
+LEELA_MODEL_DIRECTORY = "demo/leela_models"
+FIGURE_DIRECTORY = "demo/figures"

app/convert_interface.py

@@ -0,0 +1,201 @@
+"""
+Gradio interface for converting models.
+"""
+
+import os
+import uuid
+
+import gradio as gr
+
+from demo import constants, utils
+from lczerolens.model import lczero as lczero_utils
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info()
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def convert_model(
+    model_name: str,
+):
+    """
+    Convert the model.
+    """
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return list_models(), ""
+    if model_name.endswith(".onnx"):
+        gr.Warning(
+            "ONNX conversion not implemented.",
+        )
+        return list_models(), ""
+    try:
+        lczero_utils.convert_to_onnx(
+            f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}",
+            f"{constants.MODEL_DIRECTORY}/{model_name[:-6]}.onnx",
+        )
+    except RuntimeError:
+        gr.Warning(
+            f"Could not convert net at `{model_name}`.",
+        )
+        return list_models(), "Conversion failed"
+    return list_models(), "Conversion successful"
+
+
+def upload_model(
+    model_file: gr.File,
+):
+    """
+    Convert the model.
+    """
+    if model_file is None:
+        gr.Warning(
+            "File not uploaded.",
+        )
+        return list_models()
+    try:
+        id = uuid.uuid4()
+        tmp_file_path = f"{constants.LEELA_MODEL_DIRECTORY}/{id}"
+        with open(
+            tmp_file_path,
+            "wb",
+        ) as f:
+            f.write(model_file)
+        utils.save_model(tmp_file_path)
+    except RuntimeError:
+        gr.Warning(
+            "Invalid file type.",
+        )
+    finally:
+        if os.path.exists(tmp_file_path):
+            os.remove(tmp_file_path)
+    return list_models()
+
+
+def get_model_description(
+    model_name: str,
+):
+    """
+    Get the model description.
+    """
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return ""
+    if model_name.endswith(".onnx"):
+        gr.Warning(
+            "ONNX description not implemented.",
+        )
+        return ""
+    try:
+        description = lczero_utils.describenet(
+            f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}",
+        )
+    except RuntimeError:
+        raise gr.Error(
+            f"Could not describe net at `{model_name}`.",
+        )
+    return description
+
+
+def get_model_path(
+    model_name: str,
+):
+    """
+    Get the model path.
+    """
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return None
+    if model_name.endswith(".onnx"):
+        return f"{constants.MODEL_DIRECTORY}/{model_name}"
+    else:
+        return f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}"
+
+
+with gr.Blocks() as interface:
+    model_file = gr.File(type="binary")
+    upload_button = gr.Button(
+        value="Upload",
+    )
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+            conversion_status = gr.Textbox(
+                label="Conversion status",
+                lines=1,
+                interactive=False,
+            )
+
+    convert_button = gr.Button(
+        value="Convert",
+    )
+    describe_button = gr.Button(
+        value="Describe model",
+    )
+    model_description = gr.Textbox(
+        label="Model description",
+        lines=1,
+        interactive=False,
+    )
+    download_button = gr.Button(
+        value="Get download link",
+    )
+    download_file = gr.File(
+        type="filepath",
+        label="Download link",
+        interactive=False,
+    )
+
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+    upload_button.click(
+        upload_model,
+        model_file,
+        model_df,
+    )
+    convert_button.click(
+        convert_model,
+        model_name,
+        [model_df, conversion_status],
+    )
+    describe_button.click(
+        get_model_description,
+        model_name,
+        model_description,
+    )
+    download_button.click(
+        get_model_path,
+        model_name,
+        download_file,
+    )

app/crp_interface.py

@@ -0,0 +1,279 @@
+"""
+Gradio interface for plotting policy.
+"""
+
+import copy
+
+import chess
+import gradio as gr
+
+from demo import constants, utils, visualisation
+
+cache = None
+boards = None
+board_index = 0
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info(leela=False)
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def compute_cache(
+    board_fen,
+    action_seq,
+    model_name,
+    plane_index,
+    history_index,
+):
+    global cache
+    global boards
+    if model_name == "":
+        gr.Warning("No model selected.")
+        return None, None, None, None, None
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        board = chess.Board()
+        gr.Warning("Invalid FEN, using starting position.")
+    boards = [board.copy()]
+    if action_seq:
+        try:
+            if action_seq.startswith("1."):
+                for action in action_seq.split():
+                    if action.endswith("."):
+                        continue
+                    board.push_san(action)
+                    boards.append(board.copy())
+            else:
+                for action in action_seq.split():
+                    board.push_uci(action)
+                    boards.append(board.copy())
+        except ValueError:
+            gr.Warning(f"Invalid action {action} stopping before it.")
+    wrapper, lens = utils.get_wrapper_lens_from_state(model_name, "crp")
+    cache = []
+    for board in boards:
+        relevance = lens.compute_heatmap(board, wrapper)
+        cache.append(copy.deepcopy(relevance))
+    return (
+        *make_plot(
+            plane_index,
+        ),
+        *make_history_plot(
+            history_index,
+        ),
+    )
+
+
+def make_plot(
+    plane_index,
+):
+    global cache
+    global boards
+    global board_index
+
+    if cache is None:
+        gr.Warning("Cache not computed!")
+        return None, None, None
+
+    board = boards[board_index]
+    relevance_tensor = cache[board_index]
+    a_max = relevance_tensor.abs().max()
+    if a_max != 0:
+        relevance_tensor = relevance_tensor / a_max
+    vmin = -1
+    vmax = 1
+    heatmap = relevance_tensor[plane_index - 1].view(64)
+    if board.turn == chess.BLACK:
+        heatmap = heatmap.view(8, 8).flip(0).view(64)
+    svg_board, fig = visualisation.render_heatmap(board, heatmap, vmin=vmin, vmax=vmax)
+    with open(f"{constants.FIGURE_DIRECTORY}/lrp.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/lrp.svg", board.fen(), fig
+
+
+def make_history_plot(
+    history_index,
+):
+    global cache
+    global boards
+    global board_index
+
+    if cache is None:
+        gr.Warning("Cache not computed!")
+        return None, None
+
+    board = boards[board_index]
+    relevance_tensor = cache[board_index]
+    a_max = relevance_tensor.abs().max()
+    if a_max != 0:
+        relevance_tensor = relevance_tensor / a_max
+    vmin = -1
+    vmax = 1
+    heatmap = relevance_tensor[13 * (history_index - 1) : 13 * history_index - 1].sum(dim=0).view(64)
+    if board.turn == chess.BLACK:
+        heatmap = heatmap.view(8, 8).flip(0).view(64)
+    if board_index - history_index + 1 < 0:
+        history_board = chess.Board(fen=None)
+    else:
+        history_board = boards[board_index - history_index + 1]
+    svg_board, fig = visualisation.render_heatmap(history_board, heatmap, vmin=vmin, vmax=vmax)
+    with open(f"{constants.FIGURE_DIRECTORY}/lrp_history.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/lrp_history.svg", fig
+
+
+def previous_board(
+    plane_index,
+    history_index,
+):
+    global board_index
+    board_index -= 1
+    if board_index < 0:
+        gr.Warning("Already at first board.")
+        board_index = 0
+    return (
+        *make_plot(
+            plane_index,
+        ),
+        *make_history_plot(
+            history_index,
+        ),
+    )
+
+
+def next_board(
+    plane_index,
+    history_index,
+):
+    global board_index
+    board_index += 1
+    if board_index >= len(boards):
+        gr.Warning("Already at last board.")
+        board_index = len(boards) - 1
+    return (
+        *make_plot(
+            plane_index,
+        ),
+        *make_history_plot(
+            history_index,
+        ),
+    )
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board starting FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            action_seq = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                max_lines=1,
+                value=("e2e3 b8c6 d2d4 e7e5 g1f3 d8e7 " "d4d5 e5e4 f3d4 c6e5 f2f4 e5g6"),
+            )
+            compute_cache_button = gr.Button("Compute heatmaps")
+
+            with gr.Group():
+                with gr.Row():
+                    plane_index = gr.Slider(
+                        label="Plane index",
+                        minimum=1,
+                        maximum=112,
+                        step=1,
+                        value=1,
+                    )
+                with gr.Row():
+                    previous_board_button = gr.Button("Previous board")
+                    next_board_button = gr.Button("Next board")
+            current_board_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+            )
+            colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            image = gr.Image(label="Board")
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Group():
+                with gr.Row():
+                    histroy_index = gr.Slider(
+                        label="History index",
+                        minimum=1,
+                        maximum=8,
+                        step=1,
+                        value=1,
+                    )
+            history_colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            history_image = gr.Image(label="Board")
+
+    base_inputs = [
+        plane_index,
+        histroy_index,
+    ]
+    outputs = [
+        image,
+        current_board_fen,
+        colorbar,
+        history_image,
+        history_colorbar,
+    ]
+
+    compute_cache_button.click(
+        compute_cache,
+        inputs=[board_fen, action_seq, model_name] + base_inputs,
+        outputs=outputs,
+    )
+
+    previous_board_button.click(previous_board, inputs=base_inputs, outputs=outputs)
+    next_board_button.click(next_board, inputs=base_inputs, outputs=outputs)
+
+    plane_index.change(
+        make_plot,
+        inputs=plane_index,
+        outputs=[image, current_board_fen, colorbar],
+    )
+    histroy_index.change(
+        make_history_plot,
+        inputs=histroy_index,
+        outputs=[history_image, history_colorbar],
+    )

app/encoding_interface.py

@@ -0,0 +1,82 @@
+"""
+Gradio interface for plotting encodings.
+"""
+
+import chess
+import gradio as gr
+
+from demo import constants, visualisation
+from lczerolens import board_encodings
+
+
+def make_encoding_plot(
+    board_fen,
+    action_seq,
+    plane_index,
+    color_flip,
+):
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        board = chess.Board()
+        gr.Warning("Invalid FEN, using starting position.")
+    if action_seq:
+        try:
+            for action in action_seq.split():
+                board.push_uci(action)
+        except ValueError:
+            gr.Warning("Invalid action sequence, using starting position.")
+            board = chess.Board()
+    board_tensor = board_encodings.board_to_input_tensor(board)
+    heatmap = board_tensor[plane_index]
+    if color_flip and board.turn == chess.BLACK:
+        heatmap = heatmap.flip(0)
+    svg_board, fig = visualisation.render_heatmap(board, heatmap.view(64), vmin=0.0, vmax=1.0)
+    with open(f"{constants.FIGURE_DIRECTORY}/encoding.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/encoding.svg", fig
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board starting FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            action_seq = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                max_lines=1,
+                value=("e2e3 b8c6 d2d4 e7e5 g1f3 d8e7 " "d4d5 e5e4 f3d4 c6e5 f2f4 e5g6"),
+            )
+            with gr.Group():
+                with gr.Row():
+                    plane_index = gr.Slider(
+                        label="Plane index",
+                        minimum=0,
+                        maximum=111,
+                        step=1,
+                        value=0,
+                        scale=3,
+                    )
+                    color_flip = gr.Checkbox(label="Color flip", value=True, scale=1)
+
+            colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            image = gr.Image(label="Board")
+
+    policy_inputs = [
+        board_fen,
+        action_seq,
+        plane_index,
+        color_flip,
+    ]
+    policy_outputs = [image, colorbar]
+    board_fen.submit(make_encoding_plot, inputs=policy_inputs, outputs=policy_outputs)
+    action_seq.submit(make_encoding_plot, inputs=policy_inputs, outputs=policy_outputs)
+    plane_index.change(make_encoding_plot, inputs=policy_inputs, outputs=policy_outputs)
+    color_flip.change(make_encoding_plot, inputs=policy_inputs, outputs=policy_outputs)
+    interface.load(make_encoding_plot, inputs=policy_inputs, outputs=policy_outputs)

app/figures/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

app/leela_models/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

app/lrp_interface.py

@@ -0,0 +1,279 @@
+"""
+Gradio interface for plotting policy.
+"""
+
+import copy
+
+import chess
+import gradio as gr
+
+from demo import constants, utils, visualisation
+
+cache = None
+boards = None
+board_index = 0
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info(leela=False)
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def compute_cache(
+    board_fen,
+    action_seq,
+    model_name,
+    plane_index,
+    history_index,
+):
+    global cache
+    global boards
+    if model_name == "":
+        gr.Warning("No model selected.")
+        return None, None, None, None, None
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        board = chess.Board()
+        gr.Warning("Invalid FEN, using starting position.")
+    boards = [board.copy()]
+    if action_seq:
+        try:
+            if action_seq.startswith("1."):
+                for action in action_seq.split():
+                    if action.endswith("."):
+                        continue
+                    board.push_san(action)
+                    boards.append(board.copy())
+            else:
+                for action in action_seq.split():
+                    board.push_uci(action)
+                    boards.append(board.copy())
+        except ValueError:
+            gr.Warning(f"Invalid action {action} stopping before it.")
+    wrapper, lens = utils.get_wrapper_lens_from_state(model_name, "lrp")
+    cache = []
+    for board in boards:
+        relevance = lens.compute_heatmap(board, wrapper)
+        cache.append(copy.deepcopy(relevance))
+    return (
+        *make_plot(
+            plane_index,
+        ),
+        *make_history_plot(
+            history_index,
+        ),
+    )
+
+
+def make_plot(
+    plane_index,
+):
+    global cache
+    global boards
+    global board_index
+
+    if cache is None:
+        gr.Warning("Cache not computed!")
+        return None, None, None
+
+    board = boards[board_index]
+    relevance_tensor = cache[board_index]
+    a_max = relevance_tensor.abs().max()
+    if a_max != 0:
+        relevance_tensor = relevance_tensor / a_max
+    vmin = -1
+    vmax = 1
+    heatmap = relevance_tensor[plane_index - 1].view(64)
+    if board.turn == chess.BLACK:
+        heatmap = heatmap.view(8, 8).flip(0).view(64)
+    svg_board, fig = visualisation.render_heatmap(board, heatmap, vmin=vmin, vmax=vmax)
+    with open(f"{constants.FIGURE_DIRECTORY}/lrp.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/lrp.svg", board.fen(), fig
+
+
+def make_history_plot(
+    history_index,
+):
+    global cache
+    global boards
+    global board_index
+
+    if cache is None:
+        gr.Warning("Cache not computed!")
+        return None, None
+
+    board = boards[board_index]
+    relevance_tensor = cache[board_index]
+    a_max = relevance_tensor.abs().max()
+    if a_max != 0:
+        relevance_tensor = relevance_tensor / a_max
+    vmin = -1
+    vmax = 1
+    heatmap = relevance_tensor[13 * (history_index - 1) : 13 * history_index - 1].sum(dim=0).view(64)
+    if board.turn == chess.BLACK:
+        heatmap = heatmap.view(8, 8).flip(0).view(64)
+    if board_index - history_index + 1 < 0:
+        history_board = chess.Board(fen=None)
+    else:
+        history_board = boards[board_index - history_index + 1]
+    svg_board, fig = visualisation.render_heatmap(history_board, heatmap, vmin=vmin, vmax=vmax)
+    with open(f"{constants.FIGURE_DIRECTORY}/lrp_history.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/lrp_history.svg", fig
+
+
+def previous_board(
+    plane_index,
+    history_index,
+):
+    global board_index
+    board_index -= 1
+    if board_index < 0:
+        gr.Warning("Already at first board.")
+        board_index = 0
+    return (
+        *make_plot(
+            plane_index,
+        ),
+        *make_history_plot(
+            history_index,
+        ),
+    )
+
+
+def next_board(
+    plane_index,
+    history_index,
+):
+    global board_index
+    board_index += 1
+    if board_index >= len(boards):
+        gr.Warning("Already at last board.")
+        board_index = len(boards) - 1
+    return (
+        *make_plot(
+            plane_index,
+        ),
+        *make_history_plot(
+            history_index,
+        ),
+    )
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board starting FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            action_seq = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                max_lines=1,
+                value=("e2e3 b8c6 d2d4 e7e5 g1f3 d8e7 " "d4d5 e5e4 f3d4 c6e5 f2f4 e5g6"),
+            )
+            compute_cache_button = gr.Button("Compute heatmaps")
+
+            with gr.Group():
+                with gr.Row():
+                    plane_index = gr.Slider(
+                        label="Plane index",
+                        minimum=1,
+                        maximum=112,
+                        step=1,
+                        value=1,
+                    )
+                with gr.Row():
+                    previous_board_button = gr.Button("Previous board")
+                    next_board_button = gr.Button("Next board")
+            current_board_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+            )
+            colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            image = gr.Image(label="Board")
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Group():
+                with gr.Row():
+                    histroy_index = gr.Slider(
+                        label="History index",
+                        minimum=1,
+                        maximum=8,
+                        step=1,
+                        value=1,
+                    )
+            history_colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            history_image = gr.Image(label="Board")
+
+    base_inputs = [
+        plane_index,
+        histroy_index,
+    ]
+    outputs = [
+        image,
+        current_board_fen,
+        colorbar,
+        history_image,
+        history_colorbar,
+    ]
+
+    compute_cache_button.click(
+        compute_cache,
+        inputs=[board_fen, action_seq, model_name] + base_inputs,
+        outputs=outputs,
+    )
+
+    previous_board_button.click(previous_board, inputs=base_inputs, outputs=outputs)
+    next_board_button.click(next_board, inputs=base_inputs, outputs=outputs)
+
+    plane_index.change(
+        make_plot,
+        inputs=plane_index,
+        outputs=[image, current_board_fen, colorbar],
+    )
+    histroy_index.change(
+        make_history_plot,
+        inputs=histroy_index,
+        outputs=[history_image, history_colorbar],
+    )

app/main.py

@@ -0,0 +1,50 @@
+"""
+Gradio demo for lczero-easy.
+"""
+
+import gradio as gr
+
+from . import (
+    attention_interface,
+    backend_interface,
+    board_interface,
+    convert_interface,
+    crp_interface,
+    encoding_interface,
+    lrp_interface,
+    policy_interface,
+    statistics_interface,
+)
+
+demo = gr.TabbedInterface(
+    [
+        crp_interface.interface,
+        statistics_interface.interface,
+        lrp_interface.interface,
+        attention_interface.interface,
+        policy_interface.interface,
+        backend_interface.interface,
+        encoding_interface.interface,
+        board_interface.interface,
+        convert_interface.interface,
+    ],
+    [
+        "CRP",
+        "Statistics",
+        "LRP",
+        "Attention",
+        "Policy",
+        "Backend",
+        "Encoding",
+        "Board",
+        "Convert",
+    ],
+    title="LczeroLens Demo",
+    analytics_enabled=False,
+)
+
+if __name__ == "__main__":
+    demo.launch(
+        server_port=8000,
+        server_name="0.0.0.0",
+    )

app/onnx_models/.gitignore

@@ -0,0 +1,2 @@
+*
+!.gitignore

app/policy_interface.py

@@ -0,0 +1,276 @@
+"""
+Gradio interface for visualizing the policy of a model.
+"""
+
+import chess
+import chess.svg
+import gradio as gr
+import torch
+
+from demo import constants, utils, visualisation
+from lczerolens import move_encodings
+from lczerolens.xai import PolicyLens
+
+current_board = None
+current_raw_policy = None
+current_policy = None
+current_value = None
+current_outcome = None
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info(leela=False)
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def compute_policy(
+    board_fen,
+    action_seq,
+    model_name,
+):
+    global current_board
+    global current_policy
+    global current_raw_policy
+    global current_value
+    global current_outcome
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return (
+            None,
+            None,
+            "",
+        )
+    try:
+        board = chess.Board(board_fen)
+    except ValueError:
+        gr.Warning("Invalid FEN.")
+        return (None, None, "", None)
+    if action_seq:
+        try:
+            for action in action_seq.split():
+                board.push_uci(action)
+        except ValueError:
+            gr.Warning("Invalid action sequence.")
+            return (None, None, "", None)
+    wrapper = utils.get_wrapper_from_state(model_name)
+    (output,) = wrapper.predict(board)
+    current_raw_policy = output["policy"][0]
+    policy = torch.softmax(output["policy"][0], dim=-1)
+
+    filtered_policy = torch.full((1858,), 0.0)
+    legal_moves = [move_encodings.encode_move(move, (board.turn, not board.turn)) for move in board.legal_moves]
+    filtered_policy[legal_moves] = policy[legal_moves]
+    policy = filtered_policy
+
+    current_board = board
+    current_policy = policy
+    current_value = output.get("value", None)
+    current_outcome = output.get("wdl", None)
+
+
+def make_plot(
+    view,
+    aggregate_topk,
+    move_to_play,
+):
+    global current_board
+    global current_policy
+    global current_raw_policy
+    global current_value
+    global current_outcome
+
+    if current_board is None or current_policy is None:
+        gr.Warning("Please compute a policy first.")
+        return (None, None, "", None)
+
+    pickup_agg, dropoff_agg = PolicyLens.aggregate_policy(current_policy, int(aggregate_topk))
+
+    if view == "from":
+        if current_board.turn == chess.WHITE:
+            heatmap = pickup_agg
+        else:
+            heatmap = pickup_agg.view(8, 8).flip(0).view(64)
+    else:
+        if current_board.turn == chess.WHITE:
+            heatmap = dropoff_agg
+        else:
+            heatmap = dropoff_agg.view(8, 8).flip(0).view(64)
+    us_them = (current_board.turn, not current_board.turn)
+    topk_moves = torch.topk(current_policy, 50)
+    move = move_encodings.decode_move(topk_moves.indices[move_to_play - 1], us_them)
+    arrows = [(move.from_square, move.to_square)]
+    svg_board, fig = visualisation.render_heatmap(current_board, heatmap, arrows=arrows)
+    with open(f"{constants.FIGURE_DIRECTORY}/policy.svg", "w") as f:
+        f.write(svg_board)
+    fig_dist = visualisation.render_policy_distribution(
+        current_raw_policy,
+        [move_encodings.encode_move(move, us_them) for move in current_board.legal_moves],
+    )
+    return (
+        f"{constants.FIGURE_DIRECTORY}/policy.svg",
+        fig,
+        (f"Value: {current_value} - WDL: {current_outcome}"),
+        fig_dist,
+    )
+
+
+def make_policy_plot(
+    board_fen,
+    action_seq,
+    view,
+    model_name,
+    aggregate_topk,
+    move_to_play,
+):
+    compute_policy(
+        board_fen,
+        action_seq,
+        model_name,
+    )
+    return make_plot(
+        view,
+        aggregate_topk,
+        move_to_play,
+    )
+
+
+def play_move(
+    board_fen,
+    action_seq,
+    view,
+    model_name,
+    aggregate_topk,
+    move_to_play,
+):
+    global current_board
+    global current_policy
+
+    move = move_encodings.decode_move(
+        current_policy.topk(50).indices[move_to_play - 1],
+        (current_board.turn, not current_board.turn),
+    )
+    current_board.push(move)
+    action_seq = f"{action_seq} {move.uci()}"
+    compute_policy(
+        board_fen,
+        action_seq,
+        model_name,
+    )
+    return [
+        *make_plot(
+            view,
+            aggregate_topk,
+            1,
+        ),
+        action_seq,
+        1,
+    ]
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            action_seq = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                value=("e2e3 b8c6 d2d4 e7e5 g1f3 d8e7 " "d4d5 e5e4 f3d4 c6e5 f2f4 e5g6"),
+            )
+            with gr.Group():
+                with gr.Row():
+                    aggregate_topk = gr.Slider(
+                        label="Aggregate top k",
+                        minimum=1,
+                        maximum=1858,
+                        step=1,
+                        value=1858,
+                        scale=3,
+                    )
+                    view = gr.Radio(
+                        label="View",
+                        choices=["from", "to"],
+                        value="from",
+                        scale=1,
+                    )
+                with gr.Row():
+                    move_to_play = gr.Slider(
+                        label="Move to play",
+                        minimum=1,
+                        maximum=50,
+                        step=1,
+                        value=1,
+                        scale=3,
+                    )
+                    play_button = gr.Button("Play")
+
+            policy_button = gr.Button("Compute policy")
+            colorbar = gr.Plot(label="Colorbar")
+            game_info = gr.Textbox(label="Game info", lines=1, max_lines=1, value="")
+        with gr.Column():
+            image = gr.Image(label="Board")
+            density_plot = gr.Plot(label="Density")
+
+    policy_inputs = [
+        board_fen,
+        action_seq,
+        view,
+        model_name,
+        aggregate_topk,
+        move_to_play,
+    ]
+    policy_outputs = [image, colorbar, game_info, density_plot]
+    policy_button.click(make_policy_plot, inputs=policy_inputs, outputs=policy_outputs)
+    board_fen.submit(make_policy_plot, inputs=policy_inputs, outputs=policy_outputs)
+    action_seq.submit(make_policy_plot, inputs=policy_inputs, outputs=policy_outputs)
+
+    fast_inputs = [
+        view,
+        aggregate_topk,
+        move_to_play,
+    ]
+    aggregate_topk.change(make_plot, inputs=fast_inputs, outputs=policy_outputs)
+    view.change(make_plot, inputs=fast_inputs, outputs=policy_outputs)
+    move_to_play.change(make_plot, inputs=fast_inputs, outputs=policy_outputs)
+
+    play_button.click(
+        play_move,
+        inputs=policy_inputs,
+        outputs=policy_outputs + [action_seq, move_to_play],
+    )

app/state.py

@@ -0,0 +1,18 @@
+"""
+Global state for the demo application.
+"""
+
+from typing import Dict
+
+from lczerolens import Lens, ModelWrapper
+
+wrappers: Dict[str, ModelWrapper] = {}
+
+lenses: Dict[str, Dict[str, Lens]] = {
+    "activation": {},
+    "lrp": {},
+    "crp": {},
+    "policy": {},
+    "probing": {},
+    "patching": {},
+}

app/statistics_interface.py

@@ -0,0 +1,189 @@
+"""
+Gradio interface for visualizing the policy of a model.
+"""
+
+import gradio as gr
+
+from demo import utils, visualisation
+from lczerolens import GameDataset
+from lczerolens.xai import ConceptDataset, HasThreatConcept
+
+current_policy_statistics = None
+current_lrp_statistics = None
+current_probing_statistics = None
+dataset = GameDataset("assets/test_stockfish_10.jsonl")
+check_concept = HasThreatConcept("K", relative=True)
+unique_check_dataset = ConceptDataset.from_game_dataset(dataset)
+unique_check_dataset.set_concept(check_concept)
+
+
+def list_models():
+    """
+    List the models in the model directory.
+    """
+    models_info = utils.get_models_info(leela=False)
+    return sorted([[model_info[0]] for model_info in models_info])
+
+
+def on_select_model_df(
+    evt: gr.SelectData,
+):
+    """
+    When a model is selected, update the statement.
+    """
+    return evt.value
+
+
+def compute_policy_statistics(
+    model_name,
+):
+    global current_policy_statistics
+    global dataset
+
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return None
+    wrapper, lens = utils.get_wrapper_lens_from_state(model_name, "policy")
+    current_policy_statistics = lens.analyse_dataset(dataset, wrapper, 10)
+    return make_policy_plot()
+
+
+def make_policy_plot():
+    global current_policy_statistics
+
+    if current_policy_statistics is None:
+        gr.Warning(
+            "Please compute policy statistics first.",
+        )
+        return None
+    else:
+        return visualisation.render_policy_statistics(current_policy_statistics)
+
+
+def compute_lrp_statistics(
+    model_name,
+):
+    global current_lrp_statistics
+    global dataset
+
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return None, None, None
+    wrapper, lens = utils.get_wrapper_lens_from_state(model_name, "lrp")
+    current_lrp_statistics = lens.compute_statistics(dataset, wrapper, 10)
+    return make_lrp_plot()
+
+
+def make_lrp_plot():
+    global current_lrp_statistics
+
+    if current_lrp_statistics is None:
+        gr.Warning(
+            "Please compute LRP statistics first.",
+        )
+        return None, None, None
+    else:
+        return visualisation.render_relevance_proportion(current_lrp_statistics)
+
+
+def compute_probing_statistics(
+    model_name,
+):
+    global current_probing_statistics
+    global check_concept
+    global unique_check_dataset
+
+    if model_name == "":
+        gr.Warning(
+            "Please select a model.",
+        )
+        return None
+    wrapper, lens = utils.get_wrapper_lens_from_state(model_name, "probing", concept=check_concept)
+    current_probing_statistics = lens.compute_statistics(unique_check_dataset, wrapper, 10)
+    return make_probing_plot()
+
+
+def make_probing_plot():
+    global current_probing_statistics
+
+    if current_probing_statistics is None:
+        gr.Warning(
+            "Please compute probing statistics first.",
+        )
+        return None
+    else:
+        return visualisation.render_probing_statistics(current_probing_statistics)
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_df = gr.Dataframe(
+                headers=["Available models"],
+                datatype=["str"],
+                interactive=False,
+                type="array",
+                value=list_models,
+            )
+        with gr.Column(scale=1):
+            with gr.Row():
+                model_name = gr.Textbox(label="Selected model", lines=1, interactive=False, scale=7)
+    model_df.select(
+        on_select_model_df,
+        None,
+        model_name,
+    )
+
+    with gr.Row():
+        with gr.Column():
+            policy_plot = gr.Plot(label="Policy statistics")
+            policy_compute_button = gr.Button(value="Compute policy statistics")
+            policy_plot_button = gr.Button(value="Plot policy statistics")
+
+            policy_compute_button.click(
+                compute_policy_statistics,
+                inputs=[model_name],
+                outputs=[policy_plot],
+            )
+            policy_plot_button.click(make_policy_plot, outputs=[policy_plot])
+
+        with gr.Column():
+            lrp_plot_hist = gr.Plot(label="LRP history statistics")
+
+    with gr.Row():
+        with gr.Column():
+            lrp_plot_planes = gr.Plot(label="LRP planes statistics")
+
+        with gr.Column():
+            lrp_plot_pieces = gr.Plot(label="LRP pieces statistics")
+
+    with gr.Row():
+        lrp_compute_button = gr.Button(value="Compute LRP statistics")
+    with gr.Row():
+        lrp_plot_button = gr.Button(value="Plot LRP statistics")
+
+    lrp_compute_button.click(
+        compute_lrp_statistics,
+        inputs=[model_name],
+        outputs=[lrp_plot_hist, lrp_plot_planes, lrp_plot_pieces],
+    )
+    lrp_plot_button.click(
+        make_lrp_plot,
+        outputs=[lrp_plot_hist, lrp_plot_planes, lrp_plot_pieces],
+    )
+
+    with gr.Column():
+        probing_plot = gr.Plot(label="Probing statistics")
+        probing_compute_button = gr.Button(value="Compute probing statistics")
+        probing_plot_button = gr.Button(value="Plot probing statistics")
+
+        probing_compute_button.click(
+            compute_probing_statistics,
+            inputs=[model_name],
+            outputs=[probing_plot],
+        )
+        probing_plot_button.click(make_probing_plot, outputs=[probing_plot])

app/utils.py

@@ -0,0 +1,121 @@
+"""
+Utils for the demo app.
+"""
+
+import os
+import re
+import subprocess
+
+from demo import constants, state
+from lczerolens import LensFactory, LczeroModel
+from lczerolens.model import lczero as lczero_utils
+
+
+def get_models_info(onnx=True, leela=True):
+    """
+    Get the names of the models in the model directory.
+    """
+    model_df = []
+    exp = r"(?P<n_filters>\d+)x(?P<n_blocks>\d+)"
+    if onnx:
+        for filename in os.listdir(constants.MODEL_DIRECTORY):
+            if filename.endswith(".onnx"):
+                match = re.search(exp, filename)
+                if match is None:
+                    n_filters = -1
+                    n_blocks = -1
+                else:
+                    n_filters = int(match.group("n_filters"))
+                    n_blocks = int(match.group("n_blocks"))
+                model_df.append(
+                    [
+                        filename,
+                        "ONNX",
+                        n_blocks,
+                        n_filters,
+                    ]
+                )
+    if leela:
+        for filename in os.listdir(constants.LEELA_MODEL_DIRECTORY):
+            if filename.endswith(".pb.gz"):
+                match = re.search(exp, filename)
+                if match is None:
+                    n_filters = -1
+                    n_blocks = -1
+                else:
+                    n_filters = int(match.group("n_filters"))
+                    n_blocks = int(match.group("n_blocks"))
+                model_df.append(
+                    [
+                        filename,
+                        "LEELA",
+                        n_blocks,
+                        n_filters,
+                    ]
+                )
+    return model_df
+
+
+def save_model(tmp_file_path):
+    """
+    Save the model to the model directory.
+    """
+    popen = subprocess.Popen(
+        ["file", tmp_file_path],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    popen.wait()
+    if popen.returncode != 0:
+        raise RuntimeError
+    file_desc = popen.stdout.read().decode("utf-8").split(tmp_file_path)[1].strip()
+    rename_match = re.search(r"was\s\"(?P<name>.+)\"", file_desc)
+    type_match = re.search(r"\:\s(?P<type>[a-zA-Z]+)", file_desc)
+    if rename_match is None or type_match is None:
+        raise RuntimeError
+    model_name = rename_match.group("name")
+    model_type = type_match.group("type")
+    if model_type != "gzip":
+        raise RuntimeError
+    os.rename(
+        tmp_file_path,
+        f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}.gz",
+    )
+    try:
+        lczero_utils.describenet(
+            f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}.gz",
+        )
+    except RuntimeError:
+        os.remove(f"{constants.LEELA_MODEL_DIRECTORY}/{model_name}.gz")
+        raise RuntimeError
+
+
+def get_wrapper_from_state(model_name):
+    """
+    Get the model wrapper from the state.
+    """
+    if model_name in state.wrappers:
+        return state.wrappers[model_name]
+    else:
+        wrapper = LczeroModel.from_path(f"{constants.MODEL_DIRECTORY}/{model_name}")
+        state.wrappers[model_name] = wrapper
+        return wrapper
+
+
+def get_wrapper_lens_from_state(model_name, lens_type, lens_name="lens", **kwargs):
+    """
+    Get the model wrapper and lens from the state.
+    """
+    if model_name in state.wrappers:
+        wrapper = state.wrappers[model_name]
+    else:
+        wrapper = LczeroModel.from_path(f"{constants.MODEL_DIRECTORY}/{model_name}")
+        state.wrappers[model_name] = wrapper
+    if lens_name in state.lenses[lens_type]:
+        lens = state.lenses[lens_type][lens_name]
+    else:
+        lens = LensFactory.from_name(lens_type, **kwargs)
+        if not lens.is_compatible(wrapper):
+            raise ValueError(f"Lens of type {lens_type} not compatible with model.")
+        state.lenses[lens_type][lens_name] = lens
+    return wrapper, lens

app/visualisation.py

@@ -0,0 +1,303 @@
+"""
+Visualisation utils.
+"""
+
+import chess
+import chess.svg
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torchviz
+
+from . import constants
+
+COLOR_MAP = matplotlib.colormaps["RdYlBu_r"].resampled(1000)
+ALPHA = 1.0
+
+
+def render_heatmap(
+    board,
+    heatmap,
+    square=None,
+    vmin=None,
+    vmax=None,
+    arrows=None,
+    normalise="none",
+):
+    """
+    Render a heatmap on the board.
+    """
+    if normalise == "abs":
+        a_max = heatmap.abs().max()
+        if a_max != 0:
+            heatmap = heatmap / a_max
+        vmin = -1
+        vmax = 1
+    if vmin is None:
+        vmin = heatmap.min()
+    if vmax is None:
+        vmax = heatmap.max()
+    norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=False)
+
+    color_dict = {}
+    for square_index in range(64):
+        color = COLOR_MAP(norm(heatmap[square_index]))
+        color = (*color[:3], ALPHA)
+        color_dict[square_index] = matplotlib.colors.to_hex(color, keep_alpha=True)
+    fig = plt.figure(figsize=(6, 0.6))
+    ax = plt.gca()
+    ax.axis("off")
+    fig.colorbar(
+        matplotlib.cm.ScalarMappable(norm=norm, cmap=COLOR_MAP),
+        ax=ax,
+        orientation="horizontal",
+        fraction=1.0,
+    )
+    if square is not None:
+        try:
+            check = chess.parse_square(square)
+        except ValueError:
+            check = None
+    else:
+        check = None
+    if arrows is None:
+        arrows = []
+    plt.close()
+    return (
+        chess.svg.board(
+            board,
+            check=check,
+            fill=color_dict,
+            size=350,
+            arrows=arrows,
+        ),
+        fig,
+    )
+
+
+def render_architecture(model, name: str = "model", directory: str = ""):
+    """
+    Render the architecture of the model.
+    """
+    out = model(torch.zeros(1, 112, 8, 8))
+    if len(out) == 2:
+        policy, outcome_probs = out
+        value = torch.zeros(outcome_probs.shape[0], 1)
+    else:
+        policy, outcome_probs, value = out
+    torchviz.make_dot(policy, params=dict(list(model.named_parameters()))).render(
+        f"{directory}/{name}_policy", format="svg"
+    )
+    torchviz.make_dot(outcome_probs, params=dict(list(model.named_parameters()))).render(
+        f"{directory}/{name}_outcome_probs", format="svg"
+    )
+    torchviz.make_dot(value, params=dict(list(model.named_parameters()))).render(
+        f"{directory}/{name}_value", format="svg"
+    )
+
+
+def render_policy_distribution(
+    policy,
+    legal_moves,
+    n_bins=20,
+):
+    """
+    Render the policy distribution histogram.
+    """
+    legal_mask = torch.Tensor([move in legal_moves for move in range(1858)]).bool()
+    fig = plt.figure(figsize=(6, 6))
+    ax = plt.gca()
+    _, bins = np.histogram(policy, bins=n_bins)
+    ax.hist(
+        policy[~legal_mask],
+        bins=bins,
+        alpha=0.5,
+        density=True,
+        label="Illegal moves",
+    )
+    ax.hist(
+        policy[legal_mask],
+        bins=bins,
+        alpha=0.5,
+        density=True,
+        label="Legal moves",
+    )
+    plt.xlabel("Policy")
+    plt.ylabel("Density")
+    plt.legend()
+    plt.yscale("log")
+    return fig
+
+
+def render_policy_statistics(
+    statistics,
+):
+    """
+    Render the policy statistics.
+    """
+    fig = plt.figure(figsize=(6, 6))
+    ax = plt.gca()
+    move_indices = list(statistics["mean_legal_logits"].keys())
+    legal_means_avg = [np.mean(statistics["mean_legal_logits"][move_idx]) for move_idx in move_indices]
+    illegal_means_avg = [np.mean(statistics["mean_illegal_logits"][move_idx]) for move_idx in move_indices]
+    legal_means_std = [np.std(statistics["mean_legal_logits"][move_idx]) for move_idx in move_indices]
+    illegal_means_std = [np.std(statistics["mean_illegal_logits"][move_idx]) for move_idx in move_indices]
+    ax.errorbar(
+        move_indices,
+        legal_means_avg,
+        yerr=legal_means_std,
+        label="Legal moves",
+    )
+    ax.errorbar(
+        move_indices,
+        illegal_means_avg,
+        yerr=illegal_means_std,
+        label="Illegal moves",
+    )
+    plt.xlabel("Move index")
+    plt.ylabel("Mean policy logits")
+    plt.legend()
+    return fig
+
+
+def render_relevance_proportion(statistics, scaled=True):
+    """
+    Render the relevance proportion statistics.
+    """
+    norm = matplotlib.colors.Normalize(vmin=0, vmax=1, clip=False)
+    fig_hist = plt.figure(figsize=(6, 6))
+    ax = plt.gca()
+    move_indices = list(statistics["planes_relevance_proportion"].keys())
+    for h in range(8):
+        relevance_proportion_avg = [
+            np.mean([rel[13 * h : 13 * (h + 1)].sum() for rel in statistics["planes_relevance_proportion"][move_idx]])
+            for move_idx in move_indices
+        ]
+        relevance_proportion_std = [
+            np.std([rel[13 * h : 13 * (h + 1)].sum() for rel in statistics["planes_relevance_proportion"][move_idx]])
+            for move_idx in move_indices
+        ]
+        ax.errorbar(
+            move_indices[h + 1 :],
+            relevance_proportion_avg[h + 1 :],
+            yerr=relevance_proportion_std[h + 1 :],
+            label=f"History {h}",
+            c=COLOR_MAP(norm(h / 9)),
+        )
+
+    relevance_proportion_avg = [
+        np.mean([rel[104:108].sum() for rel in statistics["planes_relevance_proportion"][move_idx]])
+        for move_idx in move_indices
+    ]
+    relevance_proportion_std = [
+        np.std([rel[104:108].sum() for rel in statistics["planes_relevance_proportion"][move_idx]])
+        for move_idx in move_indices
+    ]
+    ax.errorbar(
+        move_indices,
+        relevance_proportion_avg,
+        yerr=relevance_proportion_std,
+        label="Castling rights",
+        c=COLOR_MAP(norm(8 / 9)),
+    )
+    relevance_proportion_avg = [
+        np.mean([rel[108:].sum() for rel in statistics["planes_relevance_proportion"][move_idx]])
+        for move_idx in move_indices
+    ]
+    relevance_proportion_std = [
+        np.std([rel[108:].sum() for rel in statistics["planes_relevance_proportion"][move_idx]])
+        for move_idx in move_indices
+    ]
+    ax.errorbar(
+        move_indices,
+        relevance_proportion_avg,
+        yerr=relevance_proportion_std,
+        label="Remaining planes",
+        c=COLOR_MAP(norm(9 / 9)),
+    )
+    plt.xlabel("Move index")
+    plt.ylabel("Absolute relevance proportion")
+    plt.yscale("log")
+    plt.legend()
+
+    if scaled:
+        stat_key = "planes_relevance_proportion_scaled"
+    else:
+        stat_key = "planes_relevance_proportion"
+    fig_planes = plt.figure(figsize=(6, 6))
+    ax = plt.gca()
+    move_indices = list(statistics[stat_key].keys())
+    for p in range(13):
+        relevance_proportion_avg = [
+            np.mean([rel[p].item() for rel in statistics[stat_key][move_idx]]) for move_idx in move_indices
+        ]
+        relevance_proportion_std = [
+            np.std([rel[p].item() for rel in statistics[stat_key][move_idx]]) for move_idx in move_indices
+        ]
+        ax.errorbar(
+            move_indices,
+            relevance_proportion_avg,
+            yerr=relevance_proportion_std,
+            label=constants.PLANE_NAMES[p],
+            c=COLOR_MAP(norm(p / 12)),
+        )
+
+    plt.xlabel("Move index")
+    plt.ylabel("Absolute relevance proportion")
+    plt.yscale("log")
+    plt.legend()
+
+    fig_pieces = plt.figure(figsize=(6, 6))
+    ax = plt.gca()
+    for p in range(1, 13):
+        stat_key = f"configuration_relevance_proportion_threatened_piece{p}"
+        n_attackers = list(statistics[stat_key].keys())
+        relevance_proportion_avg = [
+            np.mean(statistics[f"configuration_relevance_proportion_threatened_piece{p}"][n]) for n in n_attackers
+        ]
+        relevance_proportion_std = [np.std(statistics[stat_key][n]) for n in n_attackers]
+        ax.errorbar(
+            n_attackers,
+            relevance_proportion_avg,
+            yerr=relevance_proportion_std,
+            label="PNBRQKpnbrqk"[p - 1],
+            c=COLOR_MAP(norm(p / 12)),
+        )
+
+    plt.xlabel("Number of attackers")
+    plt.ylabel("Absolute configuration relevance proportion")
+    plt.yscale("log")
+    plt.legend()
+
+    return fig_hist, fig_planes, fig_pieces
+
+
+def render_probing_statistics(
+    statistics,
+):
+    """
+    Render the probing statistics.
+    """
+    fig = plt.figure(figsize=(6, 6))
+    ax = plt.gca()
+    n_blocks = len(statistics["metrics"])
+    for metric in statistics["metrics"]["block0"]:
+        avg = []
+        std = []
+        for block_idx in range(n_blocks):
+            metrics = statistics["metrics"]
+            block_data = metrics[f"block{block_idx}"]
+            avg.append(np.mean(block_data[metric]))
+            std.append(np.std(block_data[metric]))
+        ax.errorbar(
+            range(n_blocks),
+            avg,
+            yerr=std,
+            label=metric,
+        )
+    plt.xlabel("Block index")
+    plt.ylabel("Metric")
+    plt.yscale("log")
+    plt.legend()
+    return fig