Update app.py

app.py

@@ -13,7 +13,7 @@ def summarize_function(notes):
 st.markdown("<h1 style='text-align: center; color: #489DDB;'>GPT Clinical Notes Summarizer</h1>", unsafe_allow_html=True)
 st.markdown("<h6 style='text-align: center; color: #489DDB;'>by Bryan Mildort</h1>", unsafe_allow_html=True)
 
-from transformers import AutoTokenizer, GPTJForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 # from accelerate import infer_auto_device_map
 import torch
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -21,11 +21,15 @@ device_str = f"""Device being used: {device}"""
 st.write(device_str)
 
 # model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt-clinical-notes-summarizer", load_in_8bit=True, device_map="auto")
-# model = model.to(device)
+# 
 
-tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt-clinical-notes-summarizer")
-checkpoint = "bryanmildort/gpt-clinical-notes-summarizer"
-model = GPTJForCausalLM.from_pretrained(checkpoint)
+# tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt-clinical-notes-summarizer")
+# checkpoint = "bryanmildort/gpt-clinical-notes-summarizer"
+# model = GPTJForCausalLM.from_pretrained(checkpoint)
+
+model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt-neo-small-notes-summarizer")
+tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt-neo-small-notes-summarizer")
+model = model.to(device)
 
 # device_map = infer_auto_device_map(model, dtype="float16")
 # st.write(device_map)