Inference Playground

d3 = require("d3@7")

// Import our modules dynamically (using absolute paths from root)
Tokenizer = (await import('/playground/tokenizer.js')).Tokenizer
GPTModel = (await import('/playground/model.js')).GPTModel
tokenViz = (await import('/playground/viz/tokenViz.js')).tokenViz
embeddingViz = (await import('/playground/viz/embeddingViz.js')).embeddingViz
attentionViz = (await import('/playground/viz/attentionViz.js')).attentionViz
logitsViz = (await import('/playground/viz/logitsViz.js')).logitsViz

// Load data files
tokenizerData = FileAttachment("playground/tokenizer.json").json()
weightsData = FileAttachment("playground/weights.json").json()

// Initialize tokenizer and model
tokenizer = new Tokenizer(tokenizerData)
model = new GPTModel(weightsData)

// Model config for display
modelConfig = weightsData.config

isDarkMode = {
  const check = () => document.body.classList.contains('quarto-dark');
  let current = check();

  const observer = new MutationObserver(() => {
    const newValue = check();
    if (newValue !== current) {
      current = newValue;
    }
  });

  observer.observe(document.body, {
    attributes: true,
    attributeFilter: ['class']
  });

  return current;
}

getCSSVar = function(name, fallback = null) {
  if (typeof document === 'undefined') return fallback;
  const value = getComputedStyle(document.documentElement).getPropertyValue(name).trim();
  return value || fallback;
}

theme = {
  const lightFallbacks = {
    surfacePrimary: '#ffffff',
    surfaceSecondary: '#f8fafc',
    surfaceTertiary: '#f1f5f9',
    textPrimary: '#1e293b',
    textSecondary: '#475569',
    textMuted: '#64748b',
    borderLight: '#e2e8f0',
    borderMedium: '#cbd5e1',
    accent: '#0ea5e9',
    highlight: '#f97316'
  };

  const darkFallbacks = {
    surfacePrimary: '#0e1010',
    surfaceSecondary: '#141616',
    surfaceTertiary: '#191a19',
    textPrimary: '#ffffff',
    textSecondary: '#ffffef',
    textMuted: '#f0ddbf',
    borderLight: '#3c3e3b',
    borderMedium: '#474946',
    accent: '#38bdf8',
    highlight: '#fb923c'
  };

  const f = isDarkMode ? darkFallbacks : lightFallbacks;

  return {
    surfacePrimary: getCSSVar('--surface-primary', f.surfacePrimary),
    surfaceSecondary: getCSSVar('--surface-secondary', f.surfaceSecondary),
    surfaceTertiary: getCSSVar('--surface-tertiary', f.surfaceTertiary),
    textPrimary: getCSSVar('--text-primary', f.textPrimary),
    textSecondary: getCSSVar('--text-secondary', f.textSecondary),
    textMuted: getCSSVar('--text-muted', f.textMuted),
    borderLight: getCSSVar('--border-light', f.borderLight),
    borderMedium: getCSSVar('--border-medium', f.borderMedium),
    accent: f.accent,
    highlight: f.highlight,
    isDark: isDarkMode
  };
}

header = {
  const t = theme;
  return html`
    <div class="playground-header" style="
      display: flex;
      justify-content: space-between;
      align-items: center;
      padding: 24px 0;
      border-bottom: 1px solid ${t.borderLight};
      margin-bottom: 32px;
    ">
      <div>
        <h1 style="
          margin: 0 0 8px 0;
          font-size: 28px;
          font-weight: 700;
          color: ${t.textPrimary};
          font-family: 'Crimson Pro', 'Georgia', serif;
          letter-spacing: -0.02em;
        ">Inference Playground</h1>
        <p style="
          margin: 0;
          font-size: 14px;
          color: ${t.textMuted};
          font-family: 'IBM Plex Sans', system-ui, sans-serif;
        ">Explore how a transformer processes text, step by step</p>
      </div>
      <div class="model-badge" style="
        display: flex;
        gap: 12px;
        padding: 10px 16px;
        background: ${t.surfaceTertiary};
        border-radius: 8px;
        border: 1px solid ${t.borderLight};
        font-family: 'JetBrains Mono', 'Fira Code', monospace;
        font-size: 11px;
      ">
        <span style="color: ${t.textMuted};">Model:</span>
        <span style="color: ${t.textPrimary}; font-weight: 500;">MiniGPT</span>
        <span style="color: ${t.borderMedium};">|</span>
        <span style="color: ${t.textMuted};" title="Transformer layers">${modelConfig.num_layers} Layers</span>
        <span style="color: ${t.textMuted};" title="Attention heads per layer">${modelConfig.num_heads} Heads</span>
        <span style="color: ${t.textMuted};" title="Embedding dimension">${modelConfig.embed_dim} Dim</span>
      </div>
    </div>
  `;
}

howToUse = {
  const t = theme;
  return html`
    <details style="
      margin-bottom: 32px;
      background: ${t.surfaceSecondary};
      border-radius: 8px;
      border: 1px solid ${t.borderLight};
    ">
      <summary style="
        padding: 16px 20px;
        cursor: pointer;
        font-size: 14px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
        list-style: none;
        display: flex;
        align-items: center;
        gap: 8px;
      ">
        <span style="
          display: inline-block;
          width: 20px;
          height: 20px;
          line-height: 20px;
          text-align: center;
          background: ${t.accent};
          color: white;
          border-radius: 50%;
          font-size: 11px;
        ">?</span>
        How to Use This Playground
      </summary>
      <div style="
        padding: 0 20px 20px 20px;
        font-size: 13px;
        color: ${t.textSecondary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
        line-height: 1.7;
      ">
        <p style="margin: 0 0 12px 0;">This playground shows the four stages of transformer inference:</p>
        <ol style="margin: 0 0 16px 0; padding-left: 20px;">
          <li><strong>Tokenization</strong> — How text becomes numbers the model processes</li>
          <li><strong>Embedding</strong> — How tokens become vectors in high-dimensional space</li>
          <li><strong>Attention</strong> — How tokens share information with each other</li>
          <li><strong>Prediction</strong> — How the model decides what comes next</li>
        </ol>
        <p style="margin: 0;"><strong>Try this:</strong> Change the input text and watch each stage respond. Click tokens to select them. Adjust temperature to see how it affects predictions.</p>
      </div>
    </details>
  `;
}

inputSection = {
  const t = theme;
  return html`
    <div class="input-section" style="
      margin-bottom: 40px;
    ">
      <div style="
        display: flex;
        align-items: center;
        gap: 8px;
        margin-bottom: 12px;
      ">
        <span style="
          display: inline-block;
          width: 24px;
          height: 24px;
          line-height: 24px;
          text-align: center;
          background: ${t.accent};
          color: white;
          border-radius: 6px;
          font-size: 12px;
          font-weight: 600;
          font-family: 'JetBrains Mono', monospace;
        ">0</span>
        <span style="
          font-size: 15px;
          font-weight: 600;
          color: ${t.textPrimary};
          font-family: 'IBM Plex Sans', system-ui, sans-serif;
        ">Input Text</span>
      </div>
    </div>
  `;
}

examplePrompts = [
  { label: "Function definition", value: "def foo():", hint: "Watch how the model predicts argument names or body start" },
  { label: "Class definition", value: "class Model:", hint: "Notice inheritance and method predictions" },
  { label: "Import statement", value: "import torch", hint: "See common Python imports the model knows" },
  { label: "Variable assignment", value: "x = 42", hint: "Simple assignment—what operations follow?" },
  { label: "Print statement", value: "print(\"hello\")", hint: "String completion and formatting predictions" },
  { label: "For loop", value: "for i in range", hint: "Loop patterns are highly predictable" },
  { label: "If statement", value: "if x > 0:", hint: "Conditional logic completion" },
  { label: "Return statement", value: "return result", hint: "End-of-function patterns" }
]

viewof selectedExample = Inputs.select(examplePrompts, {
  label: "Quick examples",
  format: x => x.label,
  value: examplePrompts[0]
})

// Display hint for selected example
exampleHint = {
  const t = theme;
  return html`
    <p style="
      margin: 8px 0 0 0;
      font-size: 12px;
      color: ${t.textMuted};
      font-style: italic;
      font-family: 'IBM Plex Sans', system-ui, sans-serif;
    ">${selectedExample.hint}</p>
  `;
}

viewof inputText = Inputs.textarea({
  label: "Enter text to tokenize and process:",
  placeholder: "def hello():",
  value: selectedExample.value,
  rows: 2,
  width: "100%"
})

tokenIds = tokenizer.encode(inputText)

// Get token strings for display
tokens = tokenIds.map(id => tokenizer.idToToken(id))

// Run model forward pass and capture trace
trace = model.forward(tokenIds)

// Extract embeddings as 2D array for visualization
embeddings = {
  const emb = trace.embeddings;
  const [seqLen, embedDim] = emb.shape;
  const result = [];
  for (let i = 0; i < seqLen; i++) {
    const row = [];
    for (let j = 0; j < embedDim; j++) {
      row.push(emb.data[i * embedDim + j]);
    }
    result.push(row);
  }
  return result;
}

// Number of layers and heads from config
numLayers = modelConfig.num_layers
numHeads = modelConfig.num_heads

stage1Header = {
  const t = theme;
  return html`
    <div style="
      display: flex;
      align-items: center;
      gap: 8px;
      margin: 40px 0 16px 0;
    ">
      <span style="
        display: inline-block;
        width: 24px;
        height: 24px;
        line-height: 24px;
        text-align: center;
        background: ${t.accent};
        color: white;
        border-radius: 6px;
        font-size: 12px;
        font-weight: 600;
        font-family: 'JetBrains Mono', monospace;
      ">1</span>
      <span style="
        font-size: 15px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">Tokenization</span>
      <span style="
        font-size: 13px;
        color: ${t.textMuted};
        margin-left: 8px;
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">${tokenIds.length} tokens</span>
    </div>
    <p style="
      margin: 0 0 16px 32px;
      font-size: 13px;
      color: ${t.textSecondary};
      font-family: 'IBM Plex Sans', system-ui, sans-serif;
      max-width: 600px;
    ">BPE (Byte Pair Encoding) splits text into subword tokens. Click any token to select it.</p>
  `;
}

// Token selection state
viewof selectedTokenIdx = {
  const container = html`<div class="token-viz-container" style="width: 100%;"></div>`;

  let currentSelection = tokenIds.length - 1; // Default to last token

  const render = () => {
    tokenViz(container, tokens, tokenIds, {
      selectedIdx: currentSelection,
      onSelect: (idx) => {
        currentSelection = idx;
        container.value = idx;
        container.dispatchEvent(new CustomEvent("input"));
        render();
      }
    });
  };

  // Initial render
  setTimeout(render, 0);

  container.value = currentSelection;
  return container;
}

stage2Header = {
  const t = theme;
  return html`
    <div style="
      display: flex;
      align-items: center;
      gap: 8px;
      margin: 48px 0 16px 0;
    ">
      <span style="
        display: inline-block;
        width: 24px;
        height: 24px;
        line-height: 24px;
        text-align: center;
        background: ${t.accent};
        color: white;
        border-radius: 6px;
        font-size: 12px;
        font-weight: 600;
        font-family: 'JetBrains Mono', monospace;
      ">2</span>
      <span style="
        font-size: 15px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">Token Embeddings</span>
      <span style="
        font-size: 13px;
        color: ${t.textMuted};
        margin-left: 8px;
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">${modelConfig.embed_dim}-dimensional vectors</span>
    </div>
    <p style="
      margin: 0 0 16px 32px;
      font-size: 13px;
      color: ${t.textSecondary};
      font-family: 'IBM Plex Sans', system-ui, sans-serif;
      max-width: 600px;
    ">Each token becomes a dense vector. PCA (Principal Component Analysis) projects these vectors to 2D for visualization.</p>
  `;
}

embeddingVizContainer = {
  const container = html`<div class="embedding-viz-container" style="width: 100%; max-width: 500px;"></div>`;

  setTimeout(() => {
    embeddingViz(container, embeddings, tokens, {
      selectedIdx: selectedTokenIdx
    });
  }, 0);

  return container;
}

stage3Header = {
  const t = theme;
  return html`
    <div style="
      display: flex;
      align-items: center;
      gap: 8px;
      margin: 48px 0 16px 0;
    ">
      <span style="
        display: inline-block;
        width: 24px;
        height: 24px;
        line-height: 24px;
        text-align: center;
        background: ${t.accent};
        color: white;
        border-radius: 6px;
        font-size: 12px;
        font-weight: 600;
        font-family: 'JetBrains Mono', monospace;
      ">3</span>
      <span style="
        font-size: 15px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">Attention Patterns</span>
      <span style="
        font-size: 13px;
        color: ${t.textMuted};
        margin-left: 8px;
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">${numLayers} layers, ${numHeads} heads per layer</span>
    </div>
    <p style="
      margin: 0 0 16px 32px;
      font-size: 13px;
      color: ${t.textSecondary};
      font-family: 'IBM Plex Sans', system-ui, sans-serif;
      max-width: 600px;
    ">Each token gathers information from previous tokens via self-attention. The causal mask ensures tokens see only what came before—never ahead.</p>
  `;
}

// Attention interpretation hint
attentionHint = {
  const t = theme;
  return html`
    <div style="
      margin: 0 0 16px 32px;
      padding: 12px 16px;
      background: ${t.surfaceTertiary};
      border-radius: 6px;
      font-size: 12px;
      color: ${t.textSecondary};
      font-family: 'IBM Plex Sans', system-ui, sans-serif;
      max-width: 600px;
    ">
      <strong style="color: ${t.textPrimary};">Reading the heatmap:</strong>
      Each row shows where a token "looks" to gather context.
      Bright cells = strong attention. The diagonal often glows because tokens attend to themselves.
      Striped cells are masked (future tokens the model cannot see).
    </div>
  `;
}

viewof selectedLayer = Inputs.range([0, numLayers - 1], {
  value: 0,
  step: 1,
  label: "Layer"
})

viewof selectedHead = Inputs.range([0, numHeads - 1], {
  value: 0,
  step: 1,
  label: "Head"
})

attentionVizContainer = {
  const container = html`<div class="attention-viz-container" style="width: 100%; max-width: 500px;"></div>`;

  // Get attention weights for selected layer
  const layerTrace = trace.layers[selectedLayer];
  const attentionWeights = layerTrace.attention_weights;

  setTimeout(() => {
    attentionViz(container, attentionWeights, tokens, {
      head: selectedHead
    });
  }, 0);

  return container;
}

stage4Header = {
  const t = theme;
  return html`
    <div style="
      display: flex;
      align-items: center;
      gap: 8px;
      margin: 48px 0 16px 0;
    ">
      <span style="
        display: inline-block;
        width: 24px;
        height: 24px;
        line-height: 24px;
        text-align: center;
        background: ${t.accent};
        color: white;
        border-radius: 6px;
        font-size: 12px;
        font-weight: 600;
        font-family: 'JetBrains Mono', monospace;
      ">4</span>
      <span style="
        font-size: 15px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">Next Token Prediction</span>
    </div>
    <p style="
      margin: 0 0 16px 32px;
      font-size: 13px;
      color: ${t.textSecondary};
      font-family: 'IBM Plex Sans', system-ui, sans-serif;
      max-width: 600px;
    ">The model predicts probabilities for each vocabulary token. Lower temperature sharpens predictions; higher temperature spreads them.</p>
  `;
}

viewof temperature = Inputs.range([0.1, 2.0], {
  value: 1.0,
  step: 0.1,
  label: "Temperature",
  description: "Controls randomness: low (0.1) = focused predictions, high (2.0) = diverse predictions"
})

viewof topK = Inputs.range([1, 20], {
  value: 10,
  step: 1,
  label: "Top-k tokens",
  description: "Number of most likely tokens to display"
})

logitsVizContainer = {
  const container = html`<div class="logits-viz-container" style="width: 100%; max-width: 600px;"></div>`;

  setTimeout(() => {
    logitsViz(container, trace.logits, tokenizer, {
      temperature: temperature,
      topK: topK
    });
  }, 0);

  return container;
}

footer = {
  const t = theme;
  return html`
    <div style="
      margin-top: 64px;
      padding: 24px;
      background: ${t.surfaceSecondary};
      border-radius: 8px;
      border: 1px solid ${t.borderLight};
    ">
      <h3 style="
        margin: 0 0 12px 0;
        font-size: 14px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">About this playground</h3>
      <p style="
        margin: 0 0 16px 0;
        font-size: 13px;
        color: ${t.textSecondary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
        line-height: 1.6;
      ">
        This playground runs a small GPT model entirely in your browser using JavaScript.
        The model has ${modelConfig.num_layers} transformer layers, ${modelConfig.num_heads} attention heads,
        and ${modelConfig.embed_dim}-dimensional embeddings. It was trained on Python code.
        The architecture matches what you'll build in <a href="modules/m06_transformer/lesson.html" style="color: ${t.accent};">Module 6</a>.
      </p>
      <h4 style="
        margin: 0 0 8px 0;
        font-size: 13px;
        font-weight: 600;
        color: ${t.textPrimary};
        font-family: 'IBM Plex Sans', system-ui, sans-serif;
      ">Build it yourself</h4>
      <div style="
        display: flex;
        flex-wrap: wrap;
        gap: 8px;
      ">
        <a href="modules/m03_tokenization/lesson.html" style="
          padding: 6px 12px;
          background: ${t.surfaceTertiary};
          border-radius: 4px;
          font-size: 12px;
          color: ${t.textPrimary};
          text-decoration: none;
          font-family: 'IBM Plex Sans', system-ui, sans-serif;
        ">Module 3: Tokenization</a>
        <a href="modules/m04_embeddings/lesson.html" style="
          padding: 6px 12px;
          background: ${t.surfaceTertiary};
          border-radius: 4px;
          font-size: 12px;
          color: ${t.textPrimary};
          text-decoration: none;
          font-family: 'IBM Plex Sans', system-ui, sans-serif;
        ">Module 4: Embeddings</a>
        <a href="modules/m05_attention/lesson.html" style="
          padding: 6px 12px;
          background: ${t.surfaceTertiary};
          border-radius: 4px;
          font-size: 12px;
          color: ${t.textPrimary};
          text-decoration: none;
          font-family: 'IBM Plex Sans', system-ui, sans-serif;
        ">Module 5: Attention</a>
        <a href="modules/m06_transformer/lesson.html" style="
          padding: 6px 12px;
          background: ${t.surfaceTertiary};
          border-radius: 4px;
          font-size: 12px;
          color: ${t.textPrimary};
          text-decoration: none;
          font-family: 'IBM Plex Sans', system-ui, sans-serif;
        ">Module 6: Transformer</a>
      </div>
    </div>
  `;
}