Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import pandas as pd
 import re
 
 def summarize_function(notes):
@@ -24,78 +25,28 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 def load_model():
     model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True)
     tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    return pipe
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
    
-# model = model.to(device)
+model = model.to(device)
 
 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"""
+notes_df = pd.read_csv('notes_small.csv')
+examples_tuple = ()
+for i in range(len(notes_df)):
+    examples_tuple += (f"Patient {i+1}", )
 
+example = st.sidebar.selectbox('Example', (examples_tuple), index=0)
+st.write(example)
 
+prompt = notes_df.iloc[int(example[-1:])-1].PARSED
 input_text = st.text_area("Notes:", prompt)
 if st.button('Summarize'):
     parsed_input = re.sub(r'\n\s*\n', '\n\n', input_text)
     parsed_input = re.sub(r'\n+', '\n',parsed_input)
     final_input = f"""[Notes]:\n{parsed_input}\n[Summary]:\n"""
-    st.write(summarize_function(final_input))
+    st.write(summarize_function(final_input))
+
+
+