Update app.py

app.py

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