|
from transformers import PretrainedConfig |
|
import torch |
|
|
|
from . import utils |
|
|
|
|
|
|
|
class CybersecurityKnowledgeGraphConfig(PretrainedConfig): |
|
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 = 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 |
|
|
|
super().__init__(**kwargs) |