cpi-connect commited on
Commit
ac863c4
·
1 Parent(s): fc3ce31

Upload model

Browse files
Files changed (1) hide show
  1. model.py +1 -15
model.py CHANGED
@@ -35,21 +35,7 @@ class CybersecurityKnowledgeGraphModel(PreTrainedModel):
35
  self.event_argument_model = ArgumentModel(num_classes = 43)
36
  self.event_realis_model = RealisModel(num_classes_realis = 4)
37
 
38
- self.event_nugget_model.load_state_dict(torch.load(self.event_nugget_model_path))
39
- self.event_realis_model.load_state_dict(torch.load(self.event_realis_model_path))
40
- self.event_argument_model.load_state_dict(torch.load(self.event_argument_model_path))
41
-
42
- role_classifiers = {}
43
- folder_path = '/cybersecurity_knowledge_graph/arg_role_models'
44
-
45
- for filename in os.listdir(os.getcwd() + folder_path):
46
- if filename.endswith('.joblib'):
47
- file_path = os.getcwd() + os.path.join(folder_path, filename)
48
- clf = joblib.load(file_path)
49
- arg = filename.split(".")[0]
50
- role_classifiers[arg] = clf
51
-
52
- self.role_classifiers = role_classifiers
53
  self.embed_model = SentenceTransformer('all-MiniLM-L6-v2')
54
 
55
 
 
35
  self.event_argument_model = ArgumentModel(num_classes = 43)
36
  self.event_realis_model = RealisModel(num_classes_realis = 4)
37
 
38
+ self.role_classifiers = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  self.embed_model = SentenceTransformer('all-MiniLM-L6-v2')
40
 
41