new repo structure

.gitignore

@@ -0,0 +1,139 @@
+# 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
+
+*secret*

app.py

@@ -4,11 +4,15 @@ Main Gradio module.
 
 import gradio as gr
 
+from src.interfaces import feature_interface
+
 
 demo = gr.TabbedInterface(
     [
+        feature_interface,
     ],
     [
+        "Feature Activation",
     ],
     title="Lczero Planning Demo",
     analytics_enabled=False,

assets/.gitignore

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

figures/.gitignore

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

requirements.txt

@@ -0,0 +1,7 @@
+git+https://github.com/Xmaster6y/lczerolens
+chess
+matplotlib
+numpy
+torch
+tensordict
+einops

src/constants.py

@@ -0,0 +1,18 @@
+"""Manage constants for the app.
+"""
+
+import os
+import pathlib
+import torch
+
+
+ASSETS_FOLDER = pathlib.Path(__file__).parent.parent / "assets"
+FIGURES_FOLER = pathlib.Path(__file__).parent.parent / "figures"
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
+
+MODEL_NAME = "lc0-10-4238.onnx"
+SAE_CONFIG = "debug"
+LAYER = 9
+N_FEATURES = 7680

src/global_variables.py

@@ -0,0 +1,54 @@
+"""Manage global variables for the app.
+"""
+
+from huggingface_hub import HfApi
+
+import gradio as gr
+from lczerolens import ModelWrapper
+import torch
+
+from src.constants import HF_TOKEN, ASSETS_FOLDER, DEVICE, MODEL_NAME, SAE_CONFIG, LAYER
+from src.helpers import SparseAutoEncoder, OutputGenerator
+
+hf_api: HfApi
+wrapper: ModelWrapper
+sae: SparseAutoEncoder
+generator: OutputGenerator
+
+
+def setup():
+    global hf_api
+    global wrapper
+    global sae
+    global generator
+
+    hf_api = HfApi(token=HF_TOKEN)
+    hf_api.snapshot_download(
+        local_dir=f"{ASSETS_FOLDER}/models",
+        repo_id="Xmaster6y/lczero-planning-models",
+        repo_type="model",
+    )
+    hf_api.snapshot_download(
+        local_dir=f"{ASSETS_FOLDER}/saes",
+        repo_id="Xmaster6y/lczero-planning-saes",
+        repo_type="model",
+    )
+
+    wrapper = ModelWrapper.from_onnx_path(f"{ASSETS_FOLDER}/models/{MODEL_NAME}").to(DEVICE)
+    sae_dict = torch.load(
+        f"{ASSETS_FOLDER}/saes/{SAE_CONFIG}/model.pt",
+        map_location=DEVICE,
+        weights_only=True
+    )
+    sae = SparseAutoEncoder()
+    sae.load_state_dict(
+        sae_dict
+    )
+    generator = OutputGenerator(
+        sae=sae,
+        wrapper=wrapper,
+        module_exp=rf".*block{LAYER}/conv2/relu"
+    )
+
+if gr.NO_RELOAD:
+    setup()

src/helpers/__init__.py

@@ -0,0 +1,4 @@
+
+
+from .generator import OutputGenerator
+from .sae import SparseAutoEncoder

src/helpers/generator.py

@@ -0,0 +1,59 @@
+"""Script to generate features for a given board state.
+"""
+
+from typing import Optional
+
+from lczerolens import ModelWrapper
+from lczerolens.xai import ActivationLens
+from lczerolens.encodings import InputEncoding
+import chess
+import einops
+import torch
+
+from .sae import SparseAutoEncoder
+
+
+class OutputGenerator:
+    
+    def __init__(self, sae: SparseAutoEncoder, wrapper: ModelWrapper, module_exp: Optional[str] = None):
+        self.sae = sae
+        self.wrapper = wrapper
+        self.lens = ActivationLens(module_exp=module_exp)
+
+    def generate(
+        self,
+        root_fen: Optional[str] = None,
+        traj_fen: Optional[str] = None,
+        root_board: Optional[chess.Board] = None,
+        traj_board: Optional[chess.Board] = None,
+    ):
+        if root_board is not None and traj_board is not None:
+            input_encoding = InputEncoding.INPUT_CLASSICAL_112_PLANE
+        elif root_fen is not None and traj_fen is not None:
+            root_board = chess.Board(root_fen)
+            traj_board = chess.Board(traj_fen)
+            input_encoding = InputEncoding.INPUT_CLASSICAL_112_PLANE_REPEATED
+        else:
+            raise ValueError
+        iter_boards = iter([[root_board, traj_board]])
+        act_dict, (model_output,) = self.lens.analyse_batched_boards(
+            iter_boards,
+            self.wrapper,
+            {
+                "return_output": True,
+                "wrapper_kwargs": {
+                    "input_encoding": input_encoding,
+                }
+            }
+        )
+        if len(act_dict) == 0:
+            raise ValueError("No module matced the given expression.")
+        elif len(act_dict) > 1:
+            raise ValueError("Multiple modules matched the given expression.")
+        acts = next(iter(act_dict.values()))
+        root_acts = einops.rearrange(acts[0], "c h w -> (h w) c")
+        traj_acts = einops.rearrange(acts[1], "c h w -> (h w) c")
+        pixel_acts = torch.cat([root_acts, traj_acts], dim=1)
+        sae_output = self.sae(pixel_acts, output_features=True)
+        return model_output, pixel_acts, sae_output
+        

src/helpers/sae.py

@@ -0,0 +1,93 @@
+"""
+Defines the dictionary classes
+"""
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+
+
+class SparseAutoEncoder(nn.Module):
+    """
+    A 2-layer sparse autoencoder.
+    """
+
+    def __init__(
+        self,
+        activation_dim,
+        dict_size,
+        pre_bias=False,
+        init_normalise_dict=None,
+    ):
+        super().__init__()
+        self.activation_dim = activation_dim
+        self.dict_size = dict_size
+        self.pre_bias = pre_bias
+        self.init_normalise_dict = init_normalise_dict
+
+        self.b_enc = nn.Parameter(torch.zeros(self.dict_size))
+        self.relu = nn.ReLU()
+
+        self.W_dec = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(
+                    self.dict_size,
+                    self.activation_dim,
+                )
+            )
+        )
+        if init_normalise_dict == "l2":
+            self.normalize_dict_(less_than_1=False)
+            self.W_dec *= 0.1
+        elif init_normalise_dict == "less_than_1":
+            self.normalize_dict_(less_than_1=True)
+
+        self.W_enc = nn.Parameter(self.W_dec.t())
+        self.b_dec = nn.Parameter(
+            torch.zeros(
+                self.activation_dim,
+            )
+        )
+
+    @torch.no_grad()
+    def normalize_dict_(
+        self,
+        less_than_1=False,
+    ):
+        norm = self.W_dec.norm(dim=1)
+        positive_mask = norm != 0
+        if less_than_1:
+            greater_than_1_mask = (norm > 1) & (positive_mask)
+            self.W_dec[greater_than_1_mask] /= norm[greater_than_1_mask].unsqueeze(1)
+        else:
+            self.W_dec[positive_mask] /= norm[positive_mask].unsqueeze(1)
+
+    def encode(self, x):
+        return x @ self.W_enc + self.b_enc
+
+    def decode(self, f):
+        return f @ self.W_dec + self.b_dec
+
+    def forward(self, x, output_features=False, ghost_mask=None):
+        """
+        Forward pass of an autoencoder.
+        x : activations to be autoencoded
+        output_features : if True, return the encoded features as well
+            as the decoded x
+        ghost_mask : if not None, run this autoencoder in "ghost mode"
+            where features are masked
+        """
+        if self.pre_bias:
+            x = x - self.b_dec
+        f_pre = self.encode(x)
+        out = TensorDict({}, batch_size=x.shape[0])
+        if ghost_mask is not None:
+            f_ghost = torch.exp(f_pre) * ghost_mask.to(f_pre)
+            x_ghost = f_ghost @ self.W_dec
+            out["x_ghost"] = x_ghost
+        f = self.relu(f_pre)
+        if output_features:
+            out["features"] = f
+        x_hat = self.decode(f)
+        out["x_hat"] = x_hat
+        return out

src/interfaces/__init__.py

@@ -0,0 +1,2 @@
+
+from .feature_interface import interface as feature_interface

src/interfaces/feature_interface.py

@@ -0,0 +1,121 @@
+"""
+Gradio interface for plotting policy.
+"""
+
+import chess
+import gradio as gr
+import uuid
+
+from lczerolens.encodings import encode_move
+
+from src import constants, global_variables, visualisation
+
+
+def compute_features_fn(
+    features,
+    model_output, 
+    file_id,
+    root_fen, 
+    traj_fen, 
+    feature_index
+):
+    model_output, _, sae_output = global_variables.generator.generate(
+        root_fen=root_fen,
+        traj_fen=traj_fen
+    )
+    features = sae_output["f"]
+    first_output = render_feature_index(
+        features,
+        model_output,
+        file_id,
+        feature_index,
+        traj_fen,
+    )
+    game_info = f"WDL: {model_output.get('wdl')}"
+    return *first_output, game_info
+
+
+def render_feature_index(
+    features,
+    model_output,
+    file_id,
+    feature_index,
+    traj_fen,
+):
+    if file_id is None:
+        file_id = str(uuid.uuid4())
+    board = chess.Board(traj_fen)
+    pixel_features = features[:,feature_index]
+    if board.turn:
+        heatmap = pixel_features.view(64)
+    else:
+        heatmap = pixel_features.view(8,8).flip(0).view(64)
+
+    best_legal_logit = None
+    best_legal_move = None
+    for move in board.legal_moves:
+        move_index = encode_move(move, (board.turn, not board.turn))
+        logit = model_output["policy"][1,move_index].item()
+        if best_legal_logit is None:
+            best_legal_logit = logit
+        else:
+            best_legal_move = move
+
+    svg_board, fig = visualisation.render_heatmap(
+        board,
+        heatmap,
+        arrows=[(best_legal_move.from_square, best_legal_move.to_square)],
+    )
+    with open(f"{constants.FIGURES_FOLER}/{file_id}.svg", "w") as f:
+        f.write(svg_board)
+    return (
+        features,
+        model_output,
+        file_id,
+        f"{constants.FIGURES_FOLER}/{file_id}.svg",
+        fig
+    ) 
+
+with gr.Blocks() as interface:
+    with gr.Row():
+        with gr.Column():
+            root_fen = gr.Textbox(
+                label="Root FEN",
+                lines=1,
+                max_lines=1,
+                value=chess.STARTING_FEN,
+            )
+            traj_fen = gr.Textbox(
+                label="Trajectory FEN",
+                lines=1,
+                max_lines=1,
+                value="rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1",
+            )
+            compute_features = gr.Button("Compute features")
+
+            with gr.Group():
+                with gr.Row():
+                    feature_index = gr.Slider(
+                        label="Feature index",
+                        minimum=0,
+                        maximum=constants.N_FEATURES,
+                        step=1,
+                        value=0,
+                    )
+                    
+            with gr.Group():
+                with gr.Row():    
+                    game_info = gr.Textbox(label="Game info", lines=1, max_lines=1, value="")
+                with gr.Row():
+                    colorbar = gr.Plot(label="Colorbar")
+        with gr.Column():
+            board_image = gr.Image(label="Board")
+
+    features = gr.State(None)
+    model_output = gr.State(None)
+    file_id = gr.State(None)
+    compute_features.click(
+        compute_features_fn,
+        inputs=[features, model_output, file_id, root_fen, traj_fen, feature_index],
+        outputs=[features, model_output, file_id, board_image, colorbar, game_info],
+    )

src/visualisation.py

@@ -0,0 +1,108 @@
+"""
+Visualisation utils.
+"""
+
+import chess
+import chess.svg
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+COLOR_MAP = matplotlib.colormaps["RdYlBu_r"].resampled(1000)
+ALPHA = 1.0
+NORM = matplotlib.colors.Normalize(vmin=0, vmax=1, clip=False)
+
+
+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_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