Upload model

config.json

@@ -146,7 +146,7 @@
     "B-GPE",
     "I-GPE"
   ],
-  "event_argument_model_path": "cybersecurity_knowledge_graph/argument_model_state_dict.pth",
+  "event_argument_model_path": "argument_model_state_dict.pth",
   "event_nugget_list": [
     "O",
     "B-Ransom",
@@ -160,8 +160,8 @@
     "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",
+  "event_nugget_model_path": "nugget_model_state_dict.pth",
+  "event_realis_model_path": "realis_model_state_dict.pth",
   "model_type": "cybersecurity_knowledge_graph",
   "realis_list": [
     "O",

configuration.py

@@ -1,7 +1,7 @@
 from transformers import PretrainedConfig
 import torch
 
-from cybersecurity_knowledge_graph.utils import event_args_list, event_nugget_list, realis_list, arg_2_role
+from utils import event_args_list, event_nugget_list, realis_list, arg_2_role
 
 
 class CybersecurityKnowledgeGraphConfig(PretrainedConfig):

model.py

@@ -6,15 +6,15 @@ 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
-from cybersecurity_knowledge_graph.realis_model_utils import CustomRobertaWithPOS as RealisModel
+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 configuration import CybersecurityKnowledgeGraphConfig
 
-from cybersecurity_knowledge_graph.event_nugget_predict import create_dataloader as event_nugget_dataloader
-from cybersecurity_knowledge_graph.event_realis_predict import create_dataloader as event_realis_dataloader
-from cybersecurity_knowledge_graph.event_arg_predict import create_dataloader as event_argument_dataloader
+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
@@ -40,7 +40,7 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
         self.event_argument_model.load_state_dict(torch.load(self.event_argument_model_path))
 
         role_classifiers = {}
-        folder_path = '/cybersecurity_knowledge_graph/arg_role_models'
+        folder_path = '/arg_role_models'
 
         for filename in os.listdir(os.getcwd() + folder_path):
             if filename.endswith('.joblib'):
@@ -50,7 +50,7 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
                 role_classifiers[arg] = clf
 
         self.role_classifiers = role_classifiers
-        self.embed_model = SentenceTransformer('sentence_transformer')
+        self.embed_model = SentenceTransformer('all-MiniLM-L6-v2')
 
 
         self.event_nugget_list = config.event_nugget_list