bert_v1.py

#%%
import torch
from torch.utils.data import  DataLoader
from transformers import BertTokenizer
from sklearn.model_selection import train_test_split

import pandas as pd
# import numpy as np
from sklearn import preprocessing

from transformers import AutoModelForSequenceClassification
from transformers import AdamW


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

tokenizer = BertTokenizer.from_pretrained(
    './model/bert_uncased_L-4_H-512_A-8',
    do_lower_case = True,
    )
print(tokenizer)

#%% lebel processing

label = data['attackVector']

encoder = preprocessing.LabelEncoder()
label_val = encoder.fit_transform(label)

#%%

#%% 
import collections
label_num = len(collections.Counter(label_val))
X_train, X_test, y_train, y_test = train_test_split(data['value'], label_val, test_size=0.2, random_state=42)
X_train_data = X_train.values.tolist()
X_test_data = X_test.values.tolist()

X_train_ = tokenizer(X_train_data, return_tensors="pt", padding="max_length", max_length=128, truncation=True)
print(X_train_)

#%%
X_test_ = tokenizer(X_test_data, return_tensors="pt", padding="max_length", max_length=128, truncation=True)


# %%
# 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__)

#%%
num_of_epochs = 10
learning_rate = 5e-5

# %%

# Dataset & dataloader
train_dataset = Dataset(X_train_, y_train)
val_dataset = Dataset(X_test_, y_test)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)
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')
# device = 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 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__)
# %%
print(model)
# %%
from tqdm.auto import tqdm

tqdm.pandas()

for name, param in model.named_parameters():
    if 'classifier' not in name: # classifier layer
        param.requires_grad = False
    else:
        param.requires_grad = True
# print(model)        
   #%%     
for i in tqdm(range(num_of_epochs//2)):
    print("Epoch: #{}".format(i+1))
    train(train_loader, model, optimizer)
    test(val_loader, model)
    
print("After Unfreezing the layer......................")

for name, param in model.named_parameters():
		param.requires_grad = True
        
for i in tqdm(range(num_of_epochs//2, num_of_epochs)):
    print("Epoch: #{}".format(i+1))
    train(train_loader, model, optimizer)
    test(val_loader, model)