Update app.py

app.py

@@ -181,9 +181,11 @@ def on_btn_click():
 
 @st.cache_resource
 def load_model():
-    model = (AutoModelForCausalLM.from_pretrained(
-        model_name_or_path,
-        trust_remote_code=True).to(torch.bfloat16).cuda())
+    # model = (AutoModelForCausalLM.from_pretrained(
+    #     model_name_or_path,
+    #     trust_remote_code=True).to(torch.bfloat16).cuda())
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+    model.to_bettertransformer()
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                               trust_remote_code=True)
     return model, tokenizer