Update rag.py

rag.py

@@ -3,6 +3,7 @@ from sentence_transformers import SentenceTransformer, util
 from groq import Groq
 import datetime
 import requests
+from datasets import load_dataset, Dataset
 from io import BytesIO
 from PIL import Image, ImageDraw, ImageFont
 import numpy as np
@@ -27,6 +28,23 @@ dataset_questions = [item.get("Question", "").lower().strip() for item in datase
 dataset_answers = [item.get("Answer", "") for item in dataset]
 dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)
 
+# Save unmatched queries to Hugging Face
+def manage_unmatched_queries(query: str):
+    try:
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        try:
+            ds = load_dataset(HF_DATASET_REPO, token=HF_TOKEN)
+            df = ds["train"].to_pandas()
+        except:
+            df = pd.DataFrame(columns=["Query", "Timestamp", "Processed"])
+        if query not in df["Query"].values:
+            new_entry = {"Query": query, "Timestamp": timestamp, "Processed": False}
+            df = pd.concat([df, pd.DataFrame([new_entry])], ignore_index=True)
+            updated_ds = Dataset.from_pandas(df)
+            updated_ds.push_to_hub(HF_DATASET_REPO, token=HF_TOKEN)
+    except Exception as e:
+        print(f"Failed to save query: {e}")
+
 def query_groq_llm(prompt, model_name="llama3-70b-8192"):
     try:
         chat_completion = groq_client.chat.completions.create(
@@ -60,6 +78,9 @@ def get_best_answer(user_input):
     best_match_idx = similarities.argmax().item()
     best_score = similarities[best_match_idx].item()
 
+    if best_score < 0.65:
+        manage_unmatched_queries(user_input)
+    
     if best_score >= 0.65:
         original_answer = dataset_answers[best_match_idx]
         prompt = f"""As an official assistant for University of Education Lahore, provide a clear response: