Allow user to set confidence threshold for frames + roles, default to 0.95

sociolome/lome_webserver.py

@@ -47,7 +47,7 @@ def convert_to_seq_labels(sentence: List[str], structures: Dict[int, Dict[str, A
                 labels[i].append(f"{prefix}:{frame}:{role_label}@{struct_id:02}")
     return labels
 
-def make_prediction(sentence, spacy_model, predictor):
+def make_prediction(sentence, spacy_model, predictor, confidence_threshold):
     spacy_doc = spacy_model(sentence)
     tokens = [t.text for t in spacy_doc]
     tgt_spans, fr_labels, fr_probas = predictor.force_decode(tokens)
@@ -59,7 +59,7 @@ def make_prediction(sentence, spacy_model, predictor):
             continue
         if frm.upper() == frm:
             continue
-        if fr_proba.max() != 1.0:
+        if fr_proba.max() != confidence_threshold:
             continue
         
         arg_spans, arg_labels, label_probas = predictor.force_decode(tokens, parent_span=tgt, parent_label=frm)
@@ -70,7 +70,7 @@ def make_prediction(sentence, spacy_model, predictor):
                 "roles": [
                     {"boundary": bnd, "label": label}
                     for bnd, label, probas in zip(arg_spans, arg_labels, label_probas)
-                    if label != "Target" and max(probas) == 1.0
+                    if label != "Target" and max(probas) == confidence_threshold
                 ]
             }
 
@@ -96,9 +96,10 @@ app = Flask(__name__)
 @app.route("/analyze")
 def analyze():
     text = request.args.get("text")
+    confidence_threshold = float(request.args.get("confidence_threshold", 0.95))
     analyses = []
     for sentence in text.split("\n"):
-        analyses.append(make_prediction(sentence, nlp, predictor))
+        analyses.append(make_prediction(sentence, nlp, predictor, confidence_threshold))
 
     return jsonify({
         "result": "OK",