import streamlit as st
import torch

def summarize_function(notes):
    gen_text = pipe(notes, max_length=(len(notes.split(' '))*2*1.225), temperature=0.8, num_return_sequences=1, top_p=0.2)[0]['generated_text'][len(notes):]
    for i in range(len(gen_text)):
        if gen_text[-i-8:].startswith('[Notes]:'):
            gen_text = gen_text[:-i-8]
    st.write('Summary: ')
    return gen_text

st.markdown("<h1 style='text-align: center; color: #489DDB;'>GPT Clinical Notes Summarizer 0.1v</h1>", unsafe_allow_html=True)
st.markdown("<h6 style='text-align: center; color: #489DDB;'>by Bryan Mildort</h1>", unsafe_allow_html=True)

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
# from accelerate import infer_auto_device_map
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# device_str = f"""Device being used: {device}"""
# st.write(device_str)
# device_map = infer_auto_device_map(model, dtype="float16")
# st.write(device_map)

@st.cache(allow_output_mutation=True)
def load_model():
    model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True, load_in_8bit=True, device_map='auto')
    # model.to(device)
    tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
    return pipe

pipe = load_model()

prompt = """Admission Date:  2130-4-14     Discharge Date:  2130-4-17
Date of Birth:   2082-12-11     Sex:  M
Service:  #58
HISTORY OF PRESENT ILLNESS:  Mr. Jefferson is a 47 year-old man
with extreme obesity with a body weight of 440 pounds who is
5'7" tall and has a BMI of 69.  He has had numerous weight
loss programs in the past without significant long term
effect and also has significant venostasis ulcers in his
lower extremities.  He has no known drug allergies.
His only past medical history other then obesity is
osteoarthritis for which he takes Motrin and smoker's cough
secondary to smoking one pack per day for many years.  He has
used other narcotics, cocaine and marijuana, but has been
clean for about fourteen years.
He was admitted to the General Surgery Service status post
gastric bypass surgery on 2130-4-14.  The surgery was
uncomplicated, however, Mr. Jefferson was admitted to the Surgical
Intensive Care Unit after his gastric bypass secondary to
unable to extubate secondary to a respiratory acidosis.  The
patient had decreased urine output, but it picked up with
intravenous fluid hydration.  He was successfully extubated
on 4-15 in the evening and was transferred to the floor
on 2130-4-16 without difficulty.  He continued to have
slightly labored breathing and was requiring a face tent mask
to keep his saturations in the high 90s.  However, was
advanced according to schedule and tolerated a stage two diet
and was transferred to the appropriate pain management.  He
was out of bed without difficulty and on postoperative day
three he was advanced to a stage three diet and then slowly
was discontinued.  He continued to use a face tent overnight,
but this was discontinued during the day and he was advanced
to all of the usual changes for postoperative day three
gastric bypass patient.  He will be discharged home today
postoperative day three in stable condition status post
gastric bypass.
DISCHARGE MEDICATIONS:  Vitamin B-12 1 mg po q.d., times two
months, Zantac 150 mg po b.i.d. times two months, Actigall
300 mg po b.i.d. times six months and Roxicet elixir one to
two teaspoons q 4 hours prn and Albuterol Atrovent meter dose
inhaler one to two puffs q 4 to 6 hours prn.
He will follow up with Dr. Morrow in approximately two weeks as
well as with the Lowery Medical Center Clinic.
                          Kevin Gonzalez, M.D.  R35052373
Dictated By:Dotson
MEDQUIST36
D:  2130-4-17  08:29
T:  2130-4-18  08:31
JOB#:  Job Number 20340"""


input_text = st.text_area("Notes:", prompt)
if st.button('Summarize'):
    final_input = f"""[Notes]:\n{input_text}\n[Summary]:\n"""
    st.write(summarize_function(final_input))