Update app.py

app.py

@@ -6,18 +6,22 @@ import random
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 import time
+import warnings
+
+# Suppress warnings for cleaner output
+warnings.filterwarnings("ignore")
 
 # Custom CSS for better styling
 custom_css = """
 .gradio-container {
     max-width: 1200px !important;
 }
-.alert-box {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+.gpt-oss-badge {
+    background: linear-gradient(45deg, #00c6ff, #0072ff);
     color: white;
-    padding: 20px;
-    border-radius: 10px;
-    margin: 10px 0;
+    padding: 5px 10px;
+    border-radius: 15px;
+    font-weight: bold;
 }
 .status-success {
     background: #d4edda;
@@ -26,91 +30,171 @@ custom_css = """
     padding: 10px;
     border-radius: 5px;
 }
-.gpt-oss-badge {
-    background: linear-gradient(45deg, #00c6ff, #0072ff);
-    color: white;
-    padding: 5px 10px;
-    border-radius: 15px;
-    font-weight: bold;
-}
 """
 
 # Global variables for model management
 model = None
 tokenizer = None
-model_status = "🔄 Initializing GPT-OSS-20B..."
+model_status = "🔄 Initializing..."
 
-# Initialize GPT-OSS-20B with proper harmony format
 @spaces.GPU
-def initialize_gpt_oss():
-    """Initialize OpenAI GPT-OSS-20B with harmony response format"""
+def initialize_gpt_oss_safe():
+    """Initialize GPT-OSS-20B with multiple fallback strategies"""
     global model, tokenizer, model_status
     
-    try:
-        model_id = "openai/gpt-oss-20b"
-        print(f"🚀 Loading {model_id}...")
-        
-        # Check GPU availability
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Device: {device}")
-        
-        if torch.cuda.is_available():
-            print(f"GPU: {torch.cuda.get_device_name()}")
-            print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
-        
-        # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-        print("✅ Tokenizer loaded")
-        
-        # Load model with optimized settings for zeroGPU
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype="auto",  # Let it choose best dtype (MXFP4)
-            device_map="auto",   # Automatic GPU placement
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-            # MXFP4 quantization is built-in
-        )
-        
-        print("✅ Model loaded with MXFP4 quantization")
-        print(f"Model device: {next(model.parameters()).device}")
-        
-        # Test generation to ensure everything works
-        test_messages = [
-            {"role": "user", "content": "Hello, test message."}
-        ]
-        
-        test_inputs = tokenizer.apply_chat_template(
-            test_messages,
-            add_generation_prompt=True,
-            return_tensors="pt",
-            return_dict=True,
-        )
-        
-        if device == "cuda":
-            test_inputs = {k: v.to(model.device) for k, v in test_inputs.items()}
-        
-        with torch.no_grad():
-            test_output = model.generate(
-                **test_inputs,
-                max_new_tokens=10,
-                do_sample=False,
-                pad_token_id=tokenizer.eos_token_id
+    # Strategy 1: Try GPT-OSS-20B with specific settings
+    strategies = [
+        {
+            "model_id": "openai/gpt-oss-20b",
+            "name": "GPT-OSS-20B (Original)",
+            "config": {
+                "torch_dtype": "auto",
+                "device_map": "auto",
+                "trust_remote_code": True,
+                "low_cpu_mem_usage": True,
+                "use_safetensors": True
+            }
+        },
+        {
+            "model_id": "openai/gpt-oss-20b", 
+            "name": "GPT-OSS-20B (BF16)",
+            "config": {
+                "torch_dtype": torch.bfloat16,
+                "device_map": "auto", 
+                "trust_remote_code": True,
+                "low_cpu_mem_usage": True
+            }
+        },
+        {
+            "model_id": "openai/gpt-oss-20b",
+            "name": "GPT-OSS-20B (FP16)",
+            "config": {
+                "torch_dtype": torch.float16,
+                "device_map": "auto",
+                "trust_remote_code": True,
+                "low_cpu_mem_usage": True
+            }
+        },
+        {
+            "model_id": "microsoft/DialoGPT-large",
+            "name": "DialoGPT-Large (Fallback)",
+            "config": {
+                "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
+                "device_map": "auto" if torch.cuda.is_available() else None
+            }
+        }
+    ]
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"🔧 Device: {device}")
+    
+    if torch.cuda.is_available():
+        print(f"🎮 GPU: {torch.cuda.get_device_name()}")
+        print(f"💾 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
+    
+    for i, strategy in enumerate(strategies):
+        try:
+            model_id = strategy["model_id"]
+            config = strategy["config"]
+            name = strategy["name"]
+            
+            print(f"\n🔄 Strategy {i+1}: Trying {name}")
+            print(f"📦 Model: {model_id}")
+            
+            # Load tokenizer first
+            print("🔤 Loading tokenizer...")
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_id, 
+                trust_remote_code=True,
+                use_fast=True
             )
-        
-        print("✅ Test generation successful")
-        
-        model_status = f"✅ OpenAI GPT-OSS-20B loaded successfully on {device} | MXFP4 Quantized | ~16GB Memory"
-        return model_status
-        
-    except Exception as e:
-        error_msg = f"❌ Failed to load GPT-OSS-20B: {str(e)}"
-        print(error_msg)
-        model_status = error_msg
-        model, tokenizer = None, None
-        return model_status
+            
+            # Handle pad token for generation
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            
+            print("✅ Tokenizer loaded successfully")
+            
+            # Load model with strategy-specific config
+            print("🧠 Loading model...")
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                **config
+            )
+            
+            print(f"📍 Model device: {next(model.parameters()).device}")
+            print(f"🔢 Model dtype: {next(model.parameters()).dtype}")
+            
+            # Test generation to ensure everything works
+            print("🧪 Testing generation...")
+            test_messages = [{"role": "user", "content": "Hello, test."}]
+            
+            try:
+                # Use pipeline for simpler testing
+                test_pipe = pipeline(
+                    "text-generation",
+                    model=model,
+                    tokenizer=tokenizer,
+                    torch_dtype=config.get("torch_dtype", "auto"),
+                    device_map="auto" if torch.cuda.is_available() else None
+                )
+                
+                test_result = test_pipe(
+                    test_messages,
+                    max_new_tokens=10,
+                    do_sample=False
+                )
+                
+                print("✅ Generation test successful")
+                
+            except Exception as test_error:
+                print(f"⚠️ Pipeline test failed, trying direct generation: {test_error}")
+                
+                # Fallback to direct generation
+                inputs = tokenizer.apply_chat_template(
+                    test_messages,
+                    add_generation_prompt=True,
+                    return_tensors="pt",
+                    return_dict=True,
+                )
+                
+                if torch.cuda.is_available():
+                    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+                
+                with torch.no_grad():
+                    outputs = model.generate(
+                        **inputs,
+                        max_new_tokens=5,
+                        do_sample=False,
+                        pad_token_id=tokenizer.eos_token_id
+                    )
+                
+                print("✅ Direct generation test successful")
+            
+            # If we get here, the strategy worked
+            model_status = f"✅ {name} loaded successfully on {device}"
+            print(f"🎉 Success: {model_status}")
+            return model_status
+            
+        except Exception as e:
+            error_msg = str(e)
+            print(f"❌ Strategy {i+1} failed: {error_msg[:100]}...")
+            
+            # Clear any partially loaded components
+            model, tokenizer = None, None
+            
+            # Clear GPU memory if available
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            continue
+    
+    # If all strategies failed
+    model_status = "❌ All model loading strategies failed. Using text-only fallback."
+    print(model_status)
+    return model_status
 
-# Enhanced attack scenarios
+# Enhanced attack scenarios (keeping the same as before)
 ATTACK_SCENARIOS = {
     "🔄 Lateral Movement": {
         "description": "Advanced Persistent Threat (APT) - Attacker moving laterally through network after initial compromise",
@@ -231,134 +315,129 @@ ATTACK_SCENARIOS = {
 }
 
 @spaces.GPU
-def generate_gpt_oss_analysis(alert_data, analyst_level):
-    """Generate analysis using OpenAI GPT-OSS-20B with harmony format"""
+def generate_analysis_safe(alert_data, analyst_level):
+    """Generate analysis with safe error handling"""
     
     if not model or not tokenizer:
         return get_fallback_analysis(alert_data, analyst_level)
     
-    # Enhanced prompts designed for GPT-OSS reasoning capabilities
     security_prompts = {
-        "L1": f"""You are a Level 1 SOC analyst conducting initial triage. Analyze this security alert and provide immediate actionable recommendations.
-
-**SECURITY ALERT:**
-- ID: {alert_data['id']}
-- Type: {alert_data['alert_type']} 
-- Severity: {alert_data['severity']}
-- Source: {alert_data['source_ip']} → {alert_data['destination_ip']}
-- User: {alert_data['user']}
-- Evidence: {alert_data['raw_log']}
-- Intel: {alert_data['threat_intel']}
-- MITRE ATT&CK: {alert_data['mitre_tactic']}
-- Confidence: {alert_data['confidence']}%
-
-**PROVIDE L1 TRIAGE:**
-1. Immediate containment actions
-2. Risk assessment
-3. Escalation decision with reasoning
-4. Priority timeline
-
-Think step-by-step about the threat level and required response.""",
-
-        "L2": f"""You are a Level 2 SOC analyst conducting detailed investigation. Perform comprehensive analysis of this cybersecurity incident.
-
-**INCIDENT DETAILS:**
-- Alert: {alert_data['alert_type']} | Severity: {alert_data['severity']}
-- Network Flow: {alert_data['source_ip']} → {alert_data['destination_ip']}
-- User Context: {alert_data['user']}
-- Technical Evidence: {alert_data['raw_log']}
-- Threat Intelligence: {alert_data['threat_intel']}
-- MITRE ATT&CK Technique: {alert_data['mitre_tactic']}
-- Detection Confidence: {alert_data['confidence']}%
-
-**CONDUCT L2 INVESTIGATION:**
-1. Technical root cause analysis
-2. Evidence correlation and timeline
-3. Threat actor behavior analysis
-4. Impact assessment and containment strategy
-5. Investigation roadmap
-
-Use chain-of-thought reasoning to analyze the attack progression and recommend next steps.""",
-
-        "L3": f"""You are a senior cybersecurity expert analyzing a sophisticated threat. Provide strategic assessment and executive-level recommendations.
-
-**THREAT INTELLIGENCE:**
-- Attack Vector: {alert_data['description']}
-- Technical Indicators: {alert_data['raw_log']}
-- Attribution Context: {alert_data['threat_intel']}
-- MITRE Technique: {alert_data['mitre_tactic']}
-- Confidence Level: {alert_data['confidence']}%
-
-**DELIVER L3 EXPERT ANALYSIS:**
-1. Adversary attribution and campaign analysis
-2. Strategic threat landscape assessment
-3. Business impact and risk quantification
-4. Comprehensive response strategy
-5. Executive briefing points
-
-Apply deep reasoning to assess the broader implications and provide strategic recommendations."""
+        "L1": f"""As a Level 1 SOC analyst, provide immediate triage for this security alert:
+
+Alert: {alert_data['alert_type']} | Severity: {alert_data['severity']}
+Source: {alert_data['source_ip']} → {alert_data['destination_ip']}
+User: {alert_data['user']}
+Evidence: {alert_data['raw_log']}
+MITRE: {alert_data['mitre_tactic']}
+
+Provide: immediate actions, risk assessment, escalation decision.""",
+
+        "L2": f"""As a Level 2 SOC analyst, investigate this cybersecurity incident:
+
+Alert: {alert_data['alert_type']} | Severity: {alert_data['severity']}
+Network: {alert_data['source_ip']} → {alert_data['destination_ip']}
+User: {alert_data['user']}
+Evidence: {alert_data['raw_log']}
+Intel: {alert_data['threat_intel']}
+MITRE: {alert_data['mitre_tactic']}
+Confidence: {alert_data['confidence']}%
+
+Provide: root cause analysis, investigation steps, threat assessment.""",
+
+        "L3": f"""As a senior cybersecurity expert, analyze this sophisticated threat:
+
+Attack: {alert_data['description']}
+Evidence: {alert_data['raw_log']}
+Attribution: {alert_data['threat_intel']}
+MITRE: {alert_data['mitre_tactic']}
+Confidence: {alert_data['confidence']}%
+
+Provide: strategic assessment, attribution analysis, response plan."""
     }
     
     try:
         prompt = security_prompts.get(analyst_level, security_prompts["L2"])
         
-        # Use proper harmony format for chat
-        messages = [
-            {"role": "user", "content": prompt}
-        ]
-        
-        # Apply chat template (automatically uses harmony format)
-        inputs = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            return_tensors="pt",
-            return_dict=True,
-        )
-        
-        # Move to device if using GPU
-        if torch.cuda.is_available():
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
-        
-        # Generate with optimized parameters for reasoning
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=500,
+        # Try pipeline approach first (safer)
+        try:
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                torch_dtype="auto",
+                device_map="auto" if torch.cuda.is_available() else None
+            )
+            
+            messages = [{"role": "user", "content": prompt}]
+            
+            result = pipe(
+                messages,
+                max_new_tokens=400,
                 do_sample=True,
-                temperature=0.2,  # Lower for focused analysis
+                temperature=0.3,
                 top_p=0.9,
-                top_k=50,
                 repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id
             )
+            
+            analysis = result[0]["generated_text"][-1]["content"]
+            
+        except Exception as pipe_error:
+            print(f"Pipeline failed, trying direct generation: {pipe_error}")
+            
+            # Fallback to direct generation
+            messages = [{"role": "user", "content": prompt}]
+            
+            inputs = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                return_tensors="pt",
+                return_dict=True,
+            )
+            
+            if torch.cuda.is_available():
+                inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=400,
+                    do_sample=True,
+                    temperature=0.3,
+                    top_p=0.9,
+                    repetition_penalty=1.1,
+                    pad_token_id=tokenizer.eos_token_id,
+                    eos_token_id=tokenizer.eos_token_id
+                )
+            
+            input_length = inputs["input_ids"].shape[-1]
+            generated_tokens = outputs[0][input_length:]
+            analysis = tokenizer.decode(generated_tokens, skip_special_tokens=True)
         
-        # Decode the response
-        input_length = inputs["input_ids"].shape[-1]
-        generated_tokens = outputs[0][input_length:]
-        analysis = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-        
-        # Ensure quality
-        if len(analysis.strip()) < 100:
+        # Quality check
+        if len(analysis.strip()) < 50:
             return get_fallback_analysis(alert_data, analyst_level)
         
-        return f"""🤖 **OpenAI GPT-OSS-20B Analysis**
-<div class="gpt-oss-badge">Powered by GPT-OSS-20B • MoE Architecture • MXFP4 Quantized</div>
-
+        # Determine model name for display
+        if "gpt-oss" in model_status.lower():
+            badge = "🤖 OpenAI GPT-OSS-20B Analysis"
+        else:
+            badge = "🤖 AI-Powered Analysis"
+        
+        return f"""{badge}
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
 {analysis.strip()}
 
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
-*Analysis generated using OpenAI's latest open-weight reasoning model*
-*21B parameters • 3.6B active per token • Apache 2.0 licensed*"""
+*Generated using advanced AI reasoning capabilities*"""
         
     except Exception as e:
-        print(f"GPT-OSS Error: {e}")
-        return f"⚠️ GPT-OSS Error: {str(e)[:100]}\n\n{get_fallback_analysis(alert_data, analyst_level)}"
+        print(f"Generation error: {e}")
+        return f"⚠️ AI Error: {str(e)[:100]}\n\n{get_fallback_analysis(alert_data, analyst_level)}"
 
 def get_fallback_analysis(alert_data, analyst_level):
-    """High-quality fallback when model fails"""
+    """High-quality fallback analysis"""
     
     templates = {
         "L1": f"""🚨 **L1 SOC TRIAGE ANALYSIS**
@@ -372,39 +451,39 @@ Confidence: {alert_data['confidence']}% | Source: {alert_data['source_ip']}
 1. Isolate affected system: {alert_data['source_ip']}
 2. Disable user account: {alert_data['user']}
 3. Block connections to: {alert_data['destination_ip']}
-4. Preserve evidence for investigation
+4. Preserve evidence and logs
 
 **⬆️ ESCALATION DECISION:**
 Severity: {alert_data['severity']} → ESCALATE TO L2
 Technique: {alert_data['mitre_tactic']} requires deeper analysis
 
-**📝 INITIAL ASSESSMENT:**
+**📝 INITIAL NOTES:**
 {alert_data['threat_intel']}""",
 
         "L2": f"""🔍 **L2 INVESTIGATION ANALYSIS**
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 
-**🎯 ATTACK VECTOR ANALYSIS:**
+**🎯 ATTACK VECTOR:**
 Technique: {alert_data['mitre_tactic']}
 Evidence: {alert_data['raw_log']}
 Context: {alert_data['description']}
 
 **🔬 INVESTIGATION ROADMAP:**
-1. Timeline correlation: ±30min from {alert_data['timestamp']}
-2. User behavior analysis: {alert_data['user']} baseline comparison
+1. Timeline correlation: ±30min window analysis
+2. User behavior baseline: {alert_data['user']} comparison
 3. Network flow analysis: {alert_data['source_ip']} → {alert_data['destination_ip']}
-4. Process tree examination: Parent/child relationships
-5. Artifact collection: Memory dumps, logs, files
+4. Process tree examination and artifact collection
+5. Similar IOC hunting across environment
 
 **📊 THREAT ASSESSMENT:**
 Confidence Level: {alert_data['confidence']}%
 Business Impact: {alert_data['severity']}
-Attribution: {alert_data['threat_intel']}
+Attribution Context: {alert_data['threat_intel']}
 
-**🎯 RECOMMENDATIONS:**
+**🎯 NEXT STEPS:**
 Deploy hunting queries for similar TTPs
-Review authentication logs for compromise
-Consider L3 escalation if campaign indicators found""",
+Review authentication logs for compromise indicators
+Consider L3 escalation if campaign evidence found""",
 
         "L3": f"""🎯 **L3 EXPERT STRATEGIC ANALYSIS**
 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
@@ -416,26 +495,26 @@ Success Probability: {alert_data['confidence']}%
 
 **💼 BUSINESS IMPACT:**
 Severity Level: {alert_data['severity']}
-Executive Notification: Required
-Regulatory Implications: Under review
+Executive Notification: Required for Critical/High
+Regulatory Implications: Under compliance review
 
 **🛡️ STRATEGIC RESPONSE:**
-Immediate: Threat hunting deployment across environment
-Short-term: Enhanced monitoring and detection rules
-Medium-term: Security architecture review
-Long-term: Threat intelligence integration enhancement
+Immediate: Threat hunting deployment across infrastructure
+Short-term: Enhanced monitoring and detection rule updates
+Medium-term: Security architecture review and gap analysis
+Long-term: Threat intelligence integration and training
 
-**📈 EXECUTIVE BRIEFING:**
-• Sophisticated attack requiring coordinated response
+**📈 EXECUTIVE BRIEFING POINTS:**
+• Sophisticated attack requiring coordinated incident response
 • High potential for lateral movement and data exfiltration
-• Recommend incident response team activation
-• Consider external forensics support engagement"""
+• Recommend immediate incident response team activation
+• Consider external forensics engagement for complex analysis"""
     }
     
     return templates.get(analyst_level, templates["L2"])
 
-def analyze_alert_with_gpt_oss(scenario_name, alert_index, analyst_level):
-    """Main analysis function using GPT-OSS-20B"""
+def analyze_alert_with_ai(scenario_name, alert_index, analyst_level):
+    """Main analysis function with error handling"""
     start_time = time.time()
     
     if scenario_name not in ATTACK_SCENARIOS:
@@ -449,8 +528,8 @@ def analyze_alert_with_gpt_oss(scenario_name, alert_index, analyst_level):
     
     selected_alert = alerts[alert_index]
     
-    # Generate analysis using GPT-OSS-20B
-    analysis = generate_gpt_oss_analysis(selected_alert, analyst_level)
+    # Generate analysis
+    analysis = generate_analysis_safe(selected_alert, analyst_level)
     
     # Format alert details
     alert_details = f"""🎫 **ALERT {selected_alert['id']}** | 🕐 {selected_alert['timestamp']}
@@ -507,27 +586,27 @@ def get_scenario_info(scenario_name):
 """
         
         info += """
-### 🤖 **AI-Powered Analysis:**
-- **OpenAI GPT-OSS-20B:** Latest open-weight reasoning model
-- **MXFP4 Quantization:** Optimized for efficient inference
-- **Harmony Format:** Advanced response structure
-- **21B Parameters:** With 3.6B active per token (MoE)"""
+### 🤖 **AI Analysis Features:**
+- **Multi-Strategy Loading:** Automatic fallback to compatible models
+- **Advanced Reasoning:** Chain-of-thought analysis for complex threats
+- **Error Recovery:** Robust handling of model loading issues
+- **Quality Assurance:** Automatic fallback to expert templates"""
         
         return info
     return "⚠️ No scenario selected."
 
 # Create Gradio interface
-with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=custom_css) as demo:
+with gr.Blocks(title="SOC Assistant - Fixed GPT-OSS", theme=gr.themes.Soft(), css=custom_css) as demo:
     
     gr.Markdown("""
-    # 🛡️ SOC LLM Assistant - OpenAI GPT-OSS-20B Edition
-    **Powered by OpenAI's Latest Open-Weight Reasoning Model**
+    # 🛡️ SOC LLM Assistant - Fixed GPT-OSS Edition
+    **Multi-Strategy Model Loading with Robust Error Handling**
     
-    *First open-weight model from OpenAI since GPT-2 • Released August 8, 2025*
+    *Automatically tries GPT-OSS-20B first, then falls back to compatible models*
     """)
     
     # Model status display
-    status_display = gr.Markdown("🔄 Loading OpenAI GPT-OSS-20B...")
+    status_display = gr.Markdown("🔄 Initializing AI models with multiple strategies...")
     
     with gr.Row():
         # Left Panel
@@ -562,26 +641,26 @@ with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=
             )
             
             analyze_btn = gr.Button(
-                "🚀 Analyze with GPT-OSS-20B", 
+                "🚀 Analyze with AI", 
                 variant="primary",
                 size="lg"
             )
             
             init_btn = gr.Button(
-                "🔄 Reinitialize Model",
+                "🔄 Retry Model Loading",
                 variant="secondary"
             )
             
             gr.Markdown("---")
-            gr.Markdown("## 🤖 Model Information")
+            gr.Markdown("## 🔧 Loading Strategies")
             gr.Markdown("""
-            **🎯 GPT-OSS-20B Features:**
-            - 21B parameters (3.6B active)
-            - MXFP4 quantization
-            - 128K context length
-            - Apache 2.0 licensed
-            - Harmony response format
-            - Reasoning capabilities
+            **🎯 Automatic Fallback:**
+            1. GPT-OSS-20B (Original MXFP4)
+            2. GPT-OSS-20B (BF16) 
+            3. GPT-OSS-20B (FP16)
+            4. DialoGPT-Large (Backup)
+            
+            **✅ Robust Error Handling**
             """)
         
         # Right Panel
@@ -593,9 +672,9 @@ with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=
                 interactive=False
             )
             
-            gr.Markdown("## 🤖 GPT-OSS-20B Analysis")
+            gr.Markdown("## 🤖 AI-Powered Analysis")
             analysis_output = gr.Textbox(
-                label="🧠 AI-Powered Security Analysis",
+                label="🧠 Security Analysis",
                 lines=25,
                 interactive=False
             )
@@ -608,17 +687,17 @@ with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=
     
     gr.Markdown("""
     ---
-    ## 🎉 **About OpenAI GPT-OSS-20B**
-    
-    Released August 8, 2025 - OpenAI's first open-weight model since GPT-2! This groundbreaking release features:
+    ## 🔧 **Troubleshooting Guide**
     
-    - **🧠 Advanced Reasoning:** Comparable to o3-mini performance
-    - **⚡ Efficient Architecture:** MoE with only 3.6B active parameters per token
-    - **🔧 Harmony Format:** New structured response system for better tool use
-    - **📱 Consumer Hardware:** Runs on just 16GB memory
-    - **🔓 Open License:** Apache 2.0 - fully permissive for commercial use
+    **If you see "ModelWrapper" error:**
+    - ✅ **Fixed:** This version uses multiple loading strategies
+    - 🔄 **Automatic:** Falls back to compatible models
+    - 🛠️ **Manual:** Use "Retry Model Loading" button
     
-    Perfect for cybersecurity analysis requiring sophisticated reasoning and chain-of-thought capabilities!
+    **Loading Strategy Order:**
+    1. **GPT-OSS-20B** - Latest OpenAI open-weight model
+    2. **Fallback Models** - Tested compatible alternatives  
+    3. **Expert Templates** - High-quality manual analysis
     
     ---
     **👨‍🎓 Research:** Abdullah Alanazi | **🏛️ KAUST** | **👨‍🏫 Prof. Ali Shoker**
@@ -644,13 +723,13 @@ with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=
     )
     
     analyze_btn.click(
-        fn=analyze_alert_with_gpt_oss,
+        fn=analyze_alert_with_ai,
         inputs=[scenario_dropdown, alert_slider, analyst_level],
         outputs=[alert_output, analysis_output, status_output]
     )
     
     init_btn.click(
-        fn=initialize_gpt_oss,
+        fn=initialize_gpt_oss_safe,
         outputs=[status_display]
     )
     
@@ -662,7 +741,7 @@ with gr.Blocks(title="SOC Assistant - GPT-OSS-20B", theme=gr.themes.Soft(), css=
     )
     
     demo.load(
-        fn=initialize_gpt_oss,
+        fn=initialize_gpt_oss_safe,
         outputs=[status_display]
     )