Update app.py

app.py

@@ -1,91 +1,130 @@
-import os
+import time
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from datasets import load_dataset, Dataset
+from huggingface_hub import hf_hub_download
+from sentence_transformers import SentenceTransformer, util
+import torch
 
-# ---------- CONFIG ----------
-MODEL_ID = "YOUR_MODEL_ID_HF"  # Replace with your HF model ID
+# ---------------------------
+# CONFIGURATION
+# ---------------------------
+HF_TOKEN = "<YOUR_HF_TOKEN>"  # set your HF token
 DATASET_NAME = "guardian-ai-qna"
-SYSTEM_PROMPT = "You are Guardian AI, a cybersecurity expert. Answer concisely."
-
-# ---------- LOAD TOKENIZER & MODEL ----------
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-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)
+MAX_QUESTIONS = 5          # max questions per TIME_WINDOW
+TIME_WINDOW = 3600         # 1 hour in seconds
+EMBED_MODEL = "all-MiniLM-L6-v2"  # small but effective embedding model
 
-# ---------- LOAD DATASET ----------
+# ---------------------------
+# LOAD OR CREATE DATASET
+# ---------------------------
 try:
-    dataset = load_dataset("huggingface", DATASET_NAME, split="train")
+    dataset = load_dataset(DATASET_NAME, use_auth_token=HF_TOKEN)
+    dataset = dataset["train"]
 except:
     dataset = Dataset.from_dict({"question": [], "answer": []})
 
-# ---------- EMBEDDING HELPER ----------
-from sentence_transformers import SentenceTransformer, util
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
+# ---------------------------
+# EMBEDDING MODEL
+# ---------------------------
+embedder = SentenceTransformer(EMBED_MODEL)
 
-# Cache embeddings in memory
+# Precompute embeddings for existing Q&A
 if len(dataset) > 0:
     dataset_embeddings = embedder.encode(dataset["question"], convert_to_tensor=True)
 else:
-    dataset_embeddings = []
+    dataset_embeddings = torch.empty((0, embedder.get_sentence_embedding_dimension()))
+
+# ---------------------------
+# USER RATE LIMITING
+# ---------------------------
+user_limits = {}
+
+def check_rate_limit(session_id):
+    current_time = time.time()
+    if session_id not in user_limits:
+        user_limits[session_id] = {"count": 0, "start_time": current_time}
+    
+    user_data = user_limits[session_id]
+    if current_time - user_data["start_time"] > TIME_WINDOW:
+        user_data["count"] = 0
+        user_data["start_time"] = current_time
+    
+    if user_data["count"] >= MAX_QUESTIONS:
+        return False, f"You have reached the max of {MAX_QUESTIONS} questions. Please wait before asking more."
+    
+    user_data["count"] += 1
+    return True, None
+
+# ---------------------------
+# HELPER FUNCTIONS
+# ---------------------------
+def find_similar_answer(user_input):
+    if len(dataset) == 0:
+        return None
+    
+    query_emb = embedder.encode(user_input, convert_to_tensor=True)
+    scores = util.cos_sim(query_emb, dataset_embeddings)
+    top_idx = torch.argmax(scores)
+    top_score = scores[0][top_idx].item()
+    
+    if top_score > 0.6:  # threshold for similarity
+        return dataset["answer"][top_idx]
+    return None
 
-# ---------- 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]
+        "question": dataset["question"] + new_entry["question"],
+        "answer": dataset["answer"] + new_entry["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 ""
-    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 ----------
-def chat(history, user_input):
-    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:"
+    
+    # update embeddings incrementally
+    new_emb = embedder.encode([question], convert_to_tensor=True)
+    if len(dataset_embeddings) == 0:
+        dataset_embeddings = new_emb
+    else:
+        dataset_embeddings = torch.vstack([dataset_embeddings, new_emb])
+    
+    # save to HF dataset (push to hub)
+    dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
 
-    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))
+# ---------------------------
+# MAIN CHAT FUNCTION
+# ---------------------------
+def chat(history, user_input, session_id="default"):
+    # Rate limit check
+    allowed, message = check_rate_limit(session_id)
+    if not allowed:
+        history.append(("System", message))
+        return history, history
+    
+    # Check existing similar Q&A
+    response = find_similar_answer(user_input)
+    
+    if not response:
+        # Fallback / simple generative response
+        response = f"Guardian AI: Sorry, I don’t know the answer yet. I’m learning!"
+    
+    # Save new Q&A for incremental learning
     save_qna(user_input, response)
+    
+    # Update chat history
+    history.append((user_input, response))
     return history, history
 
-# ---------- GRADIO APP ----------
+# ---------------------------
+# GRADIO INTERFACE
+# ---------------------------
 with gr.Blocks() as app:
     chatbot = gr.Chatbot()
-    state = gr.State([])
-    with gr.Row():
-        user_msg = gr.Textbox(label="Type your message")
-        send_btn = gr.Button("Send")
-
-    send_btn.click(chat, [state, user_msg], [chatbot, state])
+    msg = gr.Textbox(label="Your question")
+    session_state = gr.State("default")  # default session
+    
+    def user_submit(message, history, session_id):
+        return chat(history, message, session_id)
+    
+    msg.submit(user_submit, inputs=[msg, chatbot, session_state], outputs=[chatbot, chatbot])
 
-app.launch(share=True)
+# Launch app
+app.launch(server_name="0.0.0.0", server_port=7860, share=True)