First push

.gitignore

@@ -0,0 +1,2 @@
+.venv
+__pycache__

app.py

@@ -0,0 +1,625 @@
+"""
+Logit Lens Explorer - Gradio Application.
+
+Interactive text generation tool that surfaces the logit lens for each
+generated token. Users input a prompt, the model generates text, and
+clicking any token reveals what the model was predicting at each
+intermediate layer.
+
+Part of E02: Logit Lens Explorer.
+"""
+
+import html as html_lib
+import json
+from typing import Generator
+
+import gradio as gr
+
+try:
+    import spaces
+    SPACES_AVAILABLE = True
+except ImportError:
+    SPACES_AVAILABLE = False
+
+from model import generate_with_logit_lens, load_model, TokenData
+
+
+def gpu_decorator(duration: int = 120):
+    """Return @spaces.GPU decorator if available, otherwise a no-op."""
+    if SPACES_AVAILABLE:
+        return spaces.GPU(duration=duration)
+    return lambda fn: fn
+
+
+def build_token_html(tokens: list[TokenData]) -> str:
+    """Build HTML output from accumulated tokens as plain clickable spans.
+
+    Each token span carries all data needed for client-side logit lens
+    rendering: the token text, probability, and per-layer predictions
+    as JSON data attributes.
+
+    Args:
+        tokens: List of TokenData objects.
+
+    Returns:
+        HTML string with clickable token spans.
+    """
+    font_family = "'Cascadia Code', 'Fira Code', Consolas, monospace"
+    style_tag = "<style>.token-span:hover { text-decoration: underline !important; }</style>"
+
+    if not tokens:
+        return (
+            f'{style_tag}<div class="token-container" '
+            f'style="font-family: {font_family}; line-height: 1.8; padding: 10px;"></div>'
+        )
+
+    spans = []
+    for i, token_data in enumerate(tokens):
+        token_text = html_lib.escape(token_data.token)
+        if "\n" in token_text:
+            token_text = token_text.replace("\n", "<br>")
+            spans.append(token_text)
+        else:
+            # Serialize layer predictions as JSON for client-side rendering
+            layers_json = html_lib.escape(json.dumps([
+                {
+                    "layer_index": lp.layer_index,
+                    "top_tokens": lp.top_tokens,
+                }
+                for lp in token_data.layer_predictions
+            ]))
+
+            span = (
+                f'<span class="token-span" data-token-index="{i}"'
+                f' data-token="{html_lib.escape(token_data.token)}"'
+                f' data-prob="{token_data.probability}"'
+                f' data-layers="{layers_json}"'
+                f' style="cursor: pointer;">{token_text}</span>'
+            )
+            spans.append(span)
+
+    html_content = "".join(spans)
+    return (
+        f'{style_tag}<div class="token-container" style="font-family: {font_family};'
+        f' line-height: 1.8; padding: 10px; white-space: pre-wrap;">{html_content}</div>'
+    )
+
+
+@gpu_decorator(duration=120)
+def run_inference(prompt: str) -> list[TokenData]:
+    """Run full text generation on GPU and return all tokens.
+
+    On HuggingFace Spaces with ZeroGPU, this function is decorated with
+    @spaces.GPU to allocate GPU resources for the duration of inference.
+
+    Args:
+        prompt: User prompt text.
+
+    Returns:
+        List of TokenData with token strings, IDs, probabilities,
+        and per-layer logit lens predictions.
+    """
+    return list(generate_with_logit_lens(prompt))
+
+
+def generate_streaming(prompt: str) -> Generator[str, None, None]:
+    """Stream token generation with progressive HTML output.
+
+    Runs full inference first (GPU-bound), then streams HTML rendering
+    from pre-computed tokens (no GPU needed). This architecture is
+    required for HuggingFace ZeroGPU compatibility.
+
+    Args:
+        prompt: User prompt text.
+
+    Yields:
+        HTML string with accumulated tokens.
+    """
+    if not prompt or not prompt.strip():
+        yield '<div style="color: #666; padding: 10px;">Please enter a prompt.</div>'
+        return
+
+    # Show loading indicator during GPU inference
+    loading = """<div style="color: #60a5fa; padding: 10px; display: flex; align-items: center; gap: 10px;">
+        <div style="width: 20px; height: 20px; border: 2px solid #60a5fa;
+             border-top-color: transparent; border-radius: 50%;
+             animation: spin 1s linear infinite;"></div>
+        <style>@keyframes spin { to { transform: rotate(360deg); } }</style>
+        Generating...
+    </div>"""
+    yield loading
+
+    # Full inference (GPU allocated here on ZeroGPU)
+    tokens = run_inference(prompt)
+
+    if not tokens:
+        yield '<div style="color: #666; padding: 10px;">No tokens generated.</div>'
+        return
+
+    # Stream HTML rendering (no GPU needed)
+    accumulated: list[TokenData] = []
+    for token_data in tokens:
+        accumulated.append(token_data)
+        yield build_token_html(accumulated)
+
+
+# JavaScript for token click handling -- reads layer data from span attributes
+# and renders the logit lens panel entirely client-side (no server round-trip).
+# Matches the pattern from the OCR app's alternatives panel.
+TOKEN_CLICK_JS = """
+(function() {
+    console.log('[logit-lens] Click handler installed');
+
+    var CARD_TOP_K = 5;   // Show top 5 in each layer card
+    var CHART_TOP_N = 20; // Track top 20 most recurring tokens in chart
+
+    // 20 distinct colors for chart lines
+    var LINE_COLORS = [
+        '#60a5fa','#f87171','#34d399','#fbbf24','#a78bfa',
+        '#fb923c','#2dd4bf','#f472b6','#818cf8','#4ade80',
+        '#e879f9','#38bdf8','#facc15','#fb7185','#a3e635',
+        '#c084fc','#22d3ee','#fdba74','#86efac','#fca5a5'
+    ];
+
+    function escapeHtml(text) {
+        var div = document.createElement('div');
+        div.textContent = text;
+        return div.innerHTML;
+    }
+
+    function renderLayerCard(layer, finalToken, nLayers, layerIdx) {
+        var lastLayer = nLayers - 1;
+        var label;
+        if (layerIdx === 0) {
+            label = 'Layer 0 (embed)';
+        } else if (layerIdx === lastLayer) {
+            label = 'Layer ' + layerIdx + ' (final)';
+        } else {
+            label = 'Layer ' + layerIdx;
+        }
+
+        var tokenCells = '';
+        var displayCount = Math.min(layer.top_tokens.length, CARD_TOP_K);
+        for (var i = 0; i < displayCount; i++) {
+            var entry = layer.top_tokens[i];
+            var tok = escapeHtml(entry.token);
+            var pct = (entry.probability * 100);
+            var barWidth = Math.max(pct, 0.5);
+            var isMatch = entry.token === finalToken;
+            var tokColor = isMatch ? '#60a5fa' : '#e5e7eb';
+            var barColor = isMatch ? '#60a5fa' : '#4b5563';
+            var fontWeight = isMatch ? '700' : '400';
+
+            tokenCells +=
+                '<div style="display:flex;align-items:center;gap:6px;margin:2px 0;">' +
+                '<span style="width:80px;overflow:hidden;text-overflow:ellipsis;' +
+                'white-space:nowrap;font-family:monospace;font-size:12px;' +
+                'color:' + tokColor + ';font-weight:' + fontWeight + ';">' + tok + '</span>' +
+                '<span style="width:44px;text-align:right;color:#9ca3af;' +
+                'font-size:11px;flex-shrink:0;">' + pct.toFixed(1) + '%</span>' +
+                '<div style="flex:1;height:8px;background:#1f2937;' +
+                'border-radius:4px;overflow:hidden;min-width:30px;">' +
+                '<div style="width:' + barWidth + '%;height:100%;' +
+                'background:' + barColor + ';border-radius:4px;"></div>' +
+                '</div></div>';
+        }
+
+        var cardBg = (layerIdx % 2 === 0) ? '#111827' : '#0d1117';
+        return '<div style="background:' + cardBg + ';border-radius:6px;padding:8px 10px;">' +
+            '<div style="color:#9ca3af;font-size:11px;font-family:monospace;' +
+            'margin-bottom:4px;font-weight:600;">' + label + '</div>' +
+            tokenCells +
+            '</div>';
+    }
+
+    function renderLineChart(layersData, finalToken, nLayers) {
+        // Collect frequency counts: how many layers each token appears in
+        var tokenFreq = {};  // token -> count of layers it appears in
+        var tokenProbs = {}; // token -> array of {layer, prob}
+        for (var li = 0; li < nLayers; li++) {
+            var tops = layersData[li].top_tokens;
+            for (var ti = 0; ti < tops.length; ti++) {
+                var tok = tops[ti].token;
+                var prob = tops[ti].probability;
+                if (!tokenFreq[tok]) {
+                    tokenFreq[tok] = 0;
+                    tokenProbs[tok] = [];
+                }
+                tokenFreq[tok]++;
+                tokenProbs[tok].push({layer: li, prob: prob});
+            }
+        }
+
+        // Sort by frequency descending, take top N
+        var allTokens = Object.keys(tokenFreq);
+        allTokens.sort(function(a, b) { return tokenFreq[b] - tokenFreq[a]; });
+        var chartTokens = allTokens.slice(0, CHART_TOP_N);
+
+        // Ensure the final token is always included
+        if (chartTokens.indexOf(finalToken) === -1 && tokenFreq[finalToken]) {
+            chartTokens.pop();
+            chartTokens.push(finalToken);
+        }
+
+        console.log('[logit-lens] Chart tokens:', chartTokens.length, chartTokens);
+
+        if (chartTokens.length === 0) return null;
+
+        // Build lookup: token -> layer -> probability (0 if absent)
+        var data = {}; // token -> array of length nLayers
+        var maxProb = 0;
+        for (var ci = 0; ci < chartTokens.length; ci++) {
+            var t = chartTokens[ci];
+            data[t] = new Array(nLayers);
+            for (var l = 0; l < nLayers; l++) { data[t][l] = 0; }
+            var entries = tokenProbs[t];
+            for (var ei = 0; ei < entries.length; ei++) {
+                var p = entries[ei].prob * 100;
+                data[t][entries[ei].layer] = p;
+                if (p > maxProb) maxProb = p;
+            }
+        }
+
+        // Build color map for each token
+        var colorMap = {};
+        for (var ci = 0; ci < chartTokens.length; ci++) {
+            var tok = chartTokens[ci];
+            colorMap[tok] = (tok === finalToken) ? '#60a5fa' : LINE_COLORS[ci % LINE_COLORS.length];
+        }
+
+        // SVG dimensions
+        var W = 700, H = 300;
+        var padL = 45, padR = 20, padT = 20, padB = 30;
+        var plotW = W - padL - padR;
+        var plotH = H - padT - padB;
+        var yMax = Math.ceil(maxProb / 10) * 10;
+        if (yMax < 10) yMax = 10;
+
+        function xPos(layer) { return padL + (layer / (nLayers - 1)) * plotW; }
+        function yPos(pct) { return padT + plotH - (pct / yMax) * plotH; }
+
+        // Start SVG
+        var svg = '<svg class="logit-chart-svg" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 ' + W + ' ' + H +
+            '" style="width:100%;max-width:' + W + 'px;height:auto;background:#111827;border-radius:8px;display:block;">';
+
+        // Y-axis gridlines and labels
+        var yTicks = 5;
+        for (var yi = 0; yi <= yTicks; yi++) {
+            var yVal = (yMax / yTicks) * yi;
+            var y = yPos(yVal);
+            svg += '<line x1="' + padL + '" y1="' + y + '" x2="' + (W - padR) + '" y2="' + y +
+                '" stroke="#374151" stroke-width="1"/>';
+            svg += '<text x="' + (padL - 6) + '" y="' + (y + 4) +
+                '" text-anchor="end" fill="#9ca3af" font-size="10" font-family="monospace">' +
+                yVal.toFixed(0) + '%</text>';
+        }
+
+        // X-axis labels (every 4 layers + first and last)
+        for (var xi = 0; xi < nLayers; xi++) {
+            if (xi === 0 || xi === nLayers - 1 || xi % 4 === 0) {
+                var x = xPos(xi);
+                svg += '<text x="' + x + '" y="' + (H - 8) +
+                    '" text-anchor="middle" fill="#9ca3af" font-size="10" font-family="monospace">' +
+                    xi + '</text>';
+            }
+        }
+
+        // Draw lines for each token
+        for (var ci = 0; ci < chartTokens.length; ci++) {
+            var tok = chartTokens[ci];
+            var color = colorMap[tok];
+            var strokeW = (tok === finalToken) ? '2.5' : '1.5';
+            var opacity = (tok === finalToken) ? '1' : '0.7';
+
+            var points = '';
+            for (var l = 0; l < nLayers; l++) {
+                if (l > 0) points += ' ';
+                points += xPos(l).toFixed(1) + ',' + yPos(data[tok][l]).toFixed(1);
+            }
+            svg += '<polyline points="' + points + '" fill="none" stroke="' + color +
+                '" stroke-width="' + strokeW + '" opacity="' + opacity + '"/>';
+        }
+
+        // Invisible overlay rect to capture mouse events across the full plot area
+        svg += '<rect class="logit-chart-overlay" x="' + padL + '" y="' + padT +
+            '" width="' + plotW + '" height="' + plotH + '" fill="transparent" pointer-events="all" style="cursor:crosshair;"/>';
+
+        // Vertical crosshair line (hidden initially)
+        svg += '<line class="logit-chart-crosshair" x1="0" y1="' + padT + '" x2="0" y2="' + (padT + plotH) +
+            '" stroke="#9ca3af" stroke-width="1" stroke-dasharray="4,3" visibility="hidden"/>';
+
+        svg += '</svg>';
+
+        // Tooltip div (hidden, positioned absolutely over the chart)
+        var tooltip = '<div class="logit-chart-tooltip" style="' +
+            'display:none;position:absolute;pointer-events:none;z-index:10;' +
+            'background:#1e293b;border:1px solid #475569;border-radius:6px;padding:8px 10px;' +
+            'font-family:monospace;font-size:11px;color:#e5e7eb;' +
+            'box-shadow:0 4px 12px rgba(0,0,0,0.4);max-width:220px;' +
+            '"></div>';
+
+        // Legend (horizontal wrapping)
+        var legend = '<div style="display:flex;flex-wrap:wrap;gap:8px 14px;margin-top:8px;">';
+        for (var ci = 0; ci < chartTokens.length; ci++) {
+            var tok = chartTokens[ci];
+            var color = colorMap[tok];
+            var weight = (tok === finalToken) ? '700' : '400';
+            legend += '<div style="display:flex;align-items:center;gap:4px;">' +
+                '<div style="width:12px;height:3px;background:' + color + ';border-radius:2px;"></div>' +
+                '<span style="font-family:monospace;font-size:11px;color:' + color +
+                ';font-weight:' + weight + ';">' + escapeHtml(tok) + '</span>' +
+                '</div>';
+        }
+        legend += '</div>';
+
+        // Return HTML + metadata object (avoids DOM attribute serialization issues)
+        var chartMeta = {
+            tokens: chartTokens,
+            data: data,
+            colors: colorMap,
+            nLayers: nLayers,
+            padL: padL,
+            padR: padR,
+            padT: padT,
+            plotW: plotW,
+            plotH: plotH,
+            W: W,
+            finalToken: finalToken
+        };
+
+        var html = '<div class="logit-chart-wrapper" style="position:relative;margin-bottom:16px;">' +
+            '<div style="color:#9ca3af;font-size:11px;font-family:monospace;margin-bottom:6px;font-weight:600;">' +
+            'Probability by Layer (top ' + chartTokens.length + ' recurring tokens)</div>' +
+            svg + tooltip + legend + '</div>';
+
+        return { html: html, meta: chartMeta };
+    }
+
+    function attachChartHover(meta) {
+        var panel = document.getElementById('logit-lens-panel');
+        if (!panel) { console.error('[logit-lens] hover: panel not found'); return; }
+
+        var wrapper = panel.querySelector('.logit-chart-wrapper');
+        if (!wrapper) { console.error('[logit-lens] hover: wrapper not found'); return; }
+
+        var svgEl = wrapper.querySelector('.logit-chart-svg');
+        var crosshair = wrapper.querySelector('.logit-chart-crosshair');
+        var tooltipEl = wrapper.querySelector('.logit-chart-tooltip');
+        if (!svgEl || !crosshair || !tooltipEl) { console.error('[logit-lens] hover: SVG elements not found', !!svgEl, !!crosshair, !!tooltipEl); return; }
+
+        // Sort tokens by probability descending at each layer for tooltip display
+        function getLayerEntries(layerIdx) {
+            var entries = [];
+            for (var i = 0; i < meta.tokens.length; i++) {
+                var tok = meta.tokens[i];
+                var pct = meta.data[tok][layerIdx];
+                if (pct > 0) {
+                    entries.push({token: tok, pct: pct, color: meta.colors[tok]});
+                }
+            }
+            entries.sort(function(a, b) { return b.pct - a.pct; });
+            return entries;
+        }
+
+        function mouseToLayer(e) {
+            var rect = svgEl.getBoundingClientRect();
+            // Map pixel position to SVG viewBox coordinates
+            var scaleX = meta.W / rect.width;
+            var svgX = (e.clientX - rect.left) * scaleX;
+            // Convert SVG X to layer index
+            var layerFrac = (svgX - meta.padL) / meta.plotW;
+            var layer = Math.round(layerFrac * (meta.nLayers - 1));
+            return Math.max(0, Math.min(meta.nLayers - 1, layer));
+        }
+
+        function svgXForLayer(layer) {
+            return meta.padL + (layer / (meta.nLayers - 1)) * meta.plotW;
+        }
+
+        svgEl.addEventListener('mousemove', function(e) {
+            var layer = mouseToLayer(e);
+            var x = svgXForLayer(layer);
+
+            // Update crosshair position
+            crosshair.setAttribute('x1', x);
+            crosshair.setAttribute('x2', x);
+            crosshair.setAttribute('visibility', 'visible');
+
+            // Build tooltip content
+            var entries = getLayerEntries(layer);
+            var label = 'Layer ' + layer;
+            if (layer === 0) label += ' (embed)';
+            else if (layer === meta.nLayers - 1) label += ' (final)';
+
+            var html = '<div style="font-weight:600;margin-bottom:4px;color:#9ca3af;">' + label + '</div>';
+            for (var i = 0; i < entries.length; i++) {
+                var entry = entries[i];
+                var isFinal = (entry.token === meta.finalToken);
+                var w = isFinal ? '700' : '400';
+                html += '<div style="display:flex;align-items:center;gap:5px;margin:1px 0;">' +
+                    '<div style="width:8px;height:8px;border-radius:50%;background:' + entry.color +
+                    ';flex-shrink:0;"></div>' +
+                    '<span style="color:' + entry.color + ';font-weight:' + w + ';overflow:hidden;' +
+                    'text-overflow:ellipsis;white-space:nowrap;max-width:120px;">' +
+                    escapeHtml(entry.token) + '</span>' +
+                    '<span style="color:#9ca3af;margin-left:auto;">' + entry.pct.toFixed(1) + '%</span></div>';
+            }
+            if (entries.length === 0) {
+                html += '<div style="color:#6b7280;font-style:italic;">No tracked tokens at this layer</div>';
+            }
+
+            tooltipEl.innerHTML = html;
+            tooltipEl.style.display = 'block';
+
+            // Position tooltip relative to wrapper
+            var wrapperRect = wrapper.getBoundingClientRect();
+            var svgRect = svgEl.getBoundingClientRect();
+            var pixelX = (x / meta.W) * svgRect.width + svgRect.left - wrapperRect.left;
+            var tooltipW = tooltipEl.offsetWidth;
+
+            // Flip to left side if tooltip would overflow right edge
+            if (pixelX + tooltipW + 12 > wrapperRect.width) {
+                tooltipEl.style.left = (pixelX - tooltipW - 12) + 'px';
+            } else {
+                tooltipEl.style.left = (pixelX + 12) + 'px';
+            }
+            tooltipEl.style.top = (svgRect.top - wrapperRect.top + meta.padT) + 'px';
+        });
+
+        svgEl.addEventListener('mouseleave', function() {
+            crosshair.setAttribute('visibility', 'hidden');
+            tooltipEl.style.display = 'none';
+        });
+
+        console.log('[logit-lens] Chart hover attached');
+    }
+
+    document.addEventListener('click', function(e) {
+        var token = e.target.closest('.token-span[data-token-index]');
+        if (!token) return;
+
+        console.log('[logit-lens] Token clicked:', token.textContent, 'index:', token.dataset.tokenIndex);
+
+        // Highlight selected token, clear previous
+        document.querySelectorAll('.token-span').forEach(function(s) {
+            s.style.background = '';
+        });
+        token.style.background = 'rgba(96, 165, 250, 0.2)';
+
+        // Read data from span attributes
+        var finalToken = token.dataset.token;
+        var prob = parseFloat(token.dataset.prob) || 0;
+        var idx = parseInt(token.dataset.tokenIndex);
+
+        var layersData;
+        try {
+            layersData = JSON.parse(token.dataset.layers);
+        } catch (err) {
+            console.error('[logit-lens] Failed to parse layers data:', err);
+            return;
+        }
+
+        console.log('[logit-lens] Layers:', layersData.length, 'Final token:', JSON.stringify(finalToken));
+
+        var nLayers = layersData.length;
+
+        // Find first layer where final token appears in top-k
+        var firstAppearance = -1;
+        for (var li = 0; li < nLayers; li++) {
+            var tops = layersData[li].top_tokens;
+            for (var ti = 0; ti < tops.length; ti++) {
+                if (tops[ti].token === finalToken) {
+                    firstAppearance = li;
+                    break;
+                }
+            }
+            if (firstAppearance >= 0) break;
+        }
+
+        // Build header
+        var appearanceNote = '';
+        if (firstAppearance >= 0) {
+            appearanceNote = ' &middot; first in top-k at layer ' + firstAppearance;
+        } else if (nLayers > 0) {
+            appearanceNote = ' &middot; <span style="color:#f87171;">never in top-k</span>';
+        }
+
+        var header = '<div style="font-weight:600;margin-bottom:12px;padding-bottom:8px;' +
+            'border-bottom:1px solid #374151;">' +
+            'Selected: "<span style="color:#60a5fa;">' + escapeHtml(finalToken) + '</span>"' +
+            ' (token ' + idx + ', ' + (prob * 100).toFixed(2) + '%)' +
+            appearanceNote + '</div>';
+
+        // Build line chart (returns {html, meta})
+        var chartResult = renderLineChart(layersData, finalToken, nLayers);
+        var chartHtml = chartResult ? chartResult.html : '';
+        var chartMeta = chartResult ? chartResult.meta : null;
+
+        // Build layer cards (reversed: final layer at top, embedding at bottom)
+        var cards = '';
+        for (var i = nLayers - 1; i >= 0; i--) {
+            cards += renderLayerCard(layersData[i], finalToken, nLayers, i);
+        }
+
+        var grid = '<div style="display:grid;grid-template-columns:repeat(auto-fill,minmax(200px,1fr));gap:6px;">' +
+            cards + '</div>';
+
+        // Update panel: header -> chart -> grid
+        var panel = document.getElementById('logit-lens-panel');
+        if (panel) {
+            panel.innerHTML = header + chartHtml + grid;
+            if (chartMeta) attachChartHover(chartMeta);
+            console.log('[logit-lens] Panel updated with chart +', nLayers, 'layers');
+        } else {
+            console.error('[logit-lens] Panel element #logit-lens-panel not found');
+        }
+    });
+})();
+"""
+
+# Initial HTML for the logit lens panel
+LOGIT_LENS_PANEL_INITIAL = """
+<div id="logit-lens-panel" style="
+    padding: 16px;
+    background: #1f2937;
+    border-radius: 8px;
+    color: #e5e7eb;
+    font-family: system-ui, -apple-system, sans-serif;
+    font-size: 14px;
+    min-height: 100px;
+    max-height: 600px;
+    overflow-y: auto;
+">
+    <div style="color: #9ca3af; font-style: italic;">
+        Click on any generated token to see per-layer predictions.
+    </div>
+</div>
+"""
+
+
+# Build Gradio interface
+with gr.Blocks(title="Logit Lens Explorer") as demo:
+    gr.Markdown("# Logit Lens Explorer")
+    gr.Markdown(
+        "Enter a prompt to generate text. Click any token to see per-layer predictions."
+    )
+
+    prompt_input = gr.Textbox(
+        label="Prompt",
+        placeholder="Enter a prompt...",
+        lines=2,
+    )
+    submit_btn = gr.Button("Generate", variant="primary")
+
+    gr.Markdown("### Generated Tokens")
+    gr.Markdown("*Click any token to inspect its per-layer predictions.*")
+    token_display = gr.HTML(
+        value='<div style="color: #666; padding: 10px;">Enter a prompt and click Generate to start.</div>',
+    )
+
+    gr.Markdown("### Logit Lens Panel")
+    logit_lens_panel = gr.HTML(
+        value=LOGIT_LENS_PANEL_INITIAL,
+    )
+
+    # Wire up generation: button click and Enter key in textbox
+    submit_btn.click(
+        fn=generate_streaming,
+        inputs=[prompt_input],
+        outputs=[token_display],
+    )
+    prompt_input.submit(
+        fn=generate_streaming,
+        inputs=[prompt_input],
+        outputs=[token_display],
+    )
+
+
+if __name__ == "__main__":
+    if not SPACES_AVAILABLE:
+        print("Preloading model (local development)...")
+        load_model()
+    else:
+        print("ZeroGPU detected - model will load on first inference request")
+    print("Starting Gradio server...")
+    demo.launch(server_port=7861, js=TOKEN_CLICK_JS)

model.py

@@ -0,0 +1,220 @@
+"""
+Model loading and inference for Logit Lens Explorer.
+
+Loads Llama-3.2-3B-Instruct and provides inference with hidden state
+capture for logit lens visualization.
+
+Part of E02: Logit Lens Explorer.
+"""
+
+from dataclasses import dataclass
+from typing import Generator
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
+
+MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
+
+_model = None
+_tokenizer = None
+_device = None
+
+
+@dataclass
+class LayerPrediction:
+    """Top-k token predictions from a single transformer layer."""
+
+    layer_index: int  # 0 = embedding, 1-28 = transformer layers
+    top_tokens: list[dict]  # [{"token": str, "probability": float}, ...]
+
+
+@dataclass
+class TokenData:
+    """Data for a single generated token with per-layer logit lens predictions."""
+
+    token: str
+    token_id: int
+    probability: float
+    layer_predictions: list[LayerPrediction]  # len = 29 (embedding + 28 layers)
+
+
+def load_model():
+    """Load the Llama model and tokenizer. Uses cached singleton."""
+    global _model, _tokenizer, _device
+
+    if _model is not None:
+        return _model, _tokenizer
+
+    _device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {_device}")
+    print(f"Loading model: {MODEL_ID}...")
+
+    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    _model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        attn_implementation="flash_attention_2",
+        torch_dtype=torch.float16,
+    ).to(_device).eval()
+
+    print("Model loaded successfully")
+    return _model, _tokenizer
+
+
+def project_hidden_states(
+    hidden_states: torch.Tensor,
+    model,
+    tokenizer,
+    top_k: int = 20,
+) -> list[LayerPrediction]:
+    """Batch-project hidden states through RMSNorm + lm_head.
+
+    Takes stacked hidden states from all layers and projects them through
+    the model's final normalization and unembedding head in a single
+    batched operation.
+
+    Args:
+        hidden_states: Stacked hidden states, shape (n_layers, 1, hidden_dim).
+        model: The causal LM model with .model.norm and .lm_head.
+        tokenizer: Tokenizer for decoding token IDs.
+        top_k: Number of top predictions per layer.
+
+    Returns:
+        List of LayerPrediction, one per layer.
+    """
+    # Reshape to (n_layers, hidden_dim), removing any size-1 middle dims, upcast to float32
+    n_layers = hidden_states.shape[0]
+    hidden_dim = hidden_states.shape[-1]
+    hs = hidden_states.reshape(n_layers, hidden_dim).float()
+
+    # Apply final RMSNorm (float32 for numerical stability)
+    normed = model.model.norm(hs)
+    # Cast back to model weight dtype for lm_head linear projection
+    logits = model.lm_head(normed.to(model.lm_head.weight.dtype))
+
+    # Softmax in float32 to avoid overflow
+    probs = torch.softmax(logits.float(), dim=-1)
+    top_probs, top_indices = torch.topk(probs, k=top_k, dim=-1)
+
+    # Move to CPU once for all layers
+    top_probs_cpu = top_probs.cpu().tolist()
+    top_indices_cpu = top_indices.cpu().tolist()
+
+    predictions = []
+    for layer_idx in range(len(top_probs_cpu)):
+        top_tokens = [
+            {"token": tokenizer.decode([int(idx)]), "probability": prob}
+            for prob, idx in zip(top_probs_cpu[layer_idx], top_indices_cpu[layer_idx])
+        ]
+        predictions.append(LayerPrediction(
+            layer_index=layer_idx,
+            top_tokens=top_tokens,
+        ))
+    return predictions
+
+
+def generate_with_logit_lens(
+    prompt: str,
+    max_new_tokens: int = 512,
+    top_k: int = 20,
+) -> Generator[TokenData, None, None]:
+    """Generate text token-by-token with per-layer logit lens predictions.
+
+    Uses greedy decoding (argmax) for deterministic text generation, but
+    records the natural softmax probabilities (temperature=1) for the logit
+    lens visualization so layer predictions reflect the model's true
+    confidence distribution.
+
+    Args:
+        prompt: User prompt text.
+        max_new_tokens: Maximum tokens to generate.
+        top_k: Number of top predictions per layer for logit lens.
+
+    Yields:
+        TokenData with token string, ID, probability, and per-layer predictions.
+    """
+    model, tokenizer = load_model()
+
+    messages = [{"role": "user", "content": prompt}]
+    prompt_full = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    inputs = tokenizer(prompt_full, return_tensors="pt").to(_device)
+    input_ids = inputs.input_ids
+    attention_mask = inputs.attention_mask
+
+    # EOS token IDs for stopping
+    eos_token_id = model.config.eos_token_id
+    if isinstance(eos_token_id, int):
+        eos_token_id = [eos_token_id]
+    elif eos_token_id is None:
+        eos_token_id = []
+
+    generated_ids = input_ids.clone()
+    past_key_values = DynamicCache()
+    seq_length = input_ids.shape[1]
+
+    with torch.no_grad():
+        for step in range(max_new_tokens):
+            if step == 0:
+                cache_position = torch.arange(seq_length, device=_device)
+                outputs = model(
+                    input_ids=generated_ids,
+                    attention_mask=attention_mask,
+                    cache_position=cache_position,
+                    past_key_values=past_key_values,
+                    output_hidden_states=True,
+                    return_dict=True,
+                    use_cache=True,
+                )
+            else:
+                cache_position = torch.tensor([seq_length], device=_device)
+                outputs = model(
+                    input_ids=generated_ids[:, -1:],
+                    attention_mask=attention_mask,
+                    cache_position=cache_position,
+                    past_key_values=past_key_values,
+                    output_hidden_states=True,
+                    return_dict=True,
+                    use_cache=True,
+                )
+
+            past_key_values = outputs.past_key_values
+
+            # Greedy decoding with natural probability recording
+            next_token_logits = outputs.logits[:, -1, :].float()
+            probs = torch.softmax(next_token_logits, dim=-1)
+            next_token_id = torch.argmax(probs, dim=-1).item()
+            next_token_prob = probs[0, next_token_id].item()
+
+            if next_token_id in eos_token_id:
+                break
+
+            # Eager logit lens: stack last-position hidden state from each layer
+            # outputs.hidden_states is a tuple of (n_layers+1) tensors,
+            # each shape (batch, seq_len, hidden_dim)
+            hidden_states = torch.stack([
+                hs[:, -1:, :] for hs in outputs.hidden_states
+            ])  # (n_layers, 1, hidden_dim)
+
+            layer_predictions = project_hidden_states(
+                hidden_states, model, tokenizer, top_k=top_k
+            )
+
+            token_str = tokenizer.decode([next_token_id])
+
+            yield TokenData(
+                token=token_str,
+                token_id=next_token_id,
+                probability=next_token_prob,
+                layer_predictions=layer_predictions,
+            )
+
+            # Update for next iteration
+            next_token_tensor = torch.tensor([[next_token_id]], device=_device)
+            generated_ids = torch.cat([generated_ids, next_token_tensor], dim=-1)
+            attention_mask = torch.cat(
+                [attention_mask, torch.ones((1, 1), device=_device, dtype=attention_mask.dtype)],
+                dim=-1,
+            )
+            seq_length += 1

requirements.txt

@@ -0,0 +1,26 @@
+# Logit Lens - Gradio Dependencies
+
+# Gradio UI
+gradio>=6.4.0
+
+# HuggingFace Spaces (ZeroGPU support)
+spaces
+
+# PyTorch + CUDA
+torch==2.6.0
+torchvision
+
+# Transformers + Qwen VL
+transformers==4.57.3
+qwen-vl-utils
+huggingface_hub
+
+# Attention + Acceleration
+flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
+git+https://github.com/huggingface/accelerate.git
+git+https://github.com/huggingface/peft.git
+transformers-stream-generator
+
+# Image processing
+Pillow
+sentencepiece