Update app.py

app.py

@@ -2,13 +2,11 @@
 import gradio as gr
 import torch
 import requests
-import re
-from datetime import datetime
 import json
-
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from qdrant_client import QdrantClient
+from datetime import datetime
 
 # === Load Models ===
 print("Loading zero-shot classifier...")
@@ -17,43 +15,32 @@ classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnl
 print("Loading embedding model...")
 embedding_model = SentenceTransformer("intfloat/e5-large")
 
-print("Loading WizardMath model...")
-tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardMath-7B-V1.1")
-model = AutoModelForCausalLM.from_pretrained(
-    "WizardLM/WizardMath-7B-V1.1", torch_dtype=torch.float16, device_map="auto"
-)
+print("Loading step-by-step generator...")
+tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
+model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")
 
 # === Qdrant Setup ===
 print("Connecting to Qdrant...")
 qdrant_client = QdrantClient(path="qdrant_data")
 collection_name = "math_problems"
 
-# === Guard Functions ===
+# === Guard Function ===
 def is_valid_math_question(text):
     candidate_labels = ["math", "not math"]
     result = classifier(text, candidate_labels)
     return result['labels'][0] == "math" and result['scores'][0] > 0.7
 
-def output_guardrails(answer):
-    if not answer or len(answer.strip()) < 10:
-        return False
-    math_keywords = ["solve", "equation", "integral", "derivative", "value", "expression", "steps", "solution"]
-    if not any(word in answer.lower() for word in math_keywords):
-        return False
-    banned_keywords = ["kill", "bomb", "hate", "politics", "violence"]
-    if any(word in answer.lower() for word in banned_keywords):
-        return False
-    if re.match(r"^\s*I'm just a model|Sorry, I can't|As an AI", answer, re.IGNORECASE):
-        return False
-    return True
-
-# === Retrieval ===
+# === Retrieval from Qdrant ===
 def retrieve_from_qdrant(query):
     query_vector = embedding_model.encode(query).tolist()
-    hits = qdrant_client.search(collection_name=collection_name, query_vector=query_vector, limit=3)
+    hits = qdrant_client.query_points(
+        collection_name=collection_name,
+        query_vector=query_vector,
+        limit=3
+    )
     return [hit.payload for hit in hits] if hits else []
 
-# === Web Search ===
+# === Web Search Fallback ===
 def web_search_tavily(query):
     TAVILY_API_KEY = "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"
     response = requests.post(
@@ -62,14 +49,10 @@ def web_search_tavily(query):
     )
     return response.json().get("answer", "No answer found from Tavily.")
 
-# === Answer Generation ===
-def generate_step_by_step_answer(question, context=""):
-    prompt = f"### Question:\n{question}\n"
-    if context:
-        prompt += f"### Context:\n{context}\n"
-    prompt += "### Let's solve it step by step:\n"
-
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+# === Generator ===
+def generate_step_by_step_answer(question, context):
+    prompt = f"Answer the following math question step-by-step:\nQuestion: {question}\nContext: {context}\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(
         **inputs,
         max_new_tokens=256,
@@ -78,67 +61,65 @@ def generate_step_by_step_answer(question, context=""):
         do_sample=True,
         pad_token_id=tokenizer.eos_token_id
     )
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = decoded.split("### Let's solve it step by step:")[-1].strip()
-    return answer
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # === Router ===
 def router(question):
     if not is_valid_math_question(question):
-        return "❌ Only math questions are accepted. Please rephrase."
-
-    context_items = retrieve_from_qdrant(question)
-    context = "\n".join([item.get("solution", "") for item in context_items])
+        return "❌ Only math questions are accepted. Please rephrase.", ""
 
+    retrieved = retrieve_from_qdrant(question)
+    context = "\n".join([item["solution"] for item in retrieved if "solution" in item])
     if context:
         answer = generate_step_by_step_answer(question, context)
-        if output_guardrails(answer):
-            return answer
-
-    answer = web_search_tavily(question)
-    return answer if output_guardrails(answer) else "⚠️ No valid math answer found."
+        return answer, context
+    else:
+        fallback = web_search_tavily(question)
+        return fallback, "Tavily Search"
 
 # === Feedback Storage ===
-def store_feedback(question, answer, feedback, correct_answer):
+def store_feedback(question, answer, correct_answer):
     entry = {
         "question": question,
         "model_answer": answer,
-        "feedback": feedback,
         "correct_answer": correct_answer,
         "timestamp": str(datetime.now())
     }
     with open("feedback.json", "a") as f:
         f.write(json.dumps(entry) + "\n")
 
-# === Gradio UI ===
+# === Gradio Functions ===
 def ask_question(question):
-    answer = router(question)
+    answer, context = router(question)
     return answer, question, answer
 
-def submit_feedback(question, model_answer, feedback):
-    store_feedback(question, model_answer, feedback, "")
+def submit_feedback(question, model_answer, correct_answer):
+    store_feedback(question, model_answer, correct_answer)
     return "✅ Feedback received. Thank you!"
 
+# === Gradio UI ===
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧮 Math Tutor with AI Guardrails + Feedback")
+    gr.Markdown("## 🧮 Math Question Answering with Retrieval + Feedback")
 
     with gr.Row():
         question_input = gr.Textbox(label="Enter your math question", lines=2)
         submit_btn = gr.Button("Get Answer")
 
-    answer_output = gr.Markdown()
+    answer_output = gr.Markdown(label="Answer")
     hidden_q = gr.Textbox(visible=False)
     hidden_a = gr.Textbox(visible=False)
 
     submit_btn.click(fn=ask_question, inputs=[question_input], outputs=[answer_output, hidden_q, hidden_a])
 
-    gr.Markdown("### 📝 Feedback")
-    fb_like = gr.Radio(["👍", "👎"], label="Was this answer helpful?")
-    fb_submit_btn = gr.Button("Submit Feedback")
+    gr.Markdown("### 📝 Submit Feedback")
+    fb_correct = gr.Textbox(label="Correct Answer (optional)")
+    fb_submit = gr.Button("Submit Feedback")
     fb_status = gr.Textbox(label="Status", interactive=False)
 
-    fb_submit_btn.click(fn=submit_feedback,
-                        inputs=[hidden_q, hidden_a, fb_like],
-                        outputs=[fb_status])
+    fb_submit.click(
+        fn=submit_feedback,
+        inputs=[hidden_q, hidden_a, fb_correct],
+        outputs=[fb_status]
+    )
 
-demo.launch(share=True, debug=True)
+demo.launch(share=True)