Update app.py

app.py

@@ -3,7 +3,7 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.llms import CTransformers
 from langchain.chains import RetrievalQA
-import chainlit as cl
+import gradio as gr
 
 DB_FAISS_PATH = "vectorstores/db_faiss"
 
@@ -18,74 +18,58 @@ Helpful answer:
 """
 
 def set_custom_prompt():
-    """
-    Prompt template for QA retrieval for each vectorstore
-    """
-    prompt = PromptTemplate(template=custom_prompt_template,
-                            input_variables=['context', 'question'])
+    prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context', 'question'])
     return prompt
 
-
 def load_llm():
-    """
-    Load the language model
-    """
     llm = CTransformers(
         model="C:/Users/sanath/Downloads/llama-2-7b-chat.ggmlv3.q8_0.bin",
-        model_type = "llama",
-        max_new_tokens = 512,
-        temperature = 0.5
+        model_type="llama",
+        max_new_tokens=512,
+        temperature=0.5
     )
     return llm
 
-def retrieval_QA_chain(llm,prompt,db):
-  qachain = 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 qachain
-
+def retrieval_QA_chain(llm, prompt, db):
+    qachain = 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 qachain
 
 def qa_bot():
-  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
-
-def final_result(query):
-  qa_result = qa_bot()
-  response = qa_result({'query':query})
-  return response
+    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-miniLM-L6-V2', model_kwargs={'device': 'cpu'})
+    db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
+    llm = load_llm()
+    qa_prompt = set_custom_prompt()
+    qa = retrieval_QA_chain(llm, qa_prompt, db)
+    return qa
 
-@cl.on_chat_start
-async def start():
-  chain = qa_bot()
-  msg = cl.Message(content = "Starting the bot...")
-  await msg.send()
-  msg.content = "Hi, Welcome to the Medical bot. What is your query?"
-  await msg.update()
+def chatbot_response(query):
+    try:
+        qa = qa_bot()
+        response = qa({'query': query})
+        answer = response["result"]
+        sources = response["source_documents"]
+        if sources:
+            answer += f"\nSources:" + str(sources)
+        else:
+            answer += "\nNo sources found"
+        return answer
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
 
-  cl.user_session.set("chain",chain)
+# Create a Gradio interface with updated API
+iface = gr.Interface(
+    fn=chatbot_response,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your question..."),
+    outputs="text",
+    title="Medical Chatbot",
+    description="Ask a medical question and get answers based on the provided context."
+)
 
-@cl.on_message
-async def main(message):
-  chain = cl.user_session.get("chain")
-  cb = cl.AsyncLangchainCallbackHandler(
-      stream_final_answer = True, answer_prefix_tokens = ["FINAL","ANSWER"]
-  )
-  cb.answer_reached = True
-  res = await chain.acall(message, callbacks=[cb])
-  answer = res["result"]
-  sources = res["source_documents"]
-  if sources:
-    answer += f"\nSources:" + str(sources)
-  else:
-    answer += "\nNo sources found"
-  
-  await cl.Message(content=answer).send()
+# Launch the Gradio app
+iface.launch()