Upload model

event_arg_predict.py

@@ -39,7 +39,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=Tru
 
 from .args_model_utils import CustomRobertaWithPOS as ArgumentModel
 model_nugget = ArgumentModel(num_classes=43)
-model_nugget.load_state_dict(torch.load("cybersecurity_knowledge_graph/argument_model_state_dict.pth", map_location=device)) 
+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()
 
 """

event_nugget_predict.py

@@ -35,7 +35,7 @@ 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("cybersecurity_knowledge_graph/nugget_model_state_dict.pth", map_location=device))  
+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()
 
 """

event_realis_predict.py

@@ -51,7 +51,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=Tru
 
 from .realis_model_utils import CustomRobertaWithPOS as RealisModel
 model_realis = RealisModel(num_classes_realis=4)
-model_realis.load_state_dict(torch.load("cybersecurity_knowledge_graph/realis_model_state_dict.pth", map_location=device))  
+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()
 
 """