forked from Eshe/competition-vd
62 lines
2.1 KiB
Python
Executable File
62 lines
2.1 KiB
Python
Executable File
import torch
|
|
from torch.utils.data import Dataset
|
|
import numpy as np
|
|
from transformers import RobertaTokenizer
|
|
|
|
class VulnerabilityDataset(Dataset):
|
|
def __init__(self, functions, labels, tokenizer, max_length=512):
|
|
self.functions = functions
|
|
self.labels = labels
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
|
|
def __len__(self):
|
|
return len(self.functions)
|
|
|
|
def __getitem__(self, idx):
|
|
function = self.functions[idx]
|
|
label = self.labels[idx]
|
|
|
|
encoding = self.tokenizer.encode_plus(
|
|
function,
|
|
add_special_tokens=True,
|
|
max_length=self.max_length,
|
|
return_token_type_ids=False,
|
|
padding='max_length',
|
|
truncation=True,
|
|
return_attention_mask=True,
|
|
return_tensors='pt',
|
|
)
|
|
|
|
return {
|
|
'input_ids': encoding['input_ids'].flatten(),
|
|
'attention_mask': encoding['attention_mask'].flatten(),
|
|
'labels': torch.tensor(label, dtype=torch.long)
|
|
}
|
|
|
|
def augment_data(functions, labels):
|
|
augmented_functions = []
|
|
augmented_labels = []
|
|
|
|
for function, label in zip(functions, labels):
|
|
augmented_functions.append(function)
|
|
augmented_labels.append(label)
|
|
|
|
augmented_function = function.replace('int', 'long').replace('float', 'double')
|
|
augmented_functions.append(augmented_function)
|
|
augmented_labels.append(label)
|
|
|
|
augmented_function = function.replace('for', 'while')
|
|
augmented_functions.append(augmented_function)
|
|
augmented_labels.append(label)
|
|
|
|
return np.array(augmented_functions), np.array(augmented_labels)
|
|
|
|
def get_data_loaders(functions, labels, batch_size=32):
|
|
tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
|
|
augmented_functions, augmented_labels = augment_data(functions, labels)
|
|
|
|
dataset = VulnerabilityDataset(augmented_functions, augmented_labels, tokenizer)
|
|
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
|
|
|
|
return dataloader |