Update app.py

app.py

@@ -19,11 +19,17 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # 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)
-tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
+@st.cache_resource
+def load_model()
+    st.write('Loading model...')
+    model = pipeline("text-generation", model="bryanmildort/gpt_neo_notes")
+    st.write('Success!')
+    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