# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/pdf/03_ner.ipynb.

# %% auto 0
__all__ = ['load_ner_model', 'load_job_model', 'proc_ner', 'job_ner', 'edu_ner', 'work_ner', 'is_valid_jner', 'is_valid_tner']

# %% ../../nbs/pdf/03_ner.ipynb 2
from ..imports import *

# %% ../../nbs/pdf/03_ner.ipynb 4
def load_ner_model(model_name="tner/deberta-v3-large-ontonotes5", device="cpu"):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForTokenClassification.from_pretrained(model_name)
    return pipeline(
        "ner",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple",
        device=device,
    )


def load_job_model(model_name="ismail-lucifer011/autotrain-job_all-903929564", device="cpu"):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForTokenClassification.from_pretrained(model_name)
    return pipeline(
        "ner",
        model=model,
        tokenizer=tokenizer,
        aggregation_strategy="simple",
        device=device,
    )


def proc_ner(txt, ner, ner_dict={"institute": "", "date": ""}, thresh=3):
    ner_dict = copy.deepcopy(ner_dict)
    org_key = "institute" if "institute" in ner_dict else "company"
    mapper = {
        "ORG": org_key,
        "FAC": org_key,
        "GPE": org_key,
        "LOC": org_key,
        "Job": "role",
        "WORK_OF_ART": "degree",
        "DATE": "date",
    }
    ner_dict["text"] = txt
    for d in ner:
        eg = d["entity_group"]
        w = " " + d["word"].strip()
        k = mapper.get(eg, None)
        if k is not None and ner_dict.get(k, None) is not None and not w.startswith("##"):
            ner_dict[k] = (ner_dict[k] + w).strip()
    res = {k: v for k, v in ner_dict.items() if len(v) > thresh}
    if res.get(org_key, None) is None:
        return {}
    return res


def job_ner(txt, tner, jner):
    return tner(txt) + jner(txt)


def edu_ner(txt, tner, ner_dict={"institute": "", "date": ""}):
    ner = tner(txt)
    return proc_ner(txt, ner, ner_dict)


def work_ner(txt, tner, jner, ner_dict={"company": "", "date": ""}):
    ner = job_ner(txt, tner, jner)
    return proc_ner(txt, ner, ner_dict)


def is_valid_jner(ner, thresh=3):
    return ner.get("company", None) is not None


def is_valid_tner(ner, thresh=3):
    return ner.get("institute", None) is not None
