Update app.py

app.py

@@ -25,7 +25,8 @@ def summarization_model():
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def generation_model():
-    generator = pipeline("text-generation")
+    model_name = "distilgpt2"
+    generator = pipeline(model=model_name, tokenizer=model_name, task="text-generation")
     return generator
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
@@ -40,10 +41,11 @@ if option == "Extractive question answering":
         with open("sample.txt", "r") as text_file:
             sample_text = text_file.read()
         context = st.text_area('Use the example below or input your own text in English (10,000 characters max)', value=sample_text, max_chars=10000, height=330)
-        question = st.text_input(label='Enter your question')
-        button = st.button('Get answer')
+        question = st.text_input(label="Enter your question")
+        button = st.button("Get answer")
         if button:
-            question_answerer = question_model()
+            with st.spinner(text="Loading question model..."):
+                question_answerer = question_model()
             with st.spinner(text="Getting answer..."):
                 answer = question_answerer(context=context, question=question)
                 answer = answer["answer"]
@@ -57,7 +59,8 @@ if option == "Extractive question answering":
             question = st.text_input(label="Enter your question")
             button = st.button("Get answer")
             if button:
-                question_answerer = question_model()
+                with st.spinner(text="Loading summarization model..."):
+                    question_answerer = question_model()
                 with st.spinner(text="Getting answer..."):
                     answer = question_answerer(context=context, question=question)
                     answer = answer["answer"]
@@ -72,7 +75,8 @@ elif option == "Text summarization":
         text = st.text_area("Input a text in English (10,000 characters max) or use the example below", value=sample_text, max_chars=10000, height=330)
         button = st.button("Get summary")
         if button:
-            summarizer = summarization_model()
+            with st.spinner(text="Loading summarization model..."):
+                summarizer = summarization_model()
             with st.spinner(text="Summarizing text..."):
                 summary = summarizer(text, max_length=130, min_length=30)
                 st.write(summary[0]["summary_text"])
@@ -90,22 +94,24 @@ elif option == "Text summarization":
                     summary = summarizer(text, max_length=130, min_length=30)
                     st.write(summary[0]["summary_text"])
                 
-elif option == 'Text generation':
+elif option == "Text generation":
     st.markdown("<h2 style='text-align: center; color:grey;'>Generate text</h2>", unsafe_allow_html=True)
-    text = st.text_input(label='Enter one line of text and let the NLP model generate the rest for you')
-    button = st.button('Generate text')
+    text = st.text_input(label="Enter one line of text and let the NLP model generate the rest for you")
+    button = st.button("Generate text")
     if button:
-        generator = generation_model()
+        with st.spinner(text="Loading text generation model..."):
+            generator = generation_model()
         with st.spinner(text="Generating text..."):
             generated_text = generator(text, max_length=50)
             st.write(generated_text[0]["generated_text"])
                 
-elif option == 'Sentiment analysis':
+elif option == "Sentiment analysis":
     st.markdown("<h2 style='text-align: center; color:grey;'>Classify review</h2>", unsafe_allow_html=True)
     text = st.text_input(label='Enter a sentence to get its sentiment analysis')
-    button = st.button('Get sentiment analysis')
+    button = st.button("Get sentiment analysis")
     if button:
-        sentiment_analysis = sentiment_model()
+        with st.spinner(text="Loading sentiment analysis model..."):
+            sentiment_analysis = sentiment_model()
         with st.spinner(text="Getting sentiment analysis..."):
             sentiment = sentiment_analysis(text)
             st.write(sentiment[0]["label"])