File size: 2,674 Bytes
231a212 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | 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()
|