Update app.py

app.py

@@ -1,34 +1,181 @@
+# app.py
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from peft import PeftModel
+import json
+import os
 
-# Load the tokenizer
-model_name = "Ogero79/threatscope-cyberthreat-analyst"  # Ensure this matches your public model name
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# --- 1. Configuration ---
+adapter_model_name = "Ogero79/threatscope-cyberthreat-analyst" 
+base_model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
 
-# Load model in FP32 on CPU (no quantization)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float32,  # Use FP32
-    low_cpu_mem_usage=True  # Use this to reduce memory usage
+# --- 2. Model Loading ---
+print("--- Loading Model and Tokenizer ---")
+# Load the tokenizer from the adapter repo
+tokenizer = AutoTokenizer.from_pretrained(adapter_model_name)
+
+# Load the base Llama 3 model.
+# We use float16 to save memory on the CPU Space.
+# device_map="auto" will intelligently place the model on the CPU.
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    token=os.environ.get("HF_TOKEN"), # Use the token from Space secrets
 )
 
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    
-    # Reduce max_length to fit CPU memory
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=100)
-    
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Create Gradio Interface
-demo = gr.Interface(
-    fn=generate_response,
-    inputs="text",
-    outputs="text",
-    examples=["Phishing email detected", "Potential DDoS attack"]
+# Load the PEFT adapter and merge it into the base model for faster inference.
+model = PeftModel.from_pretrained(base_model, adapter_model_name)
+model = model.merge_and_unload()
+model.eval()
+
+# Create the text-generation pipeline. device=-1 ensures it runs on CPU.
+generator = pipeline(
+    "text-generation", 
+    model=model, 
+    tokenizer=tokenizer, 
+    device=-1, # Explicitly set to CPU
+    torch_dtype=torch.float16
 )
 
-if __name__ == "__main__":
-    demo.launch()
+print("✅ Model and pipeline loaded successfully!")
+
+
+# --- 3. Inference Function (copied and adapted from your notebook) ---
+def generate_response(prompt_text, max_new_tokens=512, temperature=0.01):
+    # Define the safe/default JSON structure for non-threats
+    safe_default_response = {
+        "summary": "No actionable cybersecurity threat detected",
+        "threat_type": "Non-Threat",
+        "risk_score": 0,
+        "risk_level": "None",
+        "suggested_defense": "No action required",
+        "iocs": [],
+        "threat_actor": "None",
+        "geographical_scope": "None"
+    }
+    
+    messages = [
+        {
+            "role": "system", 
+            "content": (
+                "You are an expert cybersecurity analyst. Analyze input and return JSON with these fields:\n"
+                "- summary: If input describes a threat, summarize it. Otherwise, state no threat detected\n"
+                "- threat_type: Threat category if valid, otherwise 'Non-Threat'\n"
+                "- risk_score: 0-100 (0 for non-threats)\n"
+                "- risk_level: Critical/High/Medium/Low/None\n"
+                "- suggested_defense: Recommendations or 'No action required'\n"
+                "- iocs: Empty list for non-threats\n"
+                "- threat_actor: 'None' for non-threats\n"
+                "- geographical_scope: 'None' for non-threats\n"
+                "For CLEAR non-threats (e.g., 'Hello', weather queries), return the safe default format immediately."
+            )
+        },
+        {
+            "role": "user", 
+            "content": f"Analyze this input for cybersecurity threats: {prompt_text}\n"
+                      f"Return ONLY the JSON output with all fields populated."
+        }
+    ]
+
+    try:
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        outputs = generator(prompt, max_new_tokens=max_new_tokens, temperature=temperature, 
+                          top_p=0.9, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+        generated_full_text = outputs[0]["generated_text"]
+        response = generated__text[len(prompt):].strip()
+
+        # First try to find and parse JSON
+        try:
+            first_brace = response.find('{')
+            last_brace = response.rfind('}')
+            if first_brace == -1 or last_brace == -1:
+                raise ValueError("No JSON detected in response")
+            
+            parsed_json = json.loads(response[first_brace:last_brace+1])
+            
+            # Validate required fields exist
+            required_fields = {
+                'summary': str, 'threat_type': str, 'risk_score': int,
+                'risk_level': str, 'suggested_defense': str, 'iocs': list,
+                'threat_actor': str, 'geographical_scope': str
+            }
+            
+            for field, field_type in required_fields.items():
+                if field not in parsed_json or not isinstance(parsed_json.get(field), field_type):
+                    parsed_json[field] = safe_default_response[field]
+            
+            return parsed_json
+            
+        except (json.JSONDecodeError, ValueError):
+            # If JSON parsing fails, analyze the raw response for threat indicators
+            threat_keywords = ["malware", "attack", "phishing", "breach", "exploit", "hack", "ransomware"]
+            if any(keyword in response.lower() for keyword in threat_keywords):
+                # If threat keywords found but JSON invalid, return error with the raw analysis
+                return {
+                    **safe_default_response,
+                    "summary": f"Potential threat detected but invalid format. Analyst review recommended. Raw response: {response[:200]}...",
+                    "threat_type": "Unknown (Format Error)",
+                    "risk_score": 50,
+                    "risk_level": "Medium"
+                }
+            else:
+                # No threat keywords detected - definitely safe to return default
+                return safe_default_response
+                
+    except Exception as e:
+        # Critical error case - return safe format with error details
+        safe_default_response["summary"] = f"System error: {str(e)}. Default safe response returned"
+        return safe_default_response
+
+# --- 4. Gradio Interface ---
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 900px;
+}
+"""
+
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # 🤖 ThreatScope: AI Cybersecurity Analyst
+            
+            Enter a description of a potential security event below. The fine-tuned Llama 3 model will analyze it and return a structured JSON response with a risk assessment and suggested actions.
+            
+            **Note:** This is an 8B parameter model running on a CPU. The first inference may be slow, but subsequent ones will be faster.
+            """
+        )
+        
+        with gr.Row():
+            prompt_input = gr.Textbox(
+                label="Enter Threat Description", 
+                placeholder="e.g., Our DNS server is being flooded with requests from thousands of botnet IPs.",
+                lines=4
+            )
+
+        analyze_button = gr.Button("Analyze Threat")
+        
+        output_json = gr.JSON(label="Analysis Result")
+        
+        gr.Examples(
+            [
+                "A misconfigured cloud storage bucket exposed sensitive customer data online for months.",
+                "Urgent: Employee received a suspicious email with a malicious attachment claiming to be from HR.",
+                "An ex-employee's credentials were used to log into the main database at 2 AM.",
+                "What's the capital of France?",
+            ],
+            inputs=prompt_input,
+            outputs=output_json,
+            fn=generate_response,
+        )
+
+    analyze_button.click(
+        fn=generate_response,
+        inputs=prompt_input,
+        outputs=output_json
+    )
+
+demo.launch()