reset space

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.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*

README.md

@@ -0,0 +1,14 @@
+---
+title: Lczerolens Demo
+emoji: 🔬
+colorFrom: gray
+colorTo: green
+sdk: gradio
+sdk_version: 5.28.0
+app_file: app.py
+pinned: true
+license: mit
+short_description: Demo lczerolens features
+---
+
+See the documentation [here](https://lczerolens.readthedocs.io/).

app.py

@@ -0,0 +1,36 @@
+"""
+Gradio demo for lczerolens.
+"""
+
+import gradio as gr
+
+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="lczerolens Demo",
+    analytics_enabled=False,
+)
+
+
+if __name__ == "__main__":
+    demo.launch()

demo/constants.py

@@ -0,0 +1,21 @@
+"""
+Constants for the demo.
+"""
+
+import os
+from pathlib import Path
+
+DEMO_DIRECTORY = Path(__file__).parent
+
+ONNX_MODEL_DIRECTORY = DEMO_DIRECTORY / "onnx-models"
+LEELA_MODEL_DIRECTORY = DEMO_DIRECTORY / "leela-models"
+FIGURE_DIRECTORY = DEMO_DIRECTORY / "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/figures/.gitignore

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

demo/interfaces/activations.py

@@ -0,0 +1,168 @@
+"""
+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 demo import constants
+from demo.utils import get_info
+
+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)
+    info = get_info(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return model, board, activations, info, *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)
+    info = get_info(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return board, activations, info, *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)
+    info = get_info(model, board)
+    plots = render_activations(board, activations, layer_index, channel_index)
+    return model, activations, info, *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,
+                )
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=constants.ONNX_MODEL_NAMES,
+                )
+            with gr.Group():
+                info = gr.Textbox(label="Info", lines=1, value="")
+            with gr.Group(): 
+                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, info, image_board, colorbar],
+        concurrency_limit=1, 
+        concurrency_id="trace_queue"
+    )
+    game_pgn.submit(
+        on_board_change,
+        inputs=[model, game_pgn, board_fen, layer_index, channel_index],
+        outputs=[board, activations, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    board_fen.submit(
+        on_board_change,
+        inputs=[model, game_pgn, board_fen, layer_index, channel_index],
+        outputs=[board, activations, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    model_name.change(
+        on_model_change,
+        inputs=[model_name, board, layer_index, channel_index],
+        outputs=[model, activations, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    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],
+    )

demo/interfaces/board.py

@@ -0,0 +1,57 @@
+"""
+Gradio interface for plotting a board.
+"""
+
+import chess
+import chess.svg
+import gradio as gr
+
+from lczerolens.board import LczeroBoard
+
+from ..utils import create_board_figure
+
+def make_board_plot(board_fen, arrows, square):
+    try:
+        board = LczeroBoard(board_fen)
+    except ValueError:
+        board = LczeroBoard()
+        gr.Warning("Invalid FEN, using starting position.")
+    filepath = create_board_figure(board, arrows=arrows, square=square, name="board")
+    return filepath
+
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            board_fen = gr.Textbox(
+                label="Board starting FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            arrows = gr.Textbox(
+                label="Arrows",
+                lines=1,
+                max_lines=1,
+                value="",
+                placeholder="e2e4 e7e5",
+            )
+            square = gr.Textbox(
+                label="Square",
+                lines=1,
+                max_lines=1,
+                value="",
+                placeholder="e4",
+            )
+        with gr.Column():
+            image = gr.Image(label="Board", interactive=False)
+
+    inputs = [
+        board_fen,
+        arrows,
+        square,
+    ]
+    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)
+    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,186 @@
+"""
+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 demo import constants
+from demo.utils import get_info
+
+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, target: str):
+    lens = Lens.from_name("gradient")
+
+    def init_target(model):
+        if target == "best_move":
+            return getattr(model, "output/policy").output.max(dim=1).values
+        else:
+            wdl_index = {"win": 0, "draw": 1, "loss": 2}[target]
+            return getattr(model, "output/wdl").output[:, wdl_index]
+    results = lens.analyse(model, board, init_target=init_target)
+
+    return results["input_grad"]
+
+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, 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, target)
+    info = get_info(model, board)
+    plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
+    return model, board, gradients, info, *plots
+
+def on_board_change(model: LczeroModel, game_pgn: str, board_fen: str, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+    board = get_board(game_pgn, board_fen)
+    gradients = get_gradients(model, board, target)
+    info = get_info(model, board)
+    plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
+    return board, gradients, info, *plots
+
+def on_model_change(model_name: str, board: LczeroBoard, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+    model = get_model(model_name)
+    gradients = get_gradients(model, board, target)
+    info = get_info(model, board)
+    plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
+    return model, gradients, info, *plots
+
+def on_target_change(model: LczeroModel, board: LczeroBoard, target: str, average_over_planes:bool, begin_average_index:int, end_average_index:int, plane_index: int):
+    gradients = get_gradients(model, board, target)
+    plots = render_gradients(board, gradients, average_over_planes, begin_average_index, end_average_index, plane_index)
+    return gradients, *plots
+
+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,
+                )
+                model_name = gr.Dropdown(
+                    label="Model",
+                    choices=constants.ONNX_MODEL_NAMES,
+                )
+            with gr.Group():
+                info = gr.Textbox(label="Info", lines=1, value="")
+            with gr.Group():
+                target = gr.Radio(
+                    ["win", "draw", "loss", "best_move"], label="Target",
+                    value="win",
+                )
+                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, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[model, board, gradients, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    game_pgn.submit(
+        on_board_change,
+        inputs=[model, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[board, gradients, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    board_fen.submit(
+        on_board_change,
+        inputs=[model, game_pgn, board_fen, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[board, gradients, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    model_name.change(
+        on_model_change,
+        inputs=[model_name, board, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[model, gradients, info, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    target.change(
+        on_target_change,
+        inputs=[model, board, target, average_over_planes, begin_average_index, end_average_index, plane_index],
+        outputs=[gradients, image_board, colorbar],
+        concurrency_id="trace_queue"
+    )
+    for render_arg in [average_over_planes, begin_average_index, end_average_index, plane_index]:
+        render_arg.change(
+            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,139 @@
+"""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 demo import constants
+from demo.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="",
+            )
+            model_name = gr.Dropdown(
+                label="Model",
+                choices=constants.ONNX_MODEL_NAMES,
+            )
+            with gr.Column():
+                move_to_play = gr.Textbox(
+                    label="Move to play (UCI)",
+                    lines=1,
+                    max_lines=1,
+                    value="",
+                )
+                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/leela-models/.gitignore

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

demo/onnx-models/lc0-10-4238.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:096bc5b833572de5357961f8cb2d059fe4f0ff303598131f20278fb23bc58919
+size 15152179

demo/onnx-models/lc0-19-1876.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df7e2b87f0f4525d7de7201ee2e946fc10172200ef97c6786b3fdcfb8610b88a
+size 97139290

demo/onnx-models/lc0-19-4508.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcf7c13451f00fc9d1859f98a52a37056aa3d10ac7f874408ca46d0cde50c241
+size 97139290

demo/onnx-models/look-ahead-lc0.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a3dad19ed0546106e3b97d7aa7dbf0e3cc7041a6c63979f923360ee71cb4a24
+size 378669573

demo/onnx-models/maia-1100.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1a4ac9b99aee3c127285b7cdc2024d054a6de72abdfa2fe3db5e4f7078c96ad
+size 3484716

demo/onnx-models/maia-1200.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82c21e9cf65f0de9a72469eaae8845c3e92a32fe93e1e626a064cd3d24cca4cf
+size 3483901

demo/onnx-models/maia-1300.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7c3b8d9da7a2cc586f409b0264abcb9d00390276ffe5f39a09d591787555e50
+size 3483901

demo/onnx-models/maia-1400.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3651ae981d8bc05cf735c79d1a51298d91d3450592ac874d9c2c66e6c90b896b
+size 3483901

demo/onnx-models/maia-1500.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94333b25bb8d25685334645afbb7494b358ae5724ad651d735ae1683fb0a9e3f
+size 3483901

demo/onnx-models/maia-1600.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d46d32c6c1ff2a9f04df5c0652e50ae502b8870c6e0895e6abeeb02bba01d423
+size 3483901

demo/onnx-models/maia-1700.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b01ef1c6bf0136480526f2933f88623bc5b1e2bbf319fc6045be68524f942d8
+size 3483901

demo/onnx-models/maia-1800.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b9b999de7d0fe4b2efc9d6591b30f4010ec7490c36676dd920ee0420f202a67
+size 3483901

demo/onnx-models/maia-1900.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65ee89dcee614d2b7f5bf8fc5950e83050bf855ecb4d34f6e6214b09acc64572
+size 3484716

demo/onnx-models/t1-256x10-distilled-swa-2432500.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5cf50f8d5e08c1f42c9a30bdcbd3a2deeef209f24ca4104d4170eda1b193bad
+size 80896651

demo/utils.py

@@ -0,0 +1,68 @@
+
+import gradio as gr
+import chess.svg
+
+from lczerolens import LczeroBoard, LczeroModel, Lens
+
+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"
+
+
+class OutputLens(Lens):
+    def _intervene(self, model: LczeroModel, **kwargs) -> dict:
+        return model.output.save()
+
+def get_info(model: LczeroModel, board: LczeroBoard):
+    lens = OutputLens()
+    output = lens.analyse(model, board)
+    w = output["wdl"][0,0]
+    d = output["wdl"][0,1]
+    l = output["wdl"][0,2]
+    legal_indices = board.get_legal_indices()
+    best_move_idx = output["policy"].gather(dim=1, index=legal_indices.unsqueeze(0)).argmax(dim=1).item()
+    best_move = board.decode_move(legal_indices[best_move_idx])
+    info = f"w: {w:.2f}, d: {d:.2f}, l: {l:.2f}, best: {best_move}"
+    return info

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.3",
+]

requirements.txt

@@ -0,0 +1,2 @@
+lczerolens[viz]>=0.3.3
+gdown

resolve-assets.sh

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