CyberPeace-Institute/Cybersecurity-Knowledge-Graph

1from transformers import PreTrainedModel2import torch3import joblib, os4import numpy as np5from sentence_transformers import SentenceTransformer6from transformers import AutoTokenizer7 8 9from .nugget_model_utils import CustomRobertaWithPOS as NuggetModel10from .args_model_utils import CustomRobertaWithPOS as ArgumentModel11from .realis_model_utils import CustomRobertaWithPOS as RealisModel12 13from .configuration import CybersecurityKnowledgeGraphConfig14 15from .event_nugget_predict import create_dataloader as event_nugget_dataloader16from .event_realis_predict import create_dataloader as event_realis_dataloader17from .event_arg_predict import create_dataloader as event_argument_dataloader18 19class CybersecurityKnowledgeGraphModel(PreTrainedModel):20    config_class = CybersecurityKnowledgeGraphConfig21 22    def __init__(self, config):23        super().__init__(config)24        self.tokenizer = AutoTokenizer.from_pretrained("ehsanaghaei/SecureBERT")25        26        self.event_nugget_model_path = config.event_nugget_model_path27        self.event_argument_model_path = config.event_argument_model_path28        self.event_realis_model_path = config.event_realis_model_path29 30        self.event_nugget_dataloader = event_nugget_dataloader31        self.event_argument_dataloader = event_argument_dataloader32        self.event_realis_dataloader = event_realis_dataloader33 34        self.event_nugget_model = NuggetModel(num_classes = 11)35        self.event_argument_model = ArgumentModel(num_classes = 43)36        self.event_realis_model = RealisModel(num_classes_realis = 4)37 38        self.role_classifiers = {}39        self.embed_model = SentenceTransformer('all-MiniLM-L6-v2')40 41 42        self.event_nugget_list = config.event_nugget_list43        self.event_args_list = config.event_args_list44        self.realis_list = config.realis_list45        self.arg_2_role = config.arg_2_role46 47 48    def forward(self, text):49        nugget_dataloader, _ = self.event_nugget_dataloader(text)50        argument_dataloader, _ = self.event_argument_dataloader(self.event_nugget_model, text)51        realis_dataloader, _ = self.event_realis_dataloader(self.event_nugget_model, text)52 53        nugget_pred = self.forward_model(self.event_nugget_model, nugget_dataloader)54        no_nuggets = torch.all(nugget_pred == 0, dim=1)55 56        argument_preds = torch.empty(nugget_pred.size())57        realis_preds = torch.empty(nugget_pred.size())58        for idx, (batch, no_nugget) in enumerate(zip(nugget_pred, no_nuggets)):59            if no_nugget:60                argument_pred, realis_pred = torch.zeros(batch.size()), torch.zeros(batch.size())61            else:62                argument_pred = self.forward_model(self.event_argument_model, argument_dataloader)63                realis_pred = self.forward_model(self.event_realis_model, realis_dataloader)64            argument_preds[idx] = argument_pred65            realis_preds[idx] = realis_pred66        67        attention_mask = [batch["attention_mask"] for batch in nugget_dataloader]68        attention_mask = torch.cat(attention_mask, dim=-1)69 70        input_ids = [batch["input_ids"] for batch in nugget_dataloader]71        input_ids = torch.cat(input_ids, dim=-1)72        73        output = {"nugget" : nugget_pred, "argument" : argument_preds, "realis" : realis_preds, "input_ids" : input_ids, "attention_mask" : attention_mask}74        no_of_batch = output['input_ids'].shape[0]75 76        structured_output = []77        for b in range(no_of_batch):78            token_mask = [True if self.tokenizer.decode(token) not in self.tokenizer.all_special_tokens else False for token in output['input_ids'][b]]79            filtered_ids = output['input_ids'][b][token_mask]80            filtered_tokens = [self.tokenizer.decode(token) for token in filtered_ids]81 82            filtered_nuggets = output['nugget'][b][token_mask]83            filtered_args = output['argument'][b][token_mask]84            filtered_realis = output['realis'][b][token_mask]85 86            batch_output = [{"id" : id.item(), "token" : token, "nugget" : self.event_nugget_list[int(nugget.item())], "argument" : self.event_args_list[int(arg.item())], "realis" : self.realis_list[int(realis.item())]} 87                            for id, token, nugget, arg, realis in zip(filtered_ids, filtered_tokens, filtered_nuggets, filtered_args, filtered_realis)]88            structured_output.extend(batch_output)89        90        91        # args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(structured_output) if item["argument"]!= "O"]92        93        # entities = []94        # current_entity = None95        # for position, label, token in args:96        #     if label.startswith('B-'):97        #         if current_entity is not None:98        #             entities.append(current_entity)99        #         current_entity = {'label': label[2:], 'text': token.replace(" ", ""), 'start': position, 'end': position}100        #     elif label.startswith('I-'):101        #         if current_entity is not None:102        #             current_entity['text'] += ' ' + token.replace(" ", "")103        #             current_entity['end'] = position104 105        # for entity in entities:106        #     context = self.tokenizer.decode([item["id"] for item in structured_output[max(0, entity["start"] - 15) : min(len(structured_output), entity["end"] + 15)]])107        #     entity["context"] = context108        109        # for entity in entities:110        #     if len(self.arg_2_role[entity["label"]]) > 1:111        #         sent_embed = self.embed_model.encode(entity["context"])112        #         arg_embed = self.embed_model.encode(entity["text"])113        #         embed = np.concatenate((sent_embed, arg_embed))114 115        #         arg_clf = self.role_classifiers[entity["label"]]116        #         role_id = arg_clf.predict(embed.reshape(1, -1))117        #         role = self.arg_2_role[entity["label"]][role_id[0]]118 119        #         entity["role"] = role120        #     else:121        #         entity["role"] = self.arg_2_role[entity["label"]][0]122        123        # for item in structured_output:124        #     item["role"] = "O"125        # for entity in entities:126        #     for i in range(entity["start"], entity["end"] + 1):127        #         structured_output[i]["role"] = entity["role"]128        return structured_output129 130    def forward_model(self, model, dataloader):131        predicted_label = []132        for batch in dataloader:133            with torch.no_grad():134                logits = model(**batch)135            batch_predicted_label = logits.argmax(-1)136            predicted_label.append(batch_predicted_label)137        return torch.cat(predicted_label, dim=-1)