| from transformers import PretrainedConfig |
| import torch |
|
|
| from cybersecurity_knowledge_graph.utils import event_args_list, event_nugget_list, realis_list, arg_2_role |
|
|
|
|
| class CybersecurityKnowledgeGraphConfig(PretrainedConfig): |
| model_type = "cybersecurity_knowledge_graph" |
|
|
| def __init__( |
| self, |
| event_nugget_model_path : str = "nugget_model_state_dict.pth", |
| event_argument_model_path : str = "argument_model_state_dict.pth", |
| event_realis_model_path : str = "realis_model_state_dict.pth", |
| **kwargs, |
| ): |
| self.event_nugget_model_path = event_nugget_model_path |
| self.event_argument_model_path = event_argument_model_path |
| self.event_realis_model_path = event_realis_model_path |
|
|
| self.event_nugget_list = event_nugget_list |
| self.event_args_list = event_args_list |
| self.realis_list = realis_list |
| self.arg_2_role = arg_2_role |
|
|
| super().__init__(**kwargs) |
