Upload predictor_up.py

predictor_up.py

@@ -0,0 +1,412 @@
+from model import PredicateAwareSRL
+import torch, json
+from transformers import AutoTokenizer
+import spacy
+from spacy import cli as spacy_cli
+
+_CACHE = {
+    "ckpt_path": None,
+    "bert_name": None,
+    "spacy_model": None,
+    "device": None,
+    "model": None,
+    "tokenizer": None,
+    "label2id": None,
+    "id2label": None,
+    "hparams": None,
+    "nlp": None,
+}
+
+_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")
+    """
+
+    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 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:
+        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"
+
+    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
+    )
+
+def _pick_device(dev=None): 
+    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)
+
+
+@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,
+            truncation=True,
+            max_length=500,
+            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_enc = tokenizer(
+            [words[predicate_word_idx]],
+            is_split_into_words=True,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=10,
+            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):  
+    enc = tokenizer(
+        words,
+        is_split_into_words=True,
+        add_special_tokens=False,
+        truncation=True,
+        max_length=max_length,
+        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): 
+    """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, max_length=500)
+    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:
+        if p >= n_words:
+            continue
+        
+        pred_wp_ids = tokenizer(
+            [words[p]],
+            is_split_into_words=True,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=10,
+            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)
+    _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
+
+