Upload predicator.py

predicator.py

@@ -1,5 +1,7 @@
-## This is testing 
 
+from SRL_model import SRL_BERT_model
+from transformers import AutoTokenizer
+import spacy
 import torch
 
 @torch.no_grad()
@@ -64,6 +66,58 @@ def predict_srl_single(model, tokenizer, words, predicate_word_idx, id2label, de
     tags = [id2label[i] for i in pred_ids]
     return tags, logits.squeeze(0).cpu()  # [L_word, num_labels]
 
+
+
+def spacy_verb_indices(nlp, sentence: str,word_spans):
+    """
+    Map spaCy POS to each tokenizer word span by max-overlap token,
+    and return indices that are verbs (VB*, or POS in {VERB, AUX}).
+    """
+    doc = nlp(sentence)
+    verb_idxs = []
+    for i, (wb, we) in enumerate(word_spans):
+        # find spaCy token with maximum overlap
+        best_tok, best_olap = None, 0
+        for tok in doc:
+            tb, te = tok.idx, tok.idx + len(tok)
+            olap = max(0, min(we, te) - max(wb, tb))
+            if olap > best_olap:
+                best_olap = olap
+                best_tok  = tok
+        if best_tok is None: 
+            continue
+        is_verb = best_tok.tag_.startswith("VB") or best_tok.pos_ in {"VERB", "AUX"}
+        if is_verb:
+            verb_idxs.append(i)
+
+    
+    return verb_idxs
+
+def words_and_spans_from_tokenizer(sentence: str, nlp, tokenizer):
+    """
+    Returns:
+      words : list[str] (tokenizer-aligned words)
+      spans : list[(start_char, end_char)] for each word
+    """
+    enc = tokenizer(sentence, add_special_tokens=False, return_offsets_mapping=True)
+    word_ids = enc.word_ids()
+    offsets  = enc["offset_mapping"]
+    # print("Offsets:", offsets)
+    
+    words, spans, seen = [], [], set()
+    for wid, (b, e) in zip(word_ids, offsets):
+        if wid is None or wid in seen:
+            continue
+        seen.add(wid)
+        words.append(sentence[b:e])
+        spans.append((b, e))
+        
+        
+    doc = nlp(sentence)
+    words_final = [token for token in doc]
+    
+    return words, spans, words_final
+
 def bio_to_spans(tags):
     spans = []
     i = 0
@@ -84,58 +138,188 @@ def bio_to_spans(tags):
     return spans
 
 @torch.no_grad()
-def predict_srl_all_predicates(model, tokenizer, sentence, id2label, device="cuda", prob_threshold=0.50):
-    words = sentence.split()
-    # find the numeric id for "B-V"
-    b_v_id = None
-    for k, v in id2label.items():
-        if v == "B-V":
-            b_v_id = k
-            break
-    if b_v_id is None:
-        raise ValueError("Label set has no 'B-V' tag.")
+def predict_srl_allennlp_like_spacy(
+    model, tokenizer, nlp, sentence, id2label,
+    device="cuda"):
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.eval()
+
+    # 0) tokenizer-aligned words (and spans for spaCy alignment)
+    words, spans, words_final = words_and_spans_from_tokenizer(sentence, nlp,tokenizer)
+    # words = words_and_spans_from_tokenizer(sentence, nlp)
+    # words = sentence.split(' ')
+    # print(spans)
+    n = len(words)
+    if n == 0:
+        return [], []
+
+    # 1) verb candidates via spaCy POS (aligned to tokenizer words by overlap)
+    verb_idxs = spacy_verb_indices(nlp, sentence, spans)
+    # verb_idxs = spacy_verb_indices(nlp, sentence)
+    if not verb_idxs:
+        # no verbs → either return empty or consider all tokens as a fallback
+        return words, []
+
+    # 2) find the predicate label id ('B-V' or 'V')
+    pred_ids = [i for i, t in id2label.items() if t in ("B-V", "V")]
+    if not pred_ids:
+        raise ValueError("Label set has no predicate tag ('B-V' or 'V').")
+    b_v_id = pred_ids[0]
 
+    keep = verb_idxs
+
+    # 4) PASS 2: final tagging with indicator ON
     results = []
-    for p in range(len(words)):
-        tags, logits = predict_srl_single(model, tokenizer, words, p, id2label, device=device)
-        # check predicate decision at position p
-        pred_id_at_p = logits.argmax(-1)[p].item()
-        keep = (pred_id_at_p == b_v_id)
+    for p in keep:
+        tags, logits = predict_srl_single(
+            model, tokenizer, words, p, id2label, device=device
+            # , use_indicator=True
+        )
+        # if require_core_arg and not has_core_argument(tags):
+        #     continue
+        p_bv = torch.softmax(logits[p], dim=-1)[b_v_id].item()
+        spans_out = bio_to_spans(tags)
+        results.append({
+            "predicate_index": p,
+            "predicate": words[p],
+            "p_bv": p_bv,
+            "tags": tags,
+            "spans": spans_out
+        })
+
+    return words_final, results
+
+
+def create_description(words, tag_list):
+    desc_list = []
+    for tok, tag in zip(words, tag_list):
+        desc_list.append("["+tag+": "+tok+"]")
+
+    return (' ').join(desc_list)
 
-        # optional confidence gate
-        if prob_threshold is not None:
-            probs = torch.softmax(logits[p], dim=-1)
-            keep = keep and (probs[b_v_id].item() >= prob_threshold)
+def print_srl_frames_pretty(words, frames, show_grid=True, color=False):
+    """
+    Pretty-print SRL frames.
+    - Description: Token+Labels
+    - Frames: Predicate/Roles
+    - show_grid: also print a token/label grid aligned by column
+    - color: add simple ANSI colors per role (terminal only)
+    """
+    import itertools
 
-        if keep:
-            spans = bio_to_spans(tags)
-            results.append({
-                "predicate_index": p,
-                "predicate": words[p],
-                "tags": tags,
-                "spans": spans
-            })
-    return words, results
+    # tiny colorizer (terminal); safe no-op if color=False
+    ANSI = {
+        "ARG0": "\033[38;5;34m", "ARG1": "\033[38;5;33m", "ARG2": "\033[38;5;129m",
+        "ARG3": "\033[38;5;172m", "ARG4": "\033[38;5;166m", "ARGM": "\033[38;5;244m",
+        "V": "\033[1;37m", "RESET": "\033[0m"
+    }
+    def paint(txt, role):
+        if not color: return txt
+        key = "ARGM" if role.startswith("ARGM") else ("V" if role.endswith("V") or role=="V" else role)
+        return f"{ANSI.get(key, '')}{txt}{ANSI['RESET']}"
 
+    def spans_from_bio(tags):
+        spans = []
+        i = 0
+        while i < len(tags):
+            t = tags[i]
+            if t == "O":
+                i += 1; continue
+            if t.endswith("-V"):  # you can include/exclude the V span as you like
+                spans.append(("V", i, i))
+                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
+    
+    words = [word.text for word in words]
+    print("Sentence:", " ".join(words))
+    if not frames:
+        print("  (no predicates detected)")
+        return
 
+    for k, fr in enumerate(frames, 1):
+        tags = fr["tags"]
+        spans = fr.get("spans") or spans_from_bio(tags)
+        pred_idx = fr["predicate_index"]
+        pred = fr["predicate"]
+        p_bv = fr.get("p_bv", None)
 
-# words, preds = predict_srl_all_predicates(model, tokenizer, sentence, id2label, device=device)
+        print("\n" + "—"*60)
 
+        print(create_description(words, tags))
+        
+        # Aggregate phrases per role for a clean summary
+        by_role = {}
+        for role, s, e in spans:
+            phrase = " ".join(words[s:e+1])
+            by_role.setdefault(role, []).append(phrase)
 
-def predicator_srl(sentence):
-    words, preds = predict_srl_all_predicates(model, tokenizer, sentence, id2label, device=device)
+        # Put V first, then core args, then ARGM*
+        order = (
+            (("V",),),
+            tuple((r,) for r in ["ARG0","ARG1","ARG2","ARG3","ARG4"]),
+            (tuple(sorted([r for r in by_role if r.startswith("ARGM")])),)
+        )
+        ordered_roles = []
+        for group in order:
+            for r in itertools.chain.from_iterable(group):
+                if r in by_role: ordered_roles.append(r)
+        # add any leftover roles
+        for r in sorted(by_role):
+            if r not in ordered_roles: ordered_roles.append(r)
+        print("Predicate:")
+        print(f"  {r:<8}: {pred}")
+        print("Roles:")
+        for r in ordered_roles:
+            joined = "; ".join(by_role[r])
+            print(f"  {r:<8}: {paint(joined, r)}")
 
-    return words, preds
+        if show_grid:
+            # token/tag grid aligned by column width
+            colw = [max(len(w), len(t)) for w, t in zip(words, tags)]
+            tok_row = " ".join(w.ljust(colw[i]) for i, w in enumerate(words))
+            tag_row = " ".join((t if t != "O" else ".").ljust(colw[i]) for i, t in enumerate(tags))
+            print("\nTOKEN:", tok_row)
+            print("LABEL:", tag_row)
 
-if __name__ == "__main__":
-    sentence = "Hojeong decide to go to the school"
-    words, preds = predicator_srl(sentence)
-    print(words)
-    for r in preds:
-        print(f"Predicate: {r['predicate']} (idx {r['predicate_index']})")
-        print("Tags:", list(zip(words, r["tags"])))
-        print("Spans:", r["spans"])  # (ROLE, start, end) indices over words
-        print("-" * 60)
+def main_predictor(model_path, bert_name, sentence):
 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # model_path = "/blue/bonniejdorr/youms/SRL-Aware_Model/model/best_srl_Sep_29.ckpt"
+    ckpt = torch.load(model_path, map_location=device)
+    hp = ckpt["hparams"]
 
+    model = SRL_BERT_model.PredicateAwareSRL(**hp).to(device)
+    model.load_state_dict(ckpt["state_dict"])
+    model.eval()
+
+    label2id = ckpt["label2id"]
+    id2label = {v:k for k,v in label2id.items()}
+
+    # bert_name = "bert-large-cased" or "bert-based-cased"
+    bert_name = bert_name
+    tokenizer = AutoTokenizer.from_pretrained(bert_name)
+
+    nlp = spacy.load("en_core_web_md")
+
+    words, frames = predict_srl_allennlp_like_spacy(
+        model, tokenizer, nlp, sentence, id2label,
+        device=device,
+        prob_threshold=0.40,   # tune on dev; try 0.3–0.6
+        top_k=None,
+        pick_best_if_none=True
+    )
+    return words, frames
 
+if __name__ =="__main__":
+    words, frames = main_predictor(model_path, bert_namem sentence)
+    print_srl_frames_pretty(words, frames, show_grid=True, color=False)