|
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 |
|
|
|
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 = PromptTemplate(template=custom_prompt_template, |
|
input_variables=['context', 'question']) |
|
return prompt |
|
|
|
def load_llm(): |
|
llm = CTransformers( |
|
model="TheBloke/Llama-2-7B-Chat-GGML", |
|
model_type="llama", |
|
max_new_tokens=512, |
|
temperature=0.5 |
|
) |
|
return llm |
|
|
|
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} |
|
) |
|
response = qa({'query': query}) |
|
return response |
|
|
|
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__': |
|
main() |
|
|
|
|