Update app.py

app.py

@@ -1,30 +1,51 @@
 import gradio as gr
-import pandas as pd
 import torch
+import pandas as pd
 
 from core.model_loader import load_model, SUPPORTED_MODELS
 from core.profiler import ActivationProfiler
 from core.stress import stress_layer
 from core.sensitivity import compute_si, assign_policy
-from core.reevaluator import reevaluate
 from core.visualization import create_heatmap
 
-def run_cortex(model_choice, calibration_text):
-    try:
-        model, tokenizer = load_model(model_choice)
+# =====================================================
+# 🔥 LOAD MODEL ONCE AT STARTUP (IMPORTANT)
+# =====================================================
+
+print("Loading model at startup...")
+
+MODEL_CHOICE_DEFAULT = "DistilGPT2 (Fast CPU)"
+model, tokenizer = load_model(MODEL_CHOICE_DEFAULT)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model.eval()
 
-        # Fix GPT2 padding issue
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
+print("Model loaded successfully.")
 
+# =====================================================
+# MAIN FUNCTION
+# =====================================================
+
+def run_cortex(calibration_text):
+
+    try:
         texts = calibration_text.strip().split("\n")
         texts = [t.strip() for t in texts if t.strip() != ""]
 
         if len(texts) == 0:
             return {"error": "No valid input sentences provided."}, None, 0.0, None
 
-        inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
+        inputs = tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=64
+        )
 
+        # ===== Activation Profiling =====
         profiler = ActivationProfiler(model)
         profiler.register()
 
@@ -34,22 +55,56 @@ def run_cortex(model_choice, calibration_text):
         profiler.remove()
         stats = profiler.get()
 
+        # ===== SAFE LIMITED STRESS TEST =====
         stress_results = {}
+        layer_names = list(stats.keys())[:3]   # 🔥 LIMIT TO 3 LAYERS
 
-        # 🔥 LIMIT layers for stability (important for HF CPU)
-        for name in list(stats.keys())[:6]:
+        for name in layer_names:
             mse_val, kl_val = stress_layer(model, name, inputs, bits=8)
             stress_results[name] = (mse_val, kl_val)
 
+        # ===== Sensitivity Index =====
         scores = compute_si(stats, stress_results)
         policy = assign_policy(scores)
 
-        policy, degradation = reevaluate(model, tokenizer, texts, policy)
+        # Skip heavy reevaluation for HF CPU
+        degradation = 0.0
 
         heatmap = create_heatmap(scores)
 
-        return policy, heatmap, float(degradation), "layer_metrics.csv"
+        df = pd.DataFrame({
+            "Layer": list(scores.keys()),
+            "Sensitivity_Index": list(scores.values()),
+            "Policy": [policy.get(k, "FP16") for k in scores]
+        })
+
+        csv_path = "layer_metrics.csv"
+        df.to_csv(csv_path, index=False)
+
+        return policy, heatmap, degradation, csv_path
 
     except Exception as e:
         return {"error": str(e)}, None, 0.0, None
 
+
+# =====================================================
+# GRADIO UI
+# =====================================================
+
+demo = gr.Interface(
+    fn=run_cortex,
+    inputs=gr.Textbox(
+        label="Calibration Text (One sentence per line)",
+        lines=6
+    ),
+    outputs=[
+        gr.JSON(label="Precision Policy"),
+        gr.Image(label="Critical Layer Map"),
+        gr.Number(label="Model Degradation (Approx)"),
+        gr.File(label="Download CSV Metrics")
+    ],
+    title="Cortex in the Loop — HF Stable Mode",
+    description="Runtime Sensitivity Analyzer (HF CPU Safe Version)"
+)
+
+demo.launch(server_name="0.0.0.0", server_port=7860)