"""Neural layer introspection helpers for the Phase 4 Visual Layers GUI. This module provides PrivacyBERT-focused analysis with layer activations, attention summaries, math breakdown snapshots, and architecture map payloads. """ from __future__ import annotations from dataclasses import dataclass from functools import lru_cache from math import sqrt from pathlib import Path import importlib import json import tempfile import uuid from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple from prert.phase4.compliance_assessor import split_policy_clauses DEFAULT_MAX_CLAUSES = 16 DEFAULT_TOP_TOKENS = 8 BERT_LAYER_COUNT = 12 BERT_HEAD_COUNT = 12 HEATMAP_MAX_TOKENS = 64 @dataclass(frozen=True) class _ModelBundle: tokenizer: Any model: Any torch: Any model_id: str revision: str def _normalize_model_id(model_id: str) -> str: return str(model_id or "").strip() def _normalize_revision(revision: str) -> str: value = str(revision or "").strip() return value or "main" def _load_kwargs(revision: str) -> Dict[str, Any]: kwargs: Dict[str, Any] = {"revision": revision} token = ( importlib.import_module("os").getenv("HF_TOKEN") or importlib.import_module("os").getenv("HUGGINGFACEHUB_API_TOKEN") ) if token: kwargs["token"] = token return kwargs def _force_eager_attention(model: Any) -> None: config = getattr(model, "config", None) if config is not None: setattr(config, "output_hidden_states", True) setattr(config, "output_attentions", True) # Newer Transformers versions may read either public or private keys. if hasattr(config, "attn_implementation"): setattr(config, "attn_implementation", "eager") if hasattr(config, "_attn_implementation"): setattr(config, "_attn_implementation", "eager") setter = getattr(model, "set_attn_implementation", None) if callable(setter): try: setter("eager") except Exception: # Keep graceful fallback for architectures without runtime switching. pass @lru_cache(maxsize=3) def _get_model_bundle(model_id: str, revision: str) -> _ModelBundle: normalized_model_id = _normalize_model_id(model_id) if not normalized_model_id: raise ValueError("MODEL_ID is required for Visual Layers analysis.") normalized_revision = _normalize_revision(revision) torch_module = importlib.import_module("torch") transformers_module = importlib.import_module("transformers") load_kwargs = _load_kwargs(normalized_revision) tokenizer = transformers_module.AutoTokenizer.from_pretrained(normalized_model_id, **load_kwargs) eager_load_kwargs = dict(load_kwargs) eager_load_kwargs["attn_implementation"] = "eager" try: model = transformers_module.AutoModelForSequenceClassification.from_pretrained( normalized_model_id, **eager_load_kwargs, ) except TypeError: # Some older model classes do not accept attn_implementation at load time. model = transformers_module.AutoModelForSequenceClassification.from_pretrained( normalized_model_id, **load_kwargs, ) _force_eager_attention(model) model.eval() return _ModelBundle( tokenizer=tokenizer, model=model, torch=torch_module, model_id=normalized_model_id, revision=normalized_revision, ) def _coerce_indices(selected: Optional[Sequence[Any]], maximum: int, default_count: int) -> List[int]: if not selected: return list(range(min(default_count, maximum))) parsed: List[int] = [] for value in selected: try: index = int(value) except (TypeError, ValueError): continue if 0 <= index < maximum: parsed.append(index) unique = sorted(set(parsed)) if unique: return unique return list(range(min(default_count, maximum))) def _sigmoid(value: float) -> float: import math return 1.0 / (1.0 + math.exp(-value)) def _clean_tokens(tokens: Sequence[str]) -> List[str]: cleaned: List[str] = [] for token in tokens: normalized = str(token) if normalized in {"[PAD]", "[CLS]", "[SEP]"}: cleaned.append(normalized) continue cleaned.append(normalized.replace("##", "")) return cleaned def _sequence_length_from_mask(mask: Any) -> int: try: return int(mask.sum().item()) except Exception: try: return int(mask.sum()) except Exception: return int(len(mask)) def _build_architecture_map(layer_scores: Sequence[Mapping[str, Any]]) -> Dict[str, Any]: nodes: List[Dict[str, Any]] = [ {"id": "input", "label": "Input Tokens", "group": "io", "level": 0, "trigger": 1.0}, {"id": "embedding", "label": "Embeddings", "group": "embedding", "level": 1, "trigger": 1.0}, ] edges: List[Dict[str, Any]] = [ {"src": "input", "dst": "embedding", "weight": 1.0}, ] for index, score in enumerate(layer_scores, start=1): node_id = f"layer_{index}" trigger = float(score.get("mean_norm", 0.0)) nodes.append( { "id": node_id, "label": f"Layer {index}", "group": "transformer", "level": index + 1, "trigger": trigger, } ) previous = "embedding" if index == 1 else f"layer_{index - 1}" edges.append({"src": previous, "dst": node_id, "weight": trigger}) nodes.extend( [ {"id": "pooler", "label": "Pooler", "group": "head", "level": len(layer_scores) + 2, "trigger": 1.0}, {"id": "classifier", "label": "Classifier", "group": "head", "level": len(layer_scores) + 3, "trigger": 1.0}, {"id": "softmax", "label": "Softmax", "group": "head", "level": len(layer_scores) + 4, "trigger": 1.0}, ] ) last_layer = f"layer_{len(layer_scores)}" if layer_scores else "embedding" edges.extend( [ {"src": last_layer, "dst": "pooler", "weight": 1.0}, {"src": "pooler", "dst": "classifier", "weight": 1.0}, {"src": "classifier", "dst": "softmax", "weight": 1.0}, ] ) return { "nodes": nodes, "edges": edges, "triggered_node_ids": [str(node["id"]) for node in nodes], } def _summarize_layer_activations(torch_module: Any, hidden_states: Sequence[Any], token_count: int) -> List[Dict[str, Any]]: summaries: List[Dict[str, Any]] = [] usable = list(hidden_states[1:]) for index, tensor in enumerate(usable, start=1): layer_tensor = tensor[0, :token_count, :] norms = torch_module.linalg.vector_norm(layer_tensor, dim=-1) summaries.append( { "layer": index, "mean_norm": float(norms.mean().item()), "max_norm": float(norms.max().item()), "mean_abs": float(layer_tensor.abs().mean().item()), "activation_score": float(_sigmoid(float(layer_tensor.mean().item()))), } ) return summaries def _attention_entropy(row: Sequence[float]) -> float: import math epsilon = 1e-12 return float(-sum(float(p) * math.log(float(p) + epsilon) for p in row)) def _summarize_attention( torch_module: Any, attentions: Sequence[Any], tokens: Sequence[str], token_count: int, selected_layers: Sequence[int], selected_heads: Sequence[int], top_tokens: int, ) -> List[Dict[str, Any]]: summaries: List[Dict[str, Any]] = [] for layer_index in selected_layers: if layer_index < 0 or layer_index >= len(attentions): continue layer_attention = attentions[layer_index][0, :, :token_count, :token_count] heads = min(int(layer_attention.shape[0]), BERT_HEAD_COUNT) for head_index in selected_heads: if head_index < 0 or head_index >= heads: continue head_matrix = layer_attention[head_index] cls_row = head_matrix[0] values, indices = torch_module.topk(cls_row, k=min(top_tokens, token_count)) top_focus = [ { "token": str(tokens[int(idx)]), "score": float(val), } for val, idx in zip(values.detach().cpu().tolist(), indices.detach().cpu().tolist()) ] summaries.append( { "layer": layer_index + 1, "head": head_index, "mean": float(head_matrix.mean().item()), "max": float(head_matrix.max().item()), "entropy": _attention_entropy(cls_row.detach().cpu().tolist()), "cls_focus": top_focus, "tokens": [str(token) for token in tokens[: min(token_count, HEATMAP_MAX_TOKENS)]], "matrix": head_matrix[ : min(token_count, HEATMAP_MAX_TOKENS), : min(token_count, HEATMAP_MAX_TOKENS), ] .detach() .cpu() .tolist(), } ) return summaries def _math_breakdown( bundle: _ModelBundle, hidden_states: Sequence[Any], attentions: Sequence[Any], token_count: int, selected_layer_index: int, ) -> Dict[str, Any]: if selected_layer_index < 0: selected_layer_index = 0 if selected_layer_index >= len(attentions): selected_layer_index = max(0, len(attentions) - 1) details: Dict[str, Any] = { "layer": selected_layer_index + 1, "supported": False, "note": "Detailed Q/K/V extraction unavailable for this architecture.", } base_model = getattr(bundle.model, "base_model", None) encoder = getattr(base_model, "encoder", None) if base_model is not None else None layers = getattr(encoder, "layer", None) if encoder is not None else None if layers is None or selected_layer_index >= len(layers): return details layer_module = layers[selected_layer_index] attention_self = getattr(getattr(layer_module, "attention", None), "self", None) intermediate = getattr(layer_module, "intermediate", None) output_dense = getattr(getattr(layer_module, "output", None), "dense", None) if attention_self is None or intermediate is None or output_dense is None: return details torch_module = bundle.torch hidden_in = hidden_states[selected_layer_index][0, :token_count, :] query_projection = attention_self.query(hidden_in) key_projection = attention_self.key(hidden_in) value_projection = attention_self.value(hidden_in) num_heads = int(getattr(attention_self, "num_attention_heads", BERT_HEAD_COUNT)) head_dim = int(getattr(attention_self, "attention_head_size", int(query_projection.shape[-1] / max(num_heads, 1)))) q = query_projection.view(token_count, num_heads, head_dim).transpose(0, 1) k = key_projection.view(token_count, num_heads, head_dim).transpose(0, 1) v = value_projection.view(token_count, num_heads, head_dim).transpose(0, 1) scale = 1.0 / sqrt(float(head_dim)) scores = torch_module.matmul(q, k.transpose(-1, -2)) * scale probs = torch_module.softmax(scores, dim=-1) context = torch_module.matmul(probs, v) ffn_intermediate = intermediate(hidden_in) ffn_output = output_dense(ffn_intermediate) observed_attention = attentions[selected_layer_index][0, :, :token_count, :token_count] divergence = float((probs - observed_attention).abs().mean().item()) details = { "layer": selected_layer_index + 1, "supported": True, "q_mean_norm": float(torch_module.linalg.vector_norm(query_projection, dim=-1).mean().item()), "k_mean_norm": float(torch_module.linalg.vector_norm(key_projection, dim=-1).mean().item()), "v_mean_norm": float(torch_module.linalg.vector_norm(value_projection, dim=-1).mean().item()), "attention_score_mean": float(scores.mean().item()), "attention_prob_mean": float(probs.mean().item()), "context_mean_abs": float(context.abs().mean().item()), "ffn_intermediate_mean": float(ffn_intermediate.mean().item()), "ffn_output_mean": float(ffn_output.mean().item()), "attention_reconstruction_mae": divergence, } return details def _step_trace( text: str, token_count: int, predictions: Sequence[Mapping[str, Any]], selected_layers: Sequence[int], ) -> List[Dict[str, Any]]: trace: List[Dict[str, Any]] = [ { "step": "input", "detail": f"Received {len(text)} characters.", }, { "step": "tokenize", "detail": f"Generated {token_count} tokens after truncation and padding.", }, ] for layer_index in selected_layers: trace.append( { "step": f"transformer_layer_{layer_index + 1}", "detail": "Self-attention + feed-forward block executed.", } ) trace.extend( [ { "step": "classifier", "detail": "Linear classification head projected pooled representation into logits.", }, { "step": "softmax", "detail": f"Normalized logits across {len(predictions)} output classes.", }, ] ) return trace def _analyze_clause( bundle: _ModelBundle, text: str, max_length: int, selected_layers: Sequence[int], selected_heads: Sequence[int], top_tokens: int, ) -> Dict[str, Any]: if not text.strip(): raise ValueError("Input text is empty.") encoded = bundle.tokenizer( text, truncation=True, padding="max_length", max_length=max_length, return_tensors="pt", ) with bundle.torch.no_grad(): outputs = bundle.model( **encoded, output_hidden_states=True, output_attentions=True, return_dict=True, ) hidden_states = outputs.hidden_states attentions = outputs.attentions if hidden_states is None or attentions is None: raise RuntimeError("Model did not return hidden_states and attentions. Use a BERT-compatible checkpoint.") token_count = _sequence_length_from_mask(encoded["attention_mask"][0]) input_ids = encoded["input_ids"][0, :token_count] raw_tokens = bundle.tokenizer.convert_ids_to_tokens(input_ids) tokens = _clean_tokens(raw_tokens) logits = outputs.logits[0] probabilities = bundle.torch.softmax(logits, dim=-1).detach().cpu().tolist() labels = [] id2label = getattr(bundle.model.config, "id2label", None) for index in range(len(probabilities)): label = str(id2label.get(index, str(index))) if isinstance(id2label, dict) else str(index) labels.append(label) predictions = [ { "label": labels[index], "score": float(probabilities[index]), } for index in range(len(probabilities)) ] predictions.sort(key=lambda item: float(item.get("score", 0.0)), reverse=True) layer_activation = _summarize_layer_activations(bundle.torch, hidden_states, token_count) architecture = _build_architecture_map(layer_activation) attention = _summarize_attention( bundle.torch, attentions, tokens, token_count, selected_layers, selected_heads, top_tokens, ) math = _math_breakdown(bundle, hidden_states, attentions, token_count, selected_layers[0] if selected_layers else 0) trace = _step_trace(text, token_count, predictions, selected_layers) return { "text": text, "token_count": token_count, "tokens": tokens, "predictions": predictions, "layer_activation": layer_activation, "attention": attention, "math_breakdown": math, "step_trace": trace, "architecture": architecture, } def _aggregate_policy_results(results: Sequence[Mapping[str, Any]]) -> Dict[str, Any]: if not results: return { "clauses": 0, "avg_token_count": 0.0, "layer_activation": [], "attention": [], "label_distribution": {}, } clause_count = len(results) avg_tokens = sum(int(item.get("token_count", 0)) for item in results) / max(clause_count, 1) layer_table: Dict[int, Dict[str, float]] = {} for item in results: for layer in item.get("layer_activation", []): layer_idx = int(layer.get("layer", 0)) bucket = layer_table.setdefault(layer_idx, {"mean_norm": 0.0, "max_norm": 0.0, "mean_abs": 0.0, "count": 0.0}) bucket["mean_norm"] += float(layer.get("mean_norm", 0.0)) bucket["max_norm"] += float(layer.get("max_norm", 0.0)) bucket["mean_abs"] += float(layer.get("mean_abs", 0.0)) bucket["count"] += 1.0 layer_activation: List[Dict[str, Any]] = [] for layer_idx in sorted(layer_table): bucket = layer_table[layer_idx] count = max(float(bucket.get("count", 1.0)), 1.0) layer_activation.append( { "layer": layer_idx, "mean_norm": float(bucket.get("mean_norm", 0.0) / count), "max_norm": float(bucket.get("max_norm", 0.0) / count), "mean_abs": float(bucket.get("mean_abs", 0.0) / count), } ) attention_rows: List[Dict[str, Any]] = [] attention_table: Dict[Tuple[int, int], Dict[str, float]] = {} for item in results: for row in item.get("attention", []): key = (int(row.get("layer", 0)), int(row.get("head", 0))) bucket = attention_table.setdefault(key, {"mean": 0.0, "max": 0.0, "entropy": 0.0, "count": 0.0}) bucket["mean"] += float(row.get("mean", 0.0)) bucket["max"] += float(row.get("max", 0.0)) bucket["entropy"] += float(row.get("entropy", 0.0)) bucket["count"] += 1.0 for (layer, head), bucket in sorted(attention_table.items()): count = max(float(bucket.get("count", 1.0)), 1.0) attention_rows.append( { "layer": layer, "head": head, "mean": float(bucket.get("mean", 0.0) / count), "max": float(bucket.get("max", 0.0) / count), "entropy": float(bucket.get("entropy", 0.0) / count), } ) label_counts: Dict[str, int] = {} for item in results: predictions = item.get("predictions", []) if not predictions: continue top_label = str(predictions[0].get("label", "unknown")) label_counts[top_label] = label_counts.get(top_label, 0) + 1 return { "clauses": clause_count, "avg_token_count": avg_tokens, "layer_activation": layer_activation, "attention": attention_rows, "label_distribution": label_counts, } def run_visual_layers_analysis( *, mode: str, clause_text: str, policy_text: str, model_id: str, model_revision: str, max_length: int, max_clauses: int, selected_layers: Optional[Sequence[Any]] = None, selected_heads: Optional[Sequence[Any]] = None, top_tokens: int = DEFAULT_TOP_TOKENS, ) -> Dict[str, Any]: normalized_mode = str(mode or "single_clause").strip().lower() if normalized_mode not in {"single_clause", "full_policy"}: raise ValueError("Mode must be either 'single_clause' or 'full_policy'.") selected_layer_indices = _coerce_indices(selected_layers, BERT_LAYER_COUNT, default_count=4) selected_head_indices = _coerce_indices(selected_heads, BERT_HEAD_COUNT, default_count=3) requested_max_length = max(32, min(int(max_length), 512)) requested_max_clauses = max(1, min(int(max_clauses), 64)) requested_top_tokens = max(1, min(int(top_tokens), 24)) bundle = _get_model_bundle(_normalize_model_id(model_id), _normalize_revision(model_revision)) if normalized_mode == "single_clause": clause = str(clause_text or "").strip() if not clause: raise ValueError("Single-clause mode requires input text.") clause_result = _analyze_clause( bundle, clause, requested_max_length, selected_layer_indices, selected_head_indices, requested_top_tokens, ) return { "mode": "single_clause", "model": {"model_id": bundle.model_id, "revision": bundle.revision}, "selected_layers": selected_layer_indices, "selected_heads": selected_head_indices, "max_length": requested_max_length, "result": clause_result, } clauses = split_policy_clauses(policy_text) if not clauses: raise ValueError("Full-policy mode requires valid policy text.") limited_clauses = clauses[:requested_max_clauses] clause_results: List[Dict[str, Any]] = [] for clause in limited_clauses: clause_results.append( _analyze_clause( bundle, clause, requested_max_length, selected_layer_indices, selected_head_indices, requested_top_tokens, ) ) aggregate = _aggregate_policy_results(clause_results) architecture = _build_architecture_map(aggregate.get("layer_activation", [])) return { "mode": "full_policy", "model": {"model_id": bundle.model_id, "revision": bundle.revision}, "selected_layers": selected_layer_indices, "selected_heads": selected_head_indices, "max_length": requested_max_length, "requested_clauses": requested_max_clauses, "result": { "clauses": clause_results, "aggregate": aggregate, "architecture": architecture, }, } def render_visual_layers_svg(analysis: Mapping[str, Any]) -> str: mode = str(analysis.get("mode", "single_clause")) result = analysis.get("result") architecture = {} if mode == "full_policy": architecture = dict((result or {}).get("architecture", {})) if isinstance(result, dict) else {} else: architecture = dict((result or {}).get("architecture", {})) if isinstance(result, dict) else {} nodes = architecture.get("nodes", []) if isinstance(architecture, dict) else [] edges = architecture.get("edges", []) if isinstance(architecture, dict) else [] if not nodes: return "