| import os, json |
| from cybersecurity_knowledge_graph.utils import get_content, get_event_args, get_event_nugget, get_idxs_from_text, get_args_entity_from_idx, find_dict_by_overlap |
| from tqdm import tqdm |
| import spacy |
| import jsonlines |
| from sklearn.model_selection import train_test_split |
| import math |
| from transformers import pipeline |
| from sentence_transformers import SentenceTransformer |
| import numpy as np |
|
|
| embed_model = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
| pipe = pipeline("token-classification", model="CyberPeace-Institute/SecureBERT-NER") |
|
|
| nlp = spacy.load('en_core_web_sm') |
|
|
| """ |
| Class: EventArgumentRoleDataset |
| Description: This class represents a dataset for training and evaluating event argument role classifiers. |
| Attributes: |
| - path: The path to the folder containing JSON files with event data. |
| - tokenizer: A tokenizer for encoding text data. |
| - arg: The specific argument type (subtype) for which the dataset is created. |
| - data: A list to store data samples, each consisting of an embedding and a label. |
| - train_data, val_data, test_data: Lists to store the split training, validation, and test data samples. |
| - datapoint_id: An identifier for tracking data samples. |
| Methods: |
| - __len__(): Returns the total number of data samples in the dataset. |
| - __getitem__(index): Retrieves a data sample at a specified index. |
| - to_jsonlines(train_path, val_path, test_path): Writes the dataset to JSON files for train, validation, and test sets. |
| - train_val_test_split(): Splits the data into training and test sets. |
| - load_data(): Loads and preprocesses event data from JSON files, creating embeddings for argument-role classification. |
| """ |
| class EventArgumentRoleDataset(): |
| def __init__(self, path, tokenizer, arg): |
| self.path = path |
| self.tokenizer = tokenizer |
| self.arg = arg |
| self.data = [] |
| self.train_data, self.val_data, self.test_data = None, None, None |
| self.datapoint_id = 0 |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, index): |
| sample = self.data[index] |
| return sample |
| |
| def to_jsonlines(self, train_path, val_path, test_path): |
| if self.train_data is None or self.test_data is None: |
| raise ValueError("Do the train-val-test split") |
| with jsonlines.open(train_path, "w") as f: |
| f.write_all(self.train_data) |
| |
| |
| with jsonlines.open(test_path, "w") as f: |
| f.write_all(self.test_data) |
|
|
| def train_val_test_split(self): |
| self.train_data, self.test_data = train_test_split(self.data, test_size=0.1, random_state=42, shuffle=True) |
| |
| |
| def load_data(self): |
| folder_path = self.path |
| json_files = [file for file in os.listdir(folder_path) if file.endswith('.json')] |
|
|
| |
| for idx, file_path in enumerate(tqdm(json_files)): |
| try: |
| with open(self.path + file_path, "r") as f: |
| file_json = json.load(f) |
| except: |
| print("Error in ", file_path) |
| content = get_content(file_json) |
| content = content.replace("\xa0", " ") |
|
|
| event_args = get_event_args(file_json) |
| doc = nlp(content) |
| |
| sentence_indexes = [] |
| for sent in doc.sents: |
| start_index = sent[0].idx |
| end_index = sent[-1].idx + len(sent[-1].text) |
| sentence_indexes.append((start_index, end_index)) |
| |
| for idx, (start, end) in enumerate(sentence_indexes): |
| sentence = content[start:end] |
| is_arg_sentence = [event_arg["startOffset"] >= start and event_arg["endOffset"] <= end for event_arg in event_args] |
| args = [event_args[idx] for idx, boolean in enumerate(is_arg_sentence) if boolean] |
| if args != []: |
| sentence_doc = nlp(sentence) |
| sentence_embed = embed_model.encode(sentence) |
| for arg in args: |
| if arg["type"] == self.arg: |
| arg_embed = embed_model.encode(arg["text"]) |
| embedding = np.concatenate((sentence_embed, arg_embed)) |
|
|
| self.data.append({"embedding" : embedding, "label" : arg["role"]["type"]}) |
