Update app.py

app.py

@@ -1,469 +1,229 @@
 import marimo
 
-__generated_with = "0.9.2"
-app = marimo.App()
+__generated_with = "0.8.22"
+app = marimo.App(width="medium")
 
 
 @app.cell
 def __():
     import marimo as mo
-
-    mo.md("# Welcome to marimo! 🌊🍃")
     return (mo,)
 
 
 @app.cell
-def __(mo):
-    slider = mo.ui.slider(1, 22)
-    return (slider,)
-
-
-@app.cell
-def __(mo, slider):
-    mo.md(
-        f"""
-        marimo is a **reactive** Python notebook.
-
-        This means that unlike traditional notebooks, marimo notebooks **run
-        automatically** when you modify them or
-        interact with UI elements, like this slider: {slider}.
-
-        {"##" + "🍃" * slider.value}
-        """
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.accordion(
-        {
-            "Tip: disabling automatic execution": mo.md(
-                rf"""
-            marimo lets you disable automatic execution: just go into the
-            notebook settings and set
-
-            "Runtime > On Cell Change" to "lazy".
-
-            When the runtime is lazy, after running a cell, marimo marks its
-            descendants as stale instead of automatically running them. The
-            lazy runtime puts you in control over when cells are run, while
-            still giving guarantees about the notebook state.
-            """
-            )
-        }
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        Tip: This is a tutorial notebook. You can create your own notebooks
-        by entering `marimo edit` at the command line.
-        """
-    ).callout()
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        ## 1. Reactive execution
-
-        A marimo notebook is made up of small blocks of Python code called
-        cells.
-
-        marimo reads your cells and models the dependencies among them: whenever
-        a cell that defines a global variable  is run, marimo
-        **automatically runs** all cells that reference that variable.
-
-        Reactivity keeps your program state and outputs in sync with your code,
-        making for a dynamic programming environment that prevents bugs before they
-        happen.
-        """
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(changed, mo):
-    (
-        mo.md(
-            f"""
-            **✨ Nice!** The value of `changed` is now {changed}.
-
-            When you updated the value of the variable `changed`, marimo
-            **reacted** by running this cell automatically, because this cell
-            references the global variable `changed`.
-
-            Reactivity ensures that your notebook state is always
-            consistent, which is crucial for doing good science; it's also what
-            enables marimo notebooks to double as tools and  apps.
-            """
-        )
-        if changed
-        else mo.md(
-            """
-            **🌊 See it in action.** In the next cell, change the value of the
-            variable  `changed` to `True`, then click the run button.
-            """
-        )
-    )
-    return
+def __():
+    import pandas as pd 
+    df = pd.read_csv("our_visualization/datasets/test_set.csv")
+    df.head()
+    return df, pd
 
 
 @app.cell
 def __():
-    changed = False
-    return (changed,)
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.accordion(
-        {
-            "Tip: execution order": (
-                """
-                The order of cells on the page has no bearing on
-                the order in which cells are executed: marimo knows that a cell
-                reading a variable must run after the cell that  defines it. This
-                frees you to organize your code in the way that makes the most
-                sense for you.
-                """
-            )
-        }
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        **Global names must be unique.** To enable reactivity, marimo imposes a
-        constraint on how names appear in cells: no two cells may define the same
-        variable.
-        """
+    import pickle
+    from utils import ChessBoard
+    import onnxruntime as ort
+    from leela_board import _idx_to_move_bn, _idx_to_move_wn
+    import numpy as np 
+    from onnx2torch import convert
+    import onnx
+    import torch 
+    import os 
+
+    def get_models(root="/our_visualization/models"):
+        paths = os.listdir(root)
+        model_paths = []
+        for path in paths:
+            if ".onnx" in path: model_paths.append(os.path.join(root, path))
+        return model_paths
+
+    def get_activations_from_model(model_path, pattern, fen):
+        # Write hooks for selected model path
+        def register_hooks_for_capture(model, pattern):
+            activations = {}
+            def get_activation(name):
+                def hook(module, input, output):
+                    activations[name] = output.detach().numpy()
+                return hook 
+
+            handles = []
+            for n, m in model.named_modules():
+                if pattern in n:
+                    handle = m.register_forward_hook(get_activation(n))
+                    handles.append(handle)
+            return activations, handles
+
+        # Load model and register hooks for it 
+        model = convert(onnx.load(model_path))
+        act, handles = register_hooks_for_capture(model, pattern)
+
+        # Get fen and pass it through model to generate activations
+        board = ChessBoard(fen)
+        inputs = board.t
+        _, _, _ = model(inputs.unsqueeze(dim=0))
+
+        # Remove handles
+        [h.remove() for h in handles]
+        return act
+    return (
+        ChessBoard,
+        convert,
+        get_activations_from_model,
+        get_models,
+        np,
+        onnx,
+        ort,
+        os,
+        pickle,
+        torch,
     )
-    return
 
 
-@app.cell(hide_code=True)
-def __(mo):
-    mo.accordion(
-        {
-            "Tip: encapsulation": (
-                """
-                By encapsulating logic in functions, classes, or Python modules,
-                you can minimize the number of global variables in your notebook.
-                """
-            )
-        }
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.accordion(
-        {
-            "Tip: private variables": (
-                """
-                Variables prefixed with an underscore are "private" to a cell, so
-                they can be defined by multiple cells.
-                """
-            )
-        }
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        ## 2. UI elements
+@app.cell
+def __(df, mo):
+    min_elo, max_elo = df["Rating"].min() // 100 * 100, df["Rating"].max() // 100 * 100
+    elo_list = [f"{elo}" for elo in range(min_elo, max_elo + 100, 100)]
+    dropdown_elo = mo.ui.dropdown(value = "1000", options=elo_list, label=f"Select rating in range of {min_elo} - {max_elo}")
+    dropdown_elo
+    return dropdown_elo, elo_list, max_elo, min_elo
 
-        Cells can output interactive UI elements. Interacting with a UI
-        element **automatically triggers notebook execution**: when
-        you interact with a UI element, its value is sent back to Python, and
-        every cell that references that element is re-run.
 
-        marimo provides a library of UI elements to choose from under
-        `marimo.ui`.
-        """
+@app.cell
+def __(df, dropdown_elo, mo):
+    unique_themes = set()
+    df_rated = df[(df["Rating"] >= int(dropdown_elo.value)) & (df["Rating"] <= int(dropdown_elo.value) + 100)]
+    for i in range(len(df_rated)):
+        themes = df_rated.iloc[i]["Themes"].split(" ")
+        for theme in themes: unique_themes.add(theme)
+    unique_themes_list = list(unique_themes)
+    unique_themes_list.sort()
+
+    dropdown_themes = mo.ui.dropdown(value=unique_themes_list[0], options=unique_themes_list, label=f"Select puzzle theme")
+    dropdown_themes
+    return (
+        df_rated,
+        dropdown_themes,
+        i,
+        theme,
+        themes,
+        unique_themes,
+        unique_themes_list,
     )
-    return
 
 
 @app.cell
-def __(mo):
-    mo.md("""**🌊 Some UI elements.** Try interacting with the below elements.""")
-    return
+def __(df_rated, dropdown_themes):
+    themes_mask = []
+    def _(themes_mask):
+        for i in range(len(df_rated)):
+            themes_new = df_rated.iloc[i]["Themes"].split(" ")
+            if dropdown_themes.value in themes_new: themes_mask.append(i)
+    _(themes_mask)
+    fens = list(df_rated.iloc[themes_mask]["FEN"])
+    df_rated.iloc[themes_mask][["FEN", "Moves", "Themes", "Rating"]]
+    return fens, themes_mask
 
 
 @app.cell
-def __(mo):
-    icon = mo.ui.dropdown(["🍃", "🌊", "✨"], value="🍃")
-    return (icon,)
+def __(fens, mo):
+    dropdown_fen = mo.ui.dropdown(value = fens[0], options=fens, label="Select FEN")
+    dropdown_fen
+    return (dropdown_fen,)
 
 
 @app.cell
-def __(icon, mo):
-    repetitions = mo.ui.slider(1, 16, label=f"number of {icon.value}: ")
-    return (repetitions,)
+def __(df_rated, dropdown_fen, mo):
+    moves = df_rated[df_rated["FEN"] == dropdown_fen.value]["Moves"].iloc[0].split(" ")
+    player_moves = moves[1::2]
+    board_moves = []
+    def _(board_moves):
+        for i in range(len(player_moves)):
+            board_moves.append(moves[:2 * i + 1])
+    _(board_moves)
+    moves_dict = {pm: om for pm, om in zip(player_moves, board_moves)}
+    dropdown_moves = mo.ui.dropdown(options=moves_dict, value=player_moves[0], label="Select which player move to look at")
+    # print(moves)
+    dropdown_moves
+    return board_moves, dropdown_moves, moves, moves_dict, player_moves
 
 
 @app.cell
-def __(icon, repetitions):
-    icon, repetitions
-    return
+def __(dropdown_moves, mo):
+    dropdown_layer = mo.ui.dropdown(value="0", options=[f"{i}" for i in range(15)], label="Select layer (smaller - closer to input)")
+    focus_square = mo.ui.text_area(value=dropdown_moves.selected_key[:2], placeholder="Input square to look at (e.g. a1, b8, ...")
+    mo.vstack([dropdown_layer, focus_square])
+    return dropdown_layer, focus_square
 
 
 @app.cell
-def __(icon, mo, repetitions):
-    mo.md("# " + icon.value * repetitions.value)
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        ## 3. marimo is just Python
-
-        marimo cells parse Python (and only Python), and marimo notebooks are
-        stored as pure Python files — outputs are _not_ included. There's no
-        magical syntax.
-
-        The Python files generated by marimo are:
-
-        - easily versioned with git, yielding minimal diffs
-        - legible for both humans and machines
-        - formattable using your tool of choice,
-        - usable as Python  scripts, with UI  elements taking their default
-        values, and
-        - importable by other modules (more on that in the future).
-        """
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        ## 4. Running notebooks as apps
-
-        marimo notebooks can double as apps. Click the app window icon in the
-        bottom-right to see this notebook in "app view."
-
-        Serve a notebook as an app with `marimo run` at the command-line.
-        Of course, you can use marimo just to level-up your
-        notebooking, without ever making apps.
-        """
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        ## 5. The `marimo` command-line tool
-
-        **Creating and editing notebooks.** Use
-
-        ```
-        marimo edit
-        ```
-
-        in a terminal to start the marimo notebook server. From here
-        you can create a new notebook or edit existing ones.
-
-
-        **Running as apps.** Use
-
-        ```
-        marimo run notebook.py
-        ```
-
-        to start a webserver that serves your notebook as an app in read-only mode,
-        with code cells hidden.
-
-        **Convert a Jupyter notebook.** Convert a Jupyter notebook to a marimo
-        notebook using `marimo convert`:
-
-        ```
-        marimo convert your_notebook.ipynb > your_app.py
-        ```
-
-        **Tutorials.** marimo comes packaged with tutorials:
-
-        - `dataflow`: more on marimo's automatic execution
-        - `ui`: how to use UI elements
-        - `markdown`: how to write markdown, with interpolated values and
-           LaTeX
-        - `plots`: how plotting works in marimo
-        - `sql`: how to use SQL
-        - `layout`: layout elements in marimo
-        - `fileformat`: how marimo's file format works
-        - `markdown-format`: for using `.md` files in marimo
-        - `for-jupyter-users`: if you are coming from Jupyter
-
-        Start a tutorial with `marimo tutorial`; for example,
-
-        ```
-        marimo tutorial dataflow
-        ```
-
-        In addition to tutorials, we have examples in our
-        [our GitHub repo](https://www.github.com/marimo-team/marimo/tree/main/examples).
-        """
-    )
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        ## 6. The marimo editor
-
-        Here are some tips to help you get started with the marimo editor.
-        """
-    )
-    return
+def __(ChessBoard, dropdown_fen, dropdown_moves):
+    def _():
+        board = ChessBoard(dropdown_fen.value)
+        for move in dropdown_moves.value:
+            print(move)
+            # board.move(move)
+        return board.board.pc_board.fen()
+    FEN = _()
+    return (FEN,)
 
 
 @app.cell
-def __(mo, tips):
-    mo.accordion(tips)
-    return
-
-
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md("""## Finally, a fun fact""")
-    return
+def __(focus_square):
+    import chess 
+    from global_data import global_data
+
+    focus_square_ind = 8 * (int(focus_square.value[1]) - 1) + ord(focus_square.value[0]) - ord("a")
+
+    def set_plotting_parameters(act, layer_number, fen):
+        layer_key = [k for k in act.keys() if "0" in k][0].replace("0", f"{layer_number}")
+        print(act.keys())
+        global_data.model = 'test'
+        global_data.activations = act[layer_key][0, :, ::-1 , :]
+        print(global_data.activations.shape)
+        global_data.subplot_rows = 8
+        global_data.subplot_cols = 4
+        global_data.board = chess.Board(fen)
+        global_data.show_all_heads = True
+        # global_data.selected_head = 1
+        global_data.visualization_mode = 'ROW'
+        global_data.focused_square_ind = focus_square_ind
+        # global_data.heatmap_horizontal_gap = 0.001
+
+        global_data.visualization_mode_is_64x64 = False
+        global_data.colorscale_mode = "mode1"
+        global_data.show_colorscale = False
+    return chess, focus_square_ind, global_data, set_plotting_parameters
 
 
-@app.cell(hide_code=True)
-def __(mo):
-    mo.md(
-        """
-        The name "marimo" is a reference to a type of algae that, under
-        the right conditions, clumps together to form a small sphere
-        called a "marimo moss ball". Made of just strands of algae, these
-        beloved assemblages are greater than the sum of their parts.
-        """
-    )
-    return
+@app.cell
+def __(
+    FEN,
+    dropdown_layer,
+    get_activations_from_model,
+    get_models,
+    set_plotting_parameters,
+):
+    # FEN = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
+    # board = ChessBoard("r1b2rk1/pp2pp1p/6p1/3Qb2q/1P4n1/2P1BN2/P2N1PPP/R4RK1 w - - 0 14")
+    # board.move("f3e5")
+    # FEN = board.board.pc_board.fen()
+    PATTERN = "mha/QK/softmax"
+    # PATTERN = "smolgen_weights"
+    MODEL = get_models()[-1]
+    ACTIVATIONS = get_activations_from_model(MODEL, PATTERN, FEN)
+    set_plotting_parameters(ACTIVATIONS, int(dropdown_layer.value), FEN)
+    from activation_heatmap import heatmap_figure
+    fig = heatmap_figure()
+    fig.update_layout(height=1500, width=1200)
+    fig
+    return ACTIVATIONS, MODEL, PATTERN, fig, heatmap_figure
 
 
-@app.cell(hide_code=True)
+@app.cell
 def __():
-    tips = {
-        "Saving": (
-            """
-            **Saving**
-
-            - _Name_ your app using the box at the top of the screen, or
-              with `Ctrl/Cmd+s`. You can also create a named app at the
-              command line, e.g., `marimo edit app_name.py`.
-
-            - _Save_ by clicking the save icon on the bottom right, or by
-              inputting `Ctrl/Cmd+s`. By default marimo is configured
-              to autosave.
-            """
-        ),
-        "Running": (
-            """
-            1. _Run a cell_ by clicking the play ( ▷ ) button on the top
-            right of a cell, or by inputting `Ctrl/Cmd+Enter`.
-
-            2. _Run a stale cell_  by clicking the yellow run button on the
-            right of the cell, or by inputting `Ctrl/Cmd+Enter`. A cell is
-            stale when its code has been modified but not run.
-
-            3. _Run all stale cells_ by clicking the play ( ▷ ) button on
-            the bottom right of the screen, or input `Ctrl/Cmd+Shift+r`.
-            """
-        ),
-        "Console Output": (
-            """
-            Console output (e.g., `print()` statements) is shown below a
-            cell.
-            """
-        ),
-        "Creating, Moving, and Deleting Cells": (
-            """
-            1. _Create_ a new cell above or below a given one by clicking
-                the plus button to the left of the cell, which appears on
-                mouse hover.
-
-            2. _Move_ a cell up or down by dragging on the handle to the 
-                right of the cell, which appears on mouse hover.
-
-            3. _Delete_ a cell by clicking the trash bin icon. Bring it
-                back by clicking the undo button on the bottom right of the
-                screen, or with `Ctrl/Cmd+Shift+z`.
-            """
-        ),
-        "Disabling Automatic Execution": (
-            """
-            Via the notebook settings (gear icon) or footer panel, you
-            can disable automatic execution. This is helpful when
-            working with expensive notebooks or notebooks that have
-            side-effects like database transactions.
-            """
-        ),
-        "Disabling Cells": (
-            """
-            You can disable a cell via the cell context menu.
-            marimo will never run a disabled cell or any cells that depend on it.
-            This can help prevent accidental execution of expensive computations
-            when editing a notebook.
-            """
-        ),
-        "Code Folding": (
-            """
-            You can collapse or fold the code in a cell by clicking the arrow
-            icons in the line number column to the left, or by using keyboard
-            shortcuts.
-
-            Use the command palette (`Ctrl/Cmd+k`) or a keyboard shortcut to
-            quickly fold or unfold all cells.
-            """
-        ),
-        "Code Formatting": (
-            """
-            If you have [ruff](https://github.com/astral-sh/ruff) installed,
-            you can format a cell with the keyboard shortcut `Ctrl/Cmd+b`.
-            """
-        ),
-        "Command Palette": (
-            """
-            Use `Ctrl/Cmd+k` to open the command palette.
-            """
-        ),
-        "Keyboard Shortcuts": (
-            """
-            Open the notebook menu (top-right) or input `Ctrl/Cmd+Shift+h` to
-            view a list of all keyboard shortcuts.
-            """
-        ),
-        "Configuration": (
-            """
-           Configure the editor by clicking the gears icon near the top-right
-           of the screen.
-           """
-        ),
-    }
-    return (tips,)
+    # Add fens after opponents moves 
+    # Default squares of interest
+    return
 
 
 if __name__ == "__main__":