Update app.py

app.py

@@ -1,97 +1,91 @@
+import os
 import gradio as gr
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from datasets import load_dataset, Dataset, concatenate_datasets
-import os
-
-# -------------------------------
-# 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"
+from datasets import load_dataset, Dataset
 
-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.
-"""
+# ---------- CONFIG ----------
+MODEL_ID = "YOUR_MODEL_ID_HF"  # Replace with your HF model ID
+DATASET_NAME = "guardian-ai-qna"
+SYSTEM_PROMPT = "You are Guardian AI, a cybersecurity expert. Answer concisely."
 
-# -------------------------------
-# Load model & tokenizer
-# -------------------------------
+# ---------- LOAD TOKENIZER & MODEL ----------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
+device = 0 if torch.cuda.is_available() else -1
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
-# -------------------------------
-# 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
+# ---------- LOAD DATASET ----------
+try:
+    dataset = load_dataset("huggingface", DATASET_NAME, split="train")
+except:
+    dataset = Dataset.from_dict({"question": [], "answer": []})
 
-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)
+# ---------- EMBEDDING HELPER ----------
+from sentence_transformers import SentenceTransformer, util
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
-def retrieve_similar_qna(user_input, top_k=3):
-    dataset = load_qna_dataset()
+# Cache embeddings in memory
+if len(dataset) > 0:
+    dataset_embeddings = embedder.encode(dataset["question"], convert_to_tensor=True)
+else:
+    dataset_embeddings = []
+
+# ---------- SAVE QNA FUNCTION ----------
+def save_qna(question, answer):
+    global dataset, dataset_embeddings
+    new_entry = Dataset.from_dict({"question": [question], "answer": [answer]})
+    dataset = Dataset.from_dict({
+        "question": dataset["question"] + [question],
+        "answer": dataset["answer"] + [answer]
+    })
+    # update embeddings
+    dataset_embeddings.append(embedder.encode(question, convert_to_tensor=True))
+    # push to HF dataset
+    dataset.push_to_hub(DATASET_NAME, token=os.environ.get("HF_TOKEN"))
+
+# ---------- RETRIEVE SIMILAR QNA ----------
+def retrieve_similar_qna(query, top_k=3):
     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)
+    query_emb = embedder.encode(query, convert_to_tensor=True)
+    similarities = util.cos_sim(query_emb, dataset_embeddings)[0]
+    top_results = similarities.topk(k=min(top_k, len(similarities)))
+    context = ""
+    for idx in top_results.indices:
+        context += f"Q: {dataset[idx]['question']}\nA: {dataset[idx]['answer']}\n"
+    return context
 
-# -------------------------------
-# Chat function
-# -------------------------------
+# ---------- CHAT FUNCTION ----------
 def chat(history, user_input):
-    # 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"]
+    with torch.no_grad():
+        result = generator(
+            prompt,
+            max_new_tokens=150,
+            do_sample=True,
+            temperature=0.6,
+            top_p=0.85
+        )[0]["generated_text"]
 
     response = result.split("Guardian AI:")[-1].strip()
     history.append((user_input, response))
     save_qna(user_input, response)
     return history, history
 
-# -------------------------------
-# Gradio UI
-# -------------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## 🛡️ Guardian AI – Cybersecurity Educator")
-    chatbot = gr.Chatbot(type="messages")  # Updated type to avoid deprecation warning
+# ---------- GRADIO APP ----------
+with gr.Blocks() as app:
+    chatbot = gr.Chatbot()
     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")
+        user_msg = gr.Textbox(label="Type your message")
+        send_btn = gr.Button("Send")
 
-    send_btn.click(chat, [state, user_input], [chatbot, state])
-    user_input.submit(chat, [state, user_input], [chatbot, state])
+    send_btn.click(chat, [state, user_msg], [chatbot, state])
 
-demo.launch()
+app.launch(share=True)