Update rag.py

rag.py

@@ -1,11 +1,15 @@
 import json
 from sentence_transformers import SentenceTransformer, util
 from groq import Groq
-from datetime import datetime
+import datetime
+import requests
+from io import BytesIO
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+from dotenv import load_dotenv
 import os
 import pandas as pd
 from datasets import load_dataset, Dataset
-from dotenv import load_dotenv
 
 # Load environment variables
 load_dotenv()
@@ -13,12 +17,17 @@ load_dotenv()
 # Initialize Groq client
 groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 
-# Load similarity model
+# Load models and dataset
 similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
 
-# Config
-HF_DATASET_REPO = "midrees2806/unmatched_queries"
-HF_TOKEN = os.getenv("HF_TOKEN")
+# Load dataset (automatically using the path)
+with open('dataset.json', 'r') as f:
+    dataset = json.load(f)
+
+# Precompute embeddings
+dataset_questions = [item.get("input", "").lower().strip() for item in dataset]
+dataset_answers = [item.get("response", "") for item in dataset]
+dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)
 
 # Greeting list
 GREETINGS = [
@@ -27,25 +36,14 @@ GREETINGS = [
     "hey there", "greetings"
 ]
 
-# Load local dataset
-try:
-    with open('dataset.json', 'r') as f:
-        dataset = json.load(f)
-    if not all(isinstance(item, dict) and 'Question' in item and 'Answer' in item for item in dataset):
-        raise ValueError("Invalid dataset structure")
-except Exception as e:
-    print(f"Error loading dataset: {e}")
-    dataset = []
-
-# Precompute embeddings
-dataset_questions = [item.get("Question", "").lower().strip() for item in dataset]
-dataset_answers = [item.get("Answer", "") for item in dataset]
-dataset_embeddings = similarity_model.encode(dataset_questions, convert_to_tensor=True)
+# Hugging Face config
+HF_DATASET_REPO = "midrees2806/unmatched_queries"
+HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Save unmatched queries to Hugging Face
 def manage_unmatched_queries(query: str):
     try:
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         try:
             ds = load_dataset(HF_DATASET_REPO, token=HF_TOKEN)
             df = ds["train"].to_pandas()
@@ -59,7 +57,6 @@ def manage_unmatched_queries(query: str):
     except Exception as e:
         print(f"Failed to save query: {e}")
 
-# Query Groq LLM
 def query_groq_llm(prompt, model_name="llama3-70b-8192"):
     try:
         chat_completion = groq_client.chat.completions.create(
@@ -76,16 +73,10 @@ def query_groq_llm(prompt, model_name="llama3-70b-8192"):
         print(f"Error querying Groq API: {e}")
         return ""
 
-# Main logic function to be called from Gradio
 def get_best_answer(user_input):
-    if not user_input.strip():
-        return "Please enter a valid question."
-
     user_input_lower = user_input.lower().strip()
 
-    if len(user_input_lower.split()) < 3 and not any(greet in user_input_lower for greet in GREETINGS):
-        return "Please ask your question properly with at least 3 words."
-
+    # 👉 Greeting functionality
     if any(greet in user_input_lower for greet in GREETINGS):
         greeting_response = query_groq_llm(
             f"You are an official assistant for University of Education Lahore. "
@@ -93,18 +84,21 @@ def get_best_answer(user_input):
         )
         return greeting_response if greeting_response else "Hello! How can I assist you today?"
 
-    if any(keyword in user_input_lower for keyword in ["fee structure", "fees structure", "semester fees", "semester fee"]):
+    # 👉 Check if question is about fee
+    if any(keyword in user_input_lower for keyword in ["semester fee", "semester fees"]):
         return (
             "💰 For complete and up-to-date fee details for this program, we recommend visiting the official University of Education fee structure page.\n"
-            "You'll find comprehensive information regarding tuition, admission charges, and other applicable fees there.\n"
+            "You’ll find comprehensive information regarding tuition, admission charges, and other applicable fees there.\n"
             "🔗 https://ue.edu.pk/allfeestructure.php"
         )
 
+    # 🔁 Continue with normal similarity-based logic
     user_embedding = similarity_model.encode(user_input_lower, convert_to_tensor=True)
     similarities = util.pytorch_cos_sim(user_embedding, dataset_embeddings)[0]
     best_match_idx = similarities.argmax().item()
     best_score = similarities[best_match_idx].item()
 
+    # 👉 Save unmatched query if no close match
     if best_score < 0.65:
         manage_unmatched_queries(user_input)
 
@@ -124,14 +118,16 @@ def get_best_answer(user_input):
     llm_response = query_groq_llm(prompt)
 
     if llm_response:
-        for marker in ["Improved Answer:", "Official Answer:", "Rephrased Answer:"]:
+        for marker in ["Improved Answer:", "Official Answer:"]:
             if marker in llm_response:
-                return llm_response.split(marker)[-1].strip()
-        return llm_response
+                response = llm_response.split(marker)[-1].strip()
+                break
+        else:
+            response = llm_response
     else:
-        return dataset_answers[best_match_idx] if best_score >= 0.65 else (
-            "For official information:\n"
-            "📞 +92-42-99262231-33\n"
-            "✉️ info@ue.edu.pk\n"
-            "🌐 https://ue.edu.pk"
-        )
+        response = dataset_answers[best_match_idx] if best_score >= 0.65 else """For official information:
+        📞 +92-42-99262231-33
+        ✉️ info@ue.edu.pk
+        🌐 ue.edu.pk"""
+
+    return response