app.py

import streamlit as st

def summarize_function(notes):
    notes_input = tokenizer(notes, return_tensors='pt')
    output = model(**notes_input)
    max_length = len(output['logits'][0]) + 40
    # max_length = len(notes.split(' ')) + 40
    input_ids = notes_input.input_ids
    gen_tokens = model.generate(input_ids, do_sample=True, temperature = 0.5, max_length=max_length)
    gen_text = tokenizer.batch_decode(gen_tokens)[0]
    return gen_text[len(notes):]

st.markdown("<h1 style='text-align: center; color: #489DDB;'>GPT Clinical Notes Summarizer</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, GPTJForCausalLM
# from accelerate import infer_auto_device_map
import torch
device = "cuda:0" if torch.cuda.is_available() else "cpu"
device_str = f"""Device being used: {device}"""
st.write(device_str)

# model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt-clinical-notes-summarizer", load_in_8bit=True, device_map="auto")
# model = model.to(device)

tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt-clinical-notes-summarizer")
checkpoint = "bryanmildort/gpt-clinical-notes-summarizer"
model = GPTJForCausalLM.from_pretrained(checkpoint)

# device_map = infer_auto_device_map(model, dtype="float16")
# st.write(device_map)

prompt = """Edie Whelan is a 26 yo female who is here for follow-up following an ED visit for palpitations 2 weeks ago. She had a normal workup in the ED consisting of CBC, metabolic panel, cardiac enzymes and and ECG. She states that she has a 5 yr hx of palpitations with SOB, throat tightening, nause and clamminess. these have been increasing in frequency in the past few weeks to 1-2x a day until 2 weeks ago when she had associated numbness in her fingers. She states that there are no precipitating events and these stop on their own. She endorses recent life stressors having bought a condo 3 months ago and lost her job 2 months ago.  She denies fevers, chills, changes in skin/hair/nails or chest pain/tightness. She denies vision or hearing changes. 
PMH: none, Meds: none, Allergies: none, Surgeries: None, Family hx: none
Social: denies caffeine, tobacco, alcohol, and drug use. Is in a monogomous relationship and has sex with condoms"""

prefix = """[Notes]: 35 yo female,  comes in to clinic due to heavy mentrual periods. initiated 5 months ago, her last menstrual period was 6 months ago, she denias any pain while MP. she also details. weight gain of 18 pounds, she has had an unstable diet. she tried to concieve 9 years ago  with any success, she has no hayd any misscarragies. she denies any shortness of breath, diaphoresis, head flushing, palpitations, sincope, chills, fever or abdominal pain. pap smear 6 months ago.

ros: as above

fh: parents healthy, granmother cervical cancer

psh: none

sh: alcohol: beer in ocaciones, tabaco: none, drugs: none

sexual active, with boyfriend, does no use condoms

ob gyn: menarche 12 yo , regular cycles until 6 months ago, prolonged to more than a month, took contraceptive pills before (control of mentrual periods)

allergies: none, 

meds: none 

diet: no strict diet

exercise: none
[Summary]: 35yo F c/o heavy menstrual periods x5mos, weight gain, unstable diet, no conception, no pap smear in 6mos, FH grandmother cervical cancer, SH alcohol, no tabaco, no drugs,
[Notes]: 17 yr Old M comes to clinic with c/o Heart Pounding, which started 2-3 Months. Symptoms comes and goes and last 3-4 min. Pt had 5-6 episodes so far and goes away on its on. Also feeling chest pressure which is in center x 2 days and difficulty in breathing x 2 days. Denies headache, diaphoresis, loss of consciouness, flushing, diarrhea, constipation, hair and skin chages. Diet and Appetite is normal. Pt is on Health plan for diet and eats fruits, veggies and protein meal.

ROS negative except above

PMH None Meds  Takes addrall from room mate, NKDA

PSH none   Hospital none   No trauma

FH Mother had thyroid condition, Father has heart problem

SH student, No smoking, ETOH occasional, Recreational drugs Pot one time

SEX H Active with GF, Use condom and NO STD.
[Summary]: 17yo M c/o heart pounding x2-3mos, chest pressure x2days, difficulty breathing x2days, addrall, FH mother thyroid, father heart, SH student, ETOH occasional, recreational drugs
[Notes]: 67 yo F patient presents to clinic 3 weeks after the death of her son with difficulty falling asleep and waking early in the morning. She has tried ambien to help sleep without success. She also endorses having a sad or depressed mood, decreased energy, decreased interest in things she used to enjoy, and increased appetite. She endorses visual hallucinations of seeing her son at the kitchen table and auditory hallucinations of noise coming from the neighbor. She denies any suicidal or homicidal thoughts, intents, actions, or plans. She reports that she has a good support system with her family and friends. 

ROS: denies HA, vision changes, fever, chills, night sweats, N/V, changes in bowel or bladder habits, chest pain, SOB, or changes in weight. 

PMH: HTN, breast cancer 10 years in remission. Meds: HCTZ, lisinopril, ambrian for sleep. NKDA.

Sx Hx: lumpectomy, laparotomy for complicated appendicits. 

Mother suffered depression.
[Summary]: 67yo F c/o difficulty sleeping, sad/depressed mood, decreased energy, decreased interest, increased appetite, visual/auditory hallucinations, HTN, breast cancer"""

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