"""
Inside GPT-2: Embeddings
========================
An interactive teaching demo (Gradio, runs on a free CPU Hugging Face Space).

Models
- GPT-2 small (OpenAI, 2019): tokens, token/position embeddings, hidden states.
- paraphrase-multilingual-MiniLM-L12-v2: a small multilingual sentence-embedding model,
  used for the shared prompt map and as a comparison in "Compare your own lists".
Companion Space: "Inside GPT-2: Attention".
"""

import html
import os
import re
import threading
import time
import uuid

import gradio as gr
import numpy as np
import plotly.graph_objects as go
import torch
import umap
from transformers import AutoModel, AutoTokenizer

from reference_words import REFERENCE_WORDS

torch.set_grad_enabled(False)  # note: thread-local, so model calls below also use inference_mode

# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------
GPT2_ID = "openai-community/gpt2"
ST_ID = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
# Optional: link to the companion attention Space, e.g. "https://huggingface.co/spaces/<user>/inside-gpt2-attention"
COMPANION_URL = os.environ.get("COMPANION_URL", "")

tok = AutoTokenizer.from_pretrained(GPT2_ID)
gpt2 = AutoModel.from_pretrained(GPT2_ID).eval()
st_tok = AutoTokenizer.from_pretrained(ST_ID)
st_model = AutoModel.from_pretrained(ST_ID).eval()

N_LAYERS = gpt2.config.n_layer  # 12
D_MODEL = gpt2.config.n_embd  # 768
BOS = tok.bos_token_id  # 50256, "<|endoftext|>"
N_PARAMS = sum(p.numel() for p in gpt2.parameters())

MAX_TEXT_TOKENS = 128  # tokens tab
MAX_MAP_ITEMS = 60

WTE = gpt2.wte.weight.detach().float().numpy()  # (50257, 768) token embeddings
WPE = gpt2.wpe.weight.detach().float().numpy()  # (1024, 768) position embeddings
WTE_UNIT = WTE / np.linalg.norm(WTE, axis=1, keepdims=True)
VOCAB = [tok.decode([i]) for i in range(WTE.shape[0])]
WORD_RE = re.compile(r" ?[^\W\d_]{2,}")  # whole-word tokens (letters only)
WORD_MASK = np.array([bool(WORD_RE.fullmatch(v)) for v in VOCAB])
NORM_KEY = [v.strip().lower() for v in VOCAB]
@torch.inference_mode()
def hidden_states(ids):
    """All 13 residual-stream states for a list of token ids: (13, T, 768).
    Index 0 = token + position embedding, 1..12 = after each block
    (HF applies the final LayerNorm to index 12)."""
    out = gpt2(torch.tensor([ids]), output_hidden_states=True)
    return torch.stack(out.hidden_states)[:, 0].float().numpy()


# Reference statistics used to "normalize dimensions" of hidden states.
# GPT-2's residual stream has a large shared offset and a few huge outlier
# dimensions, which make every cosine similarity look like 0.9+. We z-score
# each dimension with mean/std measured on this neutral reference text.
REF_TEXT = (
    "The history of the town goes back many centuries. Today it is known for its "
    "museums, an old stone bridge and a lively music scene. Many students come here "
    "to study engineering, design and medicine, and the market sells fresh bread, "
    "cheese and vegetables every morning. In winter the river sometimes freezes, and "
    "children skate on the ice while their parents drink tea in small cafes. A new "
    "train line connects the station with the airport, which has made the region "
    "popular with tourists. Scientists at the local institute work on batteries, "
    "weather models and software for hospitals. My neighbour repairs old bicycles in "
    "her garage and sells them online. Last year she sold more than two hundred of "
    "them. When I asked her why she enjoys it, she said that every bike has a story, "
    "and that fixing something is better than throwing it away. Football is played "
    "on Saturdays, and the stadium can hold about twenty thousand people. Prices for "
    "apartments have risen quickly, so many young families now live in the villages "
    "nearby. The mayor promised more parks, cheaper buses and faster internet, but "
    "critics say the plans are too expensive. Still, most people agree that the "
    "town is a good place to live, especially in spring when the trees are green."
)
_REF = hidden_states([BOS] + tok.encode(REF_TEXT))[:, 1:]  # drop the BOS position
REF_MU = _REF.mean(axis=1)  # (13, 768)
REF_SD = _REF.std(axis=1) + 1e-5

# ---------------------------------------------------------------------------
# Visual style (light theme, validated categorical palette)
# ---------------------------------------------------------------------------
SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK2 = "#52514e"
MUTED = "#898781"
GRID = "#e1e0d9"
AXIS = "#c3c2b7"
MASKED = "#efeee9"
SERIES = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#6250d6", "#e34948"]
CHIP_TINTS = ["#d8e7fa", "#fbdccf", "#cdeee1", "#fbe7b9", "#e3dff8"]
FONT = "system-ui, -apple-system, 'Segoe UI', Roboto, sans-serif"

BLUE_STEPS = ["#fcfcfb", "#cde2fb", "#9ec5f4", "#6da7ec", "#3987e5", "#256abf", "#184f95", "#0d366b"]
SEQ_SCALE = [[i / (len(BLUE_STEPS) - 1), c] for i, c in enumerate(BLUE_STEPS)]
DIV_SCALE = [
    [0.0, "#8f2424"], [0.17, "#e34948"], [0.34, "#f2b8b5"], [0.5, "#f0efec"],
    [0.66, "#9ec5f4"], [0.83, "#3987e5"], [1.0, "#104281"],
]


def style(fig, title=None, height=420, legend=False):
    fig.update_layout(
        template="plotly_white",
        paper_bgcolor=SURFACE,
        plot_bgcolor=SURFACE,
        font=dict(family=FONT, color=INK2, size=13),
        title=dict(text=title, font=dict(color=INK, size=15), x=0.01, xanchor="left") if title else None,
        margin=dict(l=10, r=10, t=56 if title else 16, b=10),
        height=height,
        showlegend=legend,
        legend=dict(orientation="h", yanchor="bottom", y=1.0, x=1, xanchor="right", font=dict(color=INK2)),
        hoverlabel=dict(bgcolor="white", font=dict(family=FONT, color=INK), bordercolor=AXIS),
    )
    fig.update_xaxes(gridcolor=GRID, linecolor=AXIS, zeroline=False, tickfont=dict(color=MUTED))
    fig.update_yaxes(gridcolor=GRID, linecolor=AXIS, zeroline=False, tickfont=dict(color=MUTED))
    return fig


def empty_fig(message, height=300):
    fig = go.Figure()
    fig.add_annotation(text=message, showarrow=False, font=dict(size=14, color=INK2), x=0.5, y=0.5,
                       xref="paper", yref="paper")
    fig.update_xaxes(visible=False)
    fig.update_yaxes(visible=False)
    return style(fig, height=height)


def _bytes_to_unicode():
    """GPT-2's byte <-> printable-character table (from the original GPT-2 code)."""
    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return dict(zip(bs, map(chr, cs)))


_BYTE_DECODER = {v: k for k, v in _bytes_to_unicode().items()}


def token_text(i):
    """Token string. Byte fragments (parts of emoji or rare characters) become e.g. '[F0 9F]'."""
    t = VOCAB[i]
    if "\ufffd" not in t:
        return t
    raw = bytes(_BYTE_DECODER[c] for c in tok.convert_ids_to_tokens(int(i)))
    lead = " " if raw.startswith(b" ") else ""
    return lead + "[" + "\u2009".join(f"{b:02X}" for b in raw.lstrip(b" ")) + "]"  # thin space, not a token space


def show_tok(i, spaces=True, maxlen=None):
    """Readable label for token id i: leading space as '␣', BOS as [BOS], raw bytes as [F0 9F]."""
    if i == BOS:
        return "[BOS]"
    t = token_text(i).replace("\n", "↵").replace("\t", "⇥")
    t = t.replace(" ", "␣") if spaces else t.strip() or "␣"
    if maxlen and len(t) > maxlen:
        t = t[: maxlen - 1] + "…"
    return t


def short(s, n=34):
    s = s.strip()
    return s if len(s) <= n else s[: n - 1] + "…"


def parse_groups(text):
    """Lines -> [(group_name, [items])]. A line starting with '#' names a new
    group, a blank line also starts a new group."""
    groups, name, items = [], None, []

    def flush():
        nonlocal items, name
        if items:
            groups.append((name or f"group {len(groups) + 1}", items))
        items, name = [], None

    for raw in (text or "").splitlines():
        line = raw.strip()
        if not line:
            flush()
        elif line.startswith("#"):
            flush()
            name = line.lstrip("#").strip() or None
        else:
            items.append(line)
    flush()
    return groups


def cosine_matrix(X):
    Xu = X / (np.linalg.norm(X, axis=-1, keepdims=True) + 1e-9)
    return Xu @ np.swapaxes(Xu, -1, -2)


def pca(X, k=2):
    """Project L2-normalised rows of X onto their top-k principal axes."""
    Xu = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-9)
    Xc = Xu - Xu.mean(axis=0)
    U, S, _ = np.linalg.svd(Xc, full_matrices=False)
    var = S**2 / max((S**2).sum(), 1e-12)
    P = U[:, :k] * S[:k]
    if P.shape[1] < k:
        P = np.pad(P, ((0, 0), (0, k - P.shape[1])))
        var = np.pad(var, (0, k - len(var)))
    return P, var[:k]


def similarity_heatmap(C, labels, hover_labels, title, height=None):
    n = len(labels)
    numbered = max(len(l) for l in labels) > 12
    ylabels = [f"{i + 1}. {l}" if numbered and not l[:1].isdigit() else l for i, l in enumerate(labels)]
    xlabels = [str(i + 1) for i in range(n)] if numbered else labels
    hover = [[f"{html.escape(hover_labels[i])}<br>vs. {html.escape(hover_labels[j])}<br><b>cosine {C[i, j]:.3f}</b>"
              for j in range(n)] for i in range(n)]
    fig = go.Figure(go.Heatmap(
        z=C, x=list(range(n)), y=list(range(n)), zmin=-1, zmax=1, zmid=0, colorscale=DIV_SCALE,
        text=hover, hoverinfo="text", xgap=1, ygap=1,
        colorbar=dict(title=dict(text="cosine", font=dict(color=INK2)), thickness=12, tickfont=dict(color=MUTED)),
    ))
    style(fig, title, height or max(420, 150 + 30 * n))
    fig.update_xaxes(tickvals=list(range(n)), ticktext=xlabels, tickangle=0 if numbered else -45, showgrid=False,
                     constrain="domain")
    fig.update_yaxes(tickvals=list(range(n)), ticktext=ylabels, autorange="reversed", showgrid=False,
                     scaleanchor="x", scaleratio=1, constrain="domain")
    return fig


def scatter_groups(P, names, group_idx, group_names, title, var, point_labels=None):
    """2D or 3D (if P has 3 columns) scatter, one trace per group, every point labelled."""
    three_d = P.shape[1] == 3
    fig = go.Figure()
    for g, gname in enumerate(group_names):
        sel = [i for i, gi in enumerate(group_idx) if gi == g]
        if not sel:
            continue
        common = dict(
            mode="markers+text", name=gname,
            text=[point_labels[i] if point_labels else short(names[i], 28) for i in sel],
            textposition="top center", textfont=dict(color=INK, size=12),
            hovertext=[names[i] for i in sel], hoverinfo="text",
        )
        color = SERIES[g % len(SERIES)]
        if three_d:
            fig.add_trace(go.Scatter3d(x=P[sel, 0], y=P[sel, 1], z=P[sel, 2],
                                       marker=dict(size=6, color=color, line=dict(color=SURFACE, width=1)),
                                       **common))
        else:
            fig.add_trace(go.Scatter(x=P[sel, 0], y=P[sel, 1],
                                     marker=dict(size=11, color=color, line=dict(color=SURFACE, width=2)),
                                     **common))
    style(fig, title, 560 if three_d else 520, legend=len(group_names) > 1)
    labs = [f"PCA axis {i + 1} ({100 * v:.0f}%)" for i, v in enumerate(var)]
    if three_d:
        ax = dict(showticklabels=False, backgroundcolor=SURFACE, gridcolor=GRID, zeroline=False,
                  title=dict(font=dict(color=MUTED, size=11)))
        fig.update_layout(scene=dict(
            xaxis=dict(ax, title=dict(text=labs[0], font=dict(color=MUTED, size=11))),
            yaxis=dict(ax, title=dict(text=labs[1], font=dict(color=MUTED, size=11))),
            zaxis=dict(ax, title=dict(text=labs[2], font=dict(color=MUTED, size=11))),
            aspectmode="data"))
    else:
        span = np.ptp(P, axis=0) + 1e-9
        lo, hi = P.min(axis=0), P.max(axis=0)
        fig.update_xaxes(title=dict(text=labs[0], font=dict(color=MUTED)), showticklabels=False, zeroline=False,
                         range=[lo[0] - 0.22 * span[0], hi[0] + 0.22 * span[0]])
        fig.update_yaxes(title=dict(text=labs[1], font=dict(color=MUTED)), showticklabels=False, zeroline=False,
                         range=[lo[1] - 0.1 * span[1], hi[1] + 0.15 * span[1]])
    return fig


# ---------------------------------------------------------------------------
# Prompt map: a shared, live map of reference words + everybody's prompts
# ---------------------------------------------------------------------------
@torch.inference_mode()
def st_embed(texts):
    """Sentence embeddings (mean pooling over tokens, unit length), as sentence-transformers does it."""
    batch = st_tok(list(texts), padding=True, truncation=True, max_length=128, return_tensors="pt")
    out = st_model(**batch).last_hidden_state
    mask = batch["attention_mask"].unsqueeze(-1).float()
    emb = (out * mask).sum(1) / mask.sum(1)
    return torch.nn.functional.normalize(emb, dim=1).float().numpy()


MAP_GROUPS = list(REFERENCE_WORDS)
MAP_WORDS, MAP_GROUP_IDX = [], []
for _g, _words in enumerate(REFERENCE_WORDS.values()):
    MAP_WORDS += _words
    MAP_GROUP_IDX += [_g] * len(_words)
MAP_GROUP_IDX = np.array(MAP_GROUP_IDX)
MAP_X = st_embed(MAP_WORDS)  # (320, 384)

# Layout 1: PCA, a plain linear projection (new points are projected exactly)
_PCA_MU = MAP_X.mean(axis=0)
_, _S, _VT = np.linalg.svd(MAP_X - _PCA_MU, full_matrices=False)
_PCA_AXES = _VT[:2]
PCA_VAR = (_S**2 / (_S**2).sum())[:2]

# Layout 2: UMAP, non-linear, keeps neighbourhoods (chosen over t-SNE because it can place new
# points with .transform(); see README for the comparison)
_UMAP = umap.UMAP(n_neighbors=15, min_dist=0.15, metric="cosine", random_state=0, n_jobs=1).fit(MAP_X)
_UMAP.transform(MAP_X[:2])  # warm-up: compiles numba code at startup instead of on the first prompt

LAYOUTS = ["UMAP (clusters)", "PCA (linear)"]
MAP_XY = {LAYOUTS[0]: _UMAP.embedding_.astype(float), LAYOUTS[1]: (MAP_X - _PCA_MU) @ _PCA_AXES.T}


def _group_label_positions(XY, k=8):
    """For each group, the member whose k nearest 2D neighbours contain the most members of its own group."""
    d = ((XY[:, None, :] - XY[None, :, :]) ** 2).sum(-1)
    np.fill_diagonal(d, np.inf)
    nn = np.argsort(d, axis=1)[:, :k]
    same = (MAP_GROUP_IDX[nn] == MAP_GROUP_IDX[:, None]).sum(axis=1)
    pos = []
    for g in range(len(MAP_GROUPS)):
        members = np.where(MAP_GROUP_IDX == g)[0]
        pos.append(XY[members[np.argmax(same[members])]])
    return pos


GROUP_LABEL_XY = {name: _group_label_positions(xy) for name, xy in MAP_XY.items()}


def place_points(vecs):
    """2D positions of new vectors in both layouts, without changing the reference map."""
    return {LAYOUTS[0]: _UMAP.transform(vecs).astype(float), LAYOUTS[1]: (vecs - _PCA_MU) @ _PCA_AXES.T}


MAX_PROMPT_CHARS = 200
MAX_SHARED_PROMPTS = 300
MIN_SECONDS_BETWEEN = 1.5
ADMIN_KEY = os.environ.get("ADMIN_KEY", "")


class SharedPromptStore:
    """Prompts from all visitors, kept in memory only (gone when the Space restarts)."""

    def __init__(self):
        self.lock = threading.Lock()
        self.items = []
        self.version = 0
        self.last_add = {}
        self.next_id = 1

    def add(self, text, owner):
        vec = st_embed([text])
        xy = place_points(vec)
        sims = MAP_X @ vec[0]
        order = np.argsort(-sims)[:10]
        item = dict(text=text, owner=owner, vec=vec[0], xy={k: v[0] for k, v in xy.items()},
                    nbrs=[(int(i), float(sims[i])) for i in order])
        with self.lock:
            item["id"] = self.next_id
            self.next_id += 1
            self.items.append(item)
            del self.items[:-MAX_SHARED_PROMPTS]
            self.last_add[owner] = time.time()
            self.version += 1
        return item

    def too_soon(self, owner):
        return time.time() - self.last_add.get(owner, 0) < MIN_SECONDS_BETWEEN

    def remove_last_of(self, owner):
        with self.lock:
            for k in range(len(self.items) - 1, -1, -1):
                if self.items[k]["owner"] == owner:
                    del self.items[k]
                    self.version += 1
                    return True
        return False

    def remove_last(self):
        with self.lock:
            if self.items:
                self.items.pop()
                self.version += 1
                return True
        return False

    def clear(self):
        with self.lock:
            self.items = []
            self.version += 1

    def snapshot(self):
        with self.lock:
            return self.version, list(self.items)


STORE = SharedPromptStore()


def group_vote(item):
    votes = np.bincount(MAP_GROUP_IDX[[i for i, _ in item["nbrs"]]], minlength=len(MAP_GROUPS))
    return MAP_GROUPS[int(np.argmax(votes))], int(votes.max())


def prompt_map_fig(layout, show_labels, show_lines, owner, items):
    XY = MAP_XY[layout]
    fig = go.Figure()
    for g, gname in enumerate(MAP_GROUPS):
        sel = np.where(MAP_GROUP_IDX == g)[0]
        fig.add_trace(go.Scatter(
            x=XY[sel, 0], y=XY[sel, 1], mode="markers+text" if show_labels else "markers", name=gname,
            text=[MAP_WORDS[i] for i in sel], textposition="top center",
            textfont=dict(size=10, color=INK2), customdata=[gname] * len(sel),
            hovertemplate="<b>%{text}</b><br>%{customdata}<extra></extra>",
            marker=dict(size=8, color=SERIES[g % len(SERIES)], opacity=0.9, line=dict(color=SURFACE, width=1)),
        ))
    # group names as direct labels (so colour is not the only cue), placed where each group is densest
    for g, gname in enumerate(MAP_GROUPS):
        lx, ly = GROUP_LABEL_XY[layout][g]
        fig.add_annotation(x=lx, y=ly, text=f"<b>{html.escape(gname)}</b>", showarrow=False, yshift=14,
                           font=dict(size=13, color=INK), bgcolor="rgba(252,252,251,0.75)", borderpad=2)

    mine = [it for it in items if it["owner"] == owner]
    others = [it for it in items if it["owner"] != owner]
    focus = mine[-1] if mine else None
    if show_lines and focus is not None:
        p = focus["xy"][layout]
        xs, ys = [], []
        for i, _ in focus["nbrs"][:5]:
            xs += [p[0], XY[i, 0], None]
            ys += [p[1], XY[i, 1], None]
        fig.add_trace(go.Scatter(x=xs, y=ys, mode="lines", line=dict(color=INK2, width=1.5, dash="dot"),
                                 hoverinfo="skip", name="5 nearest words (in 384 dimensions)"))
        idx = [i for i, _ in focus["nbrs"][:5]]
        fig.add_trace(go.Scatter(
            x=XY[idx, 0], y=XY[idx, 1], mode="markers+text", showlegend=False,
            text=[f"<b>{MAP_WORDS[i]}</b>" for i in idx], textposition="top center",
            textfont=dict(size=12, color=INK), hoverinfo="skip",
            marker=dict(size=13, color="rgba(0,0,0,0)", line=dict(color=INK, width=1.5))))

    def prompt_trace(its, name, symbol, size, color, bold):
        if not its:
            return
        hover = []
        for it in its:
            near = ", ".join(f"{MAP_WORDS[i]} {s:.2f}" for i, s in it["nbrs"][:5])
            hover.append(f"<b>{html.escape(short(it['text'], 60))}</b><br>nearest: {html.escape(near)}")
        labels = [(f"<b>{html.escape(short(it['text'], 32))}</b>" if bold else html.escape(short(it["text"], 28)))
                  for it in its]
        fig.add_trace(go.Scatter(
            x=[it["xy"][layout][0] for it in its], y=[it["xy"][layout][1] for it in its],
            mode="markers+text", name=name, text=labels,
            textposition="bottom center", textfont=dict(size=12 if bold else 11, color=INK),
            hovertext=hover, hoverinfo="text",
            marker=dict(symbol=symbol, size=size, color=color, line=dict(color=SURFACE, width=1.5)),
        ))

    prompt_trace(others, f"prompts from others ({len(others)})", "diamond", 12, MUTED, False)
    prompt_trace(mine, f"your prompts ({len(mine)})", "star", 18, INK, True)

    if layout == LAYOUTS[1]:
        xt = f"PCA axis 1 ({100 * PCA_VAR[0]:.0f}% of variance)"
        yt = f"PCA axis 2 ({100 * PCA_VAR[1]:.0f}% of variance)"
        title = "PCA: a flat shadow of 384 dimensions. Distances are honest but groups overlap."
    else:
        xt = yt = "UMAP axes have no meaning, only neighbourhoods do"
        title = "UMAP: groups pulled apart so neighbourhoods are visible. Distances between clusters mean little."
    style(fig, title, 760, legend=True)
    fig.update_layout(uirevision=layout, legend=dict(orientation="h", y=-0.05, x=0, xanchor="left",
                                                     yanchor="top", font=dict(size=12)),
                      margin=dict(l=10, r=10, t=56, b=80), dragmode="pan")
    fig.update_xaxes(title=dict(text=xt, font=dict(color=MUTED)), showticklabels=False, zeroline=False,
                     showgrid=False)
    fig.update_yaxes(title=dict(text=yt, font=dict(color=MUTED)), showticklabels=False, zeroline=False,
                     showgrid=False)
    if layout == LAYOUTS[1]:  # PCA distances are meaningful, so keep the aspect ratio honest
        fig.update_yaxes(scaleanchor="x", scaleratio=1)
    return fig


def prompt_map_md(owner, items):
    mine = [it for it in items if it["owner"] == owner]
    owners = {it["owner"] for it in items}
    md = f"**{len(items)} prompt{'s' if len(items) != 1 else ''} on the map** from {len(owners)} visitor" \
         f"{'s' if len(owners) != 1 else ''}."
    if mine:
        it = mine[-1]
        near = ", ".join(f"{MAP_WORDS[i]} ({s:.2f}, {MAP_GROUPS[MAP_GROUP_IDX[i]]})" for i, s in it["nbrs"][:5])
        g, n = group_vote(it)
        md += (f"\n\nYour latest prompt **“{html.escape(it['text'])}”** is closest to: {near}. "
               f"{n} of its 10 nearest words are in **{g}**. The numbers are cosine similarities in the full "
               f"384-dimensional space: that is where the dotted lines point, wherever the 2D layout puts them.")
    return md


def render_prompt_map(layout, show_labels, show_lines, owner):
    version, items = STORE.snapshot()
    return prompt_map_fig(layout, show_labels, show_lines, owner, items), prompt_map_md(owner, items), version


def add_prompt(text, layout, show_labels, show_lines, owner):
    text = " ".join((text or "").split())[:MAX_PROMPT_CHARS]
    if not text:
        gr.Warning("Type a word or a prompt first.")
        return (gr.skip(),) * 3 + (gr.skip(),)
    if not owner:
        owner = uuid.uuid4().hex
    if STORE.too_soon(owner):
        gr.Warning("One moment: you can add a prompt every couple of seconds.")
        return (gr.skip(),) * 3 + (gr.skip(),)
    STORE.add(text, owner)
    return render_prompt_map(layout, show_labels, show_lines, owner) + ("",)


def undo_prompt(layout, show_labels, show_lines, owner):
    if not STORE.remove_last_of(owner):
        gr.Info("You have no prompts on the map yet.")
    return render_prompt_map(layout, show_labels, show_lines, owner)


def refresh_if_changed(layout, show_labels, show_lines, owner, seen_version):
    version, _ = STORE.snapshot()
    if version == seen_version:
        return gr.skip(), gr.skip(), gr.skip()
    return render_prompt_map(layout, show_labels, show_lines, owner)


def admin_action(key, action, layout, show_labels, show_lines, owner):
    if not ADMIN_KEY:
        return "Teacher controls are off: set a Space secret called `ADMIN_KEY` to enable them.", \
            *render_prompt_map(layout, show_labels, show_lines, owner)
    if key != ADMIN_KEY:
        return "Wrong key.", gr.skip(), gr.skip(), gr.skip()
    if action == "clear":
        STORE.clear()
        msg = "All prompts removed."
    else:
        msg = "Most recent prompt removed." if STORE.remove_last() else "The map has no prompts."
    return msg, *render_prompt_map(layout, show_labels, show_lines, owner)


def ensure_owner(owner):
    return owner or uuid.uuid4().hex


MAP_EXAMPLES = [
    "Mach das Licht in der Küche aus", "Will it snow tomorrow?", "I'm so nervous about my exam",
    "a cosy reading corner", "Ich hätte gern einen Kaffee", "neon poster for a techno party",
    "my cat sleeps all day", "Wie komme ich zum Bahnhof?", "my phone battery died",
]


# ---------------------------------------------------------------------------
# Tab 1: Tokens & vectors
# ---------------------------------------------------------------------------
VECTOR_KINDS = [
    "Token embedding (what the token is)",
    "Position embedding (where it is)",
    "Sum: what layer 1 actually receives",
]


def tokens_view(text, kind):
    text = text or ""
    ids = tok.encode(text)
    truncated = len(ids) > MAX_TEXT_TOKENS
    ids = ids[:MAX_TEXT_TOKENS]
    if not ids:
        return "<p>Type some text above.</p>", "", empty_fig("No tokens yet")

    chips = []
    for n, i in enumerate(ids):
        tint = CHIP_TINTS[n % len(CHIP_TINTS)]
        chips.append(
            f'<span class="tok-chip" style="background:{tint}" title="token #{n}, id {i}">'
            f'<span class="tok-text">{html.escape(show_tok(i))}</span>'
            f'<span class="tok-id">{i}</span></span>'
        )
    chips_html = '<div class="tok-row">' + "".join(chips) + "</div>"

    n_words = len(text.split())
    stats = (
        f"**{len(text)} characters, {n_words} words → {len(ids)} tokens**"
        + (f" (showing the first {MAX_TEXT_TOKENS})" if truncated else "")
        + f"\n\nThe vocabulary has **{WTE.shape[0]:,} tokens**. Each token id selects one row of the "
        f"embedding table: a list of **{D_MODEL} numbers**. The whole table holds "
        f"{WTE.size / 1e6:.1f} M numbers, about {100 * WTE.size / N_PARAMS:.0f}% of all "
        f"{N_PARAMS / 1e6:.0f} M parameters of GPT-2."
    )
    if any("\ufffd" in VOCAB[i] for i in ids):
        stats += ("\n\nTokens like `[F0 9F]` are **raw bytes**. Characters GPT-2 never learned as a whole "
                  "(most emoji, rare letters) are stored as several bytes, and the tokenizer falls back to them.")

    if kind == VECTOR_KINDS[0]:
        Z = WTE[ids]
    elif kind == VECTOR_KINDS[1]:
        Z = WPE[: len(ids)]
    else:
        Z = WTE[ids] + WPE[: len(ids)]
    lim = float(np.percentile(np.abs(Z), 99)) or 1.0
    labels = [f"{show_tok(i, maxlen=14)}" + (f"  (pos {n})" if kind != VECTOR_KINDS[0] else "")
              for n, i in enumerate(ids)]
    fig = go.Figure(go.Heatmap(
        z=Z, y=list(range(len(ids))), x=list(range(D_MODEL)), zmid=0, zmin=-lim, zmax=lim,
        colorscale=DIV_SCALE, customdata=np.array(labels, dtype=object)[:, None].repeat(D_MODEL, 1),
        hovertemplate="%{customdata}<br>dimension %{x}<br>value %{z:.3f}<extra></extra>",
        colorbar=dict(thickness=12, tickfont=dict(color=MUTED), title=dict(text="value", font=dict(color=INK2))),
    ))
    style(fig, f"{kind}: one row of {D_MODEL} numbers per token", max(260, 90 + 22 * len(ids)))
    fig.update_yaxes(tickvals=list(range(len(ids))), ticktext=labels, autorange="reversed", showgrid=False)
    fig.update_xaxes(title=dict(text="dimension (0 to 767)", font=dict(color=MUTED)), showgrid=False)
    return chips_html, stats, fig


# ---------------------------------------------------------------------------
# Tab 2: Nearest neighbours & vector arithmetic
# ---------------------------------------------------------------------------
def word_vector(word, lead_space):
    w = word.strip()
    ids = tok.encode((" " + w) if lead_space else w)
    if not ids:
        return None, [], ""
    note = ""
    if len(ids) > 1:
        parts = " + ".join(f"`{show_tok(i)}`" for i in ids)
        note = (f"`{w}` is not a single token, it is split into {parts}. "
                f"Using the average of those vectors (a rough approximation).")
    return WTE[ids].mean(axis=0), ids, note


def nearest(vec, k, exclude_ids=(), exclude_keys=(), hide_variants=True):
    sims = WTE_UNIT @ (vec / (np.linalg.norm(vec) + 1e-9))
    out, seen = [], set()
    for i in np.argsort(-sims):
        if not WORD_MASK[i] or i in exclude_ids or NORM_KEY[i] in exclude_keys:
            continue
        if hide_variants:
            if NORM_KEY[i] in seen:
                continue
            seen.add(NORM_KEY[i])
        out.append((int(i), float(sims[i])))
        if len(out) >= k:
            break
    return out


def bar_fig(results, title):
    labels = [show_tok(i) for i, _ in results]
    vals = [s for _, s in results]
    fig = go.Figure(go.Bar(
        x=vals, y=list(range(len(vals))), orientation="h", marker=dict(color=SERIES[0], line=dict(width=0)),
        customdata=labels, hovertemplate="%{customdata}<br>cosine similarity %{x:.3f}<extra></extra>",
    ))
    style(fig, title, max(300, 80 + 30 * len(vals)))
    fig.update_yaxes(tickvals=list(range(len(vals))), ticktext=labels, autorange="reversed", showgrid=False,
                     tickfont=dict(color=INK, size=13))
    fig.update_xaxes(title=dict(text="cosine similarity (1 = same direction)", font=dict(color=MUTED)),
                     range=[0, 1], gridcolor=GRID)
    fig.update_traces(marker_cornerradius=4)
    return fig


def neighbours_view(word, k, lead_space, hide_variants):
    if not (word or "").strip():
        return empty_fig("Type a word"), ""
    vec, ids, note = word_vector(word, lead_space)
    if vec is None:
        return empty_fig("Could not tokenize that"), ""
    res = nearest(vec, int(k), exclude_ids=set(ids), exclude_keys={word.strip().lower()},
                  hide_variants=hide_variants)
    lookup = " + ".join(f"`{show_tok(i)}`" for i in ids)
    md = f"Looked up {lookup} (token id {', '.join(map(str, ids))}). Searching {int(WORD_MASK.sum()):,} whole-word tokens."
    if note:
        md += "\n\n" + note
    return bar_fig(res, f"Closest tokens to “{word.strip()}” in embedding space"), md


def parse_expression(expr):
    terms = []
    for sign, word in re.findall(r"([+-]?)\s*([^+\-]+)", expr or ""):
        word = word.strip()
        if word:
            terms.append((-1.0 if sign == "-" else 1.0, word))
    return terms


def arithmetic_view(expr, k, lead_space, exclude_inputs, hide_variants):
    terms = parse_expression(expr)
    if not terms:
        return empty_fig("Type an expression like: king - man + woman"), ""
    vec = np.zeros(D_MODEL, dtype=np.float32)
    ids_all, notes = set(), []
    for sign, w in terms:
        v, ids, note = word_vector(w, lead_space)
        if v is None:
            continue
        vec += sign * v
        ids_all.update(ids)
        if note:
            notes.append(note)
    ex_ids = ids_all if exclude_inputs else set()
    ex_keys = {w.lower() for _, w in terms} if exclude_inputs else set()
    res = nearest(vec, int(k), exclude_ids=ex_ids, exclude_keys=ex_keys, hide_variants=hide_variants)
    pretty = " ".join(("−" if s < 0 else "+") + " " + w for s, w in terms).lstrip("+ ")
    md = ("Input words are **excluded** from the results (the standard trick in word-analogy demos). "
          "Untick the box to see what is really closest: often one of the input words itself."
          if exclude_inputs else
          "Input words are **included**. Often the closest vector is simply one of the inputs, "
          "because adding and subtracting only nudges the starting vector.")
    if notes:
        md += "\n\n" + "\n\n".join(notes)
    return bar_fig(res, f"Closest tokens to: {pretty}"), md


# ---------------------------------------------------------------------------
# Tab 3: Semantic map
# ---------------------------------------------------------------------------
SOURCES = [
    "GPT-2 token embeddings (static lookup)",
    "GPT-2 hidden states (contextual, pick a layer)",
    "Multilingual MiniLM sentence-embedding model",
]


@torch.inference_mode()
def embed_items(items, source, layer):
    if source == SOURCES[0]:
        return np.stack([WTE[tok.encode(" " + s.strip()) or [BOS]].mean(axis=0) for s in items])
    if source == SOURCES[1]:
        L = int(layer)
        vecs = []
        for s in items:
            H = hidden_states([BOS] + tok.encode(s))[L, 1:]
            vecs.append(((H - REF_MU[L]) / REF_SD[L]).mean(axis=0))
        return np.stack(vecs)
    return st_embed(items)


def semantic_map_view(text, source, layer, three_d=False):
    groups = parse_groups(text)
    items, gidx = [], []
    for g, (_, its) in enumerate(groups):
        for it in its:
            items.append(it)
            gidx.append(g)
    items, gidx = items[:MAX_MAP_ITEMS], gidx[:MAX_MAP_ITEMS]
    if len(items) < 3:
        msg = "Enter at least 3 words or sentences, one per line."
        return empty_fig(msg), empty_fig(msg), ""
    X = embed_items(items, source, layer)
    k = 3 if three_d else 2
    P, var = pca(X, k)
    gnames = [g for g, _ in groups]
    src = source.split(" (")[0] if source != SOURCES[1] else f"GPT-2 hidden states, layer {int(layer)}"
    scatter = scatter_groups(P, items, gidx, gnames, f"{src}: {X.shape[1]} dimensions flattened to {k}", var)
    C = cosine_matrix(X)
    labels = [f"{short(s, 22)}" for s in items]
    heat = similarity_heatmap(C, labels, items, "Cosine similarity between every pair")
    kept = 100 * var.sum()
    same = np.equal.outer(gidx, gidx)
    off = ~np.eye(len(items), dtype=bool)
    md = (f"**{len(items)} items, {X.shape[1]}-dimensional vectors.** The {k}D map keeps {kept:.0f}% of the "
          f"variation, so {100 - kept:.0f}% is lost in the flattening: trust the similarity matrix more than "
          f"distances on the map.")
    if len(groups) > 1 and (same & off).any() and (~same).any():
        md += (f"\n\nAverage similarity **within** groups: {C[same & off].mean():.2f}, "
               f"**between** groups: {C[~same].mean():.2f}. The bigger the gap, the better this "
               f"embedding separates your groups.")
    return scatter, heat, md


# ---------------------------------------------------------------------------
# Tab 4: Contextual embeddings
# ---------------------------------------------------------------------------
def find_word(ids, word):
    """Position of the LAST token of the first occurrence of `word` (skipping BOS)."""
    w = word.strip()
    variants = []
    for v in [w, w.lower(), w.capitalize()]:
        for c in (" " + v, v):
            enc = tok.encode(c)
            if enc and enc not in variants:
                variants.append(enc)
    best = None
    for enc in variants:
        n = len(enc)
        for i in range(1, len(ids) - n + 1):
            if ids[i:i + n] == enc:
                if best is None or i < best[0]:
                    best = (i, i + n - 1, enc)
                break
    return best


def contextual_view(word, text, normalize, layer):
    word = (word or "").strip()
    groups = parse_groups(text)
    if not word or not groups:
        msg = "Enter a word and some sentences that contain it."
        return empty_fig(msg), empty_fig(msg), empty_fig(msg), ""
    vecs, names, gidx, used, skipped = [], [], [], [], []
    for g, (_, sents) in enumerate(groups):
        for s in sents:
            ids = [BOS] + tok.encode(s)
            hit = find_word(ids, word)
            if hit is None:
                skipped.append(s)
                continue
            start, pos, enc = hit
            H = hidden_states(ids)[:, pos]  # (13, 768)
            vecs.append(H)
            names.append(s)
            gidx.append(g)
            used.append((s, enc, pos))
    if len(vecs) < 2:
        msg = f"Need at least 2 sentences containing “{word}”."
        return empty_fig(msg), empty_fig(msg), empty_fig(msg), ""
    V = np.stack(vecs, axis=1)  # (13, N, 768)
    if normalize:
        V = (V - REF_MU[:, None]) / REF_SD[:, None]
    C = cosine_matrix(V)  # (13, N, N)
    N = len(names)
    same = np.equal.outer(gidx, gidx)
    off = ~np.eye(N, dtype=bool)
    L = int(layer)
    layers = list(range(N_LAYERS + 1))
    gnames = [g for g, _ in groups]

    # Line chart across layers
    fig = go.Figure()
    multi = len(set(gidx)) > 1
    for i in range(N):
        for j in range(i + 1, N):
            sg = same[i, j]
            color = SERIES[0] if (sg or not multi) else SERIES[1]
            fig.add_trace(go.Scatter(
                x=layers, y=C[:, i, j], mode="lines", line=dict(color=color, width=1), opacity=0.28,
                hovertemplate=f"{html.escape(short(names[i], 30))}<br>vs. {html.escape(short(names[j], 30))}"
                              "<br>layer %{x}: %{y:.3f}<extra></extra>", showlegend=False,
            ))
    if multi:
        within = np.array([C[l][same & off].mean() for l in layers])
        between = np.array([C[l][~same].mean() for l in layers])
        fig.add_trace(go.Scatter(x=layers, y=within, mode="lines+markers", name="same group (average)",
                                 line=dict(color=SERIES[0], width=3), marker=dict(size=8)))
        fig.add_trace(go.Scatter(x=layers, y=between, mode="lines+markers", name="different groups (average)",
                                 line=dict(color=SERIES[1], width=3), marker=dict(size=8)))
    else:
        avg = np.array([C[l][off].mean() for l in layers])
        fig.add_trace(go.Scatter(x=layers, y=avg, mode="lines+markers", name="average of all pairs",
                                 line=dict(color=SERIES[0], width=3), marker=dict(size=8)))
    fig.add_vline(x=L, line=dict(color=MUTED, dash="dot", width=1))
    style(fig, f"How similar is “{word}” to itself across sentences, layer by layer?", 430, legend=True)
    fig.update_xaxes(title=dict(text="layer (0 = input embedding, 12 = last layer)", font=dict(color=MUTED)),
                     tickvals=layers, range=[-0.3, N_LAYERS + 0.3])
    fig.update_yaxes(title=dict(text="cosine similarity", font=dict(color=MUTED)))

    labels = [f"{n + 1}. {short(s, 22)}" for n, s in enumerate(names)]
    heat = similarity_heatmap(C[L], labels, names, f"Similarity of “{word}” between sentences, layer {L}")
    P, var = pca(V[L], 2)
    scatter = scatter_groups(P, [f"{n + 1}. {s}" for n, s in enumerate(names)], gidx, gnames,
                             f"“{word}” vectors at layer {L}, flattened to 2D", var,
                             point_labels=[str(n + 1) for n in range(N)])

    rows = []
    for n, (s, enc, pos) in enumerate(used):
        piece = " + ".join(f"`{show_tok(i)}`" for i in enc)
        rows.append(f"{n + 1}. token {piece} at position {pos}")
    md = "**Which token was used:** " + "; ".join(rows) + "."
    if any(len(e) > 1 for _, e, _ in used):
        md += " For words split into several tokens, the last piece is used (it has seen the whole word)."
    if skipped:
        md += "\n\n**Skipped (word not found):** " + "; ".join(f"“{short(s, 40)}”" for s in skipped)
    return fig, heat, scatter, md


# ---------------------------------------------------------------------------
# Defaults and copy
# ---------------------------------------------------------------------------
DEFAULT_TOKENS_TEXT = "Hello hello HELLO! Students in Nürnberg are learning how tokenization works."
DEFAULT_MAP_TEXT = """# colours
red
green
blue
yellow

# animals
dog
cat
horse
rabbit

# cities
Paris
Berlin
Tokyo
Madrid

# verbs
run
jump
swim
walk"""
SMART_HOME_MAP_TEXT = """# lights
Turn off the lights in the kitchen
Dim the living room lamp
Switch on the bedroom light

# weather
Will it rain tomorrow?
What's the weather like this weekend?
How cold is it outside?

# cooking
How long do I boil an egg?
Give me a recipe for pasta
What can I cook with potatoes?"""
DEFAULT_CTX_WORD = "bank"
DEFAULT_CTX_TEXT = """# river bank
We had a picnic on the grassy river bank
The fisherman sat on the muddy bank
Ducks swam near the river bank

# money bank
She deposited her salary at the bank
The robbers broke into the bank
He asked his bank for a loan"""
CTX_EXAMPLE_2_WORD = "light"
CTX_EXAMPLE_2_TEXT = """# light = lamp
Please turn on the light
The kitchen light is broken
He switched off the light

# light = not heavy
This suitcase is very light
My new laptop is really light
The jacket feels surprisingly light"""
INTRO = f"""
# Inside GPT-2: embeddings
How does a language model turn text into numbers it can compute with? This app opens up **GPT-2 small**
({N_PARAMS / 1e6:.0f} M parameters, {D_MODEL}-dimensional vectors) and a small **sentence-embedding model**.
GPT-2 is small and old (2019), but today's chat models start the same way: text becomes tokens, tokens become
vectors (**embeddings**), and similar meanings end up as nearby vectors.
""" + (f"Companion app: [Inside GPT-2: attention]({COMPANION_URL})" if COMPANION_URL else "")

CSS = f"""
.gradio-container {{ max-width: 1180px !important; }}
.tok-row {{ display:flex; flex-wrap:wrap; gap:6px 4px; padding:6px 0; }}
.tok-chip {{ display:inline-flex; flex-direction:column; align-items:center; padding:4px 7px 3px;
  border-radius:6px; font-family:{FONT}; line-height:1.15; }}
.tok-text {{ font-size:17px; color:{INK}; white-space:pre; }}
.tok-id {{ font-size:10.5px; color:{INK2}; font-variant-numeric:tabular-nums; }}
"""


# Plots are designed for a light background: pin the light theme. This runs in <head>, before the
# app starts, so the one-time redirect to ?__theme=light does not interrupt any requests.
FORCE_LIGHT_HEAD = """
<script>
(() => {
  const url = new URL(window.location);
  if (url.searchParams.get('__theme') !== 'light') {
    url.searchParams.set('__theme', 'light');
    window.location.replace(url.href);
  }
})();
</script>
"""

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    neutral_hue=gr.themes.colors.stone,
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
)

# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="Inside GPT-2: embeddings") as demo:
    gr.Markdown(INTRO)

    with gr.Tabs():
        # ---- Shared prompt map
        with gr.Tab("Prompt map"):
            gr.Markdown(
                "Every coloured dot is a word (hover to read it). Its position comes from its **embedding**: 384 numbers produced "
                "by a multilingual sentence-embedding model, flattened to 2D. Words with similar meanings get "
                "similar numbers, so they land near each other.  \n"
                "**Type a word or a whole prompt, in any language.** It is turned into 384 numbers the same way "
                "and placed on the map next to the words it is most similar to. Everybody using this page shares "
                "the map: your prompts are black stars, other people's are grey diamonds.  \n"
                "*Note: everything you add is visible to everyone who has this page open, until the app "
                "restarts. Please don't type anything personal.*"
            )
            with gr.Row(equal_height=True):
                pm_text = gr.Textbox(label="Your word or prompt (any language)", scale=4,
                                     placeholder="e.g. Mach das Licht aus, a cosy reading corner, dragon",
                                     max_length=MAX_PROMPT_CHARS)
                pm_add = gr.Button("Add to map", variant="primary", scale=1)
                pm_undo = gr.Button("Remove my last", scale=1)
            pm_ex = gr.Examples([[e] for e in MAP_EXAMPLES], [pm_text], label="Try one of these")
            with gr.Row():
                pm_layout = gr.Radio(LAYOUTS, value=LAYOUTS[0], label="Layout", scale=2)
                pm_labels = gr.Checkbox(value=False, label="Show all word labels", scale=1)
                pm_lines = gr.Checkbox(value=True, label="Lines to my prompt's 5 nearest words", scale=1)
                pm_live = gr.Checkbox(value=True, label="Live updates from others", scale=1)
            pm_md = gr.Markdown()
            pm_plot = gr.Plot(show_label=False)
            gr.Markdown(
                f"**UMAP or PCA?** Both squash 384 dimensions into 2, and both lose information. "
                f"*PCA* is a flat shadow: it keeps the two directions in which the {len(MAP_WORDS)} words differ "
                f"most ({100 * PCA_VAR.sum():.0f}% of the variation) and drops the rest, so the groups overlap. "
                f"*UMAP* bends and stretches the space so that every word stays close to its nearest neighbours. "
                f"The groups separate nicely, but the size of a cluster and the gaps between clusters mean little. "
                f"The dotted lines always point to the true nearest words in 384 dimensions: when they are long, "
                f"the 2D map is distorting.  \n"
                f"*Why UMAP and not t-SNE?* On these words both made equally good maps. UMAP can place new prompts "
                f"onto a finished map, while t-SNE has to recompute the whole map (or guess), and in a test UMAP "
                f"put held-out words into the right group a little more often."
            )
            with gr.Accordion("Teacher controls", open=False):
                gr.Markdown("Needs the `ADMIN_KEY` secret of this Space." if ADMIN_KEY else
                            "Off. To enable, add a secret called `ADMIN_KEY` in the Space settings.")
                with gr.Row():
                    pm_key = gr.Textbox(label="Admin key", type="password", scale=2)
                    pm_rm_last = gr.Button("Remove the most recent prompt", scale=1)
                    pm_clear = gr.Button("Clear all prompts", variant="stop", scale=1)
                pm_admin_msg = gr.Markdown()
            pm_owner = gr.BrowserState("", storage_key="inside-gpt2-prompt-map-owner")
            pm_seen = gr.State(-1)
            pm_timer = gr.Timer(4)

            pm_view = [pm_layout, pm_labels, pm_lines, pm_owner]
            pm_out = [pm_plot, pm_md, pm_seen]
            for trig in (pm_add.click, pm_text.submit):
                trig(add_prompt, [pm_text] + pm_view, pm_out + [pm_text])
            pm_ex.load_input_event.then(add_prompt, [pm_text] + pm_view, pm_out + [pm_text])
            pm_undo.click(undo_prompt, pm_view, pm_out)
            for c in (pm_layout, pm_labels, pm_lines):
                c.input(render_prompt_map, pm_view, pm_out)
            pm_timer.tick(refresh_if_changed, pm_view + [pm_seen], pm_out)
            pm_live.input(lambda v: gr.Timer(active=v), pm_live, pm_timer)
            pm_rm_last.click(lambda k, *v: admin_action(k, "last", *v), [pm_key] + pm_view, [pm_admin_msg] + pm_out)
            pm_clear.click(lambda k, *v: admin_action(k, "clear", *v), [pm_key] + pm_view, [pm_admin_msg] + pm_out)

        # ---- Tokens & vectors
        with gr.Tab("Tokens → vectors"):
            gr.Markdown(
                "**Step 1: text becomes tokens.** GPT-2 does not see letters or words. A tokenizer "
                "chops the text into pieces from a fixed vocabulary, and every piece has an id.  \n"
                "**Step 2: tokens become vectors.** Each id picks a row from a big table of learned "
                "numbers (the *embedding*). A second table adds information about the *position* in "
                "the text. Their sum is what the first layer receives.  \n"
                "*Try:* the same word with different capitalisation, a word with or without a space "
                "in front (`␣` marks a space), a German word, a long rare word, an emoji."
            )
            t_text = gr.Textbox(value=DEFAULT_TOKENS_TEXT, label="Text", lines=2)
            t_chips = gr.HTML()
            t_stats = gr.Markdown()
            t_kind = gr.Radio(VECTOR_KINDS, value=VECTOR_KINDS[0], label="Which vectors to show")
            t_plot = gr.Plot(label="Embedding vectors")
            gr.Markdown(
                "Blue = positive, red = negative. No single number means anything on its own: "
                "meaning lives in the *pattern* across all 768 dimensions. Notice that repeated "
                "tokens get identical rows in the token embedding, and that the position embeddings "
                "change smoothly from one position to the next (except position 0, which looks quite "
                "different from all the others)."
            )
            t_text.change(tokens_view, [t_text, t_kind], [t_chips, t_stats, t_plot])
            t_kind.change(tokens_view, [t_text, t_kind], [t_chips, t_stats, t_plot])

        # ---- Neighbours & arithmetic
        with gr.Tab("Neighbours & arithmetic"):
            gr.Markdown(
                "Every token is a point in a 768-dimensional space. Tokens that are used in similar "
                "ways end up pointing in similar **directions**. We measure this with *cosine "
                "similarity*: 1 = same direction, 0 = unrelated (at right angles).  \n"
                "Nobody programmed these relationships. They are a side effect of training GPT-2 to "
                "predict the next token."
            )
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Nearest neighbours")
                    n_word = gr.Textbox(value="king", label="Word")
                    n_k = gr.Slider(5, 25, value=12, step=1, label="How many neighbours")
                    n_btn = gr.Button("Find neighbours", variant="primary")
                    n_plot = gr.Plot()
                    n_md = gr.Markdown()
                with gr.Column():
                    gr.Markdown("### Vector arithmetic")
                    a_expr = gr.Textbox(value="king - man + woman", label="Expression (use + and -)")
                    a_k = gr.Slider(5, 25, value=12, step=1, label="How many results")
                    a_ex = gr.Checkbox(value=True, label="Exclude the input words from the results")
                    a_btn = gr.Button("Calculate", variant="primary")
                    a_plot = gr.Plot()
                    a_md = gr.Markdown()
            with gr.Row():
                lead = gr.Checkbox(value=True, label="Treat words as mid-sentence (adds a leading space: "
                                                     "“king” → “␣king”, which is a different token)")
                hide_var = gr.Checkbox(value=True, label="Hide spelling variants "
                                                         "(“␣King”, “King”, “king” count as one)")
            n_inputs = [n_word, n_k, lead, hide_var]
            a_inputs = [a_expr, a_k, lead, a_ex, hide_var]
            for trig in (n_btn.click, n_word.submit):
                trig(neighbours_view, n_inputs, [n_plot, n_md])
            for trig in (a_btn.click, a_expr.submit):
                trig(arithmetic_view, a_inputs, [a_plot, a_md])
            for c in (lead, hide_var):
                c.change(neighbours_view, n_inputs, [n_plot, n_md])
                c.change(arithmetic_view, a_inputs, [a_plot, a_md])
            a_ex.change(arithmetic_view, a_inputs, [a_plot, a_md])
            with gr.Row():
                n_ex = gr.Examples([["king"], ["cat"], ["Paris"], ["happy"], ["blue"], ["design"]],
                                   [n_word], label="Neighbour examples")
                a_ex_list = gr.Examples([["king - man + woman"], ["Paris - France + Germany"],
                                         ["walking - walk + swim"], ["bigger - big + small"],
                                         ["Berlin - Germany + Japan"]],
                                        [a_expr], label="Arithmetic examples")
            n_ex.load_input_event.then(neighbours_view, n_inputs, [n_plot, n_md])
            a_ex_list.load_input_event.then(arithmetic_view, a_inputs, [a_plot, a_md])

        # ---- Compare your own lists
        with gr.Tab("Compare your own lists"):
            gr.Markdown(
                "Enter words or sentences, **one per line**. Start a group with a line like "
                "`# animals` (or separate groups with an empty line); groups get colours.  \n"
                "The vectors have hundreds of dimensions, so to draw them we flatten them to 2D "
                "with PCA, which keeps the two directions along which your items differ the most.  \n"
                "*Compare the three sources:* GPT-2 was trained to predict text, not to compare "
                "sentences. A dedicated **sentence-embedding model** like MiniLM is trained exactly "
                "for that, which is why it is what search and RAG systems use."
            )
            with gr.Row():
                with gr.Column(scale=2):
                    m_text = gr.Textbox(value=DEFAULT_MAP_TEXT, lines=14, label="Items")
                with gr.Column(scale=3):
                    m_src = gr.Radio(SOURCES, value=SOURCES[0], label="Where do the vectors come from?")
                    m_layer = gr.Slider(0, N_LAYERS, value=6, step=1, label="GPT-2 layer", visible=False)
                    m_3d = gr.Checkbox(value=False, label="3D map (drag to rotate, keeps more information)")
                    m_btn = gr.Button("Draw map", variant="primary")
                    m_md = gr.Markdown()
                    m_ex = gr.Examples(
                        [[DEFAULT_MAP_TEXT, SOURCES[0]], [SMART_HOME_MAP_TEXT, SOURCES[0]],
                         [SMART_HOME_MAP_TEXT, SOURCES[2]]],
                        [m_text, m_src], label="Examples",
                        example_labels=["Words in 4 groups (GPT-2 static)",
                                        "Smart home requests (GPT-2 static)",
                                        "Smart home requests (MiniLM)"])
            with gr.Row():
                m_scatter = gr.Plot()
                m_heat = gr.Plot()
            m_inputs = [m_text, m_src, m_layer, m_3d]
            m_outputs = [m_scatter, m_heat, m_md]
            m_btn.click(semantic_map_view, m_inputs, m_outputs)
            m_src.change(lambda s: gr.update(visible=s == SOURCES[1]), m_src, m_layer)
            m_src.input(semantic_map_view, m_inputs, m_outputs)
            m_layer.release(semantic_map_view, m_inputs, m_outputs)
            m_3d.input(semantic_map_view, m_inputs, m_outputs)
            m_ex.load_input_event.then(semantic_map_view, m_inputs, m_outputs)

        # ---- Contextual embeddings
        with gr.Tab("Same word, different meaning"):
            gr.Markdown(
                "At the input (layer 0) a word gets practically the **same** vector in every "
                "sentence: the same token embedding, plus a small signal for its position. Then every "
                "layer mixes in information from the other tokens. Watch the vectors for the same "
                "word drift apart when the meaning differs.  \n"
                "**Important:** GPT-2 reads left to right. A token can only be influenced by the words "
                "*before* it, so put the clue (“river”, “money”) before the target word."
            )
            with gr.Row():
                with gr.Column(scale=2):
                    c_word = gr.Textbox(value=DEFAULT_CTX_WORD, label="Target word")
                    c_text = gr.Textbox(value=DEFAULT_CTX_TEXT, lines=10,
                                        label="Sentences (one per line, # starts a group)")
                    c_norm = gr.Checkbox(value=True, label="Normalise dimensions (recommended)")
                    c_btn = gr.Button("Compare", variant="primary")
                    c_ex = gr.Examples([[DEFAULT_CTX_WORD, DEFAULT_CTX_TEXT],
                                        [CTX_EXAMPLE_2_WORD, CTX_EXAMPLE_2_TEXT]],
                                       [c_word, c_text], label="Examples",
                                       example_labels=["bank: river vs. money",
                                                       "light: lamp vs. not heavy"])
                with gr.Column(scale=3):
                    c_line = gr.Plot()
                    c_layer = gr.Slider(0, N_LAYERS, value=8, step=1,
                                        label="Layer for the two views below")
            with gr.Row():
                c_heat = gr.Plot()
                c_scatter = gr.Plot()
            c_md = gr.Markdown()
            gr.Markdown(
                "*About “normalise dimensions”:* GPT-2's internal vectors share a big common offset "
                "and a handful of huge “outlier” dimensions. Without correction, every pair looks "
                "~0.9 similar. We subtract the typical value of each dimension and divide by its "
                "typical spread (measured on a neutral reference text). Untick it to see the raw "
                "numbers. Layer 12 includes GPT-2's final normalisation step."
            )
            c_inputs = [c_word, c_text, c_norm, c_layer]
            c_outputs = [c_line, c_heat, c_scatter, c_md]
            c_btn.click(contextual_view, c_inputs, c_outputs)
            c_word.submit(contextual_view, c_inputs, c_outputs)
            c_norm.change(contextual_view, c_inputs, c_outputs)
            c_layer.release(contextual_view, c_inputs, c_outputs)
            c_ex.load_input_event.then(contextual_view, c_inputs, c_outputs)


    gr.Markdown(
        "<small>Models: [GPT-2 small](https://huggingface.co/openai-community/gpt2) (OpenAI, 2019) and "
        "[paraphrase-multilingual-MiniLM-L12-v2](https://huggingface.co/sentence-transformers/"
        "paraphrase-multilingual-MiniLM-L12-v2). Everything runs on CPU. Prompts on the shared map are kept "
        "in memory only and disappear when the app restarts.</small>"
    )

    # Fill every view once on page load
    demo.load(ensure_owner, pm_owner, pm_owner).then(render_prompt_map, pm_view, pm_out)
    demo.load(tokens_view, [t_text, t_kind], [t_chips, t_stats, t_plot])
    demo.load(neighbours_view, n_inputs, [n_plot, n_md])
    demo.load(arithmetic_view, a_inputs, [a_plot, a_md])
    demo.load(semantic_map_view, m_inputs, m_outputs)
    demo.load(contextual_view, c_inputs, c_outputs)

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS, head=FORCE_LIGHT_HEAD)