Fix score-all classification in vectorised path

The vectorised GPU path (used for most layers) still had the old
cascade with "semantic"/"positional" catch-alls. Updated to match
the score-all-then-rank system with behaviour types + code cues.

Also fixed self reference error — code runs in a standalone async
function, not a class method. Use dict-based cache instead.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

backend/model_service.py

@@ -2254,15 +2254,15 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
 
                         # --- Code cue scores (what code tokens are attended to) ---
                         # Decode token texts for code-aware detection (cached per step)
-                        if not hasattr(self, '_step_token_texts') or self._step_token_texts_step != step:
+                        if step_token_texts_cache.get('step') != step:
                             try:
-                                self._step_token_texts = [
+                                step_token_texts_cache['texts'] = [
                                     manager.tokenizer.decode([tid]) for tid in current_ids[0, :seq_len_hw].tolist()
                                 ]
                             except Exception:
-                                self._step_token_texts = []
-                            self._step_token_texts_step = step
-                        token_texts = self._step_token_texts
+                                step_token_texts_cache['texts'] = []
+                            step_token_texts_cache['step'] = step
+                        token_texts = step_token_texts_cache.get('texts', [])
 
                         code_cues = {}
                         if len(token_texts) == seq_len_hw:
@@ -2790,6 +2790,9 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                         pass
                 return hook
 
+            # Cache for decoded token texts (reused across heads within a step)
+            step_token_texts_cache: Dict[str, Any] = {}
+
             # Detect FFN type from first layer
             ffn_type = "gelu"  # default
 
@@ -3223,35 +3226,77 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                             ent = 0.0 if math.isnan(ent) or math.isinf(ent) else ent
                             avg_ent = 0.0 if math.isnan(avg_ent) or math.isinf(avg_ent) else avg_ent
 
-                            # Data-driven head pattern classification (priority order)
-                            pattern_type = None
-                            confidence = 0.0
-                            # 1. Attention sink: >50% weight on positions 0-2
-                            if skw > 0.5:
-                                pattern_type = "attention_sink"
-                                confidence = skw
-                            # 2. Previous token: sharp focus on immediate predecessor
-                            elif mw > 0.9 and ptw > 0.85:
-                                pattern_type = "previous_token"
-                                confidence = ptw
-                            # 3. Local: >80% weight within 5 positions of query
-                            elif seq_len_attn > 5 and lcw > 0.8:
-                                pattern_type = "local"
-                                confidence = lcw
-                            # 4. Induction: attends to positions following previous occurrences of current token
-                            elif step > 0 and idw > 0.3:
-                                pattern_type = "induction"
-                                confidence = min(1.0, idw)
-                            # 5. Positional: low entropy, focused attention
-                            elif ent < 1.0:
-                                pattern_type = "positional"
-                                confidence = 1.0 - ent
-                            # 6. Semantic: broad attention (fallback)
-                            elif ent >= 1.0:
-                                pattern_type = "semantic"
-                                confidence = min(1.0, 0.5)
+                            # Score-all-then-rank head classification
+                            # Behaviour type scores (attention geometry)
+                            behaviour_scores = {
+                                "attention_sink": skw,
+                                "previous_token": ptw,
+                                "local": lcw,
+                                "induction": min(1.0, idw),
+                                "focused": max(0.0, 1.0 - ent) if ent < 1.5 else 0.0,
+                                "diffuse": min(1.0, max(0.0, (ent - 1.0) / 2.0)),
+                            }
+                            behaviour_thresholds = {
+                                "attention_sink": 0.4,
+                                "previous_token": 0.7,
+                                "local": 0.5,
+                                "induction": 0.2,
+                                "focused": 0.3,
+                                "diffuse": 0.3,
+                            }
+                            qualified = {
+                                k: v for k, v in behaviour_scores.items()
+                                if v >= behaviour_thresholds.get(k, 0.3)
+                            }
+                            sorted_behaviours = sorted(qualified.items(), key=lambda x: x[1], reverse=True)
+                            primary = sorted_behaviours[0] if sorted_behaviours else ("diffuse", behaviour_scores["diffuse"])
+                            secondary = sorted_behaviours[1] if len(sorted_behaviours) > 1 else None
+
+                            pattern_type = primary[0]
+                            confidence = primary[1]
                             confidence = 0.0 if math.isnan(confidence) or math.isinf(confidence) else confidence
 
+                            # Code cue scores (what code tokens are attended to)
+                            # Decode token texts once per step (cached via nonlocal)
+                            if step_token_texts_cache.get('step') != step:
+                                try:
+                                    step_token_texts_cache['texts'] = [
+                                        manager.tokenizer.decode([tid]) for tid in current_ids[0, :seq_len_attn].tolist()
+                                    ]
+                                except Exception:
+                                    step_token_texts_cache['texts'] = []
+                                step_token_texts_cache['step'] = step
+                            token_texts = step_token_texts_cache.get('texts', [])
+
+                            code_cues = {}
+                            if len(token_texts) == seq_len_attn:
+                                head_weights = all_last_row[head_idx].cpu()
+                                delimiters = {'(', ')', '{', '}', '[', ']', ':', ';', ','}
+                                delim_indices = [i for i, t in enumerate(token_texts) if t.strip() in delimiters]
+                                if delim_indices:
+                                    code_cues["delimiter_sensitive"] = head_weights[delim_indices].sum().item()
+
+                                keywords = {'def', 'return', 'if', 'else', 'elif', 'for', 'while', 'class',
+                                           'import', 'from', 'try', 'except', 'with', 'as', 'in', 'not',
+                                           'and', 'or', 'True', 'False', 'None', 'self', 'yield', 'async',
+                                           'await', 'lambda', 'raise', 'pass', 'break', 'continue',
+                                           'function', 'const', 'let', 'var', 'new', 'this',
+                                           'public', 'private', 'static', 'void', 'int', 'string', 'bool',
+                                           'namespace', 'using', 'class', 'interface', 'override', 'virtual'}
+                                kw_indices = [i for i, t in enumerate(token_texts) if t.strip() in keywords]
+                                if kw_indices:
+                                    code_cues["keyword_sensitive"] = head_weights[kw_indices].sum().item()
+
+                                if idw > 0.15:
+                                    code_cues["pattern_reuse"] = min(1.0, idw * 1.5)
+
+                            cue_threshold = 0.15
+                            sorted_cues = sorted(
+                                [(k, round(v, 4)) for k, v in code_cues.items() if v >= cue_threshold],
+                                key=lambda x: x[1], reverse=True
+                            )
+                            primary_cue = sorted_cues[0] if sorted_cues else None
+
                             attention_matrix = layer_attn_cpu[head_idx]
 
                             q_matrix = None
@@ -3269,14 +3314,25 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                                 "v_matrix": v_matrix
                             })
 
-                            critical_heads.append({
+                            head_entry = {
                                 "head_idx": head_idx,
                                 "entropy": ent,
                                 "avg_entropy": avg_ent,
                                 "max_weight": mw,
                                 "has_matrices": attention_matrix is not None,
-                                "pattern": {"type": pattern_type, "confidence": confidence} if pattern_type else None
-                            })
+                                "pattern": {"type": pattern_type, "confidence": round(confidence, 4)} if pattern_type else None,
+                            }
+                            if secondary:
+                                head_entry["secondary_behaviour"] = {"type": secondary[0], "score": round(secondary[1], 4)}
+                            if primary_cue:
+                                head_entry["code_cue"] = {
+                                    "type": primary_cue[0],
+                                    "score": primary_cue[1],
+                                    "evidence": f"{round(primary_cue[1] * 100)}% attention on {primary_cue[0].replace('_', ' ')} tokens",
+                                }
+                                if len(sorted_cues) > 1:
+                                    head_entry["secondary_cue"] = {"type": sorted_cues[1][0], "score": sorted_cues[1][1]}
+                            critical_heads.append(head_entry)
 
                         critical_heads.sort(key=lambda h: h["max_weight"], reverse=True)