| | from transformers import PretrainedConfig |
| | import torch |
| |
|
| |
|
| |
|
| | 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 = ['O', |
| | 'B-Ransom', |
| | 'I-Ransom', |
| | 'B-DiscoverVulnerability', |
| | 'I-DiscoverVulnerability', |
| | 'B-PatchVulnerability', |
| | 'I-PatchVulnerability', |
| | 'B-Databreach', |
| | 'I-Databreach', |
| | 'B-Phishing', |
| | 'I-Phishing' |
| | ] |
| |
|
| | self.event_args_list = ['O', |
| | 'B-System', |
| | 'I-System', |
| | 'B-Organization', |
| | 'B-Money', |
| | 'I-Money', |
| | 'B-Device', |
| | 'B-Person', |
| | 'I-Person', |
| | 'B-Vulnerability', |
| | 'I-Vulnerability', |
| | 'B-Capabilities', |
| | 'I-Capabilities', |
| | 'I-Organization', |
| | 'B-PaymentMethod', |
| | 'I-PaymentMethod', |
| | 'B-Data', |
| | 'I-Data', |
| | 'B-Number', |
| | 'I-Number', |
| | 'B-Malware', |
| | 'I-Malware', |
| | 'B-PII', |
| | 'I-PII', |
| | 'B-CVE', |
| | 'I-CVE', |
| | 'B-Purpose', |
| | 'I-Purpose', |
| | 'B-File', |
| | 'I-File', |
| | 'I-Device', |
| | 'B-Time', |
| | 'I-Time', |
| | 'B-Software', |
| | 'I-Software', |
| | 'B-Patch', |
| | 'I-Patch', |
| | 'B-Version', |
| | 'I-Version', |
| | 'B-Website', |
| | 'I-Website', |
| | 'B-GPE', |
| | 'I-GPE' |
| | ] |
| |
|
| | self.realis_list =["O", |
| | "Generic", |
| | "Other", |
| | "Actual" |
| | ] |
| | self.arg_2_role = { |
| | "File" : ['Tool', 'Trusted-Entity'], |
| | "Person" : ['Victim', 'Attacker', 'Discoverer', 'Releaser', 'Trusted-Entity', 'Vulnerable_System_Owner'], |
| | "Capabilities" : ['Attack-Pattern', 'Capabilities', 'Issues-Addressed'], |
| | "Purpose" : ['Purpose'], |
| | "Time" : ['Time'], |
| | "PII" : ['Compromised-Data', 'Trusted-Entity'], |
| | "Data" : ['Compromised-Data', 'Trusted-Entity'], |
| | "Organization" : ['Victim', 'Releaser', 'Discoverer', 'Attacker', 'Vulnerable_System_Owner', 'Trusted-Entity'], |
| | "Patch" : ['Patch'], |
| | "Software" : ['Vulnerable_System', 'Victim', 'Trusted-Entity', 'Supported_Platform'], |
| | "Vulnerability" : ['Vulnerability'], |
| | "Version" : ['Patch-Number', 'Vulnerable_System_Version'], |
| | "Device" : ['Vulnerable_System', 'Victim', 'Supported_Platform'], |
| | "CVE" : ['CVE'], |
| | "Number" : ['Number-of-Data', 'Number-of-Victim'], |
| | "System" : ['Victim', 'Supported_Platform', 'Vulnerable_System', 'Trusted-Entity'], |
| | "Malware" : ['Tool'], |
| | "Money" : ['Price', 'Damage-Amount'], |
| | "PaymentMethod" : ['Payment-Method'], |
| | "GPE" : ['Place'], |
| | "Website" : ['Trusted-Entity', 'Tool', 'Vulnerable_System', 'Victim', 'Supported_Platform'], |
| | } |
| |
|
| | super().__init__(**kwargs) |
