Update app.py

app.py

@@ -2,14 +2,14 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # Load model & tokenizer
-model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+model_id = "google/gemma-2b-it"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
 # Create pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
-# System instruction for Guardian AI
+# System instruction
 SYSTEM_PROMPT = """You are Guardian AI, a friendly cybersecurity educator. 
 Your goal is to explain cybersecurity concepts in simple, engaging language with examples. 
 Always keep answers clear, short, and focused on security awareness.
@@ -18,14 +18,19 @@ Always keep answers clear, short, and focused on security awareness.
 # Chat function
 def chat(history, user_input):
     prompt = SYSTEM_PROMPT + "\nUser: " + user_input + "\nGuardian AI:"
-    response = generator(prompt, max_length=500, do_sample=True, temperature=0.7, top_p=0.9)[0]['generated_text']
+    result = generator(
+        prompt,
+        max_new_tokens=200,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )[0]['generated_text']
     
-    # Extract only the part after Guardian AI:
-    response = response.split("Guardian AI:")[-1].strip()
+    response = result.split("Guardian AI:")[-1].strip()
     history.append((user_input, response))
     return history, history
 
-# Build Gradio UI
+# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🛡️ Guardian AI – Cybersecurity Educator")
     chatbot = gr.Chatbot()