Update app.py

app.py

@@ -18,7 +18,14 @@ embedding_model = SentenceTransformer("intfloat/e5-large")
 
 print("Loading text generation model...")
 # Use a lighter model for testing
-qa_pipeline = pipeline("text-generation", model="gpt2")
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+
+print("Loading text generation model...")
+model_name = "google/flan-t5-large"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+qa_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+
 
 # === Qdrant Setup ===
 print("Connecting to Qdrant...")
@@ -65,9 +72,11 @@ class MathRetrievalQA(dspy.Program):
         print("Context for generation:", context)
         if not context:
             return dspy.Output(answer="", retrieved_context="")
-        prompt = f"Question: {question}\nContext: {context}\nAnswer:"
+
+        prompt = f"Question: {question}\nStep-by-step solution:\n{context}\nAnswer:"
         print("Generating answer...")
-        answer = qa_pipeline(prompt, max_new_tokens=100)[0]["generated_text"]
+        answer = qa_pipeline(prompt, max_new_tokens=150)[0]["generated_text"]
+
         print("Generated answer:", answer)
         return {"answer": answer, "retrieved_context": context}