Update app.py

app.py

@@ -1,10 +1,11 @@
+import transformers
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelWithLMHead
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2-large")
 @st.cache_data
 def load_model(model_name):
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelWithLMHead.from_pretrained(model_name)
 return model
 
 model = load_model("gpt2-large")
@@ -20,6 +21,7 @@ top_p=top_p,
 do_sample=True,
 num_return_sequences=1
 )
+
 return output_sequences
 
 default_value = "You: Ask me anything!"