app.py

from __future__ import annotations

from pathlib import Path

import gradio as gr
import plotly.graph_objects as go
import torch
from data import VALUE_TOKENS, generate_selective_memory
from model import SelectiveSSM
from safetensors.torch import load_file

ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "micro-mamba"
MODEL = SelectiveSSM(selective=True)
MODEL.load_state_dict(load_file(ARTIFACT_DIR / "selective_ssm.safetensors"))
MODEL.eval()


def inspect_memory(seed: int, length: int) -> tuple[go.Figure, dict]:
    tokens, markers, targets = generate_selective_memory(
        1, int(length), int(seed)
    )
    with torch.inference_mode():
        logits, trace = MODEL(
            torch.from_numpy(tokens),
            torch.from_numpy(markers),
            return_trace=True,
        )
    probabilities = torch.softmax(logits, dim=1).numpy()[0]
    update = trace.numpy()[0]
    display = [
        f"Q{token - VALUE_TOKENS}" if token >= VALUE_TOKENS else str(token)
        for token in tokens[0]
    ]
    colors = ["#f59e0b" if marker else "#38bdf8" for marker in markers[0]]
    figure = go.Figure(
        go.Bar(
            x=list(range(len(display))),
            y=update,
            marker_color=colors,
            customdata=display,
            hovertemplate="step=%{x}<br>token=%{customdata}<br>update=%{y:.3f}",
        )
    )
    figure.update_layout(
        title="Input-dependent state update strength (orange = marked)",
        xaxis_title="Sequence position",
        yaxis_title="Mean discretization strength",
        template="plotly_dark",
    )
    prediction = int(probabilities.argmax())
    return figure, {
        "marked_values": tokens[0][markers[0] == 1].astype(int).tolist(),
        "query": int(tokens[0, -1] - VALUE_TOKENS),
        "target": int(targets[0]),
        "prediction": prediction,
        "correct": prediction == int(targets[0]),
        "confidence": round(float(probabilities[prediction]), 4),
    }


with gr.Blocks(title="MicroMamba") as demo:
    gr.Markdown(
        "# MicroMamba\n"
        "Inspect how a tiny selective state-space model updates its hidden state "
        "while retrieving one marked symbol from a field of distractors."
    )
    with gr.Row():
        seed = gr.Number(2043, precision=0, label="Sequence seed")
        length = gr.Slider(24, 96, 48, step=8, label="Sequence length")
    run = gr.Button("Run selective scan", variant="primary")
    trace = gr.Plot()
    result = gr.JSON()
    run.click(inspect_memory, [seed, length], [trace, result])
    demo.load(inspect_memory, [seed, length], [trace, result])


if __name__ == "__main__":
    demo.launch()