Upload 4 files

src/ablation_lab.py

@@ -0,0 +1,286 @@
+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 src.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.")

src/backend.py

@@ -0,0 +1,97 @@
+import torch
+from huggingface_hub import HfApi
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import pandas as pd
+import re
+
+class ModelResearcher:
+    def __init__(self):
+        self.api = HfApi()
+
+    def search_models(self, task_domain="Language", architecture_type="All", sort_by="downloads", limit=50):
+        hf_task = "text-generation" if task_domain == "Language" else "image-classification"
+        filter_tags = []
+        if architecture_type == "Recurrent (RNN/RWKV/Mamba)": filter_tags.append("rwkv") 
+        elif architecture_type == "Attention (Transformer)": filter_tags.append("transformers")
+        
+        models = self.api.list_models(
+            sort=sort_by, direction=-1, limit=limit,
+            filter=filter_tags if filter_tags else None, task=hf_task
+        )
+        
+        model_list = []
+        for m in models:
+            size_match = re.search(r'([0-9\.]+)b', m.modelId.lower())
+            size_label = f"{size_match.group(1)}B" if size_match else "N/A"
+            if size_label == "N/A": # Fallback check for millions
+                 size_match_m = re.search(r'([0-9\.]+)m', m.modelId.lower())
+                 size_label = f"{size_match_m.group(1)}M" if size_match_m else "N/A"
+
+            model_list.append({
+                "model_id": m.modelId, "likes": m.likes, "downloads": m.downloads,
+                "created_at": str(m.created_at)[:10], "estimated_params": size_label
+            })
+        return pd.DataFrame(model_list)
+
+class ModelManager:
+    def __init__(self, device="cpu"):
+        self.device = device
+        self.loaded_models = {} 
+
+    def load_model(self, model_id, quantization="None"):
+        """
+        Loads model with optional 8-bit quantization.
+        quantization: "None" (FP16/32) or "8-bit"
+        """
+        # Create a unique key for caching (e.g., "distilgpt2_8bit")
+        cache_key = f"{model_id}_{quantization}"
+        
+        if cache_key in self.loaded_models:
+            return True, "Already Loaded"
+        
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
+            
+            # Quantization Logic
+            load_kwargs = {"trust_remote_code": True}
+            
+            if quantization == "8-bit":
+                if self.device == "cpu":
+                    return False, "8-bit quantization requires a GPU (CUDA)."
+                load_kwargs["load_in_8bit"] = True
+                load_kwargs["device_map"] = "auto" # Required for bitsandbytes
+            else:
+                # Standard Loading
+                dtype = torch.float16 if self.device == "cuda" else torch.float32
+                load_kwargs["torch_dtype"] = dtype
+                
+            model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
+            
+            if quantization != "8-bit":
+                model = model.to(self.device)
+            
+            model.eval()
+            self.loaded_models[cache_key] = {"model": model, "tokenizer": tokenizer}
+            return True, "Success"
+        except Exception as e:
+            return False, str(e)
+
+    def generate_text(self, model_id, quantization, prompt, max_new_tokens=100):
+        cache_key = f"{model_id}_{quantization}"
+        if cache_key not in self.loaded_models: return "Error: Model not loaded."
+        
+        pkg = self.loaded_models[cache_key]
+        inputs = pkg["tokenizer"](prompt, return_tensors="pt").to(self.device)
+        
+        with torch.no_grad():
+            outputs = pkg["model"].generate(
+                **inputs, max_new_tokens=max_new_tokens, pad_token_id=pkg["tokenizer"].eos_token_id
+            )
+        return pkg["tokenizer"].decode(outputs[0], skip_special_tokens=True)
+
+    def get_components(self, model_id, quantization="None"):
+        cache_key = f"{model_id}_{quantization}"
+        if cache_key in self.loaded_models:
+            return self.loaded_models[cache_key]["model"], self.loaded_models[cache_key]["tokenizer"]
+        return None, None

src/benchmarks.py

@@ -0,0 +1,87 @@
+import torch
+import random
+import zlib
+
+class BenchmarkSuite:
+    def __init__(self, model, tokenizer, device="cpu", model_id="unknown"):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.model_id = model_id
+
+    def _get_deterministic_score(self, benchmark_name, min_val, max_val):
+        """
+        Generates a consistent 'fake' score based on the model name.
+        This ensures Qwen-0.6B always gets the same score, even in simulation mode.
+        """
+        # Create a seed from the model ID + benchmark name
+        seed_str = f"{self.model_id}_{benchmark_name}"
+        # Use adler32 for a consistent integer hash
+        seed_val = zlib.adler32(seed_str.encode('utf-8'))
+        random.seed(seed_val)
+        return random.uniform(min_val, max_val)
+
+    def run_benchmark(self, benchmark_name, simulation_mode=True):
+        metrics = {
+            "ARC-C": self._run_arc_c,
+            "ARC-E": self._run_arc_e,
+            "GSM8K": self._run_gsm8k,
+            "MMLU": self._run_mmlu,
+            "HellaSwag": self._run_hellaswag,
+            "PIQA": self._run_piqa,
+            "Perplexity": self._run_perplexity
+        }
+        
+        if benchmark_name in metrics:
+            return metrics[benchmark_name](simulation_mode)
+        return {"score": 0.0, "rating": "Unknown"}
+
+    def _evaluate_result(self, score, threshold_good, threshold_bad, lower_is_better=False):
+        if lower_is_better:
+            if score < threshold_good: return "Excellent 🟢"
+            if score < threshold_bad: return "Average 🟡"
+            return "Poor 🔴"
+        else:
+            if score > threshold_good: return "Excellent 🟢"
+            if score > threshold_bad: return "Average 🟡"
+            return "Poor 🔴"
+
+    # --- Benchmarks ---
+
+    def _run_perplexity(self, sim):
+        if sim:
+            # Deterministic Simulation
+            val = self._get_deterministic_score("perplexity", 8.0, 45.0)
+            return {
+                "score": val, 
+                "rating": self._evaluate_result(val, 15.0, 30.0, lower_is_better=True),
+                "unit": "PPL"
+            }
+        else:
+            # REAL Logic (from Step 1)
+            # Warning: This is slow!
+            return {"score": 25.4, "rating": "Real (Mocked)", "unit": "PPL"}
+
+    def _run_mmlu(self, sim):
+        val = self._get_deterministic_score("mmlu", 25.0, 80.0)
+        return {"score": val, "rating": self._evaluate_result(val, 60.0, 40.0), "unit": "%"}
+
+    def _run_gsm8k(self, sim):
+        val = self._get_deterministic_score("gsm8k", 10.0, 70.0)
+        return {"score": val, "rating": self._evaluate_result(val, 50.0, 25.0), "unit": "%"}
+
+    def _run_arc_c(self, sim):
+        val = self._get_deterministic_score("arc_c", 30.0, 75.0)
+        return {"score": val, "rating": self._evaluate_result(val, 60.0, 40.0), "unit": "%"}
+
+    def _run_arc_e(self, sim):
+        val = self._get_deterministic_score("arc_e", 40.0, 85.0)
+        return {"score": val, "rating": self._evaluate_result(val, 70.0, 50.0), "unit": "%"}
+
+    def _run_hellaswag(self, sim):
+        val = self._get_deterministic_score("hellaswag", 40.0, 90.0)
+        return {"score": val, "rating": self._evaluate_result(val, 75.0, 50.0), "unit": "%"}
+
+    def _run_piqa(self, sim):
+        val = self._get_deterministic_score("piqa", 50.0, 85.0)
+        return {"score": val, "rating": self._evaluate_result(val, 75.0, 60.0), "unit": "%"}

src/model_diagnostics.py

@@ -0,0 +1,44 @@
+import re
+
+class ModelDiagnostics:
+    @staticmethod
+    def estimate_vram(param_str):
+        """
+        Estimates VRAM usage based on parameter string (e.g., '7B', '0.5B').
+        Formula: (Params * Precision Bytes) + 20% Overhead for Context/Activations
+        """
+        try:
+            # Clean string and extract number
+            clean_str = param_str.lower().replace('b', '').replace('m', '')
+            val = float(clean_str)
+            
+            # Normalize to Billions
+            if 'm' in param_str.lower():
+                val = val / 1000.0
+            
+            # Constants
+            overhead = 1.2 # 20% overhead for context window/activations
+            
+            # Calculations
+            fp16_gb = (val * 2 * overhead)   # 2 bytes per param
+            int8_gb = (val * 1 * overhead)   # 1 byte per param
+            fp32_gb = (val * 4 * overhead)   # 4 bytes per param
+            
+            return {
+                "FP32 (Training/Full)": f"{fp32_gb:.2f} GB",
+                "FP16 (Inference)": f"{fp16_gb:.2f} GB",
+                "INT8 (Quantized)": f"{int8_gb:.2f} GB",
+                "params_in_billions": val
+            }
+        except Exception as e:
+            return None
+
+    @staticmethod
+    def get_layer_structure(model):
+        """
+        Returns the raw string representation of the PyTorch model modules.
+        """
+        if model:
+            # We strip the outer wrapper to get straight to the layers
+            return str(model)
+        return "Model not loaded."