Update app.py

app.py

@@ -18,6 +18,9 @@ def load_model():
 
 interpreter, text_preprocessor, label_encoder, tokenizer = load_model()
 
+def softmax(x):
+    return tf.exp(x) / tf.math.reduce_sum(tf.exp(x))
+
 def inference(text):
     tflite_pred = "Can't Predict"
     text = text_preprocessor.preprocess(pd.Series(text))[0]
@@ -32,13 +35,14 @@ def inference(text):
         interpreter.set_tensor(input_details[1]["index"], input_ids)
         interpreter.invoke()
         tflite_pred = interpreter.get_tensor(output_details["index"])[0]
+        tflite_pred = softmax(tflite_pred)
         tflite_pred_argmax = np.argmax(tflite_pred)
-        tflite_pred = f"{label_encoder.inverse_transform([tflite_pred_argmax])} ({tflite_pred[tflite_pred_argmax]})"
+        tflite_pred = f"{label_encoder.inverse_transform([tflite_pred_argmax])} ({str(np.round(tflite_pred[tflite_pred_argmax], 3))})"
     return tflite_pred
 
 
 def main():
-    st.title("Sentiment Analysis App")
+    st.title("Sentiment Analysis")
     review = st.text_area("Enter Review:", "", height=200)
     if st.button("Submit"):
         # result = "Can't Predict"