"""RoPE Explorer Gradio app. Imports only from ``src/``.""" from __future__ import annotations import numpy as np import plotly.graph_objects as go import gradio as gr import spaces from src.absolute_pe import add_positional_encoding from src.extract import ( DEFAULT_MODEL, MAX_SEQ_LEN, MODEL_CHOICES, expand_kv_heads, extract_from_model, get_model_dimensions, random_qk, select_head, ) from src.plots import ( additive_pe_heatmaps, attention_bars, attention_heatmaps, bulk_before_after_delta, frequency_strip, norm_compare_add_vs_rope, norms_and_cosine, position_sweep, rotation_2d, theta_heatmap, ) from src.rope import ( attention_scores, pair_dim_labels, pair_frequencies, pair_xy, rotate_pair, theta_grid, ) HOWTO_MD = """ # How RoPE works After each token has a vector from the **embedding table**, attention builds two extra vectors per token with linear layers (`q_proj`, `k_proj`): - **Q (query)** — what this token is looking for - **K (key)** — what this token offers as a match Attention scores are (scaled) **dot products** `Q · K`. **RoPE rotates those Q and K vectors in 2D planes before the dot product.** It does **not** add a position vector onto the raw token embeddings. This app shows embeddings only as context, then focuses on Q and K before vs after RoPE. ## Pairwise rotation For even dimension `d`, pair `i` uses frequency $$\\omega_i = 10000^{-2i/d},\\qquad \\theta(k,i) = k\\,\\omega_i$$ **Interleaved (paper-style)** pairing `(2i, 2i+1)`: $$ x'_{2i} = x_{2i}\\cos\\theta - x_{2i+1}\\sin\\theta,\\qquad x'_{2i+1} = x_{2i}\\sin\\theta + x_{2i+1}\\cos\\theta $$ Hugging Face Llama-like models use the same frequencies but pair `(i, i + d/2)` (`rotate_half`). Random-matrix mode uses interleaved pairing; real models use the Llama layout so the numpy implementation can be checksummed against `rotary_emb`. Relative positions fall out of the algebra: `R(m)^T R(n) = R(n-m)`. Shaw relative attention (learned bias `b_{m-n}` on scores) is a **different** mechanism and is not computed here. """ PLACEHOLDER = go.Figure().update_layout( title="Run **Compute** on the Setup tab first", template="plotly_white", height=320, ) @spaces.GPU def gpu_test(): return "GPU available" def _safe_slider_max(n: int) -> int: """Gradio sliders need max > min; keep a one-step range even at edge cases.""" return max(int(n), 1) def update_dimension(source: str, model_name: str): if source.startswith("Random"): return ( gr.update(minimum=4, maximum=128, value=32, step=2, interactive=True), gr.update(value=32), gr.update(value=1), gr.update(value=32), ) try: total_dim, n_heads, head_dim = get_model_dimensions(model_name) return ( gr.update( minimum=4, maximum=max(128, head_dim), value=head_dim, step=2, interactive=False, ), gr.update(value=total_dim), gr.update(value=n_heads), gr.update(value=head_dim), ) except Exception: return gr.update(), gr.update(), gr.update(), gr.update() def compute( source: str, sentence: str, model_name: str, seq_len: int, dim: int, seed: int, base: float, progress=gr.Progress(track_tqdm=False), ): try: if source.startswith("Random"): progress(0.4, desc="Sampling random Q/K") data = random_qk(int(seq_len), int(dim), seed=int(seed), base=float(base)) else: progress(0.2, desc=f"Loading {model_name} (first time downloads weights)") data = extract_from_model(model_name, sentence) seq = int(select_head(data["q_before"], 0).shape[0]) head_dim = int(select_head(data["q_before"], 0).shape[1]) n_pairs = head_dim // 2 n_heads = max(int(data["n_q_heads"]) - 1, 0) checksum = data["checksum"] if checksum is None: status = ( f"Random Q/K · seq={seq} · dim={head_dim} · base={data['base']:g} · " f"style={data['style']}" ) else: status = ( f"Model `{data['model_name']}` · {seq} tokens · head_dim={head_dim} · " f"Q heads={data['n_q_heads']} · KV heads={data['n_kv_heads']} · " f"rope_theta={data['base']:g} · " f"max |numpy RoPE − model rotary| on Q = **{checksum:.3e}**" ) token_labels = ", ".join(data["tokens"][:seq]) status = status + f"\n\nTokens: `{token_labels}`" return ( data, status, gr.update(maximum=_safe_slider_max(n_heads), value=0), gr.update(maximum=_safe_slider_max(seq - 1), value=0), gr.update(maximum=_safe_slider_max(n_pairs - 1), value=0), gr.update(maximum=_safe_slider_max(seq - 1), value=0), ) except Exception as exc: return ( None, f"**Error:** {exc}", gr.update(), gr.update(), gr.update(), gr.update(), ) def _qk_slice(data: dict, which: str, head: int): before = data["q_before"] if which == "Q" else data["k_before"] after = data["q_after"] if which == "Q" else data["k_after"] return select_head(before, head), select_head(after, head) def update_bulk(data, which, head, mod_2pi): if not data: fig = PLACEHOLDER return fig, fig, fig, fig before, after = _qk_slice(data, which, int(head)) dim = before.shape[-1] seq = before.shape[0] return ( bulk_before_after_delta(before, after, tokens=data["tokens"]), norms_and_cosine(before, after, tokens=data["tokens"]), theta_heatmap(seq, dim, data["base"], mod_2pi=bool(mod_2pi)), frequency_strip(dim, data["base"]), ) def update_individual(data, which, head, token, pair, sweep): if not data: return "Compute on the Setup tab first.", PLACEHOLDER, PLACEHOLDER before, after = _qk_slice(data, which, int(head)) token = int(np.clip(token, 0, before.shape[0] - 1)) n_pairs = before.shape[1] // 2 pair = int(np.clip(pair, 0, n_pairs - 1)) style = data["style"] xb, yb = pair_xy(before, token, pair, style=style) xa, ya = pair_xy(after, token, pair, style=style) theta = float(theta_grid(before.shape[0], before.shape[1], data["base"])[token, pair]) cos_t, sin_t = float(np.cos(theta)), float(np.sin(theta)) xe_chk, xo_chk = rotate_pair(np.array([xb]), np.array([yb]), np.array([theta])) d0, d1 = pair_dim_labels(pair, before.shape[1], style=style) table = f""" ### Token `{token}` · pair `{pair}` (`{d0}`, `{d1}`) | | {d0} | {d1} | |---|---:|---:| | before | {xb:.6f} | {yb:.6f} | | after | {xa:.6f} | {ya:.6f} | | check (`rotate_pair`) | {float(xe_chk):.6f} | {float(xo_chk):.6f} | **θ(k,i) = {theta:.6f} rad** · cos = {cos_t:.6f} · sin = {sin_t:.6f} `x'_even = x_even cos θ − x_odd sin θ` `x'_odd = x_even sin θ + x_odd cos θ` """ neighbors = [n for n in (token - 1, token + 1, token + 2) if 0 <= n < before.shape[0]] rot = rotation_2d(before, after, token, pair, style, theta, neighbor_tokens=neighbors) if sweep: omega = float(pair_frequencies(before.shape[1], data["base"])[pair]) sweep_fig = position_sweep(xb, yb, omega, before.shape[0], token) else: sweep_fig = PLACEHOLDER sweep_fig.update_layout(title="Enable “replay same pair at every k” to see position-only spin") return table, rot, sweep_fig def update_attention(data, head, query_token): if not data: return PLACEHOLDER, PLACEHOLDER, "" q_b = select_head(data["q_before"], int(head)) q_a = select_head(data["q_after"], int(head)) k_b_all = expand_kv_heads(data["k_before"], data["n_q_heads"]) k_a_all = expand_kv_heads(data["k_after"], data["n_q_heads"]) k_b = select_head(k_b_all, int(head)) k_a = select_head(k_a_all, int(head)) sb = attention_scores(q_b, k_b) sa = attention_scores(q_a, k_a) qt = int(np.clip(query_token, 0, sb.shape[0] - 1)) note = ( "Additive PE changes values by **addition**. RoPE encodes **relative** offset " "because `R(m)^T R(n) = R(n−m)`: the score depends on the position difference, " "not on absolute indices alone." ) return attention_heatmaps(sb, sa), attention_bars(sb[qt], sa[qt], qt), note def update_compare(data): if not data: return PLACEHOLDER, PLACEHOLDER, "" emb = np.asarray(data["embeddings"], dtype=np.float64) # Compare tab always uses additive PE on the embedding matrix (may be wider than a head). pe, combined = add_positional_encoding(emb, base=data["base"]) q_b = select_head(data["q_before"], 0) q_a = select_head(data["q_after"], 0) heat = additive_pe_heatmaps(emb, pe, combined) norms = norm_compare_add_vs_rope(emb, combined, q_b, q_a) copy = """ **Absolute sinusoidal PE** *adds* a position-shaped vector, so both **norm and direction** change. **RoPE** *rotates* query/key pairs: **norm stays**, and the relative angle depends on `m − n`. Shaw-style relative attention (`q_m^T k_n + b_{m-n}`) is a third, learned-bias mechanism — not shown as a plot. """ return heat, norms, copy def toggle_source(source: str): is_random = source.startswith("Random") return ( gr.update(visible=is_random), gr.update(visible=True, interactive=is_random), gr.update(visible=is_random), gr.update(visible=not is_random), gr.update(visible=not is_random), ) with gr.Blocks(title="RoPE Explorer") as demo: state = gr.State(None) gr.Markdown("# RoPE Explorer") gr.Markdown( "Interactive view of **Rotary Position Embedding**: random Q/K matrices or " "query/key vectors from a small ungated Hugging Face model." ) with gr.Tabs(): with gr.Tab("How RoPE works"): gr.Markdown(HOWTO_MD) with gr.Tab("Setup"): source = gr.Radio( ["Random matrix", "Real model"], value="Random matrix", label="Source", ) with gr.Row(): sentence = gr.Textbox( value="RoPE rotates query and key vectors.", label="Sentence (real model)", visible=False, ) model_name = gr.Dropdown( MODEL_CHOICES, value=DEFAULT_MODEL, label="Model (ungated, Llama-like)", visible=False, ) with gr.Row(): seq_len = gr.Slider(2, MAX_SEQ_LEN, value=16, step=1, label="Sequence length") dim = gr.Slider(4, 128, value=32, step=2, label="Dimension (even; per attention head)") seed = gr.Number(value=42, label="Seed", precision=0) with gr.Row(): total_dim = gr.Number(value=32, label="Total dimension", precision=0, interactive=False) attention_heads = gr.Number(value=1, label="Attention heads", precision=0, interactive=False) head_dim = gr.Number(value=32, label="Dimension per attention head", precision=0, interactive=False) gr.Markdown( "**Why these numbers differ:** `total dimension = attention heads × dimension per head`. " "RoPE rotates each query/key head separately, so its Dimension slider uses " "the per-head value, not the model's total dimension." ) base = gr.Number( value=10000, label="RoPE base (overridden by config.rope_theta for real models)", ) compute_btn = gr.Button("Compute", variant="primary") status = gr.Markdown("Choose a source and click Compute.") head = gr.Slider(minimum=0, maximum=2, step=1, value=0, label="Head index (real models)") with gr.Tab("Bulk changes"): which = gr.Radio(["Q", "K"], value="Q", label="Tensor") mod_2pi = gr.Checkbox(False, label="θ heatmap: wrap mod 2π") bulk_main = gr.Plot(label="Before / after / delta") bulk_norm = gr.Plot(label="Norms and cosine") bulk_theta = gr.Plot(label="θ(k, i)") bulk_freq = gr.Plot(label="ω_i") with gr.Tab("Individual changes"): with gr.Row(): token_k = gr.Slider(0, 15, step=1, value=0, label="Token index k") pair_i = gr.Slider(0, 15, step=1, value=0, label="Pair index i") sweep = gr.Checkbox(True, label="Replay the same content pair at every position k") pair_table = gr.Markdown() pair_plot = gr.Plot() sweep_plot = gr.Plot() with gr.Tab("Attention effect"): query_token = gr.Slider(0, 15, step=1, value=0, label="Query token") attn_heat = gr.Plot() attn_bar = gr.Plot() attn_note = gr.Markdown() with gr.Tab("Compare to additive PE"): pe_heat = gr.Plot() pe_norm = gr.Plot() pe_note = gr.Markdown() compute_btn.click( compute, inputs=[source, sentence, model_name, seq_len, dim, seed, base], outputs=[state, status, head, token_k, pair_i, query_token], ) source.change( toggle_source, inputs=[source], outputs=[seq_len, dim, seed, sentence, model_name], ) source.change( update_dimension, inputs=[source, model_name], outputs=[dim, total_dim, attention_heads, head_dim], ) model_name.change( update_dimension, inputs=[source, model_name], outputs=[dim, total_dim, attention_heads, head_dim], ) bulk_inputs = [state, which, head, mod_2pi] bulk_outputs = [bulk_main, bulk_norm, bulk_theta, bulk_freq] for ctrl in bulk_inputs: ctrl.change(update_bulk, inputs=bulk_inputs, outputs=bulk_outputs) ind_inputs = [state, which, head, token_k, pair_i, sweep] ind_outputs = [pair_table, pair_plot, sweep_plot] for ctrl in ind_inputs: ctrl.change(update_individual, inputs=ind_inputs, outputs=ind_outputs) attn_inputs = [state, head, query_token] attn_outputs = [attn_heat, attn_bar, attn_note] for ctrl in attn_inputs: ctrl.change(update_attention, inputs=attn_inputs, outputs=attn_outputs) state.change(update_compare, inputs=[state], outputs=[pe_heat, pe_norm, pe_note]) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860)