Update app.py

app.py

@@ -13,7 +13,7 @@ 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
-# device = "cuda:0" if torch.cuda.is_available() else "cpu"
+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")
@@ -21,10 +21,11 @@ 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 = model.to(device)
+    model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True)
+    model.to(device)
+    tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    return pipe
 
 pipe = load_model()