Update app.py

app.py

@@ -2,13 +2,14 @@
 import gradio as gr
 import torch
 import requests
-from transformers import pipeline
-from sentence_transformers import SentenceTransformer
-from qdrant_client import QdrantClient
+import re
 from datetime import datetime
-import dspy
 import json
 
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer
+from qdrant_client import QdrantClient
+
 # === Load Models ===
 print("Loading zero-shot classifier...")
 classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
@@ -16,17 +17,36 @@ 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"
+)
+
 # === Qdrant Setup ===
 print("Connecting to Qdrant...")
 qdrant_client = QdrantClient(path="qdrant_data")
 collection_name = "math_problems"
 
-# === Guard Function ===
+# === Guard Functions ===
 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 ===
 def retrieve_from_qdrant(query):
     query_vector = embedding_model.encode(query).tolist()
@@ -42,35 +62,41 @@ def web_search_tavily(query):
     )
     return response.json().get("answer", "No answer found from Tavily.")
 
-# === DSPy Signature ===
-class MathAnswer(dspy.Signature):
-    question = dspy.InputField()
-    retrieved_context = dspy.InputField()
-    answer = dspy.OutputField()
-
-# === DSPy Programs ===
-class MathRetrievalQA(dspy.Program):
-    def forward(self, question):
-        context_items = retrieve_from_qdrant(question)
-        context = "\n".join([item["solution"] for item in context_items if "solution" in item])
-        if not context:
-            return {"answer": "", "retrieved_context": ""}
-        prompt = f"Question: {question}\nContext: {context}\nAnswer:"
-        return {"answer": prompt, "retrieved_context": context}
-
-class WebFallbackQA(dspy.Program):
-    def forward(self, question):
-        answer = web_search_tavily(question)
-        return {"answer": answer, "retrieved_context": "Tavily"}
-
-class MathRouter(dspy.Program):
-    def forward(self, question):
-        if not is_valid_math_question(question):
-            return {"answer": "❌ Only math questions are accepted. Please rephrase.", "retrieved_context": ""}
-        result = MathRetrievalQA().forward(question)
-        return result if result["answer"] else WebFallbackQA().forward(question)
-
-router = MathRouter()
+# === 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")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        temperature=0.7,
+        top_p=0.95,
+        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
+
+# === 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])
+
+    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."
 
 # === Feedback Storage ===
 def store_feedback(question, answer, feedback, correct_answer):
@@ -84,39 +110,35 @@ def store_feedback(question, answer, feedback, correct_answer):
     with open("feedback.json", "a") as f:
         f.write(json.dumps(entry) + "\n")
 
-# === Gradio Functions ===
+# === Gradio UI ===
 def ask_question(question):
-    result = router.forward(question)
-    return result["answer"], question, result["answer"]
+    answer = router(question)
+    return answer, question, answer
 
-def submit_feedback(question, model_answer, feedback, correct_answer):
-    store_feedback(question, model_answer, feedback, correct_answer)
+def submit_feedback(question, model_answer, feedback):
+    store_feedback(question, model_answer, feedback, "")
     return "✅ Feedback received. Thank you!"
 
-# === Gradio UI ===
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧮 Math Question Answering with DSPy + Feedback")
+    gr.Markdown("## 🧮 Math Tutor with AI Guardrails + 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()
     hidden_q = gr.Textbox(visible=False)
     hidden_a = gr.Textbox(visible=False)
-    
-    submit_btn = gr.Button("Get Answer")
+
     submit_btn.click(fn=ask_question, inputs=[question_input], outputs=[answer_output, hidden_q, hidden_a])
 
-    # Feedback section on same page
-    gr.Markdown("### 💬 Give Feedback")
-    fb_correct = gr.Textbox(label="Correct Answer (optional)")
-    fb_like = gr.Radio(["👍", "👎"], label="Was the answer helpful?")
+    gr.Markdown("### 📝 Feedback")
+    fb_like = gr.Radio(["👍", "👎"], label="Was this answer helpful?")
     fb_submit_btn = gr.Button("Submit Feedback")
-    fb_status = gr.Markdown()
+    fb_status = gr.Textbox(label="Status", interactive=False)
 
     fb_submit_btn.click(fn=submit_feedback,
-                        inputs=[hidden_q, hidden_a, fb_like, fb_correct],
+                        inputs=[hidden_q, hidden_a, fb_like],
                         outputs=[fb_status])
 
-demo.launch(share=True, debug=True)
-
+demo.launch(share=True, debug=True)