Update app.py

app.py

@@ -19,11 +19,25 @@ 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
 
+# Retrieval QA Chain
+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
+
+# Loading the model
 def load_llm():
+    # Load the locally downloaded model here
     llm = CTransformers(
         model="TheBloke/Llama-2-7B-Chat-GGML",
         model_type="llama",
@@ -32,33 +46,32 @@ def load_llm():
     )
     return llm
 
+# QA Model Function
 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 = 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}
-                                     )
+    qa = retrieval_qa_chain(llm, qa_prompt, db)
+    
     response = qa({'query': query})
-    return response['result'], response['source_documents']
-
-def main(query):
-    answer, sources = qa_bot(query)
+    answer = response['result']
+    sources = response['source_documents']
     if sources:
         answer += f"\nSources: {sources}"
     else:
         answer += "\nNo sources found"
     return answer
 
-iface = gr.Interface(fn=main, 
-                     inputs=gr.inputs.Textbox(label="Enter your medical query"),
-                     outputs=gr.outputs.Textbox(label="Answer"),
+# Gradio interface
+iface = gr.Interface(fn=qa_bot,
+                     inputs=gr.Textbox(label="Enter your medical query"),
+                     outputs=gr.Textbox(label="Answer"),
                      title="Medical Bot",
                      description="Ask any medical query and get an answer with sources if available.")
-iface.launch()
+
+if __name__ == "__main__":
+    iface.launch()
+