Update app.py

app.py

@@ -15,15 +15,29 @@ def load_text_gen_model():
     generator = pipeline("text-generation", model="gpt2-medium")
     return generator 
     
-#model = load_model("gpt2-medium")
-
+@st.cache
+def get_sentiment_model():
+    sentiment_tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-imdb-sentiment")
+    sentiment_model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-imdb-sentiment")
+    return sentiment_model ,sentiment_tokenizer 
+    
+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)
+    dec = [tokenizer.decode(ids) for ids in output]
+    label = dec[0]
+    return label
+  
+sentiment_model ,sentiment_tokenizer  = get_sentiment_model()
 text_generator = load_text_gen_model()
 
-action = st.sidebar.selectbox("Pick an Action", ["Generate an Article","Create an Image"])
+action = st.sidebar.selectbox("Pick an Action", ["Analyse a Review","Generate an Article","Create an Image"])
 
-if action == "Generate an Article":
-    prompt = st.text_input("Enter a prompt")
-    if prompt:
-        res = text_generator( prompt, max_length=100, temperature=0.7)
-        st.write(res)
+if action == "Analyse a Review":
+    review = st.text_area("Paste the review here..")
+    if review:
+        #res = text_generator( prompt, max_length=100, temperature=0.7)
+        #st.write(res)
+        sentiment = get_sentiment(review)
+        st.write(sentiment)