Upload model

event_arg_predict.py

@@ -37,10 +37,10 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 model_checkpoint = "ehsanaghaei/SecureBERT"
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)
 
-from .args_model_utils import CustomRobertaWithPOS as ArgumentModel
-model_nugget = ArgumentModel(num_classes=43)
-model_nugget.load_state_dict(torch.load(f"{os.path.dirname(os.path.abspath(__file__))}/argument_model_state_dict.pth", map_location=device)) 
-model_nugget.eval()
+# from .args_model_utils import CustomRobertaWithPOS as ArgumentModel
+# model_nugget = ArgumentModel(num_classes=43)
+# model_nugget.load_state_dict(torch.load(f"{os.path.dirname(os.path.abspath(__file__))}/argument_model_state_dict.pth", map_location=device)) 
+# model_nugget.eval()
 
 """
 Function: create_dataloader(text_input)
@@ -51,9 +51,9 @@ Output:
     - dataloader: A DataLoader for the tokenized and batched text data.
     - tokenized_dataset_ner: The tokenized dataset used for training.
 """
-def create_dataloader(text_input):
+def create_dataloader(model_nugget, text_input):
 
-    event_nuggets = get_event_nuggets(text_input)
+    event_nuggets = get_event_nuggets(model_nugget, text_input)
     doc = nlp(text_input)
 
     content_as_words_emdash = [tok.text for tok in doc]

event_nugget_predict.py

@@ -34,9 +34,9 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 model_checkpoint = "ehsanaghaei/SecureBERT"
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)
 
-model_nugget = NuggetModel(num_classes = 11)
-model_nugget.load_state_dict(torch.load(f"{os.path.dirname(os.path.abspath(__file__))}/nugget_model_state_dict.pth", map_location=device))  
-model_nugget.eval()
+# model_nugget = NuggetModel(num_classes = 11)
+# model_nugget.load_state_dict(torch.load(f"{os.path.dirname(os.path.abspath(__file__))}/nugget_model_state_dict.pth", map_location=device))  
+# model_nugget.eval()
 
 """
 Function: create_dataloader(text_input)
@@ -133,7 +133,7 @@ Inputs:
 Output:
     - predicted_label: A tensor containing the predicted labels for the input data.
 """
-def predict(dataloader):
+def predict(model_nugget, dataloader):
     predicted_label = []
     for batch in dataloader:
         with torch.no_grad():
@@ -202,9 +202,9 @@ Output:
     - predicted_event_nuggets: A list of dictionaries, each representing an extracted event nugget with start and end offsets,
       subtype, and text content.
 """
-def get_event_nuggets(text_input):
+def get_event_nuggets(model_nugget, text_input):
     dataloader, tokenized_dataset_ner = create_dataloader(text_input)
-    predicted_label = predict(dataloader)
+    predicted_label = predict(model_nugget, dataloader)
 
     predicted_event_nuggets = []
     text_length = 0 

event_realis_predict.py

@@ -49,10 +49,10 @@ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cp
 model_checkpoint = "ehsanaghaei/SecureBERT"
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)
 
-from .realis_model_utils import CustomRobertaWithPOS as RealisModel
-model_realis = RealisModel(num_classes_realis=4)
-model_realis.load_state_dict(torch.load(f"{os.path.dirname(os.path.abspath(__file__))}/realis_model_state_dict.pth", map_location=device))  
-model_realis.eval()
+# from .realis_model_utils import CustomRobertaWithPOS as RealisModel
+# model_realis = RealisModel(num_classes_realis=4)
+# model_realis.load_state_dict(torch.load(f"{os.path.dirname(os.path.abspath(__file__))}/realis_model_state_dict.pth", map_location=device))  
+# model_realis.eval()
 
 """
 Function: create_dataloader(text_input)
@@ -63,9 +63,9 @@ Output:
     - dataloader: A DataLoader for the tokenized and batched text data.
     - tokenized_dataset_ner: The tokenized dataset used for training.
 """
-def create_dataloader(text_input):
+def create_dataloader(model_nugget, text_input):
 
-    event_nuggets = get_event_nuggets(text_input)
+    event_nuggets = get_event_nuggets(model_nugget, text_input)
     doc = nlp(text_input)
 
     content_as_words_emdash = [tok.text for tok in doc]

model.py

@@ -61,8 +61,8 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
 
     def forward(self, text):
         nugget_dataloader, _ = self.event_nugget_dataloader(text)
-        argument_dataloader, _ = self.event_argument_dataloader(text)
-        realis_dataloader, _ = self.event_realis_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)