from transformers import PreTrainedModel import torch import joblib, os import numpy as np from sentence_transformers import SentenceTransformer from transformers import AutoTokenizer from .nugget_model_utils import CustomRobertaWithPOS as NuggetModel from .args_model_utils import CustomRobertaWithPOS as ArgumentModel from .realis_model_utils import CustomRobertaWithPOS as RealisModel from .configuration import CybersecurityKnowledgeGraphConfig from .event_nugget_predict import create_dataloader as event_nugget_dataloader from .event_realis_predict import create_dataloader as event_realis_dataloader from .event_arg_predict import create_dataloader as event_argument_dataloader class CybersecurityKnowledgeGraphModel(PreTrainedModel): config_class = CybersecurityKnowledgeGraphConfig def __init__(self, config): super().__init__(config) self.tokenizer = AutoTokenizer.from_pretrained("ehsanaghaei/SecureBERT") self.event_nugget_model_path = config.event_nugget_model_path self.event_argument_model_path = config.event_argument_model_path self.event_realis_model_path = config.event_realis_model_path self.event_nugget_dataloader = event_nugget_dataloader self.event_argument_dataloader = event_argument_dataloader self.event_realis_dataloader = event_realis_dataloader self.event_nugget_model = NuggetModel(num_classes = 11) self.event_argument_model = ArgumentModel(num_classes = 43) self.event_realis_model = RealisModel(num_classes_realis = 4) self.role_classifiers = {} self.embed_model = SentenceTransformer('all-MiniLM-L6-v2') self.event_nugget_list = config.event_nugget_list self.event_args_list = config.event_args_list self.realis_list = config.realis_list self.arg_2_role = config.arg_2_role def forward(self, text): nugget_dataloader, _ = self.event_nugget_dataloader(text) argument_dataloader, _ = self.event_argument_dataloader(self.event_nugget_model, text) realis_dataloader, _ = self.event_realis_dataloader(self.event_nugget_model, text) nugget_pred = self.forward_model(self.event_nugget_model, nugget_dataloader) no_nuggets = torch.all(nugget_pred == 0, dim=1) argument_preds = torch.empty(nugget_pred.size()) realis_preds = torch.empty(nugget_pred.size()) for idx, (batch, no_nugget) in enumerate(zip(nugget_pred, no_nuggets)): if no_nugget: argument_pred, realis_pred = torch.zeros(batch.size()), torch.zeros(batch.size()) else: argument_pred = self.forward_model(self.event_argument_model, argument_dataloader) realis_pred = self.forward_model(self.event_realis_model, realis_dataloader) argument_preds[idx] = argument_pred realis_preds[idx] = realis_pred attention_mask = [batch["attention_mask"] for batch in nugget_dataloader] attention_mask = torch.cat(attention_mask, dim=-1) input_ids = [batch["input_ids"] for batch in nugget_dataloader] input_ids = torch.cat(input_ids, dim=-1) output = {"nugget" : nugget_pred, "argument" : argument_preds, "realis" : realis_preds, "input_ids" : input_ids, "attention_mask" : attention_mask} no_of_batch = output['input_ids'].shape[0] structured_output = [] for b in range(no_of_batch): token_mask = [True if self.tokenizer.decode(token) not in self.tokenizer.all_special_tokens else False for token in output['input_ids'][b]] filtered_ids = output['input_ids'][b][token_mask] filtered_tokens = [self.tokenizer.decode(token) for token in filtered_ids] filtered_nuggets = output['nugget'][b][token_mask] filtered_args = output['argument'][b][token_mask] filtered_realis = output['realis'][b][token_mask] 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())]} for id, token, nugget, arg, realis in zip(filtered_ids, filtered_tokens, filtered_nuggets, filtered_args, filtered_realis)] structured_output.extend(batch_output) # args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(structured_output) if item["argument"]!= "O"] # entities = [] # current_entity = None # for position, label, token in args: # if label.startswith('B-'): # if current_entity is not None: # entities.append(current_entity) # current_entity = {'label': label[2:], 'text': token.replace(" ", ""), 'start': position, 'end': position} # elif label.startswith('I-'): # if current_entity is not None: # current_entity['text'] += ' ' + token.replace(" ", "") # current_entity['end'] = position # for entity in entities: # context = self.tokenizer.decode([item["id"] for item in structured_output[max(0, entity["start"] - 15) : min(len(structured_output), entity["end"] + 15)]]) # entity["context"] = context # for entity in entities: # if len(self.arg_2_role[entity["label"]]) > 1: # sent_embed = self.embed_model.encode(entity["context"]) # arg_embed = self.embed_model.encode(entity["text"]) # embed = np.concatenate((sent_embed, arg_embed)) # arg_clf = self.role_classifiers[entity["label"]] # role_id = arg_clf.predict(embed.reshape(1, -1)) # role = self.arg_2_role[entity["label"]][role_id[0]] # entity["role"] = role # else: # entity["role"] = self.arg_2_role[entity["label"]][0] # for item in structured_output: # item["role"] = "O" # for entity in entities: # for i in range(entity["start"], entity["end"] + 1): # structured_output[i]["role"] = entity["role"] return structured_output def forward_model(self, model, dataloader): predicted_label = [] for batch in dataloader: with torch.no_grad(): logits = model(**batch) batch_predicted_label = logits.argmax(-1) predicted_label.append(batch_predicted_label) return torch.cat(predicted_label, dim=-1)