meh

app.py

@@ -28,8 +28,11 @@ def get_tokenizer(model_name):
 
 @st.cache(allow_output_mutation=True)
 def get_model(model_name):
-    return AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
+    import torch
+    from transformers import AutoModelForCausalLM
+    mymodel = AutoModelForCausalLM.from_pretrained(model, device_map='auto', torch_dtype=torch.bfloat16)
+    print(f"Loaded model, {mymodel.num_parameters():,d} parameters.")
+    return mymodel
 original_text = st.text_area("Original Document", "This is my text that I want to change.")
 new_text = st.text_area("New Document", help="Leave this blank, this will be updated with your new document.")