Update app.py

app.py

@@ -1,10 +1,10 @@
+import streamlit as st
 from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
 from langchain.prompts import PromptTemplate
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain_community.llms import CTransformers
 from langchain.chains import RetrievalQA
-import gradio as gr
 
 DB_FAISS_PATH = 'vectorstore/db_faiss'
 
@@ -19,20 +19,16 @@ Helpful answer:
 """
 
 def set_custom_prompt():
-    """
-    Prompt template for QA retrieval for each vectorstore
-    """
     prompt = PromptTemplate(template=custom_prompt_template,
                             input_variables=['context', 'question'])
     return prompt
 
 def load_llm():
-    # Load the locally downloaded model here
     llm = CTransformers(
-        model = "TheBloke/Llama-2-7B-Chat-GGML",
+        model="TheBloke/Llama-2-7B-Chat-GGML",
         model_type="llama",
-        max_new_tokens = 512,
-        temperature = 0.5
+        max_new_tokens=512,
+        temperature=0.5
     )
     return llm
 
@@ -42,25 +38,26 @@ def qa_bot(query):
     db = FAISS.load_local(DB_FAISS_PATH, embeddings)
     llm = load_llm()
     qa_prompt = set_custom_prompt()
-    qa_chain = RetrievalQA.from_chain_type(llm=llm,
-                                           chain_type='stuff',
-                                           retriever=db.as_retriever(search_kwargs={'k': 2}),
-                                           return_source_documents=True,
-                                           chain_type_kwargs={'prompt': qa_prompt}
-                                          )
-    result = qa_chain({'query': query})
-    response = result['answers'][0]['text'] if result['answers'] else "Sorry, I don't have an answer for that."
-
+    qa = RetrievalQA.from_chain_type(llm=llm,
+                                     chain_type='stuff',
+                                     retriever=db.as_retriever(search_kwargs={'k': 2}),
+                                     return_source_documents=True,
+                                     chain_type_kwargs={'prompt': qa_prompt}
+                                     )
+    response = qa({'query': query})
     return response
 
-iface = gr.Interface(
-    fn=qa_bot,
-    inputs="text",
-    outputs="text",
-    title="Medical Query Bot",
-    description="Enter your medical query to get an answer."
-)
+def main():
+    st.title('Medical Bot')
+    query = st.text_input('Enter your medical query:')
+    if st.button('Submit'):
+        response = qa_bot(query)
+        st.write('Answer:', response['result'])
+        if response['source_documents']:
+            st.write('Sources:', response['source_documents'])
+        else:
+            st.write('No sources found')
 
 if __name__ == '__main__':
-    iface.launch()
+    main()