Update the threshold with variance and mean

functions/model_infer.py

@@ -37,6 +37,9 @@ def preprocess(sentences):
 def predict_from_document(sentences):
     preprop = preprocess(sentences)
     prediction = model.predict(preprop)
-    # Set the prediction threshold to 0.8
-    output = (prediction.flatten()[: len(sentences)] >= 0.85).astype(int)
+    # Set the prediction threshold to 0.8 instead of 0.5, now use mean
+    output = (
+        prediction.flatten()[: len(sentences)]
+        >= np.mean(prediction) + np.var(prediction) * 2
+    ).astype(int)
     return output, prediction.flatten()[: len(sentences)]