bert2vul_cuda_cv.py

# -*- coding: utf-8 -*-
"""
Created on Thu May 12 19:51:39 2022

@author: iliaskaloup
"""

import seaborn as sn
import pandas as pd
import json

import tensorflow as tf
import numpy as np
import csv
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification #, BertModel, BertTokenizer, TFBertForSequenceClassification
import matplotlib.pyplot as plt
import random
from tensorflow.keras.callbacks import CSVLogger
from tensorflow.keras.callbacks import EarlyStopping
import tensorflow.keras.backend as K
from collections import OrderedDict
import time
from sklearn.metrics import accuracy_score, recall_score, f1_score, precision_score, \
roc_auc_score, confusion_matrix, classification_report
from sklearn.model_selection import StratifiedKFold


# user parameters
n_folds = 5
n_epochs = 4
batch_size = 6
lr = 1e-5
seed = 123
model_variation = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_variation) #Tokenizer
#bert-base-uncased # roberta-base # distilbert-base-uncased # microsoft/codebert-base-mlm # microsoft/graphcodebert-base
# TFAutoModelForSequenceClassification # BertModel


def recall_metric(y_true, y_pred):
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
        recall = (true_positives + K.epsilon()) / (possible_positives + K.epsilon())
        return recall

def precision_metric(y_true, y_pred):
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
        precision = (true_positives + K.epsilon()) / (predicted_positives + K.epsilon())
        return precision

def f1_metric(y_true, y_pred):

    prec = precision_metric(y_true, y_pred)
    rec = recall_metric(y_true, y_pred)
    f1 = 2*((prec*rec)/(prec+rec+K.epsilon()))
    return f1

def f2_metric(y_true, y_pred):

    prec = precision_metric(y_true, y_pred)
    rec = recall_metric(y_true, y_pred)
    f2 = 5*((prec*rec)/(4*prec+rec+K.epsilon()))
    return f2

def f2_loss(y_true, y_pred):

    y_true = tf.cast(y_true, tf.float32)
    tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)
    #tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)
    fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)
    fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)

    p = tp / (tp + fp + K.epsilon())
    r = tp / (tp + fn + K.epsilon())

    f2 = 5*p*r / (4*p+r+K.epsilon())
    f2 = tf.where(tf.math.is_nan(f2), tf.zeros_like(f2), f2)
    
    return 1 - K.mean(f2)

def get_max_str(lst):
    return max(lst, key=len)

def indicize_labels(labels):
    """Transforms string labels into indices"""
    indices=[]
    for j in range(len(labels)):
        for i in range(n_categories):
            if labels[j]==categories[i]:
                indices.append(i)
    return indices

def dropEmpty(tokens0):
    tokens = []
    for i in range(0, len(tokens0)):
        temp = tokens0[i]
        if temp != []:
            tokens.append(temp)
    return tokens

def listToString(s): 
    
    # initialize an empty string
    str1 = "" 
    
    # traverse in the string 
    count = 0
    for ele in s: 
        if count==0:
            str1 = str1 + ele
        else:
            str1 = str1 + ' ' + ele
        count = count + 1
        #str1 += ele  
    
    # return string  
    return str1

def prepareData(data):
        
    # lowercase
    lines = []
    labels = []
    headlines = []
    for i in range(0, len(data)):
        labels.append(int(data[i][1]))
        headlines.append(data[i][0])
        line = data[i][2:]
        lows = [w.lower() for w in line]
        lines.append(lows)
    
    texts = []
    for i in range(0, len(lines)):
        texts.append(listToString(lines[i]))
    
    return texts, labels, headlines 

def makeTfData(texts, indices):
    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='tf') #Tokenized text
    dataset = tf.data.Dataset.from_tensor_slices((dict(inputs), indices)) #Create a tensorflow dataset
    return dataset
    

# Main
#Load data from json file
# Read data
with open('data_reduced_bert.csv', newline='', encoding='utf-8') as f:
    reader = csv.reader(f)
    data = list(reader)
data = dropEmpty(data)
random.shuffle(data) # shuffle the data
#data = data[0:1000]

# Data preparation
texts, labels, headlines = prepareData(data)

# sequence lenght
#max_length = len(get_max_str(texts)) 
#print(max_length)  

# explore data
n_elements=len(headlines)
print('Elements in dataset:', n_elements)
categories=sorted(list(set(labels))) #set will return the unique different entries
n_categories=len(categories)
print("{} categories found:".format(n_categories))
for category in categories:
    print(category)
    
'''fig=plt.figure(figsize=(20,8))
lbl, counts = np.unique(labels,return_counts=True)
ticks = range(len(counts))
plt.bar(ticks,counts, align='center')
plt.xticks(ticks,lbl)
plt.xticks(rotation=90)
plt.ylabel('counts')
plt.show()'''
    
# prepare the dataset to feed it to the model
indices=indicize_labels(labels) #Integer label indices

# visualize how the tokenizer splits the text into 
'''tokens=tokenizer(headlines[0:3], padding=True, truncation=True)['input_ids']
for i in range(3):
  print('Input:',headlines[i])
  print('Subword tokenization:',tokenizer.tokenize(headlines[i]))
  print('Indices:', tokens[i])'''
  
# We see that the character 'Ġ' is used to indicate the start of a new word in the text, while the other parts of the words that get
# split don't contain it. For example 'reevaluate' gets split into 'Ġre' and 'evaluate'. We further notice that the tokenize uses the
# index 0 to indicate the beginning of the text and 2 to indicate the ending. Index 1 is reserved for padding.

# shuffle dataset 10 times with different seeder in a cross-validation-like approach
texts = np.array(texts)
indices = np.array(indices)
scores=['accuracy', 'precision', 'recall', 'f1', 'roc_auc', 'f2', 'fpr']
values = [np.array([]) for i in range(0, len(scores))]
score_dict = OrderedDict(zip(scores, values))
print("Cross-Validation...")
milli_sec1 = int(round(time.time() * 1000))
kfold = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=seed)
f=0
for train_index, test_index in kfold.split(texts, indices): # for seeder in range(0, n_folds):
    f = f + 1
    print('fold number= ', f)
    '''dataset = dataset.shuffle(len(dataset), seed=seeder, reshuffle_each_iteration=False)
    
    # train test split, we use 10% of the data for validation
    val_data_size = int(0.1*n_elements)
    val_ds = dataset.take(val_data_size).batch(batch_size, drop_remainder=True) 
    train_ds = dataset.skip(val_data_size).batch(batch_size, drop_remainder=True)
    #train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)'''
    
    texts_train, texts_test = texts[train_index], texts[test_index]
    indices_train, indices_test = indices[train_index], indices[test_index]
    
    texts_train = texts_train.tolist()
    indices_train = indices_train.tolist()
    texts_test = texts_test.tolist()
    indices_test = indices_test.tolist()
    
    # tokenize the input text
    trainset = makeTfData(texts_train, indices_train)
    testset = makeTfData(texts_test, indices_test)
    train_data_size = len(texts_train)
    train_ds = trainset.take(train_data_size).batch(batch_size, drop_remainder=True)
    val_data_size = len(texts_test)
    val_ds = testset.take(val_data_size).batch(batch_size, drop_remainder=True)
    
    # train model
    model = TFAutoModelForSequenceClassification.from_pretrained(model_variation, num_labels=n_categories)  
    
    model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=lr, epsilon=1e-08, clipnorm=1.),
            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), #f2_loss,
            metrics=[f2_metric] #tf.metrics.SparseCategoricalAccuracy()
            )
    
    #csv_logger = CSVLogger('log.csv', append=True, separator=',')
    #es = EarlyStopping(monitor='val_f2_metric', mode='max', verbose=1, patience=2)
    #history = model.fit(train_ds, validation_data=val_ds, epochs=n_epochs, verbose=1, callbacks=[csv_logger,es])
    history = model.fit(train_ds, validation_data=val_ds, epochs=n_epochs, verbose=1)
    
    # save model's weights
    model.save_weights('./saved_weights.h5')
    
    '''plt.plot(history.history['f2_metric'])
    plt.plot(history.history['val_f2_metric'])
    plt.ylabel('model loss')
    plt.xlabel('epoch')
    plt.legend(['train', 'test'], loc='best')
    plt.savefig('train_history.png')
    plt.show()'''
    
    # load trained model's weights
    trained_model = TFAutoModelForSequenceClassification.from_pretrained(model_variation, num_labels=n_categories)
    trained_model.load_weights('./saved_weights.h5')
    
    # evaluation
    y_test=[]
    x_test=[]
    for i in range(val_data_size):
        y_test.append(indices_test[i])
        x_test.append(texts_test[i])
    #y_test=indicize_labels(y_test)
    tokens=tokenizer(x_test, padding=True,truncation=True, return_tensors='tf') 
    logits=trained_model.predict(dict(tokens), verbose=1).logits
    prob=tf.nn.softmax(logits, axis=1).numpy()   
    predictions=np.argmax(prob, axis=1)    
    
    print("Evaluation on Validation set")
    accuracy=accuracy_score(y_test, predictions)
    precision=precision_score(y_test, predictions)
    recall=recall_score(y_test, predictions)
    f1=f1_score(y_test, predictions)
    f2 = 5*precision*recall / (4*precision+recall)
    roc_auc=roc_auc_score(y_test, predictions)
    print(confusion_matrix(y_test, predictions, labels=[0, 1]))
    tn, fp, fn, tp = confusion_matrix(y_test, predictions).ravel()
    fpr = fp / (fp+tn)
    acc = ((tp+tn)/(tp+tn+fp+fn))
    print("Accuracy:%.2f%%"%(acc*100))
    print("Precision:%.2f%%"%(precision*100))
    print("Recall:%.2f%%"%(recall*100))
    print("F1 score:%.2f%%"%(f1*100))
    print("F2 score:%.2f%%"%(f2*100))
    print("Roc_Auc score:%.2f%%"%(roc_auc*100))
    print("FPR score:%.2f%%"%(fpr*100))
    print(classification_report(y_test, predictions))
    del model
    del trained_model
    score_dict['accuracy'] = np.append(score_dict['accuracy'], accuracy)
    score_dict['precision'] = np.append(score_dict['precision'], precision)
    score_dict['recall'] = np.append(score_dict['recall'], recall)
    score_dict['f1'] = np.append(score_dict['f1'], f1)
    score_dict['roc_auc'] = np.append(score_dict['roc_auc'], roc_auc)
    score_dict['f2'] = np.append(score_dict['f2'], f2)
    score_dict['fpr'] = np.append(score_dict['fpr'], fpr)
  
milli_sec2 = int(round(time.time() * 1000))
print("Training is completed after", milli_sec2-milli_sec1)
print("accuracy: %.2f%% (%.2f%%)" % (score_dict['accuracy'].mean()*100, score_dict['accuracy'].std()*100))
print("precision: %.2f%% (%.2f%%)" % (score_dict['precision'].mean()*100, score_dict['precision'].std()*100))
print("recall: %.2f%% (%.2f%%)" % (score_dict['recall'].mean()*100, score_dict['recall'].std()*100))
print("f1: %.2f%% (%.2f%%)" % (score_dict['f1'].mean()*100, score_dict['f1'].std()*100))
print("roc_auc: %.2f%% (%.2f%%)" % (score_dict['roc_auc'].mean()*100, score_dict['roc_auc'].std()*100))
print("f2: %.2f%% (%.2f%%)" % (score_dict['f2'].mean()*100, score_dict['f2'].std()*100))
print("fpr: %.2f%% (%.2f%%)" % (score_dict['fpr'].mean()*100, score_dict['fpr'].std()*100))