new working demo

.gitignore

@@ -0,0 +1,140 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Pickle files
+*.pkl
+
+# Various files
+ignored
+debug
+*.zip
+lc0
+!bin/lc0
+wandb
+**/.DS_Store
+
+*secret*

.pre-commit-config.yaml

@@ -1,18 +0,0 @@
-repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
-    hooks:
-    -   id: check-added-large-files
-        args: ['--maxkb=600']
-    -   id: check-yaml
-    -   id: check-json
-    -   id: check-toml
-    -   id: end-of-file-fixer
-    -   id: trailing-whitespace
-    -   id: check-docstring-first
-- repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.4.2
-  hooks:
-    - id: ruff
-      args: [ --fix ]
-    - id: ruff-format

README.md

@@ -4,11 +4,11 @@ emoji: 🔬
 colorFrom: gray
 colorTo: green
 sdk: gradio
-sdk_version: 5.11.0
-app_file: app/main.py
-pinned: false
+sdk_version: 5.28.0
+app_file: main.py
+pinned: true
 license: mit
 short_description: Demo lczerolens features
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+See the documentation [here](https://lczerolens.readthedocs.io/).

app/attention_interface.py

@@ -1,293 +0,0 @@
-"""
-Gradio interface for plotting attention.
-"""
-
-import copy
-
-import chess
-import gradio as gr
-
-from lczerolens.board import LczeroBoard
-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 = LczeroBoard(board_fen)
-    except ValueError:
-        board = LczeroBoard()
-        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

@@ -1,208 +0,0 @@
-"""
-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
-from lczerolens.board import LczeroBoard
-
-
-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 = LczeroBoard(board_fen)
-    except ValueError:
-        board = LczeroBoard()
-        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 = LczeroBoard()
-    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/constants.py

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

app/convert_interface.py

@@ -1,201 +0,0 @@
-"""
-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

@@ -1,281 +0,0 @@
-"""
-Gradio interface for plotting policy.
-"""
-
-import copy
-
-import chess
-import gradio as gr
-
-from demo import constants, utils, visualisation
-from lczerolens.board import LczeroBoard
-
-
-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 = LczeroBoard(board_fen)
-    except ValueError:
-        board = LczeroBoard()
-        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 = LczeroBoard(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

@@ -1,83 +0,0 @@
-"""
-Gradio interface for plotting encodings.
-"""
-
-import chess
-import gradio as gr
-
-from demo import constants, visualisation
-from lczerolens import board_encodings
-from lczerolens.board import LczeroBoard
-
-
-def make_encoding_plot(
-    board_fen,
-    action_seq,
-    plane_index,
-    color_flip,
-):
-    try:
-        board = LczeroBoard(board_fen)
-    except ValueError:
-        board = LczeroBoard()
-        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 = LczeroBoard()
-    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/lrp_interface.py

@@ -1,280 +0,0 @@
-"""
-Gradio interface for plotting policy.
-"""
-
-import copy
-
-import chess
-import gradio as gr
-
-from demo import constants, utils, visualisation
-from lczerolens.board import LczeroBoard
-
-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 = LczeroBoard(board_fen)
-    except ValueError:
-        board = LczeroBoard()
-        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 = LczeroBoard(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

@@ -1,50 +0,0 @@
-"""
-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/policy_interface.py

@@ -1,278 +0,0 @@
-"""
-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.board import LczeroBoard
-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 = LczeroBoard(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

@@ -1,18 +0,0 @@
-"""
-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

@@ -1,189 +0,0 @@
-"""
-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

@@ -1,121 +0,0 @@
-"""
-Utils for the demo app.
-"""
-
-import os
-import re
-import subprocess
-
-from demo import constants, state
-from lczerolens import Lens, 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 = Lens.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

@@ -1,303 +0,0 @@
-"""
-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

demo/constants.py

@@ -0,0 +1,18 @@
+"""
+Constants for the demo.
+"""
+
+import os
+
+ONNX_MODEL_DIRECTORY = "demo/onnx-models"
+LEELA_MODEL_DIRECTORY = "demo/leela-models"
+FIGURE_DIRECTORY = "demo/figures"
+
+ONNX_MODEL_NAMES = [
+    f for f in os.listdir(ONNX_MODEL_DIRECTORY)
+    if f.endswith(".onnx")
+]
+LEELA_MODEL_NAMES = [
+    f for f in os.listdir(LEELA_MODEL_DIRECTORY)
+    if f.endswith(".pb.gz")
+]

demo/interfaces/activations.py

@@ -0,0 +1,157 @@
+"""
+Gradio interface for plotting attention.
+"""
+
+import chess
+import chess.pgn
+import io
+import gradio as gr
+import os
+import torch
+
+from lczerolens import LczeroBoard, LczeroModel, Lens
+
+from .. import constants
+
+def get_model(model_name: str):
+    return LczeroModel.from_onnx_path(os.path.join(constants.ONNX_MODEL_DIRECTORY, model_name))
+
+def get_activations(model: LczeroModel, board: LczeroBoard):
+    lens = Lens.from_name("activation", "block\d/conv2/relu")
+    with torch.no_grad():
+        results = lens.analyse(model, board)
+    return [results[f"block{i}/conv2/relu_output"][0] for i in range(len(results))]
+
+def get_board(game_pgn:str, board_fen:str):
+    if game_pgn:
+        try:
+            board = LczeroBoard()
+            pgn = io.StringIO(game_pgn)
+            game = chess.pgn.read_game(pgn)
+            for move in game.mainline_moves():
+                board.push(move)
+        except Exception as e:
+            print(e)
+            gr.Warning("Error parsing PGN, using starting position.")
+            board = LczeroBoard()
+    else:
+        try:
+            board = LczeroBoard(board_fen)
+        except Exception as e:
+            print(e)
+            gr.Warning("Invalid FEN, using starting position.")
+            board = LczeroBoard()
+    return board
+
+def render_activations(board: LczeroBoard, activations, layer_index:int, channel_index:int):
+    if layer_index >= len(activations):
+        safe_layer_index = len(activations) - 1
+        gr.Warning(f"Layer index {layer_index} out of range, using last layer ({safe_layer_index}).")
+    else:
+        safe_layer_index = layer_index
+    if channel_index >= activations[safe_layer_index].shape[0]:
+        safe_channel_index = activations[safe_layer_index].shape[0] - 1
+        gr.Warning(f"Channel index {channel_index} out of range, using last channel ({safe_channel_index}).")
+    else:
+        safe_channel_index = channel_index
+    heatmap = activations[safe_layer_index][safe_channel_index].view(64)
+    board.render_heatmap(
+        heatmap,
+        save_to=f"{constants.FIGURE_DIRECTORY}/activations.svg",
+    )
+    return  f"{constants.FIGURE_DIRECTORY}/activations_board.svg", f"{constants.FIGURE_DIRECTORY}/activations_colorbar.svg"
+
+def initial_load(model_name: str, board_fen: str, game_pgn: str, layer_index: int, channel_index: int):
+    model = get_model(model_name)
+    board = get_board(game_pgn, board_fen)
+    activations = get_activations(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return model, board, activations, *plots
+
+def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, layer_index: int, channel_index: int):
+    board = get_board(game_pgn, board_fen)
+    activations = get_activations(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return board, activations, *plots
+
+def on_model_change(model_name: str, board: LczeroBoard, layer_index: int, channel_index: int):
+    model = get_model(model_name)
+    activations = get_activations(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return model, activations, *plots
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            with gr.Group():
+                gr.Markdown(
+                    "Specify the game PGN or FEN string that you want to analyse (PGN overrides FEN)."
+                )
+                game_pgn = gr.Textbox(
+                    label="Game PGN",
+                    lines=1,
+                    value="",
+                )
+                board_fen = gr.Textbox(
+                    label="Board FEN",
+                    lines=1,
+                    max_lines=1,
+                    value=chess.STARTING_FEN,
+                )
+            with gr.Group():
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=constants.ONNX_MODEL_NAMES,
+                )
+                layer_index = gr.Slider(
+                    label="Layer index",
+                    minimum=0,
+                    maximum=19,
+                    step=1,
+                    value=0,
+                )
+                channel_index = gr.Slider(
+                    label="Channel index",
+                    minimum=0,
+                    maximum=200,
+                    step=1,
+                    value=0,
+                )
+        with gr.Column():
+            image_board = gr.Image(label="Board", interactive=False)
+            colorbar = gr.Image(label="Colorbar", interactive=False)
+
+    model = gr.State(value=None)
+    board = gr.State(value=None)
+    activations = gr.State(value=None)
+
+    interface.load(
+        initial_load,
+        inputs=[model_name, game_pgn, board_fen, layer_index, channel_index],
+        outputs=[model, board, activations, image_board, colorbar],
+    )
+    game_pgn.submit(
+        on_board_change,
+        inputs=[model, game_pgn, board_fen, layer_index, channel_index],
+        outputs=[board, activations, image_board, colorbar],
+    )
+    board_fen.submit(
+        on_board_change,
+        inputs=[model, game_pgn, board_fen, layer_index, channel_index],
+        outputs=[board, activations, image_board, colorbar],
+    )
+    model_name.change(
+        on_model_change,
+        inputs=[model_name, board, layer_index, channel_index],
+        outputs=[model, activations, image_board, colorbar],
+    )
+    layer_index.change(
+        render_activations,
+        inputs=[board, activations, layer_index, channel_index],
+        outputs=[image_board, colorbar],
+    )
+    channel_index.change(
+        render_activations,
+        inputs=[board, activations, layer_index, channel_index],
+        outputs=[image_board, colorbar],
+    )

app/board_interface.py → demo/interfaces/board.py

@@ -3,11 +3,12 @@ Gradio interface for plotting a board.
 """
 
 import chess
+import chess.svg
 import gradio as gr
 
-from demo import constants
 from lczerolens.board import LczeroBoard
 
+from ..utils import create_board_figure
 
 def make_board_plot(board_fen, arrows, square):
     try:
@@ -15,34 +16,8 @@ def make_board_plot(board_fen, arrows, square):
     except ValueError:
         board = LczeroBoard()
         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"
+    filepath = create_board_figure(board, arrows=arrows, square=square, name="board")
+    return filepath
 
 
 with gr.Blocks() as interface:
@@ -76,6 +51,7 @@ with gr.Blocks() as interface:
         arrows,
         square,
     ]
+    interface.load(make_board_plot, inputs=inputs, outputs=image)
     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)
+    square.submit(make_board_plot, inputs=inputs, outputs=image)

demo/interfaces/encodings.py

@@ -0,0 +1,99 @@
+"""
+Gradio interface for plotting attention.
+"""
+
+import chess
+import chess.pgn
+import io
+import gradio as gr
+
+from lczerolens.board import LczeroBoard
+
+from ..constants import FIGURE_DIRECTORY
+
+def make_render(game_pgn:str, board_fen:str, plane_index:int):
+    if game_pgn:
+        try:
+            board = LczeroBoard()
+            pgn = io.StringIO(game_pgn)
+            game = chess.pgn.read_game(pgn)
+            for move in game.mainline_moves():
+                board.push(move)
+        except Exception as e:
+            print(e)
+            gr.Warning("Error parsing PGN, using starting position.")
+            board = LczeroBoard()
+    else:
+        try:
+            board = LczeroBoard(board_fen)
+        except Exception as e:
+            print(e)
+            gr.Warning("Invalid FEN, using starting position.")
+            board = LczeroBoard()
+    return board, *make_board_plot(board, plane_index)
+
+def make_board_plot(board:LczeroBoard, plane_index:int):
+    input_tensor = board.to_input_tensor()
+    board.render_heatmap(
+        input_tensor[plane_index].view(64),
+        save_to=f"{FIGURE_DIRECTORY}/encodings.svg",
+        vmin=0,
+        vmax=1,
+    )
+    return  f"{FIGURE_DIRECTORY}/encodings_board.svg", f"{FIGURE_DIRECTORY}/encodings_colorbar.svg"
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            with gr.Group():
+                gr.Markdown(
+                    "Specify the game PGN or FEN string that you want to analyse (PGN overrides FEN)."
+                )
+                game_pgn = gr.Textbox(
+                    label="Game PGN",
+                    lines=1,
+                    value="",
+                )
+                board_fen = gr.Textbox(
+                    label="Board FEN",
+                    lines=1,
+                    max_lines=1,
+                    value=chess.STARTING_FEN,
+                )
+            with gr.Group():
+                with gr.Row():
+                    plane_index = gr.Slider(
+                        label="Plane index",
+                        minimum=0,
+                        maximum=111,
+                        step=1,
+                        value=0,
+                    )
+        with gr.Column():
+            image_board = gr.Image(label="Board", interactive=False)
+            colorbar = gr.Image(label="Colorbar", interactive=False)
+
+    state_board = gr.State(value=LczeroBoard())
+
+    render_inputs = [game_pgn, board_fen, plane_index]
+    render_outputs = [state_board, image_board, colorbar]
+    interface.load(
+        make_render,
+        inputs=render_inputs,
+        outputs=render_outputs,
+    )
+    game_pgn.submit(
+        make_render,
+        inputs=render_inputs,
+        outputs=render_outputs,
+    )
+    board_fen.submit(
+        make_render,
+        inputs=render_inputs,
+        outputs=render_outputs,
+    )
+    plane_index.change(
+        make_board_plot,
+        inputs=[state_board, plane_index],
+        outputs=[image_board, colorbar],
+    )

demo/interfaces/gradients.py

@@ -0,0 +1,174 @@
+"""
+Gradio interface for plotting attention.
+"""
+
+import chess
+import chess.pgn
+import io
+import gradio as gr
+import os
+
+from lczerolens import LczeroBoard, LczeroModel, Lens
+
+from .. 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):
+    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]
+    results = lens.analyse(model, board, init_target=init_target)
+
+    return results["input_grad"]
+
+def get_board(game_pgn:str, board_fen:str):
+    if game_pgn:
+        try:
+            board = LczeroBoard()
+            pgn = io.StringIO(game_pgn)
+            game = chess.pgn.read_game(pgn)
+            for move in game.mainline_moves():
+                board.push(move)
+        except Exception as e:
+            print(e)
+            gr.Warning("Error parsing PGN, using starting position.")
+            board = LczeroBoard()
+    else:
+        try:
+            board = LczeroBoard(board_fen)
+        except Exception as e:
+            print(e)
+            gr.Warning("Invalid FEN, using starting position.")
+            board = LczeroBoard()
+    return board
+
+def render_gradients(board: LczeroBoard, gradients, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index:int):
+    if average_over_planes:
+        heatmap = gradients[0, begin_average_index:end_average_index].mean(dim=0).view(64)
+    else:
+        heatmap = gradients[0, plane_index].view(64)
+    board.render_heatmap(
+        heatmap,
+        save_to=f"{constants.FIGURE_DIRECTORY}/gradients.svg",
+    )
+    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):
+    model = get_model(model_name)
+    board = get_board(game_pgn, board_fen)
+    gradients = get_gradients(model, board, wdl_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):
+    board = get_board(game_pgn, board_fen)
+    gradients = get_gradients(model, board, wdl_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):
+    model = get_model(model_name)
+    gradients = get_gradients(model, board, wdl_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)
+    plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
+    return gradients, *plots
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            with gr.Group():
+                gr.Markdown(
+                    "Specify the game PGN or FEN string that you want to analyse (PGN overrides FEN)."
+                )
+                game_pgn = gr.Textbox(
+                    label="Game PGN",
+                    lines=1,
+                    value="",
+                )
+                board_fen = gr.Textbox(
+                    label="Board FEN",
+                    lines=1,
+                    max_lines=1,
+                    value=chess.STARTING_FEN,
+                )
+            with gr.Group():
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=constants.ONNX_MODEL_NAMES,
+                )
+                wdl_target = gr.Radio(
+                    ["win", "draw", "loss"], label="WDL target",
+                    value="win",
+                )
+            with gr.Group():
+                average_over_planes = gr.Checkbox(label="Average over Planes", value=False)
+                with gr.Accordion("Average over planes", open=False):
+                    begin_average_index = gr.Slider(
+                        label="Begin average index",
+                        minimum=0,
+                        maximum=111,
+                        step=1,
+                        value=0,
+                    )
+                    end_average_index = gr.Slider(
+                        label="End average index",
+                        minimum=0,
+                        maximum=111,
+                        step=1,
+                        value=111,
+                    )
+                plane_index = gr.Slider(
+                    label="Plane index",
+                    minimum=0,
+                    maximum=111,
+                    step=1,
+                    value=0,
+                )
+
+        with gr.Column():
+            image_board = gr.Image(label="Board", interactive=False)
+            colorbar = gr.Image(label="Colorbar", interactive=False)
+
+    model = gr.State(value=None)
+    board = gr.State(value=None)
+    gradients = gr.State(value=None)
+
+    interface.load(
+        initial_load,
+        inputs=[model_name, game_pgn, board_fen, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[model, board, gradients, image_board, colorbar],
+    )
+    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],
+        outputs=[board, gradients, image_board, colorbar],
+    )
+    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],
+        outputs=[board, gradients, image_board, colorbar],
+    )
+    model_name.change(
+        on_model_change,
+        inputs=[model_name, board, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[model, gradients, image_board, colorbar],
+    )
+    wdl_target.change(
+        on_wdl_target_change,
+        inputs=[model, board, wdl_target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[gradients, image_board, colorbar],
+    )
+    for render_arg in [average_over_planes, begin_average_index, end_average_index, plane_index]:
+        render_arg.change(
+            render_gradients,
+            inputs=[board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index],
+            outputs=[image_board, colorbar],
+        )

demo/interfaces/play.py

@@ -0,0 +1,140 @@
+"""Interface to play against the model.
+"""
+
+import os
+
+import chess
+import chess.pgn
+import random
+import gradio as gr
+
+from lczerolens import LczeroBoard, LczeroModel
+from lczerolens.play import PolicySampler
+
+from .. import constants
+from ..utils import create_board_figure
+
+
+def get_sampler(model_name: str):
+    model = LczeroModel.from_onnx_path(os.path.join(constants.ONNX_MODEL_DIRECTORY, model_name))
+    return PolicySampler(model)
+
+def get_pgn(board: LczeroBoard):
+    game = chess.pgn.Game()
+    for move in board.move_stack:
+        game.add_variation(move)
+    return str(game).split("\n")[-1]
+
+def render_board(
+    board: LczeroBoard,
+):
+    player = board.turn
+    if len(board.move_stack) > 0:
+        last_move_uci = board.peek().uci()
+    else:
+        last_move_uci = None
+    
+    if board.is_check():
+        check = board.king(board.turn)
+    else:
+        check = None
+    filepath = create_board_figure(
+        board, 
+        orientation=player,
+        arrows=last_move_uci,
+        square=check,
+        name="play_board",
+    )
+    return filepath
+
+def gather_outputs(board: LczeroBoard, sampler: PolicySampler):
+    return sampler, board, board.fen(), get_pgn(board), render_board(board), ""
+
+def get_init(model_name: str):
+    sampler = get_sampler(model_name)
+    is_ai_white = random.choice([True, False])
+    init_board = LczeroBoard()
+    if is_ai_white:
+        play_ai_move(init_board, sampler)
+    return gather_outputs(init_board, sampler)
+
+def play_user_move_then_ai_move(
+    uci_move: str,
+    board: LczeroBoard,
+    sampler: PolicySampler,
+):
+    board.push_uci(uci_move)
+    play_ai_move(board, sampler)
+    return gather_outputs(board, sampler)
+    
+
+def play_ai_move(
+    board: LczeroBoard,
+    sampler: PolicySampler,
+):
+    move, _ = next(iter(sampler.get_next_moves([board])))
+    board.push(move)
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            current_fen = gr.Textbox(
+                label="Board FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            current_pgn = gr.Textbox(
+                label="Action sequence",
+                lines=1,
+                value="",
+            )
+            with gr.Row():
+                move_to_play = gr.Textbox(
+                    label="Move to play (UCI)",
+                    lines=1,
+                    max_lines=1,
+                    value="",
+                )
+                with gr.Column():
+                    model_name = gr.Dropdown(
+                        label="Model",
+                        choices=constants.ONNX_MODEL_NAMES,
+                    )
+                    play_button = gr.Button("Play")
+            reset_button = gr.Button("Reset")
+        with gr.Column():
+            image_board = gr.Image(label="Board", interactive=False)
+
+    sampler = gr.State(value=None)
+    board = gr.State(value=None)
+
+    outputs = [sampler, board, current_fen, current_pgn, image_board, move_to_play]
+
+    play_button.click(
+        play_user_move_then_ai_move,
+        inputs=[move_to_play, board, sampler],
+        outputs=outputs,
+    )
+    move_to_play.submit(
+        play_user_move_then_ai_move,
+        inputs=[move_to_play, board, sampler],
+        outputs=outputs,
+    )
+
+    model_name.change(
+        get_sampler,
+        inputs=[model_name],
+        outputs=[sampler],
+    )
+
+    reset_button.click(
+        get_init,
+        inputs=[model_name],
+        outputs=outputs,
+    )
+    interface.load(
+        get_init,
+        inputs=[model_name],
+        outputs=outputs,
+    )

demo/utils.py

@@ -0,0 +1,52 @@
+
+import gradio as gr
+import chess.svg
+from typing import Optional
+
+from lczerolens.board import LczeroBoard
+
+from . import constants
+
+
+def create_board_figure(
+    board: LczeroBoard,
+    *,
+    orientation: bool = chess.WHITE,
+    arrows: str = "",
+    square: str = "",
+    name: str = "board",
+):
+    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.")
+
+    try:
+        color_dict = {chess.parse_square(square): "#FF0000"} if square else {}
+    except ValueError:
+        color_dict = {}
+        gr.Warning("Invalid square, using none.")
+
+    svg_board = chess.svg.board(
+        board,
+        size=350,
+        orientation=orientation,
+        arrows=chess_arrows,
+        fill=color_dict,
+    )
+    with open(f"{constants.FIGURE_DIRECTORY}/{name}.svg", "w") as f:
+        f.write(svg_board)
+    return f"{constants.FIGURE_DIRECTORY}/{name}.svg"

main.py

@@ -0,0 +1,40 @@
+"""
+Gradio demo for chess project.
+"""
+
+import gradio as gr
+import subprocess
+
+from demo.interfaces import (
+    board,
+    encodings,
+    gradients,
+    play,
+    activations,
+)
+
+demo = gr.TabbedInterface(
+    [
+        board.interface,
+        play.interface,
+        encodings.interface,
+        activations.interface,
+        gradients.interface,
+    ],
+    [
+        "Board",
+        "Play",
+        "Encodings",
+        "Activations",
+        "Gradients",
+    ],
+    title="Chess Project Demo",
+    analytics_enabled=False,
+)
+
+if __name__ == "__main__":
+    subprocess.run(["bash", "resolve-assets.sh"])
+    demo.launch(
+        server_port=8000,
+        server_name="0.0.0.0",
+    )

pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "lczerolens-demo"
+version = "0.1.0"
+description = "Demo lczerolens features."
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "gdown>=5.2.0",
+    "gradio>=5.20.1",
+    "lczerolens[viz]>=0.3.1",
+]

resolve-assets.sh

@@ -0,0 +1,5 @@
+uv run gdown 1cxC8_8vw7akfPyc9cZxwaAbLG2Zl4XiT -O demo/onnx-models/lc0-10-4238.onnx
+uv run gdown 15__7FHvIR5-JbJvDg2eGUhIPZpkYyM7X -O demo/onnx-models/lc0-19-1876.onnx
+uv run gdown 1CvMyX3KuYxCJUKz9kOb9VX8zIkfISALd -O demo/onnx-models/lc0-19-4508.onnx
+uv run gdown 1TI429e9mr2de7LjHp2IIl7ouMoUaDjjZ -O demo/onnx-models/maia-1100.onnx
+uv run gdown 1-8IJ5WYMPpcxOsHfIKY8xKskwk2z_yrY -O demo/onnx-models/maia-1900.onnx