Update app.py

app.py

@@ -13,38 +13,35 @@ st.markdown("<h6 style='text-align: center; color: #489DDB;'>by Bryan Mildort</h
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # 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)
 # device_map = infer_auto_device_map(model, dtype="float16")
 # st.write(device_map)
 
-model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True).cuda()
-tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
+@st.cache(allow_output_mutation=True)
+def load_model():
+    model = pipeline("text-generation", model="bryanmildort/gpt_neo_notes")
+    return model
+   
 # model = model.to(device)
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+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
@@ -66,23 +63,16 @@ 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"""
@@ -91,5 +81,4 @@ 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))
-
+    st.write(summarize_function(final_input))