Upload model

config.json

@@ -97,6 +97,10 @@
       "Supported_Platform"
     ]
   },
+  "auto_map": {
+    "AutoConfig": "configuration.CybersecurityKnowledgeGraphConfig",
+    "AutoModelForTokenClassification": "model.CybersecurityKnowledgeGraphModel"
+  },
   "event_args_list": [
     "O",
     "B-System",
@@ -158,7 +162,6 @@
   ],
   "event_nugget_model_path": "cybersecurity_knowledge_graph/nugget_model_state_dict.pth",
   "event_realis_model_path": "cybersecurity_knowledge_graph/realis_model_state_dict.pth",
-  "model_type": "cybersecurity_knowledge_graph",
   "realis_list": [
     "O",
     "Generic",

configuration.py

@@ -1,13 +1,11 @@
 from transformers import PretrainedConfig
 import torch
 
-import utils
+import cybersecurity_knowledge_graph.utils
 
 
 
 class CybersecurityKnowledgeGraphConfig(PretrainedConfig):
-    model_type = "cybersecurity_knowledge_graph"
-
     def __init__(
         self,
         event_nugget_model_path : str = "nugget_model_state_dict.pth",
@@ -19,9 +17,9 @@ class CybersecurityKnowledgeGraphConfig(PretrainedConfig):
         self.event_argument_model_path = event_argument_model_path
         self.event_realis_model_path = event_realis_model_path
 
-        self.event_nugget_list = utils.event_nugget_list
-        self.event_args_list = utils.event_args_list
-        self.realis_list = utils.realis_list
-        self.arg_2_role = utils.arg_2_role
+        self.event_nugget_list = cybersecurity_knowledge_graph.utils.event_nugget_list
+        self.event_args_list = cybersecurity_knowledge_graph.utils.event_args_list
+        self.realis_list = cybersecurity_knowledge_graph.utils.realis_list
+        self.arg_2_role = cybersecurity_knowledge_graph.utils.arg_2_role
 
         super().__init__(**kwargs)

model.py

@@ -6,15 +6,15 @@ 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 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 configuration import CybersecurityKnowledgeGraphConfig
+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
+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
 
 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 = '/arg_role_models'
+        folder_path = '/cybersecurity_knowledge_graph/arg_role_models'
 
         for filename in os.listdir(os.getcwd() + folder_path):
             if filename.endswith('.joblib'):