Upload 2 files

predictor_dev.py

@@ -0,0 +1,385 @@
+from model import PredicateAwareSRL
+import torch, json
+from transformers import AutoTokenizer
+import spacy
+from spacy import cli as spacy_cli
+
+# ============ GLOBAL CACHE (NEW) ============
+_CACHE = {
+    "ckpt_path": None,
+    "bert_name": None,
+    "spacy_model": None,
+    "device": None,
+    "model": None,
+    "tokenizer": None,
+    "label2id": None,
+    "id2label": None,
+    "hparams": None,
+    "nlp": None,
+}
+
+# --- Add near the top of your module (where _CACHE lives) ---
+_CACHE = {
+    "model": None, "tokenizer": None, "id2label": None, "nlp": None, "device": None,
+    "ckpt_path": None, "bert_name": None, "spacy_model": None,
+}
+
+def srl_init(model_path, bert_name="bert-base-cased", spacy_model="en_core_web_md"):
+    """
+    Call ONCE per session to load and cache model/tokenizer/spaCy.
+    After this, you can call: prediction("your sentence here")
+    """
+    # reuse your existing loader logic
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    ckpt = torch.load(model_path, map_location=device)
+    hp = ckpt.get("hparams", ckpt.get("hyper_parameters", {}))
+    if "bert_name" not in hp:
+        hp["bert_name"] = bert_name
+    if "num_labels" not in hp:
+        label2id = ckpt.get("label2id") or {v:k for k,v in ckpt["id2label"].items()}
+        hp["num_labels"] = len(label2id)
+
+    model = PredicateAwareSRL(**hp).to(device).eval()
+    state = ckpt.get("model_state") or ckpt.get("state_dict") or ckpt
+    model.load_state_dict(state)
+
+    tokenizer = AutoTokenizer.from_pretrained(hp.get("bert_name", bert_name), use_fast=True)
+
+    try:
+        nlp = spacy.load(spacy_model, disable=["parser","ner","lemmatizer"])
+    except OSError:
+        spacy_cli.download(spacy_model)
+        nlp = spacy.load(spacy_model, disable=["parser","ner","lemmatizer"])
+
+    label2id = ckpt.get("label2id") or {v:k for k,v in ckpt["id2label"].items()}
+    id2label = {int(v): k for k, v in label2id.items()}
+
+    _CACHE.update({
+        "model": model, "tokenizer": tokenizer, "id2label": id2label,
+        "nlp": nlp, "device": device, "ckpt_path": model_path,
+        "bert_name": hp.get("bert_name", bert_name), "spacy_model": spacy_model,
+    })
+    torch.set_grad_enabled(False)
+
+def _predict_cached(sentence):
+    """Internal: uses cached objects set by srl_init()."""
+    if _CACHE["model"] is None:
+        raise RuntimeError("Model not loaded. Call srl_init(ckpt_path, bert_name) once first.")
+    model     = _CACHE["model"]
+    tokenizer = _CACHE["tokenizer"]
+    id2label  = _CACHE["id2label"]
+    nlp       = _CACHE["nlp"]
+    device    = "cuda" if (_CACHE["device"].type == "cuda") else "cpu"
+
+    # your existing function name stays the same:
+    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
+    )
+
+def _pick_device(dev=None):  # NEW
+    if dev == "cpu":
+        return torch.device("cpu")
+    if dev and dev.startswith("cuda") and torch.cuda.is_available():
+        return torch.device(dev)
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def _ensure_loaded(model_path, bert_name, spacy_model, model_cls):  # NEW
+    """Load model/tokenizer/spaCy once and reuse."""
+    must_reload = (
+        _CACHE["model"] is None
+        or _CACHE["ckpt_path"] != model_path
+        or _CACHE["bert_name"] != bert_name
+        or _CACHE["spacy_model"] != spacy_model
+    )
+    if not must_reload:
+        return  # already warm
+
+    device = _pick_device()
+    ckpt = torch.load(model_path, map_location=device)
+    h = ckpt.get("hparams", ckpt.get("hyper_parameters", {}))
+
+    # defaults if not present in ckpt
+    if "bert_name" not in h: h["bert_name"] = bert_name
+    if "num_labels" not in h: 
+        label2id = ckpt.get("label2id")
+        if label2id is None and "id2label" in ckpt:
+            label2id = {v:k for k,v in ckpt["id2label"].items()}
+        h["num_labels"] = len(label2id) if label2id else 0
+
+    model = model_cls(**h).to(device).eval()
+    state = ckpt.get("model_state") or ckpt.get("state_dict") or ckpt
+    model.load_state_dict(state)
+
+    tok = AutoTokenizer.from_pretrained(h.get("bert_name", bert_name), use_fast=True)
+
+    try:
+        nlp = spacy.load(spacy_model, disable=["parser","ner","lemmatizer"])
+    except OSError:
+        spacy_cli.download(spacy_model)
+        nlp = spacy.load(spacy_model, disable=["parser","ner","lemmatizer"])
+
+    label2id = ckpt.get("label2id")
+    if label2id is None and "id2label" in ckpt:
+        label2id = {v:k for k,v in ckpt["id2label"].items()}
+    id2label = {int(v): k for k, v in label2id.items()}
+
+    _CACHE.update({
+        "ckpt_path": model_path,
+        "bert_name": h.get("bert_name", bert_name),
+        "spacy_model": spacy_model,
+        "device": device,
+        "model": model,
+        "tokenizer": tok,
+        "label2id": label2id,
+        "id2label": id2label,
+        "hparams": h,
+        "nlp": nlp,
+    })
+    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()
+    sent_enc = tokenizer(
+        words, is_split_into_words=True, add_special_tokens=False,
+        return_attention_mask=False, return_token_type_ids=False,
+    )
+    # word ids
+    try:
+        sent_word_ids = sent_enc.word_ids()
+    except Exception:
+        raise ValueError("Tokenizer must be fast (use_fast=True).")
+
+    sent_wp_ids = sent_enc["input_ids"]
+    if isinstance(sent_wp_ids[0], list):
+        sent_wp_ids = sent_wp_ids[0]
+
+    first_pos_by_wid = {}
+    for pos, wid in enumerate(sent_word_ids):
+        if wid is not None and wid not in first_pos_by_wid:
+            first_pos_by_wid[wid] = pos + 1
+
+    n_words = len(words)
+    word_first_wp_fullidx = torch.tensor(
+        [first_pos_by_wid[i] for i in range(n_words)], dtype=torch.long
+    ).unsqueeze(0)
+
+    pred_enc = tokenizer(
+        [words[predicate_word_idx]], is_split_into_words=True, add_special_tokens=False,
+        return_attention_mask=False, return_token_type_ids=False,
+    )
+    pred_wp_ids = pred_enc["input_ids"]
+    if isinstance(pred_wp_ids[0], list):
+        pred_wp_ids = pred_wp_ids[0]
+
+    cls_id, sep_id = tokenizer.cls_token_id, tokenizer.sep_token_id
+    input_ids      = [cls_id] + sent_wp_ids + [sep_id] + pred_wp_ids + [sep_id]
+    token_type_ids = [0] * (1 + len(sent_wp_ids) + 1) + [1] * (len(pred_wp_ids) + 1)
+    attention_mask = [1] * len(input_ids)
+
+    device = _pick_device(device)
+    input_ids      = torch.tensor(input_ids).unsqueeze(0).to(device)
+    token_type_ids = torch.tensor(token_type_ids).unsqueeze(0).to(device)
+    attention_mask = torch.tensor(attention_mask).unsqueeze(0).to(device)
+
+    sent_len       = torch.tensor([n_words], dtype=torch.long).to(device)
+    sentence_mask  = torch.ones(1, n_words, dtype=torch.bool).to(device)
+    pred_word_idx  = torch.tensor([predicate_word_idx], dtype=torch.long).to(device)
+    indicator      = torch.zeros(1, n_words, dtype=torch.long).to(device)
+    indicator[0, predicate_word_idx] = 1
+    word_first_wp_fullidx = word_first_wp_fullidx.to(device)
+
+    logits, _ = model(
+        input_ids=input_ids,
+        token_type_ids=token_type_ids,
+        attention_mask=attention_mask,
+        word_first_wp_fullidx=word_first_wp_fullidx,
+        sentence_mask=sentence_mask,
+        sent_lens=sent_len,
+        pred_word_idx=pred_word_idx,
+        indicator=indicator,
+        labels=None,
+    )
+    pred_ids = logits.argmax(-1).squeeze(0).tolist()
+    tags = [id2label[i] for i in pred_ids]
+    return tags, logits.squeeze(0).cpu()
+
+
+def _encode_sentence_once(words, tokenizer):  # NEW
+    enc = tokenizer(
+        words, is_split_into_words=True, add_special_tokens=False,
+        return_attention_mask=False, return_token_type_ids=False,
+    )
+    sent_wp_ids = enc["input_ids"]
+    if isinstance(sent_wp_ids[0], list):
+        sent_wp_ids = sent_wp_ids[0]
+    wid = enc.word_ids()
+    first_pos = {}
+    for pos, w in enumerate(wid):
+        if w is not None and w not in first_pos:
+            first_pos[w] = pos + 1  # +1 for [CLS]
+    n_words = len(words)
+    word_first = torch.tensor([first_pos[i] for i in range(n_words)], dtype=torch.long)
+    return sent_wp_ids, word_first, n_words
+
+@torch.no_grad()
+def _batch_predict_verbs(model, tokenizer, words, verb_idxs, id2label, device):  # NEW
+    """One forward pass for all verbs in the sentence."""
+    device = _pick_device(device)
+    sent_wp_ids, word_first_1, n_words = _encode_sentence_once(words, tokenizer)
+    cls_id, sep_id = tokenizer.cls_token_id, tokenizer.sep_token_id
+
+    ids_list, tt_list, am_list = [], [], []
+    pred_idx_list, ind_list, wf_list = [], [], []
+
+    for p in verb_idxs:
+        pred_wp_ids = tokenizer(
+            [words[p]], is_split_into_words=True, add_special_tokens=False,
+            return_attention_mask=False, return_token_type_ids=False,
+        )["input_ids"]
+        if isinstance(pred_wp_ids[0], list):
+            pred_wp_ids = pred_wp_ids[0]
+
+        ids = [cls_id] + sent_wp_ids + [sep_id] + pred_wp_ids + [sep_id]
+        tt  = [0]*(1 + len(sent_wp_ids) + 1) + [1]*(len(pred_wp_ids) + 1)
+        am  = [1]*len(ids)
+
+        ids_list.append(torch.tensor(ids, dtype=torch.long))
+        tt_list.append(torch.tensor(tt, dtype=torch.long))
+        am_list.append(torch.tensor(am, dtype=torch.long))
+        pred_idx_list.append(torch.tensor(p, dtype=torch.long))
+        ind = torch.zeros(n_words, dtype=torch.long); ind[p] = 1
+        ind_list.append(ind)
+        wf_list.append(word_first_1.clone())
+
+    # pad
+    def pad_1d(seq, pad_id=0):
+        L = max(x.numel() for x in seq)
+        out = torch.full((len(seq), L), pad_id, dtype=seq[0].dtype)
+        for i, x in enumerate(seq):
+            out[i, :x.numel()] = x
+        return out
+
+    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
+    input_ids      = pad_1d(ids_list, pad_id).to(device)
+    token_type_ids = pad_1d(tt_list, 0).to(device)
+    attention_mask = pad_1d(am_list, 0).to(device)
+
+    B = len(verb_idxs)
+    sent_lens      = torch.full((B,), n_words, dtype=torch.long, device=device)
+    sentence_mask  = torch.ones(B, n_words, dtype=torch.bool, device=device)
+    pred_word_idx  = torch.stack(pred_idx_list).to(device)
+    indicator      = torch.stack(ind_list).to(device)
+    word_first_wp_fullidx = torch.stack(wf_list).to(device)
+
+    logits, _ = model(
+        input_ids=input_ids,
+        token_type_ids=token_type_ids,
+        attention_mask=attention_mask,
+        word_first_wp_fullidx=word_first_wp_fullidx,
+        sentence_mask=sentence_mask,
+        sent_lens=sent_lens,
+        pred_word_idx=pred_word_idx,
+        indicator=indicator,
+        labels=None,
+    )  # [B, n_words, C]
+
+    results = []
+    for row, p in enumerate(verb_idxs):
+        row_logits = logits[row]
+        tags = [id2label[i] for i in row_logits.argmax(-1).tolist()]
+        results.append((p, tags, row_logits))
+    return results
+
+
+@torch.no_grad()
+def predict_srl_allennlp_like_spacy(
+    model, tokenizer, nlp, sentence, id2label,
+    device="cuda",
+    prob_threshold=0.50,
+    top_k=None,
+    pick_best_if_none=True
+):
+    model.eval()
+    words, _ = words_and_spans_spacy(sentence, nlp)
+    if not words:
+        return [], []
+
+    verb_idxs = spacy_verb_indices(nlp, sentence)
+    if not verb_idxs:
+        return words, []
+
+    # one forward for all verbs (fast path)
+    batch_out = _batch_predict_verbs(model, tokenizer, words, verb_idxs, id2label, device)
+    b_v_id = next((i for i,t in id2label.items() if t in ("B-V","V")), None)
+
+    frames = []
+    for p, tags, row_logits in batch_out:
+        p_bv = float(torch.softmax(row_logits[p], dim=-1)[b_v_id].item()) if b_v_id is not None else 1.0
+        frames.append({
+            "predicate_index": p,
+            "predicate": words[p],
+            "p_bv": p_bv,
+            "tags": tags,
+            "spans": bio_to_spans(tags)
+        })
+
+    # optional thresholding / top-k
+    if prob_threshold is not None:
+        keep = [f for f in frames if f["p_bv"] >= prob_threshold]
+        if not keep and pick_best_if_none and frames:
+            keep = [max(frames, key=lambda r: r["p_bv"])]
+        frames = keep
+    if top_k is not None and len(frames) > top_k:
+        frames = sorted(frames, key=lambda r: r["p_bv"], reverse=True)[:top_k]
+
+    return words, frames
+
+def main_predictor(model_path, bert_name, sentence, spacy_model="en_core_web_md"):
+    sentence = normalize_whitespace(sentence)
+    from model import PredicateAwareSRL  # keep your import style
+    _ensure_loaded(model_path, bert_name, spacy_model, PredicateAwareSRL)  # NEW: cache/warm
+    model     = _CACHE["model"]
+    tokenizer = _CACHE["tokenizer"]
+    id2label  = _CACHE["id2label"]
+    nlp       = _CACHE["nlp"]
+    device    = _CACHE["device"]
+
+    words, frames = predict_srl_allennlp_like_spacy(
+        model, tokenizer, nlp, sentence, id2label,
+        device=str(device), prob_threshold=0.40, top_k=None, pick_best_if_none=True
+    )
+    return words, frames
+
+

visualizer_dev.py

@@ -0,0 +1,191 @@
+from predictor import main_predictor
+import re
+import itertools
+
+def bio_brackets_to_spans(text: str) -> str:
+    """
+    Collapse BIO bracket chunks into non-BIO spans.
+    Example:
+        [B-ARG2: of] [I-ARG2: the] [I-ARG2: orchards] → [ARG2: of the orchards]
+        [B-V: take] → [V: take]
+    Non-bracket text (spaces, punctuation, quotes) is preserved.
+    """
+
+    BIO_RE = re.compile(r"\[(B|I)-([A-Za-z0-9\-]+):\s*([^\]]+?)\]")
+
+    out = []
+    i = 0
+    matches = list(BIO_RE.finditer(text))
+
+    m = 0
+    cursor = 0
+    while m < len(matches):
+        # plain text before next BIO chunk
+        out.append(text[cursor:matches[m].start()])
+
+        # start a run
+        prefix, role, tok = matches[m].groups()
+        tokens = [tok]
+        cursor = matches[m].end()
+        m += 1
+
+        # absorb subsequent I-<same role> chunks if only whitespace between
+        while m < len(matches):
+            between = text[cursor:matches[m].start()]
+            p2, role2, tok2 = matches[m].groups()
+            if role2 == role and p2 == "I" and between.strip() == "":
+                tokens.append(tok2)
+                cursor = matches[m].end()
+                m += 1
+            else:
+                break
+
+        # output merged span (drop B-/I-), keep V as just "V"
+        out.append(f"[{role}: {' '.join(tokens)}]")
+
+    # trailing text
+    out.append(text[cursor:])
+    return "".join(out)
+
+def create_description(words, tag_list):
+    desc_list = []
+    for tok, tag in zip(words, tag_list):
+        if tag != 'O' : 
+            desc_list.append("["+tag+": "+tok+"]")
+        else: 
+            desc_list.append(tok)
+    desc_str_temp = (' ').join(desc_list)
+
+    return bio_brackets_to_spans(desc_str_temp)
+
+def create_dict(words, frames):
+    final_dict = {}
+    verb = []
+    for f in frames: 
+        temp_dict = {}
+        temp_dict['verb'] = f['predicate']
+        temp_dict['description'] = create_description(words, f['tags'])
+        temp_dict['tags'] = f['tags']
+        verb.append(temp_dict)
+    final_dict['verbs'] = verb
+    final_dict['words'] = words
+
+    return final_dict
+
+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)
+    """
+
+
+    # tiny colorizer (terminal-only); 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)
+
+        print("\n" + "—"*60)
+        # head = f"Frame {k} — predicate: {pred} (idx {pred_idx})"
+        # if p_bv is not None:
+        #     head += f"   P(B-V)={p_bv:.3f}"
+        # print(head)
+        
+        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)
+
+        # 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)}")
+
+        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)
+
+def prediction(model_path, bert_name, sentence):
+    words, frames = main_predictor(model_path, bert_name, sentence)
+    # assumes you have print_srl_frames_pretty already defined somewhere
+    try:
+        print_srl_frames_pretty(words, frames, show_grid=True, color=False)
+    except NameError:
+        # fallback
+        print("Sentence:", " ".join(words))
+        for fr in frames:
+            print(f"\nPredicate: {fr['predicate']}  P(B-V)={fr['p_bv']:.3f}")
+            print("Tags:", list(zip(words, fr['tags'])))
+            print("Spans:", fr['spans'])
+
+def prediction_formatted(model_path, bert_name, sentence):
+    words, frames = main_predictor(model_path, bert_name, sentence)
+    # assumes you have create_dict, otherwise return raw
+    try:
+        return create_dict(words, frames)
+    except NameError:
+        return {"words": words, "frames": frames}