Update app.py

app.py

@@ -21,8 +21,10 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 @st.cache(allow_output_mutation=True)
 def load_model():
-    model = pipeline("text-generation", model="bryanmildort/gpt_neo_notes")
-    return 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
    
 # model = model.to(device)
 
@@ -93,5 +95,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))