Add conditional

functions/model_infer.py

@@ -39,8 +39,11 @@ def predict_from_document(sentences):
     preprop = preprocess(sentences)
     prediction = model.predict(preprop)
     # Set the prediction threshold to 0.8 instead of 0.5, now use mean
-    output = (
-        prediction.flatten()[: len(sentences)]
-        >= np.mean(prediction) * 1.20  # + np.std(prediction)
-    ).astype(int)
+    if np.mean(prediction) < 0.5:
+        output = (prediction.flatten()[: len(sentences)] >= 0.5).astype(int)
+    else:
+        output = (
+            prediction.flatten()[: len(sentences)]
+            >= np.mean(prediction) * 1.20  # + np.std(prediction)
+        ).astype(int)
     return output, prediction.flatten()[: len(sentences)]