update

visualization/html_generator.py

@@ -127,28 +127,6 @@ def get_token_info_for_text(text: str) -> dict:
     }
 
 
-def delta_to_color(delta: float, avg_delta: float, max_deviation: float) -> Tuple[int, int, int]:
-    """Map a delta value to an RGB color based on deviation from average."""
-    if max_deviation == 0:
-        return (255, 255, 255)
-
-    deviation = delta - avg_delta
-    normalized = max(-1, min(1, deviation / max_deviation))
-
-    if normalized < 0:
-        intensity = -normalized
-        r = int(255 * (1 - intensity * 0.7))
-        g = 255
-        b = int(255 * (1 - intensity * 0.7))
-    else:
-        intensity = normalized
-        r = 255
-        g = int(255 * (1 - intensity * 0.7))
-        b = int(255 * (1 - intensity * 0.7))
-
-    return (r, g, b)
-
-
 def generate_comparison_html(
     text: str,
     byte_losses_a: List[float],
@@ -253,12 +231,6 @@ def generate_comparison_html(
     deltas = [a - b for a, b in zip(byte_losses_a, byte_losses_b)]
     avg_delta = sum(deltas) / len(deltas) if deltas else 0
 
-    # Calculate max deviation
-    deviations = [d - avg_delta for d in deltas]
-    abs_deviations = [abs(dev) for dev in deviations]
-    max_deviation = float(np.percentile(abs_deviations, 100)) if abs_deviations else 0
-    max_deviation = max(max_deviation, 1e-6)
-
     # Calculate average compression rates
     avg_compression_a = sum(byte_losses_a) / len(byte_losses_a) * COMPRESSION_RATE_FACTOR if byte_losses_a else 0
     avg_compression_b = sum(byte_losses_b) / len(byte_losses_b) * COMPRESSION_RATE_FACTOR if byte_losses_b else 0
@@ -375,6 +347,10 @@ def generate_comparison_html(
         compression_a_str = " ".join([f"{l * COMPRESSION_RATE_FACTOR:.2f}%" for l in losses_a])
         compression_b_str = " ".join([f"{l * COMPRESSION_RATE_FACTOR:.2f}%" for l in losses_b])
 
+        # Calculate average compression rate for this token
+        avg_compression_a_token = sum(losses_a) / len(losses_a) * COMPRESSION_RATE_FACTOR if losses_a else 0
+        avg_compression_b_token = sum(losses_b) / len(losses_b) * COMPRESSION_RATE_FACTOR if losses_b else 0
+
         topk_a_json = ""
         topk_b_json = ""
         if topk_predictions_a is not None and model_a_token_ranges:
@@ -410,15 +386,10 @@ def generate_comparison_html(
 
         token_deltas = deltas[byte_start:byte_end]
         avg_token_delta = sum(token_deltas) / len(token_deltas) if token_deltas else 0
+        tuned_delta = avg_token_delta - avg_delta
 
-        # Apply power transformation to enhance color differentiation
-        # Preserve the sign and apply power to the absolute value
-        power_n = 3
-        sign = 1 if avg_token_delta >= 0 else -1
-        avg_token_delta_powered = sign * (abs(avg_token_delta) ** power_n)
-
-        color = delta_to_color(avg_token_delta_powered, avg_delta, max_deviation)
-        r, g, b = color
+        # Initial rendering uses white color, JavaScript will apply colors based on slider
+        r, g, b = 255, 255, 255
 
         token_html_parts = []
         for char in token_text:
@@ -445,7 +416,9 @@ def generate_comparison_html(
             f'data-bytes="{escape_for_attr(bytes_str)}" '
             f'data-compression-a="{escape_for_attr(compression_a_str)}" '
             f'data-compression-b="{escape_for_attr(compression_b_str)}" '
-            f'data-delta="{avg_token_delta * COMPRESSION_RATE_FACTOR:.4f}" '
+            f'data-avg-compression-a="{avg_compression_a_token:.2f}" '
+            f'data-avg-compression-b="{avg_compression_b_token:.2f}" '
+            f'data-tuned-delta="{tuned_delta:.6f}" '
             f'data-topk-a="{escape_for_attr(topk_a_json)}" '
             f'data-topk-b="{escape_for_attr(topk_b_json)}"'
         )
@@ -689,9 +662,9 @@ def generate_comparison_html(
                 <span>RWKV worse than avg</span>
             </div>
             <div class="legend-item" style="margin-left: 20px;">
-                <span style="color: #aaa;">Saturation:</span>
-                <input type="range" id="saturation-slider" min="500" max="1000" value="1000" step="1" style="width: 200px; vertical-align: middle;">
-                <span id="saturation-value" style="color: #fff; min-width: 45px; display: inline-block;">100.0%</span>
+                <span style="color: #aaa;">Color Range:</span>
+                <input type="range" id="color-range-slider" min="0" max="100" value="50" step="1" style="width: 200px; vertical-align: middle;">
+                <span id="color-range-value" style="color: #fff; min-width: 45px; display: inline-block;">50%</span>
             </div>
         </div>
     </div>
@@ -782,6 +755,8 @@ def generate_comparison_html(
                 const bytes = token.getAttribute('data-bytes') || '';
                 const compressionA = token.getAttribute('data-compression-a') || '';
                 const compressionB = token.getAttribute('data-compression-b') || '';
+                const avgCompressionA = token.getAttribute('data-avg-compression-a') || '';
+                const avgCompressionB = token.getAttribute('data-avg-compression-b') || '';
                 const top5A = token.getAttribute('data-topk-a') || '';
                 const top5B = token.getAttribute('data-topk-b') || '';
 
@@ -827,8 +802,8 @@ def generate_comparison_html(
 
                 let tooltipHtml = `
                     <div><span class="label">Bytes:</span> <span class="bytes">${{bytes || '(empty)'}}</span></div>
-                    <div><span class="label">RWKV Compression Rate:</span> <span class="loss-a">${{compressionA || '(empty)'}}</span></div>
-                    <div><span class="label">Qwen Compression Rate:</span> <span class="loss-b">${{compressionB || '(empty)'}}</span></div>
+                    <div><span class="label">RWKV Compression Rate:</span> <span class="loss-a">${{compressionA || '(empty)'}}${{avgCompressionA ? ' (avg: ' + avgCompressionA + '%)' : ''}}</span></div>
+                    <div><span class="label">Qwen Compression Rate:</span> <span class="loss-b">${{compressionB || '(empty)'}}${{avgCompressionB ? ' (avg: ' + avgCompressionB + '%)' : ''}}</span></div>
                     <hr style="border-color: #555; margin: 6px 0;">
                     <div><span class="label">RWKV:</span> <span class="model-a">${{modelA || '(empty)'}}</span></div>
                     <div><span class="label">Qwen:</span> <span class="model-b">${{modelB || '(empty)'}}</span></div>
@@ -869,85 +844,78 @@ def generate_comparison_html(
             }});
         }});
 
-        const avgDelta = {avg_delta_compression};
-        const slider = document.getElementById('saturation-slider');
-        const saturationValue = document.getElementById('saturation-value');
-        const powerN = 3;  // Must match Python's power_n
+        const slider = document.getElementById('color-range-slider');
+        const rangeValue = document.getElementById('color-range-value');
 
-        const allDeltas = [];
-        tokenSpans.forEach(token => {{
-            const delta = parseFloat(token.getAttribute('data-delta'));
-            if (!isNaN(delta)) allDeltas.push(delta);
+        // Collect all tuned_delta values
+        const tokenData = [];
+        tokenSpans.forEach((token, idx) => {{
+            const tunedDelta = parseFloat(token.getAttribute('data-tuned-delta'));
+            if (!isNaN(tunedDelta)) {{
+                tokenData.push({{ token, tunedDelta, absDelta: Math.abs(tunedDelta) }});
+            }}
         }});
 
-        // Apply power transformation to delta values (matching Python's logic)
-        function applyPower(delta) {{
-            const sign = delta >= 0 ? 1 : -1;
-            return sign * Math.pow(Math.abs(delta), powerN);
-        }}
+        // Calculate max_abs_tuned_delta for normalization
+        const maxAbsDelta = Math.max(...tokenData.map(d => d.absDelta), 1e-9);
 
-        function percentile(arr, p) {{
-            const sorted = [...arr].sort((a, b) => a - b);
-            const idx = (p / 100) * (sorted.length - 1);
-            const lower = Math.floor(idx);
-            const upper = Math.ceil(idx);
-            if (lower === upper) return sorted[lower];
-            return sorted[lower] + (sorted[upper] - sorted[lower]) * (idx - lower);
-        }}
+        // Sort by |tuned_delta| to get rankings
+        const sortedByAbs = [...tokenData].sort((a, b) => b.absDelta - a.absDelta);
+        sortedByAbs.forEach((item, rank) => {{
+            item.rank = rank;  // rank 0 = largest deviation
+        }});
 
-        function deltaToColor(deltaPowered, avgDeltaPowered, maxDeviation) {{
-            if (maxDeviation === 0) return 'rgb(255, 255, 255)';
-            const deviation = deltaPowered - avgDeltaPowered;
-            let normalized = Math.max(-1, Math.min(1, deviation / maxDeviation));
+        function tunedDeltaToColor(tunedDelta, maxAbsDelta) {{
+            // 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 = -normalized;
-                r = Math.round(255 * (1 - intensity * 0.7));
+                r = Math.round(255 * (1 - intensity * 0.85));
                 g = 255;
-                b = Math.round(255 * (1 - intensity * 0.7));
+                b = Math.round(255 * (1 - intensity * 0.85));
             }} else {{
+                // Red (RWKV worse)
                 const intensity = normalized;
                 r = 255;
-                g = Math.round(255 * (1 - intensity * 0.7));
-                b = Math.round(255 * (1 - intensity * 0.7));
+                g = Math.round(255 * (1 - intensity * 0.85));
+                b = Math.round(255 * (1 - intensity * 0.85));
             }}
             return `rgb(${{r}}, ${{g}}, ${{b}})`;
         }}
 
-        // Pre-compute powered deltas and avgDeltaPowered
-        const allDeltasPowered = allDeltas.map(d => applyPower(d));
-        const avgDeltaPowered = applyPower(avgDelta);
-
-        // Pre-compute min and max deviations for logarithmic interpolation
-        const allDeviations = allDeltasPowered.map(d => Math.abs(d - avgDeltaPowered));
-        const minDeviation = Math.max(percentile(allDeviations, 1), 1e-9);  // Use 1st percentile to avoid extreme outliers
-        const maxDeviationFull = Math.max(percentile(allDeviations, 100), 1e-6);
-
-        function updateColors(sliderValue) {{
-            // Use logarithmic interpolation for smoother perceptual control
-            // sliderValue: 50-100, maps to maxDeviation from minDeviation to maxDeviationFull
-            // Lower slider value = lower maxDeviation = more saturation (more colors hit the clamp)
-            // Higher slider value = higher maxDeviation = less saturation (fewer colors hit the clamp)
-            const t = (sliderValue - 50) / 50;  // Normalize to 0-1
-            // Logarithmic interpolation: exp(lerp(log(min), log(max), t))
-            const logMin = Math.log(minDeviation);
-            const logMax = Math.log(maxDeviationFull);
-            const maxDeviation = Math.exp(logMin + t * (logMax - logMin));
-
-            tokenSpans.forEach((token, idx) => {{
-                const delta = parseFloat(token.getAttribute('data-delta'));
-                if (!isNaN(delta)) {{
-                    const deltaPowered = applyPower(delta);
-                    token.style.backgroundColor = deltaToColor(deltaPowered, avgDeltaPowered, maxDeviation);
+        function updateColors(colorRangePercent) {{
+            // colorRangePercent: 0-100, represents the proportion of tokens to color
+            const colorCount = Math.round(tokenData.length * colorRangePercent / 100);
+
+            // Calculate max deviation within the colored range
+            let maxAbsDeltaInRange = 1e-9;
+            tokenData.forEach(item => {{
+                if (item.rank < colorCount) {{
+                    maxAbsDeltaInRange = Math.max(maxAbsDeltaInRange, item.absDelta);
+                }}
+            }});
+
+            tokenData.forEach(item => {{
+                if (item.rank < colorCount) {{
+                    // Use dynamic normalization based on colored range
+                    item.token.style.backgroundColor = tunedDeltaToColor(item.tunedDelta, maxAbsDeltaInRange);
+                }} else {{
+                    // Outside color range, white
+                    item.token.style.backgroundColor = 'rgb(255, 255, 255)';
                 }}
             }});
         }}
 
         slider.addEventListener('input', (e) => {{
-            const val = parseInt(e.target.value) / 10;
-            saturationValue.textContent = val.toFixed(1) + '%';
+            const val = parseInt(e.target.value);
+            rangeValue.textContent = val + '%';
             updateColors(val);
         }});
+
+        // Apply default color range on page load
+        updateColors(50);
     </script>
 </body>
 </html>