bert_pipeline.py

# -*- coding: utf-8 -*-
"""
Created on Sun Aug  7 17:31:01 2022

@author: IIT
"""


#%% library import
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from transformers import BertTokenizer, BertForSequenceClassification
from sklearn.model_selection import train_test_split

import pandas as pd
import numpy as np
from sklearn import preprocessing
from sklearn.model_selection import train_test_split

from transformers import DataCollatorWithPadding
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

#%% variables
#%%
num_of_epochs = 10
learning_rate = 5e-5
# %% defining custom dataset, training and testing class
import torch

class Dataset(torch.utils.data.Dataset):
    """
    Class to store the data as PyTorch Dataset
    """
    
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels
        
    def __getitem__(self, idx):
        # an encoding can have keys such as input_ids and attention_mask
        # item is a dictionary which has the same keys as the encoding has
        # and the values are the idxth value of the corresponding key (in PyTorch's tensor format)
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item
    
    def __len__(self):
        return len(self.labels)
        

print(Dataset.__doc__)


def train(dataloader, model, optimizer):
    """Method to train the model"""
    
    model.train()
    
    epoch_loss = 0
    size = len(dataloader.dataset)
    
    for i, batch in enumerate(dataloader):  
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].type(torch.LongTensor).to(device)
        
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)

        optimizer.zero_grad()
        loss = outputs.loss
        epoch_loss += loss.item()
        loss.backward()
        optimizer.step()

        
    print('Training loss: {:.3f}'.format(epoch_loss / size))
            
print(train.__doc__)

def test(dataloader, model):
    """Method to test the model's accuracy and loss on the validation set"""
    
    model.eval()
    
    size = len(dataloader.dataset)
    test_loss, accuracy = 0, 0
    
    with torch.no_grad():
        for batch in dataloader:
            X, y = batch['input_ids'].to(device), batch['labels'].type(torch.LongTensor).to(device)
            pred = model(X, labels=y)
            
            test_loss += pred.loss
            accuracy += (pred.logits.softmax(1).argmax(1) == y).type(torch.float).sum().item()
            
        test_loss /= size
        accuracy /= size
        
        print("Test loss: {:.3f}, accuracy: {:.3f}%".format(test_loss, accuracy * 100))
        
print(test.__doc__)

# %% reading data
data = pd.read_csv('../data/output/data_2019.csv')

allcols = ['description','attackVector','attackComplexity','privilegesRequired','userInteraction','scope',
        'confidentialityImpact','integrityImpact','availabilityImpact','baseScore','baseSeverity','exploitabilityScore','impactScore']
metc_clmn = ['attackVector','attackComplexity','privilegesRequired','userInteraction','scope',
        'confidentialityImpact','integrityImpact','availabilityImpact','baseScore','baseSeverity','exploitabilityScore','impactScore']
data = data[allcols].dropna()

#%% dataset splitiing for test, validate and train
import collections
X_main, X_test_res, y_main, y_test_res = train_test_split(data['description'], data[metc_clmn], test_size=0.2, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X_main, y_main, test_size=0.2, random_state=42)

#%% converting to list
X_train = X_train.values.tolist()
X_test = X_test.values.tolist()
X_test_res = X_test_res.values.tolist()

#%% tokenizing test data
tokenizer = BertTokenizer.from_pretrained('./model/bert_uncased_L-4_H-512_A-8',do_lower_case = True)

X_train = tokenizer(X_train, return_tensors="pt", padding="max_length", max_length=128, truncation=True)
X_test = tokenizer(X_test, return_tensors="pt", padding="max_length", max_length=128, truncation=True)
X_test_res = tokenizer(X_test_res, return_tensors="pt", padding="max_length", max_length=128, truncation=True)

# %% preparing label data


#%% creating dataSet and dataLoaders
from torch.utils.data import DataLoader


train_loader = DataLoader(Dataset(X_train, y_train), batch_size=64, shuffle=True)
test_loader = DataLoader(Dataset(X_test, y_test), batch_size=64, shuffle=True)
test_res_loader = DataLoader(Dataset(X_test_res, y_test_res), batch_size=64, shuffle=True)

# %%

from numpy import mean
from torch import nn

from transformers import AdamW


print('Created train & val datasets.')

# device (turn on GPU acceleration for faster execution)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# model
model = AutoModelForSequenceClassification.from_pretrained('./model/bert_uncased_L-4_H-512_A-8', num_labels=4)
model.to(device)

#optimizer
optimizer = AdamW(model.parameters(), lr=learning_rate)



#%%


def pri():
    '''doc'''
    number = 0
    print(number)

number = {'a':12,'v':432}
pri()
print(number)