Update predictor.py

predictor.py

@@ -59,6 +59,35 @@ def srl_init(model_path, bert_name="bert-base-cased", spacy_model="en_core_web_m
     })
     torch.set_grad_enabled(False)
 
+def normalize_whitespace(s: str) -> str:
+    if s is None: return ""
+    return s.replace("\u00A0", " ").replace("\u2009", " ").strip()
+
+def spacy_verb_indices(nlp, sentence: str):
+    doc = nlp(sentence)
+    return [i for i, t in enumerate(doc) if t.pos_ in ("VERB","AUX") or t.tag_.startswith("VB")]
+
+def words_and_spans_spacy(sentence: str, nlp):
+    doc = nlp(sentence)
+    words = [t.text for t in doc]
+    spans = [(t.idx, t.idx + len(t.text)) for t in doc]
+    return words, spans
+
+def bio_to_spans(tags):
+    spans = []; i = 0
+    while i < len(tags):
+        t = tags[i]
+        if t == "O" or t.endswith("-V"):
+            i += 1; continue
+        if t.startswith("B-"):
+            role = t[2:]; j = i+1
+            while j < len(tags) and tags[j] == f"I-{role}": j += 1
+            spans.append((role, i, j-1)); i = j
+        else:
+            i += 1
+    return spans
+
+
 def _predict_cached(sentence):
     """Internal: uses cached objects set by srl_init()."""
     if _CACHE["model"] is None:
@@ -69,7 +98,8 @@ def _predict_cached(sentence):
     nlp       = _CACHE["nlp"]
     device    = "cuda" if (_CACHE["device"].type == "cuda") else "cpu"
 
-
+    sentence = normalize_whitespace(sentence)
+    
     return predict_srl_allennlp_like_spacy(
         model, tokenizer, nlp, sentence, id2label,
         device=device, prob_threshold=0.40, top_k=None, pick_best_if_none=True
@@ -137,35 +167,6 @@ def _ensure_loaded(model_path, bert_name, spacy_model, model_cls):  # NEW
     torch.set_grad_enabled(False)
 
 
-def normalize_whitespace(s: str) -> str:
-    if s is None: return ""
-    return s.replace("\u00A0", " ").replace("\u2009", " ").strip()
-
-def spacy_verb_indices(nlp, sentence: str):
-    doc = nlp(sentence)
-    return [i for i, t in enumerate(doc) if t.pos_ in ("VERB","AUX") or t.tag_.startswith("VB")]
-
-def words_and_spans_spacy(sentence: str, nlp):
-    doc = nlp(sentence)
-    words = [t.text for t in doc]
-    spans = [(t.idx, t.idx + len(t.text)) for t in doc]
-    return words, spans
-
-def bio_to_spans(tags):
-    spans = []; i = 0
-    while i < len(tags):
-        t = tags[i]
-        if t == "O" or t.endswith("-V"):
-            i += 1; continue
-        if t.startswith("B-"):
-            role = t[2:]; j = i+1
-            while j < len(tags) and tags[j] == f"I-{role}": j += 1
-            spans.append((role, i, j-1)); i = j
-        else:
-            i += 1
-    return spans
-
-
 @torch.no_grad()
 def predict_srl_single(model, tokenizer, words, predicate_word_idx, id2label, device="cuda"):
     model.eval()