Practice / rag.py
midrees2806's picture
Update rag.py
6dd717f verified
raw
history blame
4.6 kB
import json
from sentence_transformers import SentenceTransformer, util
from groq import Groq
import os
import csv
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Initialize Groq client
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
# Load similarity model
similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
# Load dataset
with open('dataset.json', 'r', encoding='utf-8') 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)
# Use absolute path for unmatched_queries.csv
base_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(base_dir, "unmatched_queries.csv")
print(f"[DEBUG] Writing to absolute path: {file_path}")
def query_groq_llm(prompt, model_name="llama3-70b-8192"):
try:
chat_completion = groq_client.chat.completions.create(
messages=[{
"role": "user",
"content": prompt
}],
model=model_name,
temperature=0.7,
max_tokens=500
)
return chat_completion.choices[0].message.content.strip()
except Exception as e:
print(f"Error querying Groq API: {e}")
return ""
def get_best_answer(user_input):
user_input_lower = user_input.lower().strip()
# πŸ‘‰ Fee-specific shortcut
if any(keyword in user_input_lower for keyword in ["fee", "fees", "charges", "semester fee"]):
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"
"πŸ”— https://ue.edu.pk/allfeestructure.php"
)
# πŸ” Similarity matching
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()
# ✏️ Log to CSV if similarity is low
if best_score < 0.65:
print(f"[DEBUG] Similarity score too low: {best_score}. Logging query to: {file_path}")
# Create CSV with header if it doesn't exist
if not os.path.exists(file_path):
print(f"[DEBUG] File {file_path} does not exist. Creating file with header.")
try:
with open(file_path, mode="w", newline="", encoding="utf-8") as file:
writer = csv.writer(file)
writer.writerow(["Unmatched Queries"])
print(f"[DEBUG] Header written successfully.")
except Exception as e:
print(f"[ERROR] Failed to create file: {e}")
# Append unmatched query
try:
with open(file_path, mode="a", newline="", encoding="utf-8") as file:
writer = csv.writer(file)
writer.writerow([user_input])
print(f"[DEBUG] Query logged: {user_input}")
except Exception as e:
print(f"[ERROR] Failed to write query to CSV: {e}")
# 🧠 Construct prompt
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:
Question: {user_input}
Original Answer: {original_answer}
Improved Answer:"""
else:
prompt = f"""As an official assistant for University of Education Lahore, provide a helpful response:
Include relevant details about university policies.
If unsure, direct to official channels.
Question: {user_input}
Official Answer:"""
# 🧠 Query LLM
llm_response = query_groq_llm(prompt)
# 🧾 Process LLM output
if llm_response:
for marker in ["Improved Answer:", "Official Answer:"]:
if marker in llm_response:
response = llm_response.split(marker)[-1].strip()
break
else:
response = llm_response
else:
response = dataset_answers[best_match_idx] if best_score >= 0.65 else """For official information:
πŸ“ž +92-42-99262231-33
βœ‰οΈ [email protected]
🌐 ue.edu.pk"""
return response