Update app.py

app.py

@@ -4,11 +4,14 @@ from datetime import datetime, timedelta
 import gradio as gr
 from datasets import load_dataset, Dataset, DatasetDict
 from huggingface_hub import HfFolder
+from sentence_transformers import SentenceTransformer, util
+import torch
+import requests
 
 # ================================
 # CONFIG
 # ================================
-MODEL_TOKEN = os.environ.get("HF_TOKEN")         # for model usage
+MODEL_TOKEN = os.environ.get("HF_TOKEN")            # for model usage
 DATASET_TOKEN = os.environ.get("dataset_HF_TOKEN")  # for dataset updates
 DATASET_NAME = "guardian-ai-qna"
 
@@ -25,7 +28,10 @@ HfFolder.save_token(DATASET_TOKEN)
 try:
     dataset = load_dataset(DATASET_NAME, use_auth_token=DATASET_TOKEN)
 except:
-    dataset = DatasetDict({"train": Dataset.from_dict({"question": [], "answer": []})})
+    dataset = DatasetDict({"train": Dataset.from_dict({"question": [], "answer": [], "embedding": []})})
+
+# Load embedding model
+embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
 # ================================
 # HELPER FUNCTIONS
@@ -48,17 +54,27 @@ def log_query(user_id):
     now = datetime.now()
     user_queries.setdefault(user_id, []).append(now)
 
-def find_in_dataset(question):
+def find_in_dataset(question, threshold=0.75):
     if len(dataset["train"]) == 0:
         return None
-    for entry in dataset["train"]:
-        if question.strip().lower() == entry["question"].strip().lower():
-            return entry["answer"]
+    # Compute embedding for input
+    question_emb = embed_model.encode(question, convert_to_tensor=True)
+    # Load existing embeddings
+    existing_embs = torch.tensor(dataset["train"]["embedding"]) if dataset["train"]["embedding"] else None
+    if existing_embs is None or len(existing_embs) == 0:
+        return None
+    # Compute cosine similarities
+    similarities = util.cos_sim(question_emb, existing_embs)[0]
+    max_score, idx = torch.max(similarities, dim=0)
+    if max_score >= threshold:
+        return dataset["train"]["answer"][idx.item()]
     return None
 
 def save_qna(question, answer):
     global dataset
-    new_entry = {"question": [question], "answer": [answer]}
+    # Compute embedding
+    emb = embed_model.encode(question).tolist()
+    new_entry = {"question": [question], "answer": [answer], "embedding": [emb]}
     new_ds = Dataset.from_dict(new_entry)
     dataset["train"] = dataset["train"].concatenate(new_ds)
     dataset["train"].push_to_hub(DATASET_NAME, token=DATASET_TOKEN)
@@ -68,7 +84,6 @@ def call_render(question):
     Replace this with your actual Render API call logic
     that fetches the answer from the internet.
     """
-    import requests
     RENDER_API_URL = os.environ.get("RENDER_API_URL")
     if not RENDER_API_URL:
         return "Render API not configured."
@@ -89,7 +104,7 @@ def chat(history, message, session_id):
 
     log_query(session_id)
 
-    # Check dataset first
+    # Check dataset first (embedding-based)
     response = find_in_dataset(message)
     if response is None:
         # Call Render API fallback