Spaces:
Sleeping
Sleeping
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 chainlit as cl | |
import streamlit as st | |
DB_FAISS_PATH = 'vectorstore/db_faiss' | |
custom_prompt_template = """Use the following pieces of information to answer the user's question. | |
If you don't know the answer, just say that you don't know, don't try to make up an answer. | |
Context: {context} | |
Question: {question} | |
Only return the helpful answer below and nothing else. | |
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", | |
max_new_tokens = 512, | |
temperature = 0.5 | |
) | |
return llm | |
#QA Model Function | |
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 | |
#output function | |
def final_result(query): | |
qa_result = qa_bot() | |
response = qa_result({'query': query}) | |
return response | |
#chainlit code | |
async def start(): | |
chain = qa_bot() | |
msg = cl.Message(content="Starting the bot...") | |
await msg.send() | |
msg.content = "Hi, Welcome to Medical Bot. What is your query?" | |
await msg.update() | |
cl.user_session.set("chain", chain) | |
async def main(message: cl.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.content, 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() | |
def main(): | |
st.title("Medical Bot") | |
st.text_input("Enter your query:", key="query") | |
if st.button("Submit"): | |
query = st.session_state['query'] | |
with st.spinner("Thinking..."): | |
response = final_result(query) | |
st.markdown(response) | |
if __name__ == "__main__": | |
main() |