Upload model

config.json

@@ -2,13 +2,172 @@
   "architectures": [
     "CybersecurityKnowledgeGraphModel"
   ],
+  "arg_2_role": {
+    "CVE": [
+      "CVE"
+    ],
+    "Capabilities": [
+      "Attack-Pattern",
+      "Capabilities",
+      "Issues-Addressed"
+    ],
+    "Data": [
+      "Compromised-Data",
+      "Trusted-Entity"
+    ],
+    "Device": [
+      "Vulnerable_System",
+      "Victim",
+      "Supported_Platform"
+    ],
+    "File": [
+      "Tool",
+      "Trusted-Entity"
+    ],
+    "GPE": [
+      "Place"
+    ],
+    "Malware": [
+      "Tool"
+    ],
+    "Money": [
+      "Price",
+      "Damage-Amount"
+    ],
+    "Number": [
+      "Number-of-Data",
+      "Number-of-Victim"
+    ],
+    "Organization": [
+      "Victim",
+      "Releaser",
+      "Discoverer",
+      "Attacker",
+      "Vulnerable_System_Owner",
+      "Trusted-Entity"
+    ],
+    "PII": [
+      "Compromised-Data",
+      "Trusted-Entity"
+    ],
+    "Patch": [
+      "Patch"
+    ],
+    "PaymentMethod": [
+      "Payment-Method"
+    ],
+    "Person": [
+      "Victim",
+      "Attacker",
+      "Discoverer",
+      "Releaser",
+      "Trusted-Entity",
+      "Vulnerable_System_Owner"
+    ],
+    "Purpose": [
+      "Purpose"
+    ],
+    "Software": [
+      "Vulnerable_System",
+      "Victim",
+      "Trusted-Entity",
+      "Supported_Platform"
+    ],
+    "System": [
+      "Victim",
+      "Supported_Platform",
+      "Vulnerable_System",
+      "Trusted-Entity"
+    ],
+    "Time": [
+      "Time"
+    ],
+    "Version": [
+      "Patch-Number",
+      "Vulnerable_System_Version"
+    ],
+    "Vulnerability": [
+      "Vulnerability"
+    ],
+    "Website": [
+      "Trusted-Entity",
+      "Tool",
+      "Vulnerable_System",
+      "Victim",
+      "Supported_Platform"
+    ]
+  },
   "auto_map": {
     "AutoConfig": "configuration.CybersecurityKnowledgeGraphConfig",
     "AutoModelForTokenClassification": "model.CybersecurityKnowledgeGraphModel"
   },
+  "event_args_list": [
+    "O",
+    "B-System",
+    "I-System",
+    "B-Organization",
+    "B-Money",
+    "I-Money",
+    "B-Device",
+    "B-Person",
+    "I-Person",
+    "B-Vulnerability",
+    "I-Vulnerability",
+    "B-Capabilities",
+    "I-Capabilities",
+    "I-Organization",
+    "B-PaymentMethod",
+    "I-PaymentMethod",
+    "B-Data",
+    "I-Data",
+    "B-Number",
+    "I-Number",
+    "B-Malware",
+    "I-Malware",
+    "B-PII",
+    "I-PII",
+    "B-CVE",
+    "I-CVE",
+    "B-Purpose",
+    "I-Purpose",
+    "B-File",
+    "I-File",
+    "I-Device",
+    "B-Time",
+    "I-Time",
+    "B-Software",
+    "I-Software",
+    "B-Patch",
+    "I-Patch",
+    "B-Version",
+    "I-Version",
+    "B-Website",
+    "I-Website",
+    "B-GPE",
+    "I-GPE"
+  ],
   "event_argument_model_path": "cybersecurity_knowledge_graph/argument_model_state_dict.pth",
+  "event_nugget_list": [
+    "O",
+    "B-Ransom",
+    "I-Ransom",
+    "B-DiscoverVulnerability",
+    "I-DiscoverVulnerability",
+    "B-PatchVulnerability",
+    "I-PatchVulnerability",
+    "B-Databreach",
+    "I-Databreach",
+    "B-Phishing",
+    "I-Phishing"
+  ],
   "event_nugget_model_path": "cybersecurity_knowledge_graph/nugget_model_state_dict.pth",
   "event_realis_model_path": "cybersecurity_knowledge_graph/realis_model_state_dict.pth",
+  "realis_list": [
+    "O",
+    "Generic",
+    "Other",
+    "Actual"
+  ],
   "torch_dtype": "float32",
   "transformers_version": "4.33.2"
 }

configuration.py

@@ -1,6 +1,9 @@
 from transformers import PretrainedConfig
 import torch
 
+from cybersecurity_knowledge_graph.utils import event_args_list, event_nugget_list, realis_list, arg_2_role
+
+
 class CybersecurityKnowledgeGraphConfig(PretrainedConfig):
 
     def __init__(
@@ -13,4 +16,10 @@ class CybersecurityKnowledgeGraphConfig(PretrainedConfig):
         self.event_nugget_model_path = event_nugget_model_path
         self.event_argument_model_path = event_argument_model_path
         self.event_realis_model_path = event_realis_model_path
+
+        self.event_nugget_list = event_nugget_list
+        self.event_args_list = event_args_list
+        self.realis_list = realis_list
+        self.arg_2_role = arg_2_role
+
         super().__init__(**kwargs)

model.py

@@ -1,5 +1,10 @@
 from transformers import PreTrainedModel
 import torch
+import joblib, os
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer
+
 
 from cybersecurity_knowledge_graph.nugget_model_utils import CustomRobertaWithPOS as NuggetModel
 from cybersecurity_knowledge_graph.args_model_utils import CustomRobertaWithPOS as ArgumentModel
@@ -16,6 +21,8 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
 
     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
@@ -32,6 +39,25 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
         self.event_realis_model.load_state_dict(torch.load(self.event_realis_model_path))
         self.event_argument_model.load_state_dict(torch.load(self.event_argument_model_path))
 
+        role_classifiers = {}
+        folder_path = '/cybersecurity_knowledge_graph/arg_role_models'
+
+        for filename in os.listdir(os.getcwd() + folder_path):
+            if filename.endswith('.joblib'):
+                file_path = os.getcwd() + os.path.join(folder_path, filename)
+                clf = joblib.load(file_path)
+                arg = filename.split(".")[0]
+                role_classifiers[arg] = clf
+
+        self.role_classifiers = role_classifiers
+        self.embed_model = SentenceTransformer('sentence_transformer')
+
+
+        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)
@@ -51,15 +77,75 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
                 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]
 
-        return {"nugget" : nugget_pred, "argument" : argument_preds, "realis" : realis_preds}
+        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)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fef48b6b9271dd45d7102c4efd5a90a3e2897daeb2393dcbd6e4fc3aa94494c5
-size 1496163441
+oid sha256:5b14783587eb16be4dac27e2bc3b5d738ee1772b44cf48d0240edef88aaee6e9
+size 1587052173