LLM basics: one token at a time

import { test, expect } from '@playwright/test';

const URL = 'https://app.thetestingacademy.com/ai/blueprint/learn/llm-basics.html';

test('temperature 0 returns the top token on every run', async ({ page }) => {
  await page.goto(URL);
  await page.getByTestId('llm-prompt').selectOption('auth401');
  await page.getByTestId('llm-temp').fill('0');
  await expect(page.getByTestId('llm-temp-value')).toHaveText('T = 0.0');

  await page.getByTestId('llm-run').click();
  const runs = page.getByTestId('llm-runs').getByRole('listitem');
  await expect(runs).toHaveCount(1);
  await expect(runs.first()).toContainText('20 of 20 runs');
  await expect(page.getByTestId('llm-status')).toContainText('1 distinct token');
});

test('changing one token in the prompt moves the top candidate', async ({ page }) => {
  await page.goto(URL);
  await expect(page.getByTestId('llm-top')).toHaveText('token (46.6%)');

  await page.getByTestId('llm-prompt').selectOption('server500');
  await expect(page.getByTestId('llm-top')).toHaveText('server (43.2%)');
  const changed = page.getByTestId('llm-tok-5');
  await expect(changed).toContainText('500');
  await expect(changed).toHaveClass(/is-on/);
});

test('a measurable prompt keeps its top token as temperature rises', async ({ page }) => {
  await page.goto(URL);
  await page.getByTestId('llm-prompt').selectOption('specific');
  await expect(page.getByTestId('llm-prob-1')).toHaveText('78.1%');

  await page.getByTestId('llm-temp').fill('2');
  await expect(page.getByTestId('llm-prob-1')).toHaveText('48.9%');

  await page.getByTestId('llm-prompt').selectOption('vague');
  await expect(page.getByTestId('llm-prob-1')).toHaveText('20.4%');
  await expect(page.getByTestId('llm-spread')).toHaveText('6 candidates cover 90% of the probability');
});

function runAttention() {
  const sentence = document.getElementById("sentenceInput").value;
  const dModel = parseInt(document.getElementById("dModel").value);
  const nHeads = parseInt(document.getElementById("nHeads").value);
  const seed = parseInt(document.getElementById("seed").value);

  const tokens = tokenize(sentence);
  if (tokens.length < 2) {
    alert("Please enter at least 2 words.");
    return;
  }

  const dK = Math.max(1, Math.floor(dModel / nHeads));

  // 1. Tokenize
  renderTokens(tokens);

  // 2. Embed
  const E = embed(tokens, dModel);
  document.getElementById("embedMatrix").innerHTML = renderMatrix(
    E, tokens,
    Array.from({length: dModel}, (_, i) => `d${i}`)
  );

  // Multi-head setup
  const heads = [];
  for (let h = 0; h < nHeads; h++) {
    const rng = mulberry32(seed + h * 1000);
    const WQ = makeMatrix(dModel, dK, rng);
    const WK = makeMatrix(dModel, dK, rng);
    const WV = makeMatrix(dModel, dK, rng);
    const Q = matmul(E, WQ);
    const K = matmul(E, WK);
    const V = matmul(E, WV);
    const scoresRaw = matmul(Q, transpose(K));
    const scoresScaled = scale(scoresRaw, 1 / Math.sqrt(dK));
    const weights = softmaxRows(scoresScaled);
    const out = matmul(weights, V);
    heads.push({ WQ, WK, WV, Q, K, V, scoresRaw, scoresScaled, weights, out });
  }

  // Display HEAD 0 as the primary example through steps 3-7
  const h0 = heads[0];
  const kCols = Array.from({length: dK}, (_, i) => `d${i}`);

  // 3. Q, K, V
  document.getElementById("qShape").textContent = `(${tokens.length} × ${dK})`;
  document.getElementById("kShape").textContent = `(${tokens.length} × ${dK})`;
  document.getElementById("vShape").textContent = `(${tokens.length} × ${dK})`;
  document.getElementById("qMatrix").innerHTML = renderMatrix(h0.Q, tokens, kCols, {class: "q"});
  document.getElementById("kMatrix").innerHTML = renderMatrix(h0.K, tokens, kCols, {class: "k"});
  document.getElementById("vMatrix").innerHTML = renderMatrix(h0.V, tokens, kCols, {class: "v"});

  // 4. Raw scores
  document.getElementById("scoreMatrix").innerHTML = renderMatrix(h0.scoresRaw, tokens, tokens);

  // 5. Scaled
  document.getElementById("scaleFactor").textContent = `√${dK} ≈ ${Math.sqrt(dK).toFixed(3)}`;
  document.getElementById("scaledMatrix").innerHTML = renderMatrix(h0.scoresScaled, tokens, tokens);

  // 6. Weights heatmap
  document.getElementById("weightsMatrix").innerHTML = renderHeatmap(h0.weights, tokens, tokens);
  document.getElementById("readToken").textContent = tokens[Math.min(1, tokens.length - 1)];
  document.getElementById("readToken2").textContent = tokens[Math.min(1, tokens.length - 1)];

  // 7. Output + focus calc
  document.getElementById("outputMatrix").innerHTML = renderMatrix(h0.out, tokens, kCols, {class: "v"});
  populateFocus(tokens);
  renderFocusCalc(h0, tokens, 0);

  // 8. Multi-head
  renderMultiHead(heads, tokens, dK, kCols);

  // Token callout
  document.getElementById("tokenCallout").innerHTML =
    `Got <strong>${tokens.length}</strong> tokens. Vocabulary size in real models: 30,000 – 200,000. We're using a deterministic per-string embedding so identical words map to identical vectors.`;

  STATE = { tokens, E, heads, dModel, dK, nHeads, kCols };
}

/* ─────────────────────────────────────────────────────────────
   Deterministic PRNG (mulberry32)
─────────────────────────────────────────────────────────────── */
function mulberry32(seed) {
  return function() {
    seed |= 0; seed = seed + 0x6D2B79F5 | 0;
    let t = seed;
    t = Math.imul(t ^ t >>> 15, t | 1);
    t ^= t + Math.imul(t ^ t >>> 7, t | 61);
    return ((t ^ t >>> 14) >>> 0) / 4294967296;
  };
}

function gaussian(rng) {
  // Box-Muller
  const u = 1 - rng(); const v = rng();
  return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * v) * 0.5;
}

// Stable string-to-seed
function strHash(s) {
  let h = 0;
  for (let i = 0; i < s.length; i++) {
    h = ((h << 5) - h) + s.charCodeAt(i);
    h |= 0;
  }
  return Math.abs(h) + 1;
}

/* ─────────────────────────────────────────────────────────────
   Matrix ops
─────────────────────────────────────────────────────────────── */
function makeMatrix(rows, cols, rng) {
  const m = [];
  for (let i = 0; i < rows; i++) {
    const row = [];
    for (let j = 0; j < cols; j++) row.push(gaussian(rng));
    m.push(row);
  }
  return m;
}

function matmul(a, b) {
  const n = a.length, k = a[0].length, m = b[0].length;
  const out = Array.from({length: n}, () => new Array(m).fill(0));
  for (let i = 0; i < n; i++) {
    for (let j = 0; j < m; j++) {
      let s = 0;
      for (let p = 0; p < k; p++) s += a[i][p] * b[p][j];
      out[i][j] = s;
    }
  }
  return out;
}

function transpose(m) {
  const rows = m.length, cols = m[0].length;
  const out = Array.from({length: cols}, () => new Array(rows));
  for (let i = 0; i < rows; i++) for (let j = 0; j < cols; j++) out[j][i] = m[i][j];
  return out;
}

function scale(m, s) {
  return m.map(row => row.map(v => v * s));
}

function softmaxRows(m) {
  return m.map(row => {
    const max = Math.max(...row);
    const exp = row.map(v => Math.exp(v - max));
    const sum = exp.reduce((a, b) => a + b, 0);
    return exp.map(v => v / sum);
  });
}

function concatCols(matrices) {
  const n = matrices[0].length;
  const out = [];
  for (let i = 0; i < n; i++) {
    out.push(matrices.flatMap(m => m[i]));
  }
  return out;
}

const slides = document.querySelectorAll(".slide");
const progress = document.getElementById("progress");
const counter = document.getElementById("counter");
let idx = 0;

function show(n) {
  if (n < 0 || n >= slides.length) return;
  slides[idx].classList.remove("active");
  idx = n;
  slides[idx].classList.add("active");
  counter.textContent = (idx + 1) + " / " + slides.length;
  progress.style.width = ((idx + 1) / slides.length * 100) + "%";
  if (slides[idx].querySelector("#heatmapTarget")) renderHeatmap();
}

document.getElementById("prevBtn").addEventListener("click", () => show(idx - 1));
document.getElementById("nextBtn").addEventListener("click", () => show(idx + 1));

document.addEventListener("keydown", (e) => {
  if (e.key === "ArrowRight" || e.key === " ") { show(idx + 1); e.preventDefault(); }
  if (e.key === "ArrowLeft") show(idx - 1);
  if (e.key.toLowerCase() === "f") {
    if (!document.fullscreenElement) document.documentElement.requestFullscreen();
    else document.exitFullscreen();
  }
});

// Attention heatmap (simulated weights for "the cat will sit on the mat")
const tokens = ["The", "cat", "will", "sit", "on", "the", "mat"];
const W = [
  [0.40, 0.35, 0.05, 0.05, 0.03, 0.04, 0.08],
  [0.08, 0.20, 0.05, 0.42, 0.05, 0.05, 0.15],
  [0.05, 0.20, 0.18, 0.45, 0.04, 0.03, 0.05],
  [0.04, 0.45, 0.06, 0.10, 0.05, 0.05, 0.25],
  [0.05, 0.18, 0.04, 0.30, 0.08, 0.10, 0.25],
  [0.10, 0.12, 0.03, 0.10, 0.05, 0.20, 0.40],
  [0.05, 0.30, 0.03, 0.20, 0.10, 0.05, 0.27],
];

function renderHeatmap() {
  const t = document.getElementById("heatmapTarget");
  if (!t || t.dataset.rendered) return;
  for (let r = 0; r < tokens.length; r++) {
    const rowLabel = document.createElement("div");
    rowLabel.className = "label";
    rowLabel.style.textAlign = "right";
    rowLabel.textContent = tokens[r];
    t.appendChild(rowLabel);
    for (let c = 0; c < tokens.length; c++) {
      const cell = document.createElement("div");
      cell.className = "cell";
      const v = W[r][c];
      const alpha = Math.min(1, 0.15 + v * 1.8);
      cell.style.background = `rgba(201, 100, 66, ${alpha})`;
      cell.style.color = v > 0.25 ? "white" : "#1f1e1d";
      cell.textContent = v.toFixed(2);
      t.appendChild(cell);
    }
  }
  t.dataset.rendered = "1";
}
renderHeatmap();
show(0);
</script>

// The demo's math, in the same style as the repo page (this page's code, not the repo's)
function probabilities(logits, T) {
  if (T === 0) {                                   // greedy: all weight on the top logit
    const top = Math.max(...logits);
    return logits.map((l) => (l === top ? 1 : 0));
  }
  const scaled = logits.map((l) => l / T);         // temperature divides the logits
  const max = Math.max(...scaled);
  const exp = scaled.map((v) => Math.exp(v - max)); // same trick as softmaxRows()
  const sum = exp.reduce((a, b) => a + b, 0);
  return exp.map((v) => v / sum);
}

// one seeded stream for all 20 runs: mulberry32(seed) from attention_interactive.html
function sample(probs, rng) {
  const u = rng();
  let acc = 0;
  for (let i = 0; i < probs.length; i++) { acc += probs[i]; if (u < acc) return i; }
  return probs.length - 1;
}