Update app.py

app.py

@@ -29,10 +29,21 @@ def load_cnn():
 
   return model
 
-def predict_sentiment(text, model):
-    processor.transform(text)
-    prediction = model.predict([text])
-    return prediction
+def predict_sentiment(text, model, torch=False):
+    if torch == True:
+        processed_text = processor.transform(text)
+        with torch.no_grad():  # Ensure no gradients are computed
+            prediction = model(processed_text)  # Get raw model output
+            # Convert output to probabilities
+            probs = torch.softmax(prediction, dim=1)
+            # Get the predicted class
+            pred_class = torch.argmax(probs, dim=1)
+        return pred_class.item()  # Return the predicted class as a Python int
+    else:
+        processor.transform(text)
+        prediction = model.predict([text])
+        return prediction
+        
 
 
 model_1 = load_model()
@@ -62,7 +73,7 @@ with st.expander("Model 2: CNN Sentiment analysis"):
     user_input = st.text_area("Enter text here...", key='model2_input')
     if st.button('Analyze', key='model2_button'):
         # Displaying output
-        result = predict_sentiment(user_input, model_2)
+        result = predict_sentiment(user_input, model_2, torch=True)
         if result >= 0.5:
             st.write('The sentiment is: Positive 😀', key='model2_poswrite')
         else: