train_fake_classifier.py

import re
import os
import glob
import pandas as pd
import torch
import torch.optim as optim
import numpy as np
import random
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
from datasets import Dataset
from transformers import DataCollatorWithPadding
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, average_precision_score, roc_auc_score

data_dir = '/cluster/scratch/goezsoy/nlp_lss_datasets'

# fix the seed
seed = 42
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)

max_word_count = 100

# load quality_fake_df to save quality fake samples
quality_fake_df = pd.read_csv(os.path.join(data_dir,'quality_fake_df.csv'))

# create dataframe which contains the dataset
real_fake_df = pd.DataFrame(columns=['speech','label','perplexity'])

# put real samples
valid_df = pd.read_csv(os.path.join(data_dir,'processed_df_valid.csv'))
valid_df = valid_df.iloc[:5000]

# add perplexity score for each real sentence
real_fake_df['speech'] = valid_df['speech'].map(lambda row: ' '.join(row.split()[:max_word_count]))
real_fake_df['label'] = 1

# put fake samples
results_path = '/cluster/home/goezsoy/conditioned_speech_gen/results'
folder_name = 'finetunedgptmed_lr2e5_epoch2'
experiment_name = 'Result_w_5.0_nBeams_1_nGenSent_128_nWordsPerSent_1_topP_0.9_WC_Guar_True_glove_maxSENTENCES.txt'

print(max_word_count,folder_name,experiment_name)
shard_list = glob.glob(os.path.join(results_path,'shard*',folder_name))


for temp_shard in shard_list:
    
    generated_texts_path = os.path.join(temp_shard,experiment_name)

    if os.path.exists(generated_texts_path):
        file = open(generated_texts_path, 'r')

        temp_speech = None
        temp_perplexity = None

        flag_speech = False
        flag_perplexity = False

        for line in file:
            if line != '\n' and re.search("\ASuccess_rate:", line) is None and re.search("#.:", line) is None:
                if re.search("\APerplexity:", line) is not None:
                    temp_perplexity = line.split()[-1]
                    flag_perplexity = True
                else:
                    # remove <|endoftext|> tokens generated by k2t
                    temp_speech = line.replace('<|endoftext|>','')

                    # if initial text is <|endoftext|> removing it leads to
                    # extra space at the start of sentence, so remove it
                    if temp_speech[0] == ' ':
                        temp_speech = temp_speech[1:]
                    flag_speech = True
            
            if flag_speech and flag_perplexity:
                temp_speech = ' '.join(temp_speech.split()[:max_word_count])
                temp_df = pd.DataFrame.from_dict({'speech':[temp_speech],'label':[0],'perplexity':[temp_perplexity]})
                real_fake_df = pd.concat([real_fake_df,temp_df], ignore_index=True)

                flag_speech = False
                flag_perplexity = False

        file.close()
    
real_fake_df.to_csv(os.path.join(data_dir,'real_fake_df.csv'), index=False)

print('compiled generated texts from all shards, created real-fake dataset.\n')
print(f'dataset size: {len(real_fake_df)}.\n')

# setting device on GPU if available, else CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')


def tokenize_function(row):

    tokenizer_dict = tokenizer(row['speech'])
    tokenizer_dict['labels'] = row['label']

    return {**tokenizer_dict}


def prepare_dataloader(df, shuffle=True):

    dataset = Dataset.from_pandas(df)

    dataset = dataset.map(tokenize_function, batched=True)

    dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

    # for dynamic padding
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer,return_tensors='pt')

    dataloader = torch.utils.data.DataLoader(dataset,collate_fn=data_collator, batch_size=16, shuffle=shuffle, drop_last=False)

    return dataloader


real_fake_df = pd.read_csv(os.path.join(data_dir,'real_fake_df.csv'))

real_fake_df['speech'] = real_fake_df['speech'].map(lambda row: str(row))

#os.environ["CURL_CA_BUNDLE"]="" # only comment out if you have https connection error with huggingface
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

# train vs test split
kf = StratifiedKFold(n_splits=4, shuffle=True, random_state=0)

X = np.arange(0,len(real_fake_df))
y = real_fake_df['label'].values

for train_index, test_index in kf.split(X,y):

    # initialize model every fold
    net = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")

    real_fake_train_valid_df = real_fake_df.iloc[train_index].reset_index(drop=True)
    real_fake_test_df = real_fake_df.iloc[test_index].reset_index(drop=True)

    # train vs valid split
    X_train, X_valid, _, _ = train_test_split(real_fake_train_valid_df.index, real_fake_train_valid_df['label'], test_size=0.1, random_state=0, stratify= real_fake_train_valid_df['label'])

    real_fake_train_df = real_fake_train_valid_df.iloc[X_train].reset_index(drop=True)
    real_fake_valid_df = real_fake_train_valid_df.iloc[X_valid].reset_index(drop=True)


    train_dataloader = prepare_dataloader(real_fake_train_df)

    valid_dataloader = prepare_dataloader(real_fake_valid_df)

    test_dataloader = prepare_dataloader(real_fake_test_df, shuffle=False)

    print('data loaded, starting model training.\n')


    net = net.to(device)

    optimizer = optim.Adam(net.parameters(), lr=1e-5)

    # zero the parameters' gradients
    optimizer.zero_grad()

    epochs = 10
    for epoch in range(epochs):  # loop over dataset

        net.train()

        batch_loss_array=[]

        batch_train_gt = []
        batch_train_preds = []

        # training
        for batch_idx, batch_data in enumerate(train_dataloader): # loop over train batches
            
            batch_data = batch_data.to(device)

            # forward pass
            outputs = net(**batch_data)
            loss = outputs.loss

            batch_train_gt += batch_data['labels'].tolist()
            batch_train_preds += torch.argmax(outputs.logits,axis=1).tolist()

            # backpropagation
            loss.backward()

            # update parameters
            optimizer.step()
            optimizer.zero_grad()

            # save batch metrics
            detached_loss = loss.detach().cpu()
            batch_loss_array.append(detached_loss.item())

        # validation
        net.eval()
        with torch.no_grad():

            batch_loss_array_valid=[]

            batch_valid_gt = []
            batch_valid_preds = []

            for _, valid_batch_data in enumerate(valid_dataloader): # loop over valid batches
                
                valid_batch_data = valid_batch_data.to(device)

                # forward pass
                val_outputs = net(**valid_batch_data)
                val_loss = val_outputs.loss

                batch_valid_gt += valid_batch_data['labels'].tolist()
                batch_valid_preds += torch.argmax(val_outputs.logits,axis=1).tolist()

                # save batch metrics
                detached_val_loss = val_loss.detach().cpu()
                batch_loss_array_valid.append(detached_val_loss.item())


        # display metrics at end of epoch
        epoch_train_loss, epoch_val_loss = np.mean(batch_loss_array), np.mean(batch_loss_array_valid)

        train_acc = accuracy_score(batch_train_gt, batch_train_preds)
        valid_acc = accuracy_score(batch_valid_gt, batch_valid_preds)

        print(f'epoch: {epoch+1} / {epochs}, train_loss: {epoch_train_loss:.4f}, val_loss: {epoch_val_loss:.4f}, train_acc: {train_acc:.4f}, val_acc: {valid_acc:.4f}')

    print('training done, starting test set evaluation.\n')

    # evaluation
    net.eval()
    with torch.no_grad():

        batch_loss_array_test=[]
        
        batch_test_gt = []
        batch_test_preds = []
        batch_real_probs = []

        for _, test_batch_data in enumerate(test_dataloader): # loop over valid batches
            
            test_batch_data = test_batch_data.to(device)

            # forward pass
            test_outputs = net(**test_batch_data)
            test_loss = test_outputs.loss

            batch_test_gt += test_batch_data['labels'].tolist()
            batch_test_preds += torch.argmax(test_outputs.logits,axis=1).tolist()
            batch_real_probs += torch.nn.functional.softmax(test_outputs.logits, dim=1)[:,1].tolist()

            # save batch metrics
            detached_test_loss = test_loss.detach().cpu()
            batch_loss_array_test.append(detached_test_loss.item())

        test_loss = np.mean(batch_loss_array_test)

        test_acc = accuracy_score(batch_test_gt, batch_test_preds)
        test_f1 = f1_score(batch_test_gt, batch_test_preds)
        conf_matrix = confusion_matrix(batch_test_gt, batch_test_preds)
        auroc = roc_auc_score(batch_test_gt, batch_real_probs)
        auprc = average_precision_score(batch_test_gt, batch_real_probs)
        tn, fp, fn, tp = confusion_matrix(batch_test_gt, batch_test_preds).ravel()
        test_fpr = fp / (fp + tn)

        print(f'test_fpr: {test_fpr:.4f}, test_f1: {test_f1:.4f}, test_acc: {test_acc:.4f}, test_auroc: {auroc:.4f}, test_auprc: {auprc:.4f}')
        print('labeling -> 0: fake, 1: real')
        print(conf_matrix)
        print('------')

        # extract quality examples for saving
        real_fake_test_df['ground_truth'] = batch_test_gt
        real_fake_test_df['real_prob'] = batch_real_probs
        real_fake_test_df['prediction'] = batch_test_preds

        temp_df = real_fake_test_df[(real_fake_test_df['label']==0) & (real_fake_test_df['prediction']==1)][['speech','perplexity']]
        quality_fake_df = pd.concat([quality_fake_df,temp_df], ignore_index=True)

# save quality examples to disk
quality_fake_df.to_csv(os.path.join(data_dir,'quality_fake_df.csv'), index=False)