Update app.py

app.py

@@ -1,76 +1,96 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from datasets import load_dataset, Dataset
-from huggingface_hub import login
+from datasets import load_dataset, Dataset, concatenate_datasets
 import os
 
-# --- Hugging Face Dataset Setup ---
-HF_TOKEN = os.environ.get("dataset_HF_TOKEN")  # Secret in your HF Space
-login(token=HF_TOKEN)
+# -------------------------------
+# Config
+# -------------------------------
+HF_TOKEN = os.environ["dataset_HF_TOKEN"]
+DATASET_ID = "your-username/guardian-ai-qna"  # replace with your HF username
+MODEL_ID = "google/gemma-2b-it"
 
-dataset_name = "YOUR_USERNAME/guardian-ai-qna"  # Replace YOUR_USERNAME
-try:
-    dataset = load_dataset(dataset_name)
-except:
-    # If dataset is empty or not yet created, create an empty one
-    dataset = Dataset.from_dict({"question": [], "answer": []})
-
-# --- Load model & tokenizer ---
-model_id = "google/gemma-2b-it"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-
-generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=-1  # CPU, change to 0 if GPU available
-)
-
-# --- 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.
+Use the examples from the Q&A memory to improve your answers.
 """
 
-# --- Save Q&A to dataset ---
-def save_qna(question, answer):
-    global dataset
-    new_entry = Dataset.from_dict({"question": [question], "answer": [answer]})
-    dataset = dataset.concat(new_entry)
-    dataset.push_to_hub(dataset_name, private=False)  # push updates
+# -------------------------------
+# Load model & tokenizer
+# -------------------------------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+
+# -------------------------------
+# Dataset functions
+# -------------------------------
+def load_qna_dataset():
+    try:
+        dataset = load_dataset(DATASET_ID, use_auth_token=HF_TOKEN)["train"]
+    except:
+        dataset = Dataset.from_dict({"question": [], "answer": []})
+    return dataset
+
+def save_qna(user_input, response):
+    dataset = load_qna_dataset()
+    new_entry = Dataset.from_dict({"question": [user_input], "answer": [response]})
+    dataset = concatenate_datasets([dataset, new_entry])
+    dataset.push_to_hub(DATASET_ID, token=HF_TOKEN)
 
-# --- Chat function ---
+def retrieve_similar_qna(user_input, top_k=3):
+    dataset = load_qna_dataset()
+    if len(dataset) == 0:
+        return ""
+    # Simple keyword-based retrieval
+    # You can upgrade to semantic search later
+    relevant = []
+    for q, a in zip(dataset["question"], dataset["answer"]):
+        if any(word in user_input.lower() for word in q.lower().split()):
+            relevant.append(f"Q: {q}\nA: {a}")
+        if len(relevant) >= top_k:
+            break
+    return "\n".join(relevant)
+
+# -------------------------------
+# Chat function
+# -------------------------------
 def chat(history, user_input):
-    prompt = SYSTEM_PROMPT + "\nUser: " + user_input + "\nGuardian AI:"
+    # Retrieve past Q&A for context
+    context = retrieve_similar_qna(user_input)
+    prompt = SYSTEM_PROMPT
+    if context:
+        prompt += f"\n\nMemory of past Q&A:\n{context}"
+    prompt += f"\n\nUser: {user_input}\nGuardian AI:"
+
     result = generator(
         prompt,
         max_new_tokens=200,
         do_sample=True,
         temperature=0.7,
         top_p=0.9
-    )[0]['generated_text']
-    
+    )[0]["generated_text"]
+
     response = result.split("Guardian AI:")[-1].strip()
     history.append((user_input, response))
-    
-    # Save to dataset
     save_qna(user_input, response)
-    
     return history, history
 
-# --- Gradio UI ---
+# -------------------------------
+# Gradio UI
+# -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("## 🛡️ Guardian AI – Cybersecurity Educator")
-    chatbot = gr.Chatbot()
+    chatbot = gr.Chatbot(type="messages")  # Updated type to avoid deprecation warning
     state = gr.State([])
-    
+
     with gr.Row():
         with gr.Column(scale=8):
             user_input = gr.Textbox(show_label=False, placeholder="Ask me about cybersecurity...")
         with gr.Column(scale=2):
             send_btn = gr.Button("Send")
-    
+
     send_btn.click(chat, [state, user_input], [chatbot, state])
     user_input.submit(chat, [state, user_input], [chatbot, state])