applications/information_extraction/HugIE/api_test.py

import sys
import os
sys.path.append("./")
sys.path.append("../")
sys.path.append("../../")
sys.path.append("../../../")
from models import SPAN_EXTRACTION_MODEL_CLASSES
from models import TOKENIZER_CLASSES
import numpy as np
import torch


class HugIEAPI:
    def __init__(self, model_type, hugie_model_name_or_path) -> None:
        if model_type not in SPAN_EXTRACTION_MODEL_CLASSES[
                "global_pointer"].keys():
            raise KeyError(
                "You must choose one of the following model: {}".format(
                    ", ".join(
                        list(SPAN_EXTRACTION_MODEL_CLASSES["global_pointer"].
                             keys()))))
        self.model_type = model_type
        self.model = SPAN_EXTRACTION_MODEL_CLASSES["global_pointer"][
            self.model_type].from_pretrained(hugie_model_name_or_path)
        self.tokenizer = TOKENIZER_CLASSES[self.model_type].from_pretrained(
            hugie_model_name_or_path)
        self.max_seq_length = 512

    def fush_multi_answer(self, has_answer, new_answer):
        # 对于某个id测试集，出现多个example时（例如同一个测试样本使用了多个模板而生成了多个example），此时将预测的topk结果进行合并
        # has为已经合并的结果，new为当前新产生的结果，
        # has格式为 {"ans": {"prob": float(prob[index_ids[ei]]), "pos": (s, e)}, ...}
        # new {"ans": {"prob": float(prob[index_ids[ei]]), "pos": (s, e)}, ...}
        # print("has_answer=", has_answer)
        for ans, value in new_answer.items():
            if ans not in has_answer.keys():
                has_answer[ans] = value
            else:
                has_answer[ans]["prob"] += value["prob"]
                has_answer[ans]["pos"].extend(value["pos"])
        return has_answer

    def get_predict_result(self, probs, indices, examples):
        probs = probs.squeeze(1)  # topk结果的概率
        indices = indices.squeeze(1)  # topk结果的索引
        # print("probs=", probs) # [n, m]
        # print("indices=", indices) # [n, m]
        predictions = {}
        topk_predictions = {}
        idx = 0
        for prob, index in zip(probs, indices):
            index_ids = torch.Tensor([i for i in range(len(index))]).long()
            topk_answer = list()
            answer = []
            topk_answer_dict = dict()
            # TODO 1. 调节阈值 2. 处理输出实体重叠问题
            entity_index = index[prob > 0.1]
            index_ids = index_ids[prob > 0.1]
            for ei, entity in enumerate(entity_index):
                # 1D index转2D index
                start_end = np.unravel_index(
                    entity, (self.max_seq_length, self.max_seq_length))
                s = examples["offset_mapping"][idx][start_end[0]][0]
                e = examples["offset_mapping"][idx][start_end[1]][1]
                ans = examples["content"][idx][s:e]
                if ans not in answer:
                    answer.append(ans)
                    # topk_answer.append({"answer": ans, "prob": float(prob[index_ids[ei]]), "pos": (s, e)})
                    topk_answer_dict[ans] = {
                        "prob":
                        float(prob[index_ids[ei]]),
                        "pos": [(s.detach().cpu().numpy().tolist(),
                                 e.detach().cpu().numpy().tolist())]
                    }

            predictions[idx] = answer
            if idx not in topk_predictions.keys():
                # print("topk_answer_dict=", topk_answer_dict)
                topk_predictions[idx] = topk_answer_dict
            else:
                # print("topk_predictions[id_]=", topk_predictions[id_])
                topk_predictions[idx] = self.fush_multi_answer(
                    topk_predictions[idx], topk_answer_dict)
            idx += 1

        for idx, values in topk_predictions.items():
            # values {"ans": {}, ...}
            answer_list = list()
            for ans, value in values.items():
                answer_list.append({
                    "answer": ans,
                    "prob": value["prob"],
                    "pos": value["pos"]
                })
            topk_predictions[idx] = answer_list

        return predictions, topk_predictions

    def request(self, text: str, entity_type: str, relation: str = None):
        assert text is not None and entity_type is not None
        if relation is None:
            instruction = "找到文章中所有【{}】类型的实体？文章：【{}】".format(entity_type, text)
            pre_len = 21 - 2 + len(entity_type)
        else:
            instruction = "找到文章中【{}】的【{}】？文章：【{}】".format(
                entity_type, relation, text)
            pre_len = 19 - 4 + len(entity_type) + len(relation)

        inputs = self.tokenizer(instruction,
                                max_length=self.max_seq_length,
                                padding="max_length",
                                return_tensors="pt",
                                return_offsets_mapping=True)

        examples = {
            "content": [instruction],
            "offset_mapping": inputs["offset_mapping"]
        }

        batch_input = {
            "input_ids": inputs["input_ids"],
            "token_type_ids": inputs["token_type_ids"],
            "attention_mask": inputs["attention_mask"],
        }

        outputs = self.model(**batch_input)

        probs, indices = outputs["topk_probs"], outputs["topk_indices"]
        predictions, topk_predictions = self.get_predict_result(
            probs, indices, examples=examples)

        return predictions, topk_predictions


if __name__ == "__main__":
    from applications.information_extraction.HugIE.api_test import HugIEAPI
    model_type = "bert"
    hugie_model_name_or_path = "wjn1996/wjn1996-hugnlp-hugie-large-zh"
    hugie = HugIEAPI("bert", hugie_model_name_or_path)
    text = "央广网北京2月23日消息 据中国地震台网正式测定，2月23日8时37分在塔吉克斯坦发生7.2级地震，震源深度10公里，震中位于北纬37.98度，东经73.29度，距我国边境线最近约82公里，地震造成新疆喀什等地震感强烈。"

    ## named entity recognition
    entity_type = "国家"
    predictions, topk_predictions = hugie.request(text, entity_type)
    print("entity_type:{}".format(entity_type))
    print("predictions:\n{}".format(predictions))
    print("topk_predictions:\n{}".format(topk_predictions))
    print("\n\n")

    ## event extraction
    entity = "塔吉克斯坦地震"
    relation = "震源深度"
    predictions, topk_predictions = hugie.request(text,
                                                  entity,
                                                  relation=relation)
    print("entity:{}, relation:{}".format(entity, relation))
    print("predictions:\n{}".format(predictions))
    print("topk_predictions:\n{}".format(topk_predictions))
    print("\n\n")

    ## event extraction
    entity = "塔吉克斯坦地震"
    relation = "震源位置"
    predictions, topk_predictions = hugie.request(text,
                                                  entity,
                                                  relation=relation)
    print("entity:{}, relation:{}".format(entity, relation))
    print("predictions:\n{}".format(predictions))
    print("topk_predictions:\n{}".format(topk_predictions))
    print("\n\n")

    ## event extraction
    entity = "塔吉克斯坦地震"
    relation = "时间"
    predictions, topk_predictions = hugie.request(text,
                                                  entity,
                                                  relation=relation)
    print("entity:{}, relation:{}".format(entity, relation))
    print("predictions:\n{}".format(predictions))
    print("topk_predictions:\n{}".format(topk_predictions))
    print("\n\n")

    ## event extraction
    entity = "塔吉克斯坦地震"
    relation = "影响"
    predictions, topk_predictions = hugie.request(text,
                                                  entity,
                                                  relation=relation)
    print("entity:{}, relation:{}".format(entity, relation))
    print("predictions:\n{}".format(predictions))
    print("topk_predictions:\n{}".format(topk_predictions))
    print("\n\n")
    """
    Output results:

    entity_type:国家
predictions:
{0: ["塔吉克斯坦"]}
predictions:
{0: [{"answer": "塔吉克斯坦", "prob": 0.9999997615814209, "pos": [(tensor(57), tensor(62))]}]}



entity:塔吉克斯坦地震, relation:震源深度
predictions:
{0: ["10公里"]}
predictions:
{0: [{"answer": "10公里", "prob": 0.999994158744812, "pos": [(tensor(80), tensor(84))]}]}



entity:塔吉克斯坦地震, relation:震源位置
predictions:
{0: ["10公里", "距我国边境线最近约82公里", "北纬37.98度，东经73.29度", "北纬37.98度，东经73.29度，距我国边境线最近约82公里"]}
predictions:
{0: [{"answer": "10公里", "prob": 0.9895901083946228, "pos": [(tensor(80), tensor(84))]}, {"answer": "距我国边境线最近约82公里", "prob": 0.8584909439086914, "pos": [(tensor(107), tensor(120))]}, {"answer": "北纬37.98度，东经73.29度", "prob": 0.7202121615409851, "pos": [(tensor(89), tensor(106))]}, {"answer": "北纬37.98度，东经73.29度，距我国边境线最近约82公里", "prob": 0.11628123372793198, "pos": [(tensor(89), tensor(120))]}]}



entity:塔吉克斯坦地震, relation:时间
predictions:
{0: ["2月23日8时37分"]}
predictions:
{0: [{"answer": "2月23日8时37分", "prob": 0.9999995231628418, "pos": [(tensor(49), tensor(59))]}]}



entity:塔吉克斯坦地震, relation:影响
predictions:
{0: ["新疆喀什等地震感强烈"]}
predictions:
{0: [{"answer": "新疆喀什等地震感强烈", "prob": 0.9525265693664551, "pos": [(tensor(123), tensor(133))]}]}

    """