feat: Add top-5 predictions with deltas, certainty meter, and bar charts in accordion panels

app.py

@@ -506,12 +506,14 @@ def update_check_token_graph(check_token_data):
     [State('model-dropdown', 'value')]
 )
 def create_layer_accordions(activation_data, model_name):
-    """Create accordion panels for each layer."""
+    """Create accordion panels for each layer with top-5 bar charts, deltas, and certainty."""
     if not activation_data or not model_name:
         return html.P("Run analysis to see layer-by-layer predictions.", className="placeholder-text")
     
     try:
         from transformers import AutoModelForCausalLM, AutoTokenizer
+        import plotly.graph_objs as go
+        
         model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         
@@ -521,38 +523,98 @@ def create_layer_accordions(activation_data, model_name):
         if not layer_data:
             return html.P("No layer data available.", className="placeholder-text")
         
-        # Create accordion panels
+        # Create accordion panels (reversed to show final layer first)
         accordions = []
-        for i, layer in enumerate(layer_data):
+        for i, layer in enumerate(reversed(layer_data)):
             layer_num = layer['layer_num']
             top_token = layer.get('top_token', 'N/A')
             top_prob = layer.get('top_prob', 0.0)
-            top_3 = layer.get('top_3_tokens', [])
+            top_5 = layer.get('top_5_tokens', [])
+            deltas = layer.get('deltas', {})
+            certainty = layer.get('certainty', 0.0)
             
-            # Create summary header
+            # Create summary header with certainty
             if top_token:
-                summary_text = f"Layer L{layer_num}: '{top_token}' (p={top_prob:.3f})"
+                summary_text = f"Layer L{layer_num}: '{top_token}' (p={top_prob:.3f}, certainty={certainty:.2f})"
             else:
                 summary_text = f"Layer L{layer_num}: (no prediction)"
             
-            # Create accordion panel
+            # Create accordion panel content
+            content_items = []
+            
+            if top_5:
+                # Create horizontal bar chart for top-5 tokens
+                tokens = [tok for tok, _ in top_5]
+                probs = [prob for _, prob in top_5]
+                
+                # Create delta annotations (▲/▼ with color)
+                annotations = []
+                for idx, (token, prob) in enumerate(top_5):
+                    delta = deltas.get(token, 0.0)
+                    if abs(delta) > 0.001:  # Only show meaningful deltas
+                        symbol = '▲' if delta > 0 else '▼'
+                        color = '#28a745' if delta > 0 else '#dc3545'
+                        annotations.append({
+                            'x': prob,
+                            'y': idx,
+                            'text': f'{symbol} {abs(delta):.3f}',
+                            'showarrow': False,
+                            'xanchor': 'left',
+                            'xshift': 10,
+                            'font': {'size': 10, 'color': color}
+                        })
+                
+                # Create Plotly figure
+                fig = go.Figure(data=[
+                    go.Bar(
+                        x=probs,
+                        y=tokens,
+                        orientation='h',
+                        marker={'color': '#667eea'},
+                        text=[f'{p:.3f}' for p in probs],
+                        textposition='auto',
+                        hovertemplate='%{y}: %{x:.4f}<extra></extra>'
+                    )
+                ])
+                
+                fig.update_layout(
+                    title={
+                        'text': f'Top 5 Predictions (Certainty: {certainty:.2f})',
+                        'font': {'size': 14}
+                    },
+                    xaxis={'title': 'Probability', 'range': [0, max(probs) * 1.15]},
+                    yaxis={'title': '', 'autorange': 'reversed'},
+                    height=250,
+                    margin={'l': 100, 'r': 80, 't': 50, 'b': 40},
+                    annotations=annotations,
+                    hovermode='closest'
+                )
+                
+                content_items.append(
+                    dcc.Graph(
+                        figure=fig,
+                        config={'displayModeBar': False},
+                        style={'marginBottom': '10px'}
+                    )
+                )
+                
+                # Add certainty tooltip explanation
+                content_items.append(html.Div([
+                    html.Small([
+                        html.I(className="fas fa-info-circle", style={'marginRight': '5px', 'color': '#667eea'}),
+                        f"Certainty = 1 − H(p_top5)/log(5), where H is Shannon entropy. ",
+                        "Higher values indicate more confident predictions."
+                    ], style={'color': '#6c757d', 'fontStyle': 'italic'})
+                ], style={'marginTop': '5px'}))
+            else:
+                content_items.append(html.P("No predictions available"))
+            
             panel = html.Details([
                 html.Summary(summary_text, className="layer-summary"),
-                html.Div([
-                    html.P(f"Layer {layer_num} details (placeholder for future content)")
-                ], className="layer-content")
+                html.Div(content_items, className="layer-content")
             ], className="layer-accordion")
             
             accordions.append(panel)
-            
-            # Add token chips between adjacent layers (not after last layer)
-            if i < len(layer_data) - 1 and top_3:
-                chips = html.Div([
-                    html.Span("→", className="token-arrow"),
-                    *[html.Span(f"{tok} ({prob:.2f})", className="token-chip") 
-                      for tok, prob in top_3]
-                ], className="token-chips-row")
-                accordions.append(chips)
         
         return html.Div(accordions)
         

components/main_panel.py

@@ -45,10 +45,19 @@ def create_main_panel():
             ], id='check-token-graph-container', style={'flex': '1', 'minWidth': '300px', 'display': 'none'})
         ], className="input-container", style={"marginBottom": "1.5rem", "display": "flex", "gap": "1.5rem", "alignItems": "flex-start"}),
         
-        # Layer-based visualization section
+        # Layer-based visualization section with loading spinner
         html.Div([
             html.H3("Layer-by-Layer Predictions", className="section-title"),
-            html.Div(id='layer-accordions-container', className="layer-accordions")
+            dcc.Loading(
+                id="layer-accordions-loading",
+                type="default",
+                children=html.Div(id='layer-accordions-container', className="layer-accordions"),
+                overlay_style={"visibility":"visible", "opacity": .7, "backgroundColor": "white"},
+                custom_spinner=html.Div([
+                    html.I(className="fas fa-spinner fa-spin", style={'fontSize': '24px', 'color': '#667eea', 'marginRight': '10px'}),
+                    html.Span("Loading visuals...", style={'fontSize': '16px', 'color': '#495057'})
+                ], style={'display': 'flex', 'alignItems': 'center', 'justifyContent': 'center', 'padding': '2rem'})
+            )
         ], className="visualization-section"),
         
         # Two-Prompt Comparison section (shown when comparing)

todo.md

@@ -15,12 +15,13 @@ Note: Minimal-change approach. Reuse existing files (`app.py`, `components/main_
 - [ ] Add CSS utility classes for compact header + tokens chips row
 
 ## Feature: Per-layer predictions (top-5), deltas, certainty meter
-- [ ] Extend forward pass outputs to include per-layer top-5 tokens + probs (reusing logit lens) in `utils/model_patterns.py`
-- [ ] Compute delta vs previous layer for overlapping tokens (prob change, signed)
-- [ ] Compute certainty meter using normalized entropy over top-5 probs (0–1)
-- [ ] Render a `dcc.Graph` horizontal bar chart (top-5) inside each panel body
-- [ ] Show per-token delta as small ▲/▼ with color next to bars
-- [ ] Add tooltip explaining certainty: "certainty = 1 − H(p_top5)/log(5)"
+- [x] Extend forward pass outputs to include per-layer top-5 tokens + probs (reusing logit lens) in `utils/model_patterns.py`
+- [x] Compute delta vs previous layer for overlapping tokens (prob change, signed)
+- [x] Compute certainty meter using normalized entropy over top-5 probs (0–1)
+- [x] Render a `dcc.Graph` horizontal bar chart (top-5) inside each panel body
+- [x] Show per-token delta as small ▲/▼ with color next to bars
+- [x] Add tooltip explaining certainty: "certainty = 1 − H(p_top5)/log(5)"
+- [x] Add a spinning "Loading visuals..." after loading data until all the visualizations are loaded
 
 ## Feature: Simplified attention view + open full interactive view
 - [ ] From `activation_data['attention_outputs']`, compute top-3 attended input tokens for current position (per layer)

utils/model_patterns.py

@@ -203,9 +203,9 @@ def execute_forward_pass(model, tokenizer, prompt: str, config: Dict[str, Any])
     return result
 
 
-def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, logit_lens_parameter: str, tokenizer, norm_parameter: Optional[str] = None) -> List[Tuple[str, float]]:
+def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, logit_lens_parameter: str, tokenizer, norm_parameter: Optional[str] = None, top_k: int = 5) -> List[Tuple[str, float]]:
     """
-    Transform layer output to top 3 token probabilities using logit lens.
+    Transform layer output to top K token probabilities using logit lens.
     
     For standard logit lens, use block/layer outputs (residual stream), not component outputs.
     The residual stream contains the full hidden state with all accumulated information.
@@ -220,9 +220,10 @@ def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, lo
         logit_lens_parameter: Not used (deprecated)
         tokenizer: Tokenizer for decoding
         norm_parameter: Parameter path for final norm layer (e.g., "model.norm.weight")
+        top_k: Number of top tokens to return (default: 5)
     
     Returns:
-        List of (token_string, probability) tuples for top 3 tokens
+        List of (token_string, probability) tuples for top K tokens
     """
     with torch.no_grad():
         # Convert to tensor and ensure proper shape [batch, seq_len, hidden_dim]
@@ -242,8 +243,8 @@ def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, lo
         # Step 3: Get probabilities via softmax
         probs = F.softmax(logits[0, -1, :], dim=-1)
         
-        # Step 4: Extract top 3 tokens
-        top_probs, top_indices = torch.topk(probs, k=3)
+        # Step 4: Extract top K tokens
+        top_probs, top_indices = torch.topk(probs, k=top_k)
         
         return [
             (tokenizer.decode([idx.item()], skip_special_tokens=False), prob.item())
@@ -288,9 +289,9 @@ def get_norm_layer_from_parameter(model, norm_parameter: Optional[str]) -> Optio
     return None
 
 
-def _get_top_tokens(activation_data: Dict[str, Any], module_name: str, model, tokenizer) -> Optional[List[Tuple[str, float]]]:
+def _get_top_tokens(activation_data: Dict[str, Any], module_name: str, model, tokenizer, top_k: int = 5) -> Optional[List[Tuple[str, float]]]:
     """
-    Helper: Get top 3 tokens for a layer's block output.
+    Helper: Get top K tokens for a layer's block output.
     
     Uses block outputs (residual stream) which represent the full hidden state
     after all layer computations (attention + feedforward + residuals).
@@ -306,7 +307,7 @@ def _get_top_tokens(activation_data: Dict[str, Any], module_name: str, model, to
         norm_params = activation_data.get('norm_parameters', [])
         norm_parameter = norm_params[0] if norm_params else None
         
-        return logit_lens_transformation(layer_output, [], model, None, tokenizer, norm_parameter)
+        return logit_lens_transformation(layer_output, [], model, None, tokenizer, norm_parameter, top_k=top_k)
     except Exception as e:
         print(f"Warning: Could not compute logit lens for {module_name}: {e}")
         return None
@@ -388,12 +389,43 @@ def get_check_token_probabilities(activation_data: Dict[str, Any], model, tokeni
         return None
 
 
+def _compute_certainty(probs: List[float]) -> float:
+    """
+    Compute normalized certainty from probability distribution.
+    Formula: certainty = 1 - H(p)/log(K) where H is Shannon entropy.
+    
+    Args:
+        probs: List of probabilities (top-K)
+    
+    Returns:
+        Certainty score in [0, 1] where 1 = completely certain
+    """
+    import math
+    if not probs or len(probs) == 0:
+        return 0.0
+    
+    # Compute Shannon entropy: H = -Σ(p_i * log(p_i))
+    entropy = 0.0
+    for p in probs:
+        if p > 0:
+            entropy -= p * math.log(p)
+    
+    # Normalize by max entropy (log(K))
+    max_entropy = math.log(len(probs))
+    if max_entropy == 0:
+        return 1.0
+    
+    # Certainty = 1 - normalized_entropy
+    certainty = 1.0 - (entropy / max_entropy)
+    return max(0.0, min(1.0, certainty))  # Clamp to [0, 1]
+
+
 def extract_layer_data(activation_data: Dict[str, Any], model, tokenizer) -> List[Dict[str, Any]]:
     """
-    Extract layer-by-layer data for accordion display.
+    Extract layer-by-layer data for accordion display with top-5, deltas, and certainty.
     
     Returns:
-        List of dicts with: layer_num, top_token, top_prob, top_3_tokens (list of (token, prob))
+        List of dicts with: layer_num, top_token, top_prob, top_5_tokens, deltas, certainty
     """
     layer_modules = activation_data.get('block_modules', [])
     if not layer_modules:
@@ -407,26 +439,47 @@ def extract_layer_data(activation_data: Dict[str, Any], model, tokenizer) -> Lis
     
     logit_lens_enabled = activation_data.get('logit_lens_parameter') is not None
     layer_data = []
+    prev_token_probs = {}  # Track previous layer's token probabilities
     
     for layer_num, module_name in layer_info:
-        top_tokens = _get_top_tokens(activation_data, module_name, model, tokenizer) if logit_lens_enabled else None
+        top_tokens = _get_top_tokens(activation_data, module_name, model, tokenizer, top_k=5) if logit_lens_enabled else None
         
         if top_tokens:
             top_token, top_prob = top_tokens[0]
+            
+            # Compute deltas vs previous layer
+            deltas = {}
+            for token, prob in top_tokens:
+                prev_prob = prev_token_probs.get(token, 0.0)
+                deltas[token] = prob - prev_prob
+            
+            # Compute certainty from top-5 probabilities
+            probs = [prob for _, prob in top_tokens]
+            certainty = _compute_certainty(probs)
+            
             layer_data.append({
                 'layer_num': layer_num,
                 'module_name': module_name,
                 'top_token': top_token,
                 'top_prob': top_prob,
-                'top_3_tokens': top_tokens[:3]  # Get top 3 for chips
+                'top_3_tokens': top_tokens[:3],  # Keep for backward compatibility
+                'top_5_tokens': top_tokens[:5],  # New: top-5 for bar chart
+                'deltas': deltas,
+                'certainty': certainty
             })
+            
+            # Update previous layer probabilities
+            prev_token_probs = {token: prob for token, prob in top_tokens}
         else:
             layer_data.append({
                 'layer_num': layer_num,
                 'module_name': module_name,
                 'top_token': None,
                 'top_prob': None,
-                'top_3_tokens': []
+                'top_3_tokens': [],
+                'top_5_tokens': [],
+                'deltas': {},
+                'certainty': 0.0
             })
     
     return layer_data