Delete ablation_lab.py

ablation_lab.py

@@ -1,286 +0,0 @@
-import streamlit as st
-import torch
-import torch.nn as nn
-import pandas as pd
-import plotly.express as px
-import plotly.graph_objects as go
-import networkx as nx
-import copy
-from backend import ModelManager
-
-class AblationEngine:
-    """
-    Handles the 'Virtual Surgery' of models using PyTorch Hooks.
-    Instead of deleting code, we intercept signals during inference.
-    """
-    def __init__(self, model_manager):
-        self.manager = model_manager
-        self.active_hooks = []
-        self.ablation_log = []
-
-    def clear_hooks(self):
-        """Removes all active ablations (restores model to baseline)."""
-        for handle in self.active_hooks:
-            handle.remove()
-        self.active_hooks = []
-
-    def register_ablation(self, model, layer_name, ablation_type="zero_out", noise_level=0.1):
-        """
-        Injects a hook into a specific layer to modify its output.
-        """
-        target_module = dict(model.named_modules())[layer_name]
-        
-        def hook_fn(module, input, output):
-            if ablation_type == "zero_out":
-                # Structural Ablation: Kill the signal
-                return output * 0.0
-            
-            elif ablation_type == "add_noise":
-                # Robustness Test: Inject Gaussian noise
-                noise = torch.randn_like(output) * noise_level
-                return output + noise
-            
-            elif ablation_type == "freeze_mean":
-                # Information Bottleneck: Replace with batch mean
-                return torch.mean(output, dim=0, keepdim=True).expand_as(output)
-                
-            return output
-
-        # Register the hook
-        handle = target_module.register_forward_hook(hook_fn)
-        self.active_hooks.append(handle)
-        return f"Ablated {layer_name} ({ablation_type})"
-
-class ArchitectureVisualizer:
-    """
-    Builds a Netron-style interactive graph of the model layers using NetworkX + Plotly.
-    """
-    @staticmethod
-    def build_layer_graph(model):
-        G = nx.DiGraph()
-        prev_node = "Input"
-        G.add_node("Input", type="Input")
-
-        # Walk through modules (simplified for visualization)
-        # We limit depth to avoid 10,000 node graphs for LLMs
-        for name, module in model.named_modules():
-            # Filter for high-level blocks only (Layers, Attention, MLP)
-            if any(k in name for k in ["layer", "block", "attn", "mlp"]) and "." not in name.split(".")[-1]:
-                # Heuristic: Connect sequential blocks
-                G.add_node(name, type=module.__class__.__name__, params=sum(p.numel() for p in module.parameters()))
-                G.add_edge(prev_node, name)
-                prev_node = name
-        
-        G.add_node("Output", type="Output")
-        G.add_edge(prev_node, "Output")
-        return G
-
-    @staticmethod
-    def plot_interactive_graph(G):
-        pos = nx.spring_layout(G, seed=42, k=0.5)
-        
-        edge_x, edge_y = [], []
-        for edge in G.edges():
-            x0, y0 = pos[edge[0]]
-            x1, y1 = pos[edge[1]]
-            edge_x.extend([x0, x1, None])
-            edge_y.extend([y0, y1, None])
-
-        edge_trace = go.Scatter(
-            x=edge_x, y=edge_y,
-            line=dict(width=0.5, color='#888'),
-            hoverinfo='none', mode='lines'
-        )
-
-        node_x, node_y, node_text, node_color = [], [], [], []
-        for node in G.nodes():
-            x, y = pos[node]
-            node_x.append(x)
-            node_y.append(y)
-            info = G.nodes[node]
-            node_text.append(f"{node}<br>{info.get('type', 'Unknown')}<br>Params: {info.get('params', 'N/A')}")
-            
-            # Color coding
-            if "attn" in node.lower(): node_color.append("#FF0055") # Attention
-            elif "mlp" in node.lower(): node_color.append("#00CC96") # MLP
-            elif "layer" in node.lower(): node_color.append("#AB63FA") # Blocks
-            else: node_color.append("#FFFFFF")
-
-        node_trace = go.Scatter(
-            x=node_x, y=node_y,
-            mode='markers',
-            hoverinfo='text',
-            text=node_text,
-            marker=dict(showscale=False, color=node_color, size=15, line_width=2)
-        )
-
-        fig = go.Figure(data=[edge_trace, node_trace],
-                        layout=go.Layout(
-                            showlegend=False,
-                            hovermode='closest',
-                            margin=dict(b=0,l=0,r=0,t=0),
-                            paper_bgcolor='rgba(0,0,0,0)',
-                            plot_bgcolor='rgba(0,0,0,0)',
-                            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
-                            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False))
-                        )
-        return fig
-
-def render_ablation_dashboard():
-    # --- Custom CSS for the Dashboard Feel ---
-    st.markdown("""
-    <style>
-        .ablation-header { 
-            background: linear-gradient(90deg, #FF4B4B 0%, #FF9068 100%);
-            -webkit-background-clip: text;
-            -webkit-text-fill-color: transparent;
-            font-size: 30px; font-weight: 900;
-        }
-        .stat-box {
-            background-color: #1E1E1E; border: 1px solid #333;
-            padding: 15px; border-radius: 5px; text-align: center;
-        }
-        .risk-high { border-left: 5px solid #FF4B4B; }
-        .risk-med { border-left: 5px solid #FFAA00; }
-        .risk-low { border-left: 5px solid #00FF00; }
-    </style>
-    """, unsafe_allow_html=True)
-
-    st.markdown('<div class="ablation-header">🧪 SYSTEMATIC ABLATION LAB</div>', unsafe_allow_html=True)
-    st.caption("Surgically alter model components to measure contribution and robustness.")
-
-    if 'models' not in st.session_state:
-        st.warning("Please load models in the Discovery tab first.")
-        return
-
-    # 1. Select Subject
-    col_sel, col_viz = st.columns([1, 3])
-    
-    with col_sel:
-        st.subheader("1. Subject")
-        all_ids = st.session_state['models']['model_id'].tolist()
-        target_model_id = st.selectbox("Select Model for Surgery", all_ids)
-        
-        # Load Model Button
-        if st.button("Initialize Surgery Table"):
-            with st.spinner("Preparing model for hooks..."):
-                succ, msg = st.session_state['manager'].load_model(target_model_id)
-                if succ:
-                    st.success("Ready.")
-                    st.session_state['ablation_target'] = target_model_id
-                    # Initialize engine
-                    st.session_state['ablation_engine'] = AblationEngine(st.session_state['manager'])
-                else:
-                    st.error(msg)
-
-    # 2. Main Workspace
-    if 'ablation_target' in st.session_state:
-        target_id = st.session_state['ablation_target']
-        model_pkg = st.session_state['manager'].loaded_models.get(f"{target_id}_None") # Default FP32/16 key
-        
-        if not model_pkg:
-            st.error("Model lost from memory. Please reload.")
-            return
-
-        model = model_pkg['model']
-        
-        # --- TAB LAYOUT FOR ABLATION ---
-        t1, t2, t3 = st.tabs(["🧬 Structural Map", "🔪 Ablation Controls", "📊 Impact Report"])
-
-        # === TAB 1: ARCHITECTURE GRAPH ===
-        with t1:
-            st.markdown("### Interactive Architecture Map")
-            st.markdown("Visualize the flow to decide where to cut.")
-            
-            if st.button("Generate Graph (Heavy Compute)"):
-                with st.spinner("Tracing neural pathways..."):
-                    G = ArchitectureVisualizer.build_layer_graph(model)
-                    fig = ArchitectureVisualizer.plot_interactive_graph(G)
-                    st.plotly_chart(fig, use_container_width=True)
-
-        # === TAB 2: CONTROLS ===
-        with t2:
-            st.subheader("Configure Ablation Experiment")
-            
-            c1, c2 = st.columns(2)
-            with c1:
-                # Get all layers
-                all_layers = [n for n, _ in model.named_modules() if len(n) > 0]
-                target_layers = st.multiselect("Select Target Layers", all_layers, max_selections=5)
-            
-            with c2:
-                method = st.selectbox("Ablation Method", 
-                                      ["Zero-Out (Remove)", "Add Noise (Corrupt)", "Freeze Mean (Bottleneck)"])
-                if method == "Add Noise (Corrupt)":
-                    noise_val = st.slider("Noise Level (Std Dev)", 0.0, 2.0, 0.1)
-                else:
-                    noise_val = 0.0
-
-            if st.button("🔴 RUN ABLATION TEST"):
-                engine = st.session_state['ablation_engine']
-                engine.clear_hooks() # Reset previous
-                
-                results_log = []
-                
-                # 1. Establish Baseline
-                st.write("Measuring Baseline Performance...")
-                # We simply use a generation prompt length as a proxy for "Performance" 
-                # or run a quick perplexity check if integrated with benchmarks.
-                # For this dashboard, we run the "Prompt Integrity Test"
-                
-                prompt = "The capital of France is"
-                base_out = st.session_state['manager'].generate_text(target_id, "None", prompt)
-                results_log.append({"State": "Baseline", "Output": base_out, "Integrity": 100})
-                
-                # 2. Apply Hooks
-                for layer in target_layers:
-                    msg = engine.register_ablation(model, layer, method.lower().split()[0].replace("-","_"), noise_val)
-                    st.toast(msg)
-                
-                # 3. Measure Ablated Performance
-                st.write("Running Ablated Inference...")
-                ablated_out = st.session_state['manager'].generate_text(target_id, "None", prompt)
-                
-                # Simple heuristic: String similarity or length retention
-                integrity = (len(ablated_out) / len(base_out)) * 100 if len(base_out) > 0 else 0
-                results_log.append({"State": "Ablated", "Output": ablated_out, "Integrity": integrity})
-                
-                st.session_state['ablation_results'] = results_log
-                
-                # Cleanup
-                engine.clear_hooks()
-                st.success("Experiment Complete. Hooks Removed.")
-
-        # === TAB 3: RESULTS ===
-        with t3:
-            if 'ablation_results' in st.session_state:
-                res = st.session_state['ablation_results']
-                
-                # Visual Diff
-                st.markdown("### 📝 Output Degradation Analysis")
-                
-                col_base, col_abl = st.columns(2)
-                with col_base:
-                    st.info(f"**Baseline:** {res[0]['Output']}")
-                with col_abl:
-                    st.warning(f"**Ablated:** {res[1]['Output']}")
-                
-                # Metrics
-                deg = 100 - res[1]['Integrity']
-                st.metric("Model Degradation", f"{deg:.1f}%", delta=f"-{deg:.1f}%", delta_color="inverse")
-                
-                # Sensitivity Chart (Mocked for single run, would need loop for real sensitivity analysis)
-                st.markdown("### 🔥 Layer Sensitivity Heatmap")
-                
-                # Creating dummy data to show what the "full suite" would look like
-                sens_data = pd.DataFrame({
-                    "Layer": ["embed", "layer.0", "layer.1", "layer.2", "head"],
-                    "Sensitivity Score": [95, 10, 15, 80, 100]
-                })
-                
-                fig = px.bar(sens_data, x="Layer", y="Sensitivity Score", 
-                             color="Sensitivity Score", color_continuous_scale="RdYlGn_r",
-                             title="Estimated Contribution to Output (Simulated)")
-                st.plotly_chart(fig, use_container_width=True)
-            else:
-                st.info("Run an experiment in Tab 2 to see results.")