from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain import HuggingFaceHub
from langchain.chains import RetrievalQA
import streamlit as st

st.set_page_config(page_title = "Hospital Regulatory Chat", page_icon=":hospital:")


DB_FAISS_PATH = '.'

def get_vectorstore():
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
                                       model_kwargs={'device': 'cpu'})
    vector_store = FAISS.load_local(DB_FAISS_PATH, embeddings)
    return vector_store

vector_store = get_vectorstore()

llm = HuggingFaceHub(repo_id = "meta-llama/Llama-2-7b-chat-hf",model_kwargs={"temperature":0.5}) #, "max_length":512})

qa_chain = RetrievalQA.from_chain_type(llm=llm,
                                       chain_type='stuff',
                                       retriever=vector_store.as_retriever(search_kwargs={'k': 10}),
                                       #retriever=vector_store.as_retriever(search_kwargs={"score_threshold": .01}),
                                       return_source_documents = True
                                       )


source_dictionary = {"pdf_files\\CMS_SOMA.pdf":"[CMS State Operations Manual Appendix A](https://www.cms.gov/regulations-and-guidance/guidance/manuals/downloads/som107ap_a_hospitals.pdf)",
'pdf_files\\RCW-18-64.pdf': '[Pharmacists](https://app.leg.wa.gov/rcw/default.aspx?cite=18.64&full=true&pdf=true)', 
'pdf_files\\RCW-18-64A.pdf': '[Pharmacy Assistants](https://app.leg.wa.gov/rcw/default.aspx?cite=18.64A&full=true&pdf=true)', 
'pdf_files\\RCW-18-130.pdf': '[Regulation of Health Professions—Uniform Disciplinary Act](https://app.leg.wa.gov/rcw/default.aspx?cite=18.130&full=true&pdf=true)', 
'pdf_files\\RCW-26-44-030.pdf': '[Abuse of Children and Adult Dependent Persons](https://app.leg.wa.gov/rcw/default.aspx?cite=26.44.030&full=true&pdf=true)', 
'pdf_files\\RCW-34-05.pdf': '[Administrative Procedures Act](https://app.leg.wa.gov/rcw/default.aspx?cite=34.05&full=true&pdf=true)', 
'pdf_files\\RCW-42-56.pdf': '[Public Disclosure](https://app.leg.wa.gov/rcw/default.aspx?cite=42.56&full=true&pdf=true)', 
'pdf_files\\RCW-43-70.pdf': '[Department of Health](https://app.leg.wa.gov/rcw/default.aspx?cite=43.70&full=true&pdf=true)', 
'pdf_files\\RCW-69-04.pdf': '[Uniform Food, Drug and Cosmetic Act](https://app.leg.wa.gov/rcw/default.aspx?cite=69.04&full=true&pdf=true)', 
'pdf_files\\RCW-69-36.pdf': '[Washington Caustic Poison Act of 1929](https://app.leg.wa.gov/rcw/default.aspx?cite=69.36&full=true&pdf=true)', 
'pdf_files\\RCW-69-38.pdf': '[Poisons - Sales and Manufacturing](https://app.leg.wa.gov/rcw/default.aspx?cite=69.38&full=true&pdf=true)', 
'pdf_files\\RCW-69-40.pdf': '[Poisons and Dangerous Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.40&full=true&pdf=true)', 
'pdf_files\\RCW-69-41.pdf': '[Legend Drugs . . . Prescription Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.41&full=true&pdf=true)', 
'pdf_files\\RCW-69-43.pdf': '[Precursor Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.43&full=true&pdf=true)', 
'pdf_files\\RCW-69-45.pdf': '[Drug Samples](https://app.leg.wa.gov/rcw/default.aspx?cite=69.45&full=true&pdf=true)', 
'pdf_files\\RCW-69-48.pdf': '[Drug Take-Back Program\nSecure Medication Return webpage](https://app.leg.wa.gov/rcw/default.aspx?cite=69.48&full=true&pdf=true)',
 'pdf_files\\RCW-69-50.pdf': '[Uniform Controlled Substances Act](https://app.leg.wa.gov/rcw/default.aspx?cite=69.50&full=true&pdf=true)', 
 'pdf_files\\RCW-69-51.pdf': '[Controlled Substances Therapeutic Research Act](https://app.leg.wa.gov/rcw/default.aspx?cite=69.51&full=true&pdf=true)', 
'pdf_files\\RCW-69-51A.pdf': '[Medical Cannabis](https://app.leg.wa.gov/rcw/default.aspx?cite=69.51A&full=true&pdf=true)', 
'pdf_files\\RCW-69-52.pdf': '[Imitation Controlled Substances](https://app.leg.wa.gov/rcw/default.aspx?cite=69.52&full=true&pdf=true)', 
'pdf_files\\RCW-69-53.pdf': '[Use of Buildings for Unlawful Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.53&full=true&pdf=true)',
 'pdf_files\\RCW-69-60.pdf': '[Over-the-Counter Medications - Imprinting](https://app.leg.wa.gov/rcw/default.aspx?cite=69.60&full=true&pdf=true)',
 'pdf_files\\RCW-69-70.pdf': '[Access to prescription drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.70&full=true&pdf=true)', 
 'pdf_files\\RCW-69-75.pdf': '[Dextromethorphan](https://app.leg.wa.gov/rcw/default.aspx?cite=69.75&full=true&pdf=true)', 
'pdf_files\\RCW-70-02.pdf': '[Medical Records-Health Care Information Access and Disclosure](https://app.leg.wa.gov/rcw/default.aspx?cite=70.02&full=true&pdf=true)', 
'pdf_files\\RCW-70-54.pdf': '[Miscellaneous Health and Safety Provisions\n\n70.54.400 Retail restroom access – Customers with medical conditions – Penalty\n\n\n70.54.440 Epinephrine autoinjectors – Prescribing to certain entities-training-Liability-incident reporting](https://app.leg.wa.gov/rcw/default.aspx?cite=70.54&full=true&pdf=true)', 
'pdf_files\\RCW-70-115.pdf': '[Drug Injection Devices](https://app.leg.wa.gov/rcw/default.aspx?cite=70.115&full=true&pdf=true)', 
'pdf_files\\RCW-70-225.pdf': '[Prescription Monitoring Program](https://app.leg.wa.gov/rcw/default.aspx?cite=70.225&full=true&pdf=true)', 
'pdf_files\\RCW-70-245.pdf': '[The Washington Death With Dignity Act](https://app.leg.wa.gov/rcw/default.aspx?cite=70.245&full=true&pdf=true)', 
'pdf_files\\RCW-74-34.pdf': '[Abuse of Vulnerable Adults](https://app.leg.wa.gov/rcw/default.aspx?cite=74.34&full=true&pdf=true)', 
'pdf_files\\WAC-246-338.pdf': '[Medical Test Site Rules](https://app.leg.wa.gov/wac/default.aspx?cite=246-338&full=true&pdf=true)', 
'pdf_files\\WAC-246-320.pdf': '[Hospital Licensing Regulations](https://app.leg.wa.gov/wac/default.aspx?cite=246-320&full=true&pdf=true)', 
'pdf_files\\RCW-70-41.pdf': '[Hospital Licensing and Regulation](https://app.leg.wa.gov/rcw/default.aspx?cite=70.41&full=true&pdf=true)', 
'pdf_files\\RCW-70-42.pdf': '[Medical Test Sites](https://app.leg.wa.gov/rcw/default.aspx?cite=70.42&full=true&pdf=true)', 
'pdf_files\\WAC-246-240.pdf': '[Medical Use of Radioactive Material](https://app.leg.wa.gov/wac/default.aspx?cite=246-240&full=true&pdf=true)', 
'pdf_files\\WAC-246-221.pdf': '[Radiation Protection Standards](https://app.leg.wa.gov/wac/default.aspx?cite=246-221&full=true&pdf=true)', 
'pdf_files\\RCW-18-71.pdf': '[Physicians](https://app.leg.wa.gov/rcw/default.aspx?cite=18.71&full=true&pdf=true)', 
'pdf_files\\RCW-18-71A.pdf': '[Physician Assistants](https://app.leg.wa.gov/rcw/default.aspx?cite=18.71A&full=true&pdf=true)', 
'pdf_files\\RCW-18-79.pdf': '[Nursing Care](https://app.leg.wa.gov/rcw/default.aspx?cite=18.79&full=true&pdf=true)', 
'pdf_files\\RCW-18-84.pdf': '[Radiologic Technologies](https://app.leg.wa.gov/rcw/default.aspx?cite=18.84&full=true&pdf=true)', 
'pdf_files\\RCW-18-370.pdf': '[Medical Assistants](https://app.leg.wa.gov/rcw/default.aspx?cite=18.370&full=true&pdf=true)'}

with st.container():
    st.title("Hospital Regulation Chat")

with st.sidebar:    
    st.subheader("Find regulations for hospitals in the state of Washington.")
    st.markdown("""
    We look into these sources to find top ten most relevant excerpts:
    - [CMS State Operations Manual Appendix A](https://www.cms.gov/regulations-and-guidance/guidance/manuals/downloads/som107ap_a_hospitals.pdf)
    - [Pharmacists](https://app.leg.wa.gov/rcw/default.aspx?cite=18.64&full=true&pdf=true)
    - [Pharmacy Assistants](https://app.leg.wa.gov/rcw/default.aspx?cite=18.64A&full=true&pdf=true)
    - [Regulation of Health Professions—Uniform Disciplinary Act](https://app.leg.wa.gov/rcw/default.aspx?cite=18.130&full=true&pdf=true)
    - [Abuse of Children and Adult Dependent Persons](https://app.leg.wa.gov/rcw/default.aspx?cite=26.44.030&full=true&pdf=true)
    - [Administrative Procedures Act](https://app.leg.wa.gov/rcw/default.aspx?cite=34.05&full=true&pdf=true)
    - [Public Disclosure](https://app.leg.wa.gov/rcw/default.aspx?cite=42.56&full=true&pdf=true)
    - [Department of Health](https://app.leg.wa.gov/rcw/default.aspx?cite=43.70&full=true&pdf=true)
    - [Uniform Food, Drug and Cosmetic Act](https://app.leg.wa.gov/rcw/default.aspx?cite=69.04&full=true&pdf=true)
    - [Washington Caustic Poison Act of 1929](https://app.leg.wa.gov/rcw/default.aspx?cite=69.36&full=true&pdf=true)
    - [Poisons - Sales and Manufacturing](https://app.leg.wa.gov/rcw/default.aspx?cite=69.38&full=true&pdf=true)
    - [Poisons and Dangerous Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.40&full=true&pdf=true)
    - [Legend Drugs . . . Prescription Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.41&full=true&pdf=true)
    - [Precursor Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.43&full=true&pdf=true)
    - [Drug Samples](https://app.leg.wa.gov/rcw/default.aspx?cite=69.45&full=true&pdf=true)
    - [Drug Take-Back Program Secure Medication Return webpage](https://app.leg.wa.gov/rcw/default.aspx?cite=69.48&full=true&pdf=true)
    - [Uniform Controlled Substances Act](https://app.leg.wa.gov/rcw/default.aspx?cite=69.50&full=true&pdf=true)
    - [Controlled Substances Therapeutic Research Act](https://app.leg.wa.gov/rcw/default.aspx?cite=69.51&full=true&pdf=true)
    - [Medical Cannabis](https://app.leg.wa.gov/rcw/default.aspx?cite=69.51A&full=true&pdf=true)
    - [Imitation Controlled Substances](https://app.leg.wa.gov/rcw/default.aspx?cite=69.52&full=true&pdf=true)
    - [Use of Buildings for Unlawful Drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.53&full=true&pdf=true)
    - [Over-the-Counter Medications - Imprinting](https://app.leg.wa.gov/rcw/default.aspx?cite=69.60&full=true&pdf=true)
    - [Access to prescription drugs](https://app.leg.wa.gov/rcw/default.aspx?cite=69.70&full=true&pdf=true)
    - [Dextromethorphan](https://app.leg.wa.gov/rcw/default.aspx?cite=69.75&full=true&pdf=true)
    - [Medical Records-Health Care Information Access and Disclosure](https://app.leg.wa.gov/rcw/default.aspx?cite=70.02&full=true&pdf=true)
    - [Miscellaneous Health and Safety Provisions 70.54.400 Retail restroom access – Customers with medical conditions – Penalty 70.54.440 Epinephrine autoinjectors – Prescribing to certain entities-training-Liability-incident reporting](https://app.leg.wa.gov/rcw/default.aspx?cite=70.54&full=true&pdf=true)
    - [Drug Injection Devices](https://app.leg.wa.gov/rcw/default.aspx?cite=70.115&full=true&pdf=true)
    - [Prescription Monitoring Program](https://app.leg.wa.gov/rcw/default.aspx?cite=70.225&full=true&pdf=true)
    - [The Washington Death With Dignity Act](https://app.leg.wa.gov/rcw/default.aspx?cite=70.245&full=true&pdf=true)
    - [Abuse of Vulnerable Adults](https://app.leg.wa.gov/rcw/default.aspx?cite=74.34&full=true&pdf=true)
    - [Medical Test Site Rules](https://app.leg.wa.gov/wac/default.aspx?cite=246-338&full=true&pdf=true)
    - [Hospital Licensing Regulations](https://app.leg.wa.gov/wac/default.aspx?cite=246-320&full=true&pdf=true)
    - [Hospital Licensing and Regulation](https://app.leg.wa.gov/rcw/default.aspx?cite=70.41&full=true&pdf=true)
    - [Medical Test Sites](https://app.leg.wa.gov/rcw/default.aspx?cite=70.42&full=true&pdf=true)
    - [Medical Use of Radioactive Material](https://app.leg.wa.gov/wac/default.aspx?cite=246-240&full=true&pdf=true)
    - [Radiation Protection Standards](https://app.leg.wa.gov/wac/default.aspx?cite=246-221&full=true&pdf=true)
    - [Physicians](https://app.leg.wa.gov/rcw/default.aspx?cite=18.71&full=true&pdf=true)
    - [Physician Assistants](https://app.leg.wa.gov/rcw/default.aspx?cite=18.71A&full=true&pdf=true)
    - [Nursing Care](https://app.leg.wa.gov/rcw/default.aspx?cite=18.79&full=true&pdf=true)
    - [Radiologic Technologies](https://app.leg.wa.gov/rcw/default.aspx?cite=18.84&full=true&pdf=true)
    - [Medical Assistants](https://app.leg.wa.gov/rcw/default.aspx?cite=18.370&full=true&pdf=true)
    """) #, unsafe_allow_html=True)
    st.write("This is tool is meant to assist healthcare workers to the extent it can. Please note that the page numbers may be occasionally slightly off, use the matching excerpts to find the reference if this happens.")

st.markdown("**Ask your question and :red[click 'Find excerpts'.]**")
prompt = st.text_input("e.g. What are the rules regarding a Quality Improvement, or QAPI program?")

if (st.button("Find excerpts")):
    answer = qa_chain({"query":prompt})

    n = len(answer['source_documents'])

    for i in range(n):
        with st.container():
            page = str(answer['source_documents'][i].metadata['page'])
            page_no = "#page=" + page + ")"
            st.subheader(source_dictionary[answer['source_documents'][i].metadata['source']].replace(")",page_no))
            page_no = "**Page: " + page + "**"
            st.markdown(page_no)
            st.write("...")
            st.write(answer['source_documents'][i].page_content)
            st.write("...")
            st.write('---------------------------------\n\n')