Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 st.set_page_config(page_title="Common NLP Tasks")
 st.title("Common NLP Tasks")
@@ -25,8 +25,10 @@ def summarization_model():
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def generation_model():
-    model_name = "distilgpt2"
-    generator = pipeline(model=model_name, tokenizer=model_name, task="text-generation")
+    # model_name = "distilgpt2"
+    tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+    model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+    generator = pipeline(model=model, tokenizer=tokenizer, task="text-generation")
     return generator
 
 if option == "Extractive question answering":