| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import gradio as gr |
| import plotly.graph_objects as go |
| import torch |
| from data import VALUES, generate_bindings |
| from model import FastWeightProgrammer |
| from safetensors.torch import load_file |
|
|
| ARTIFACT_DIR = ( |
| Path(__file__).resolve().parent / "artifacts" / "fast-weight-time-machine" |
| ) |
| MODEL = FastWeightProgrammer() |
| MODEL.load_state_dict(load_file(ARTIFACT_DIR / "fast_weight.safetensors")) |
| MODEL.eval() |
|
|
|
|
| def inspect_binding( |
| seed: int, pairs: int, distractors: int |
| ) -> tuple[go.Figure, dict]: |
| keys, values, writes, targets = generate_bindings( |
| 1, int(pairs), int(distractors), int(seed) |
| ) |
| with torch.inference_mode(): |
| logits, strengths, contributions = MODEL( |
| torch.from_numpy(keys), |
| torch.from_numpy(values), |
| torch.from_numpy(writes), |
| return_trace=True, |
| ) |
| probabilities = torch.softmax(logits, dim=1).numpy()[0] |
| contribution = contributions.numpy()[0] |
| labels = [ |
| f"K{key}:V{value}" if value < VALUES else f"QUERY K{key}" |
| for key, value in zip(keys[0], values[0], strict=True) |
| ] |
| colors = ["#f59e0b" if write else "#334155" for write in writes[0]] |
| figure = go.Figure( |
| go.Bar( |
| x=list(range(len(labels))), |
| y=contribution, |
| marker_color=colors, |
| customdata=labels, |
| hovertemplate="%{customdata}<br>read contribution=%{y:.3f}", |
| ) |
| ) |
| figure.update_layout( |
| title="Query-key contribution to the fast weight matrix", |
| xaxis_title="Sequence event", |
| yaxis_title="Contribution", |
| template="plotly_dark", |
| ) |
| prediction = int(probabilities.argmax()) |
| return figure, { |
| "query": labels[-1], |
| "target_value": int(targets[0]), |
| "predicted_value": prediction, |
| "correct": prediction == int(targets[0]), |
| "confidence": round(float(probabilities[prediction]), 4), |
| "mean_active_write_strength": round( |
| float(strengths.numpy()[0][writes[0] == 1].mean()), 4 |
| ), |
| } |
|
|
|
|
| with gr.Blocks(title="Fast-Weight Time Machine") as demo: |
| gr.Markdown( |
| "# Fast-Weight Time Machine\n" |
| "Watch a learned controller write temporary key/value bindings into a " |
| "sequence-local weight matrix, then retrieve one binding." |
| ) |
| with gr.Row(): |
| seed = gr.Number(2043, precision=0, label="Sequence seed") |
| pairs = gr.Slider(2, 12, 4, step=1, label="Stored bindings") |
| distractors = gr.Slider(0, 64, 12, step=4, label="Distractors") |
| run = gr.Button("Program fast memory", variant="primary") |
| trace = gr.Plot() |
| result = gr.JSON() |
| run.click(inspect_binding, [seed, pairs, distractors], [trace, result]) |
| demo.load(inspect_binding, [seed, pairs, distractors], [trace, result]) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
