| 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) |