model.py

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)