Update app.py

app.py

@@ -1,10 +1,10 @@
-import gradio as gr
 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,59 +19,48 @@ 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 retrieval_qa_chain(llm, prompt, db):
-    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': prompt}
-                                          )
-    return qa_chain
-
 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
 
-def qa_bot():
+def qa_bot(query):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
                                        model_kwargs={'device': 'cpu'})
     db = FAISS.load_local(DB_FAISS_PATH, embeddings)
     llm = load_llm()
     qa_prompt = set_custom_prompt()
-    qa = retrieval_qa_chain(llm, qa_prompt, db)
-    return qa
+    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."
 
-# Define a function to respond to messages using your QA model
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    qa_result = qa_bot()
-    response = qa_result({'query': message})
     return response
 
-# Create a Gradio interface using the respond function
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+iface = gr.Interface(
+    fn=qa_bot,
+    inputs="text",
+    outputs="text",
+    title="Medical Query Bot",
+    description="Enter your medical query to get an answer."
 )
 
-if __name__ == "__main__":
-    demo.launch()
+if __name__ == '__main__':
+    iface.launch()
+