Capture complete attention patterns after generation

- Modified attention capture to occur after all tokens are generated
- Added tokens field to TraceData model for token display
- Increased sampling threshold from 20 to 100 tokens
- Capture all layers instead of sampling every Nth layer
- Include full token list (prompt + generated) in attention traces
- This ensures complete attention matrices that match token count

backend/model_service.py

@@ -75,6 +75,7 @@ class TraceData(BaseModel):
     type: str
     layer: Optional[str] = None
     weights: Optional[List[List[float]]] = None
+    tokens: Optional[List[str]] = None  # Add tokens field
     max_weight: Optional[float] = None
     entropy: Optional[float] = None
     mean: Optional[float] = None
@@ -128,7 +129,7 @@ class ModelManager:
             logger.error(f"Failed to load model: {e}")
             raise
     
-    def extract_attention_trace(self, layer_idx: int, attention_weights) -> TraceData:
+    def extract_attention_trace(self, layer_idx: int, attention_weights, tokens: Optional[List[str]] = None) -> TraceData:
         """Extract attention pattern trace from a layer"""
         # attention_weights is a tuple of tensors, one for each layer
         # Each tensor has shape (batch_size, num_heads, seq_len, seq_len)
@@ -138,10 +139,13 @@ class ModelManager:
         # Shape: (batch_size, num_heads, seq_len, seq_len) -> (seq_len, seq_len)
         avg_attention = layer_attention[0].mean(dim=0).detach().cpu().numpy()
         
-        # Sample the weights for efficiency
-        if avg_attention.shape[0] > 20:
-            indices = np.random.choice(avg_attention.shape[0], 20, replace=False)
+        # Don't sample if we have complete attention - we want the full matrix
+        # Only sample if the matrix is very large (>100x100)
+        if avg_attention.shape[0] > 100:
+            indices = np.random.choice(avg_attention.shape[0], 100, replace=False)
             avg_attention = avg_attention[indices][:, indices]
+            if tokens:
+                tokens = [tokens[i] for i in indices]
         
         # Ensure values are finite
         avg_attention = np.nan_to_num(avg_attention, nan=0.0, posinf=1.0, neginf=0.0)
@@ -163,6 +167,7 @@ class ModelManager:
             type="attention",
             layer=f"layer.{layer_idx}",
             weights=avg_attention.tolist(),
+            tokens=tokens,  # Include tokens in the trace
             max_weight=max_weight,
             entropy=entropy,
             timestamp=datetime.now().timestamp()
@@ -521,20 +526,9 @@ class ModelManager:
                         output_hidden_states=True
                     )
                     
-                    # Sample traces based on sampling rate
-                    if np.random.random() < sampling_rate:
-                        # Extract attention traces from multiple layers
-                        if outputs.attentions and len(outputs.attentions) > 0:
-                            # Sample every Nth layer to get good coverage
-                            num_layers = len(outputs.attentions)
-                            step = max(1, num_layers // 10)  # Get ~10 layers sampled
-                            for layer_idx in range(0, num_layers, step):
-                                try:
-                                    trace = self.extract_attention_trace(layer_idx, outputs.attentions)
-                                    traces.append(trace)
-                                    await self.broadcast_trace(trace)
-                                except Exception as e:
-                                    logger.warning(f"Failed to extract attention trace from layer {layer_idx}: {e}")
+                    # Skip mid-generation attention capture - we'll capture complete attention at the end
+                    # This ensures we get the full attention matrix for all generated tokens
+                    pass  # Removed mid-generation attention capture
                     
                     # Extract activation traces periodically (not every token to avoid overflow)
                     if outputs.hidden_states and len(outputs.hidden_states) > 0 and np.random.random() < 0.3:
@@ -633,8 +627,40 @@ class ModelManager:
                     if next_token.item() == self.tokenizer.eos_token_id:
                         break
                 
+                # After generation is complete, capture final attention patterns for all tokens
+                # Do a final forward pass with the complete sequence to get full attention
+                with torch.no_grad():
+                    final_outputs = self.model(
+                        **inputs,
+                        output_attentions=True,
+                        output_hidden_states=True
+                    )
+                    
+                    # Extract complete attention patterns from all layers
+                    if final_outputs.attentions and len(final_outputs.attentions) > 0:
+                        num_layers = len(final_outputs.attentions)
+                        
+                        # Clear previous partial traces and add complete ones
+                        traces = []  # Reset traces to only include complete attention patterns
+                        
+                        # Capture ALL layers for complete visualization
+                        for layer_idx in range(num_layers):
+                            try:
+                                # Get all token IDs (prompt + generated)
+                                all_token_ids = inputs["input_ids"][0].tolist()
+                                
+                                # Decode each token individually to preserve token boundaries
+                                all_tokens = [self.tokenizer.decode([token_id], skip_special_tokens=False) for token_id in all_token_ids]
+                                
+                                # Pass tokens to the extraction method
+                                trace = self.extract_attention_trace(layer_idx, final_outputs.attentions, all_tokens)
+                                traces.append(trace)
+                                await self.broadcast_trace(trace)
+                            except Exception as e:
+                                logger.warning(f"Failed to extract final attention trace from layer {layer_idx}: {e}")
+                
                 # Calculate final confidence
-                confidence_trace = self.calculate_confidence(logits)
+                confidence_trace = self.calculate_confidence(final_outputs.logits)
                 traces.append(confidence_trace)
                 await self.broadcast_trace(confidence_trace)