visualization/assets/main.js

        const svgOverlay = document.getElementById('svg-overlay');
        const content = document.querySelector('.content');
        const renderModelEl = document.getElementById('render-model');
        let renderModel = {};
        let renderTokens = [];
        try {
            renderModel = JSON.parse(renderModelEl ? renderModelEl.textContent : '{}') || {};
            renderTokens = Array.isArray(renderModel.tokens) ? renderModel.tokens : [];
        } catch (e) {
            console.warn('Failed to parse render model JSON:', e);
            renderModel = {};
            renderTokens = [];
        }

        function escapeControlChars(text) {
            if (!text) return text;
            let out = '';
            for (let i = 0; i < text.length; i++) {
                const ch = text[i];
                const code = text.charCodeAt(i);
                if (ch === '\\') {
                    out += '\\\\';
                } else if (ch === '\n') {
                    out += '\\n';
                } else if (ch === '\r') {
                    out += '\\r';
                } else if (ch === '\t') {
                    out += '\\t';
                } else if (code < 32 || code === 127) {
                    out += '\\x' + code.toString(16).padStart(2, '0');
                } else {
                    out += ch;
                }
            }
            return out;
        }

        if (content) {
            while (content.firstChild) {
                content.removeChild(content.firstChild);
            }
            const wordCounts = {};
            renderTokens.forEach((token) => {
                if (token && token.is_word && token.word_key) {
                    wordCounts[token.word_key] = (wordCounts[token.word_key] || 0) + 1;
                }
            });

            renderTokens.forEach((token, idx) => {
                const span = document.createElement('span');
                span.className = 'token';
                span.dataset.tokenIdx = String(idx);
                span.dataset.tunedDelta = (token && typeof token.tuned_delta === 'number') ? String(token.tuned_delta) : '0';
                span.dataset.rawDelta = (token && typeof token.raw_delta === 'number') ? String(token.raw_delta) : '';
                const kind = (token && token.display && token.display.kind) ? token.display.kind : 'normal';
                const text = (token && token.display && typeof token.display.text === 'string') ? token.display.text : '';
                const hasVisible = (() => {
                    if (!text) return false;
                    for (let i = 0; i < text.length; i++) {
                        const code = text.charCodeAt(i);
                        if (code >= 32 && code !== 127) {
                            return true;
                        }
                    }
                    return false;
                })();
                const mainKind = (kind === 'control' && hasVisible) ? 'normal' : kind;
                if (text.includes('\n') || text.includes('\r')) {
                    span.dataset.hasLinebreak = '1';
                }
                if (mainKind === 'control') {
                    span.classList.add('token-kind-control');
                    span.textContent = text;
                } else if (mainKind === 'raw') {
                    span.classList.add('token-kind-raw');
                    span.textContent = text;
                } else {
                    span.textContent = text;
                }

                if (token && token.is_word && token.word_key && wordCounts[token.word_key] > 1) {
                    span.classList.add('word');
                    span.dataset.word = token.word_key;
                    if (token.word_id !== undefined && token.word_id !== null) {
                        span.dataset.wordId = String(token.word_id);
                    }
                }

                content.appendChild(span);
            });
        }

        const words = document.querySelectorAll('.word');
        const wordGroups = {};
        words.forEach(word => {
            const wordText = word.getAttribute('data-word');
            if (!wordGroups[wordText]) {
                wordGroups[wordText] = [];
            }
            wordGroups[wordText].push(word);
        });

        function clearLines() {
            while (svgOverlay.firstChild) {
                svgOverlay.removeChild(svgOverlay.firstChild);
            }
            words.forEach(w => w.classList.remove('highlighted'));
        }

        function pickRectByY(rects, targetY) {
            if (!rects || rects.length === 0) return null;
            let best = rects[0];
            let bestDist = Infinity;
            rects.forEach(r => {
                const cy = r.top + r.height / 2;
                const dist = Math.abs(cy - targetY);
                if (dist < bestDist) {
                    best = r;
                    bestDist = dist;
                }
            });
            return best;
        }

        function getAnchorRect(element, targetY) {
            const rects = Array.from(element.getClientRects());
            if (rects.length === 0) return element.getBoundingClientRect();
            if (rects.length === 1) return rects[0];
            const picked = pickRectByY(rects, targetY);
            return picked || rects[0];
        }

        function drawLines(hoveredWord, evt) {
            clearLines();

            const wordText = hoveredWord.getAttribute('data-word');
            const wordId = parseInt(hoveredWord.getAttribute('data-word-id'));
            const sameWords = wordGroups[wordText] || [];

            const previousWords = sameWords.filter(w => {
                const id = parseInt(w.getAttribute('data-word-id'));
                return id < wordId;
            });

            if (previousWords.length === 0) return;

            sameWords.forEach(w => w.classList.add('highlighted'));

            const targetY = evt ? evt.clientY : (hoveredWord.getBoundingClientRect().top + hoveredWord.getBoundingClientRect().height / 2);
            const hoveredRect = getAnchorRect(hoveredWord, targetY);
            const hoveredX = hoveredRect.left + hoveredRect.width / 2;
            const hoveredY = hoveredRect.top + hoveredRect.height / 2;

            previousWords.forEach(prevWord => {
                const prevRect = getAnchorRect(prevWord, hoveredY);
                const prevX = prevRect.left + prevRect.width / 2;
                const prevY = prevRect.top + prevRect.height / 2;

                const midX = (hoveredX + prevX) / 2;
                const midY = Math.min(hoveredY, prevY) - 30;

                const path = document.createElementNS('http://www.w3.org/2000/svg', 'path');
                path.setAttribute('class', 'link-line');
                path.setAttribute('d', `M ${prevX} ${prevY} Q ${midX} ${midY} ${hoveredX} ${hoveredY}`);
                svgOverlay.appendChild(path);

                const dot1 = document.createElementNS('http://www.w3.org/2000/svg', 'circle');
                dot1.setAttribute('class', 'link-dot');
                dot1.setAttribute('cx', prevX);
                dot1.setAttribute('cy', prevY);
                dot1.setAttribute('r', 4);
                svgOverlay.appendChild(dot1);

                const dot2 = document.createElementNS('http://www.w3.org/2000/svg', 'circle');
                dot2.setAttribute('class', 'link-dot');
                dot2.setAttribute('cx', hoveredX);
                dot2.setAttribute('cy', hoveredY);
                dot2.setAttribute('r', 4);
                svgOverlay.appendChild(dot2);
            });
        }

        words.forEach(word => {
            word.addEventListener('mouseenter', (e) => drawLines(word, e));
            word.addEventListener('mouseleave', clearLines);
        });

        window.addEventListener('scroll', clearLines);

        const tooltip = document.getElementById('tooltip');
        const tokenSpans = document.querySelectorAll('.token');

        tokenSpans.forEach(token => {
            token.addEventListener('mouseenter', (e) => {
                const tokenIdx = parseInt(token.dataset.tokenIdx);
                const tokenInfo = (!isNaN(tokenIdx) && renderTokens[tokenIdx]) ? renderTokens[tokenIdx] : null;
                const bytes = (tokenInfo && tokenInfo.bytes_hex) ? tokenInfo.bytes_hex : '';
                const compressionA = (tokenInfo && tokenInfo.compression && tokenInfo.compression.rwkv) ? tokenInfo.compression.rwkv : '';
                const compressionB = (tokenInfo && tokenInfo.compression && tokenInfo.compression.qwen) ? tokenInfo.compression.qwen : '';
                const avgCompressionA = (tokenInfo && tokenInfo.loss && typeof tokenInfo.loss.rwkv === 'number') ? tokenInfo.loss.rwkv.toFixed(2) : '';
                const avgCompressionB = (tokenInfo && tokenInfo.loss && typeof tokenInfo.loss.qwen === 'number') ? tokenInfo.loss.qwen.toFixed(2) : '';
                const modelA = (tokenInfo && tokenInfo.model_tokens && tokenInfo.model_tokens.rwkv) ? tokenInfo.model_tokens.rwkv : null;
                const modelB = (tokenInfo && tokenInfo.model_tokens && tokenInfo.model_tokens.qwen) ? tokenInfo.model_tokens.qwen : null;
                const top5A = (tokenInfo && tokenInfo.topk && tokenInfo.topk.rwkv) ? tokenInfo.topk.rwkv : null;
                const top5B = (tokenInfo && tokenInfo.topk && tokenInfo.topk.qwen) ? tokenInfo.topk.qwen : null;

                function hasControlChars(text) {
                    if (!text) return false;
                    for (let i = 0; i < text.length; i++) {
                        const code = text.charCodeAt(i);
                        if (code < 32 || code === 127) {
                            return true;
                        }
                    }
                    return false;
                }

                function resolveKind(text, kindHint) {
                    if (kindHint === 'raw' || kindHint === 'control' || kindHint === 'normal') {
                        return kindHint;
                    }
                    if (kindHint === true) {
                        return 'raw';
                    }
                    if (hasControlChars(text)) {
                        return 'control';
                    }
                    return 'normal';
                }

                function appendEscapedWithControlColor(container, text) {
                    if (text === undefined || text === null) return;
                    let buffer = '';
                    const flush = () => {
                        if (buffer) {
                            container.appendChild(document.createTextNode(buffer));
                            buffer = '';
                        }
                    };
                    for (let i = 0; i < text.length; i++) {
                        const ch = text[i];
                        const code = text.charCodeAt(i);
                        if (ch === '\\') {
                            buffer += '\\\\';
                            continue;
                        }
                        if (ch === '\n' || ch === '\r' || ch === '\t' || code < 32 || code === 127) {
                            flush();
                            const span = document.createElement('span');
                            span.className = 'esc-control';
                            if (ch === '\n') {
                                span.textContent = '\\n';
                            } else if (ch === '\r') {
                                span.textContent = '\\r';
                            } else if (ch === '\t') {
                                span.textContent = '\\t';
                            } else {
                                span.textContent = '\\x' + code.toString(16).padStart(2, '0');
                            }
                            container.appendChild(span);
                            continue;
                        }
                        buffer += ch;
                    }
                    flush();
                }

                function appendTokenText(container, text, kindHint) {
                    const display = (text !== undefined && text !== null) ? text : '';
                    const kind = resolveKind(display, kindHint);
                    while (container.firstChild) {
                        container.removeChild(container.firstChild);
                    }
                    if (kind === 'raw') {
                        const span = document.createElement('span');
                        span.className = 'esc-raw';
                        span.textContent = display;
                        container.appendChild(span);
                        return;
                    }
                    if (kind === 'control') {
                        appendEscapedWithControlColor(container, display);
                        return;
                    }
                    container.textContent = display;
                }

                function formatTopkColumn(topkData, modelName, titleClass) {
                    const column = document.createElement('div');
                    column.className = 'topk-column';
                    const title = document.createElement('div');
                    title.className = 'topk-title ' + titleClass;
                    title.textContent = modelName;
                    column.appendChild(title);
                    const list = document.createElement('div');
                    list.className = 'topk-list';
                    column.appendChild(list);

                    if (!topkData) {
                        list.textContent = 'N/A';
                        return column;
                    }
                    try {
                        const data = topkData;
                        let actualId = null;
                        let rank = null;
                        let actualProb = null;
                        let topkList = [];
                        if (data.length >= 4) {
                            [actualId, rank, actualProb, topkList] = data;
                        } else {
                            [actualId, rank, topkList] = data;
                        }
                        topkList.forEach((item, idx) => {
                            const tokenId = item[0];
                            const prob = item[1];
                            const tokenText = item[2];
                            const isRaw = item.length > 3 ? item[3] : false;
                            const isHit = tokenId === actualId;
                            const rankClass = isHit ? 'topk-rank hit' : 'topk-rank';
                            const rawText = (tokenText !== undefined && tokenText !== null) ? tokenText : '';
                            const displayText = (rawText !== '') ? rawText : ('[' + tokenId + ']');

                            const row = document.createElement('div');
                            row.className = 'topk-item';

                            const rankSpan = document.createElement('span');
                            rankSpan.className = rankClass;
                            rankSpan.textContent = (idx + 1) + '.';
                            row.appendChild(rankSpan);

                            const tokenSpan = document.createElement('span');
                            tokenSpan.className = 'topk-token';
                            tokenSpan.title = 'ID: ' + tokenId;
                            appendTokenText(tokenSpan, displayText, isRaw);
                            row.appendChild(tokenSpan);

                            const probSpan = document.createElement('span');
                            probSpan.className = 'topk-prob';
                            probSpan.textContent = (prob * 100).toFixed(2) + '%';
                            row.appendChild(probSpan);

                            if (isHit) {
                                const hit = document.createElement('span');
                                hit.className = 'topk-hit';
                                hit.textContent = '✓';
                                row.appendChild(hit);
                            }

                            list.appendChild(row);
                        });

                        if (rank > 10) {
                            let probSuffix = '';
                            const probVal = parseFloat(actualProb);
                            if (!isNaN(probVal)) {
                                probSuffix = ' (' + (probVal * 100).toFixed(4) + '%)';
                            }
                            const miss = document.createElement('div');
                            miss.className = 'topk-item topk-miss';
                            miss.textContent = 'Actual rank: ' + rank + probSuffix;
                            list.appendChild(miss);
                        }
                        return column;
                    } catch (e) {
                        console.error('Error in formatTopkColumn for ' + modelName + ':', e);
                        console.error('topkData:', topkData);
                        list.textContent = 'Error: ' + e.message;
                        return column;
                    }
                }

                function formatTokenChips(modelData, label, labelClass) {
                    const block = document.createElement('div');
                    block.className = 'token-block';
                    const labelSpan = document.createElement('span');
                    labelSpan.className = 'label ' + labelClass;
                    labelSpan.textContent = label + ':';
                    block.appendChild(labelSpan);

                    const chips = document.createElement('div');
                    chips.className = 'token-chips';
                    block.appendChild(chips);

                    if (!modelData) {
                        const na = document.createElement('span');
                        na.className = 'topk-token token-chip';
                        na.textContent = 'N/A';
                        chips.appendChild(na);
                        return block;
                    }
                    try {
                        const tokenList = modelData;
                        tokenList.forEach((item) => {
                            const tokenId = item[0];
                            const tokenText = item[1];
                            const kindHint = item.length > 2 ? item[2] : false;
                            const probVal = item.length > 3 ? item[3] : null;
                            const displayText = (tokenText !== undefined && tokenText !== null) ? tokenText : '';

                            const group = document.createElement('span');
                            group.className = 'token-chip-group';
                            group.title = 'ID: ' + tokenId;

                            const idSpan = document.createElement('span');
                            idSpan.className = 'token-id';
                            idSpan.textContent = '[' + tokenId + ']';
                            group.appendChild(idSpan);

                            const chipSpan = document.createElement('span');
                            chipSpan.className = 'topk-token token-chip';
                            appendTokenText(chipSpan, displayText, kindHint);
                            group.appendChild(chipSpan);

                            if (probVal !== null && probVal !== undefined) {
                                const probSpan = document.createElement('span');
                                probSpan.className = 'token-prob';
                                const numericProb = typeof probVal === 'number' ? probVal : parseFloat(probVal);
                                if (!isNaN(numericProb)) {
                                    probSpan.textContent = (numericProb * 100).toFixed(2) + '%';
                                } else {
                                    probSpan.textContent = String(probVal);
                                }
                                group.appendChild(probSpan);
                            }

                            chips.appendChild(group);
                        });
                        return block;
                    } catch (e) {
                        console.error('Error in formatTokenChips for ' + label + ':', e);
                        console.error('modelData:', modelData);
                        const err = document.createElement('span');
                        err.className = 'topk-token token-chip';
                        err.textContent = 'Error: ' + e.message;
                        chips.appendChild(err);
                        return block;
                    }
                }

                tooltip.replaceChildren();

                const bytesRow = document.createElement('div');
                const bytesLabel = document.createElement('span');
                bytesLabel.className = 'label';
                bytesLabel.textContent = 'Bytes:';
                const bytesValue = document.createElement('span');
                bytesValue.className = 'bytes';
                bytesValue.textContent = bytes || '(empty)';
                bytesRow.appendChild(bytesLabel);
                bytesRow.appendChild(document.createTextNode(' '));
                bytesRow.appendChild(bytesValue);
                tooltip.appendChild(bytesRow);

                const rwkvRow = document.createElement('div');
                const rwkvLabel = document.createElement('span');
                rwkvLabel.className = 'label';
                rwkvLabel.textContent = 'RWKV Compression Rate:';
                const rwkvValue = document.createElement('span');
                rwkvValue.className = 'loss-a';
                rwkvValue.textContent = compressionA || '(empty)';
                if (avgCompressionA) {
                    rwkvValue.textContent += ' (avg: ' + avgCompressionA + '%)';
                }
                rwkvRow.appendChild(rwkvLabel);
                rwkvRow.appendChild(document.createTextNode(' '));
                rwkvRow.appendChild(rwkvValue);
                tooltip.appendChild(rwkvRow);

                const qwenRow = document.createElement('div');
                const qwenLabel = document.createElement('span');
                qwenLabel.className = 'label';
                qwenLabel.textContent = 'Qwen Compression Rate:';
                const qwenValue = document.createElement('span');
                qwenValue.className = 'loss-b';
                qwenValue.textContent = compressionB || '(empty)';
                if (avgCompressionB) {
                    qwenValue.textContent += ' (avg: ' + avgCompressionB + '%)';
                }
                qwenRow.appendChild(qwenLabel);
                qwenRow.appendChild(document.createTextNode(' '));
                qwenRow.appendChild(qwenValue);
                tooltip.appendChild(qwenRow);

                const hr = document.createElement('hr');
                hr.style.borderColor = '#555';
                hr.style.margin = '6px 0';
                tooltip.appendChild(hr);

                tooltip.appendChild(formatTokenChips(modelA, 'RWKV', 'model-a'));
                tooltip.appendChild(formatTokenChips(modelB, 'Qwen', 'model-b'));

                if (top5A || top5B) {
                    const topkSection = document.createElement('div');
                    topkSection.className = 'topk-section';
                    const topkContainer = document.createElement('div');
                    topkContainer.className = 'topk-container';
                    topkContainer.appendChild(formatTopkColumn(top5A, 'RWKV Top10', 'model-a'));
                    topkContainer.appendChild(formatTopkColumn(top5B, 'Qwen Top10', 'model-b'));
                    topkSection.appendChild(topkContainer);
                    tooltip.appendChild(topkSection);
                }

                tooltip.style.display = 'block';
            });

            token.addEventListener('mousemove', (e) => {
                const tooltipRect = tooltip.getBoundingClientRect();
                const viewportWidth = window.innerWidth;
                const viewportHeight = window.innerHeight;

                let x = e.clientX + 15;
                let y = e.clientY + 15;

                if (x + tooltipRect.width > viewportWidth - 10) {
                    x = e.clientX - tooltipRect.width - 15;
                }
                if (y + tooltipRect.height > viewportHeight - 10) {
                    y = e.clientY - tooltipRect.height - 15;
                }
                if (x < 10) x = 10;
                if (y < 10) y = 10;

                tooltip.style.left = x + 'px';
                tooltip.style.top = y + 'px';
            });

            token.addEventListener('mouseleave', () => {
                tooltip.style.display = 'none';
            });
        });

        const slider = document.getElementById('color-range-slider');
        const rangeValue = document.getElementById('color-range-value');
        const deltaModeInputs = document.querySelectorAll('input[name="delta-mode"]');
        const legendBetter = document.getElementById('legend-better');
        const legendEqual = document.getElementById('legend-equal');
        const legendWorse = document.getElementById('legend-worse');
        let deltaMode = 'relative';

        // Collect all tuned_delta values
        const tokenData = [];
        const linebreakTokens = [];
        tokenSpans.forEach((token, idx) => {
            if (token.dataset.hasLinebreak === '1') {
                linebreakTokens.push(token);
                return;
            }
            const tunedDelta = parseFloat(token.dataset.tunedDelta);
            let rawDelta = parseFloat(token.dataset.rawDelta);
            if (isNaN(rawDelta)) {
                rawDelta = tunedDelta;
            }
            if (!isNaN(tunedDelta)) {
                tokenData.push({ token, tunedDelta, rawDelta });
            }
        });
        function getDeltaValue(item) {
            return deltaMode === 'absolute' ? item.rawDelta : item.tunedDelta;
        }

        function tunedDeltaToColor(tunedDelta, maxAbsDelta, exponent) {
            // Normalize to [-1, 1]
            const normalized = Math.max(-1, Math.min(1, tunedDelta / maxAbsDelta));
            let r, g, b;
            if (normalized < 0) {
                // Green (RWKV better)
                const intensity = Math.pow(-normalized, exponent);
                r = Math.round(255 * (1 - intensity * 0.85));
                g = 255;
                b = Math.round(255 * (1 - intensity * 0.85));
            } else {
                // Red (RWKV worse)
                const intensity = Math.pow(normalized, exponent);
                r = 255;
                g = Math.round(255 * (1 - intensity * 0.85));
                b = Math.round(255 * (1 - intensity * 0.85));
            }
            return `rgb(${r}, ${g}, ${b})`;
        }

        function updateColors(colorRangePercent) {
            // colorRangePercent: 0-100, represents the proportion of tokens to color
            const absValues = tokenData.map(item => Math.abs(getDeltaValue(item)));
            const maxAbsDelta = Math.max(...absValues, 1e-9);

            const sortedByAbs = [...tokenData].sort((a, b) => Math.abs(getDeltaValue(b)) - Math.abs(getDeltaValue(a)));
            sortedByAbs.forEach((item, rank) => {
                item.rank = rank;
            });
            const colorCount = Math.round(tokenData.length * colorRangePercent / 100);

            // Calculate exponent: 100% -> 0.5, 0% -> 1.0
            const exponent = 1 - (colorRangePercent / 100) * 0.5;

            // Calculate max deviation within the colored range
            let maxAbsDeltaInRange = 1e-9;
            tokenData.forEach(item => {
                if (item.rank < colorCount) {
                    maxAbsDeltaInRange = Math.max(maxAbsDeltaInRange, Math.abs(getDeltaValue(item)));
                }
            });

            tokenData.forEach(item => {
                if (item.rank < colorCount) {
                    // Use dynamic normalization based on colored range
                    const deltaValue = getDeltaValue(item);
                    item.token.style.backgroundColor = tunedDeltaToColor(deltaValue, maxAbsDeltaInRange, exponent);
                } else {
                    // Outside color range, white
                    item.token.style.backgroundColor = 'rgb(255, 255, 255)';
                }
            });
            linebreakTokens.forEach(token => {
                token.style.backgroundColor = 'rgb(255, 255, 255)';
            });
        }

        function updateLegendText() {
            if (!legendBetter || !legendEqual || !legendWorse) return;
            if (deltaMode === 'absolute') {
                legendBetter.textContent = 'RWKV better';
                legendEqual.textContent = 'Equal';
                legendWorse.textContent = 'RWKV worse';
            } else {
                legendBetter.textContent = 'RWKV better than avg delta';
                legendEqual.textContent = 'Equal to avg delta';
                legendWorse.textContent = 'RWKV worse than avg delta';
            }
        }

        slider.addEventListener('input', (e) => {
            const val = parseFloat(e.target.value);
            rangeValue.textContent = val.toFixed(1) + '%';
            updateColors(val);
        });

        deltaModeInputs.forEach(input => {
            input.addEventListener('change', (e) => {
                const target = e.target;
                if (target && target.checked) {
                    deltaMode = target.value === 'absolute' ? 'absolute' : 'relative';
                    updateLegendText();
                    updateColors(parseFloat(slider.value));
                }
            });
        });

        updateLegendText();
        // Apply default color range on page load
        updateColors(10);