Update app.py

app.py

@@ -14,6 +14,8 @@ with open(os.path.join("models", "toxic_comment_preprocessor_classnames.bin"), "
         text_preprocessor, class_names = cloudpickle.load(model_file_obj)    
 interpreter = tf.lite.Interpreter(model_path=os.path.join("models", "toxic_comment_classifier_hf_distilbert.tflite"))
 
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
 
 def inference(text):
     text = text_preprocessor.preprocess(pd.Series(text))[0]
@@ -29,7 +31,9 @@ def inference(text):
     interpreter.set_tensor(input_details[0]["index"], attention_mask)
     interpreter.set_tensor(input_details[1]["index"], input_ids)
     interpreter.invoke()
-    tflite_pred = interpreter.get_tensor(output_details["index"])[0]
+    tflite_logits = interpreter.get_tensor(output_details["index"])[0]
+    tflite_pred = sigmoid(tflite_logits)
+    
     result_df = pd.DataFrame({'class': class_names, 'prob': tflite_pred})
     result_df.sort_values(by='prob', ascending=True, inplace=True)
     return result_df