Update app.py

app.py

@@ -1,9 +1,23 @@
 import streamlit as st
 import sympy as sp
+import chromadb
 from transformers import pipeline
+from langchain.chains import RetrievalQA
+from langchain.vectorstores import FAISS
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.llms import OpenAI
+
+# Initialize ChromaDB for Retrieval-Augmented Generation (RAG)
+chroma_client = chromadb.PersistentClient(path="./chroma_db")
+embedding_model = SentenceTransformerEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# Load the RAG-based Retrieval System
+vectorstore = FAISS.load_local("faiss_index", embedding_model)
+retriever = vectorstore.as_retriever()
+qa_chain = RetrievalQA.from_chain_type(llm=OpenAI(), retriever=retriever)
 
 # Load NLP Model
-model = pipeline("text2text-generation", model="facebook/bart-large")
+model = pipeline("text2text-generation", model="google/flan-t5-small")
 
 st.title("🤖 AI-Driven Mathematical Model Generator")
 st.write("Enter a problem statement in natural language to get a mathematical model.")
@@ -11,8 +25,13 @@ st.write("Enter a problem statement in natural language to get a mathematical mo
 user_input = st.text_area("✍️ Enter your problem:")
 
 if st.button("🚀 Generate Model"):
-    response = model(user_input, max_length=200)
-    equation = sp.sympify(response[0]['generated_text'])
+    retrieved_context = qa_chain.run(user_input)  # RAG retrieval
+    response = model(f"Generate a mathematical model for: {user_input}\nContext: {retrieved_context}", max_length=200)
+    
+    try:
+        equation = sp.sympify(response[0]['generated_text'])
+    except:
+        equation = response[0]['generated_text']  # If parsing fails, return text
     
     st.subheader("📌 Mathematical Model:")
     st.latex(sp.latex(equation))