Update app.py

app.py

@@ -4,13 +4,7 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelWithLMHead
 from transformers import pipeline
 
-#tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
 sentiment_tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-imdb-sentiment")
-
-@st.cache
-def load_model(model_name):
-    model = AutoModelWithLMHead.from_pretrained(model_name)
-    return model
     
 def load_text_gen_model():
     generator = pipeline("text-generation", model="gpt2-medium")
@@ -20,18 +14,19 @@ def load_text_gen_model():
 def get_sentiment_model():
     sentiment_model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-imdb-sentiment")
     return sentiment_model 
-    
+
+def get_summarizer_model():
+    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+    return summarizer
+
+      
 def get_sentiment(text):
     input_ids = sentiment_tokenizer .encode(text + '</s>', return_tensors='pt')
-    output = sentiment_model.generate(input_ids=input_ids,max_length=2)
+    output = sentiment_extractor.generate(input_ids=input_ids,max_length=2)
     dec = [sentiment_tokenizer.decode(ids) for ids in output]
     label = dec[0]
     return label
     
-#@st.cache(allow_output_mutation=True) 
-def get_summarizer():
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    return summarizer
 
 def get_qa_model():
     model_name = "deepset/roberta-base-squad2"
@@ -39,11 +34,9 @@ def get_qa_model():
     qa_pipeline = pipeline('question-answering', model=model_name, tokenizer=model_name)
     return qa_pipeline
         
-sentiment_model   = get_sentiment_model()
-summarizer = get_summarizer()
-answer_geerator = get_qa_model()
-
-#text_generator = load_text_gen_model()
+sentiment_extractor   = get_sentiment_model()
+summarizer = get_summarizer_model()
+answer_generator = get_qa_model()
 
 action = st.sidebar.selectbox("Pick an Action", ["Analyse a Review","Generate an Article","Create an Image"])
 
@@ -63,7 +56,7 @@ if action == "Analyse a Review":
         if st.button("Find the key topic"):
             QA_input = {'question': 'what is the review about?',
                         'context': review}
-            answer = answer_geerator (QA_input)
+            answer = answer_generator(QA_input)
             st.write(answer)