bert_intent_classifier.py

import torch
import numpy as np
from transformers import BertTokenizer
from torch import nn
from transformers import BertModel
import pandas as pd
from torch.optim import Adam
from tqdm import tqdm
from transformers import RobertaTokenizerFast, RobertaForSequenceClassification,Trainer, TrainingArguments
from sklearn.metrics import precision_recall_fscore_support as score

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
#tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base', max_length = 512)
buyer_intent = {
          'Greet-Ask':0,
          'Negotiate-Price-Decrease':1,
          'Negotiate-Remove-X':2,
          'Ask_Clarification-Y':3,
          'Negotiate-Price-NoChange':4,
          'Negotiate-Add-X':5,
          'Ask_Price':6,
          'Accept':7,
          'Reject':8,
          'Greet-Ask_Negotiate-Price-Decrease':9,
          'Negotiate-Remove-X_Negotiate-Price-Decrease':10,
          'Negotiate-Remove-delivery':11,
      }
seller_intent = {
        'Greet-Inform':0,
        'Negotiate-Price-NoChange':1,
        'Negotiate-Price-Increase':2,
        'Tell_Price':3,
        'Provide_Clarification-Y':4,
        'Negotiate-Add-X':5,
        'Accept':6,
        'Greet-Inform_Negotiate-Price-Increase':7,
        'Acknowledge acceptance':8,
        'Greet-Inform_Negotiate-Price-NoChange':9,
        'Negotiate-Remove-delivery':10,
        'Negotiate-Price-Remove-X':11,
        'avoid_rejection':12
    }


class Dataset(torch.utils.data.Dataset):

    def __init__(self, df):

        '''self.labels = [buyer_intent[df['intent'][i]] for i in range(len(df['intent'])) if int(df['speaker'][i])==0]
        #self.labels = [seller_intent[label] for label in df['intent'] if df['speaker']==1]

        self.texts = [tokenizer(df['utterance'][i], 
                               padding='max_length', max_length = 512, truncation=True,
                                return_tensors="pt") for i in range(len(df['intent'])) if int(df['speaker'][i])==0]'''
        self.labels = []
        self.texts = []
        for i in range(len(df)):
          #print(buyer_intent[df['intent'][i]])
          if df['speaker'][i]==0:
            #print(df['intent'][i])
            self.labels.append(buyer_intent[df['intent'][i].strip()])
            
            #print(buyer_intent[df['intent'][i].strip()])
            self.texts.append(tokenizer(df['utterance'][i], 
                                padding='max_length', max_length = 512, truncation=True,
                                  return_tensors="pt"))

    def classes(self):
        return self.labels

    def __len__(self):
        return len(self.labels)

    def get_batch_labels(self, idx):
        # Fetch a batch of labels
        return np.array(self.labels[idx])

    def get_batch_texts(self, idx):
        # Fetch a batch of inputs
        return self.texts[idx]

    def __getitem__(self, idx):

        batch_texts = self.get_batch_texts(idx)
        batch_y = self.get_batch_labels(idx)

        return batch_texts, batch_y
    
np.random.seed(112)
data_df = pd.read_csv('dialogue_data_v2.csv')

df_train, df_val, df_test = np.split(data_df.sample(frac=1, random_state=42), 
                                     [int(.8*len(data_df)), int(.9*len(data_df))])
df_train = df_train.reset_index()
df_val = df_val.reset_index()
df_test = df_test.reset_index()

class BertClassifier(nn.Module):

    def __init__(self, dropout=0.5):

        super(BertClassifier, self).__init__()

        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(768, 12)
        self.relu = nn.ReLU()

    def forward(self, input_id, mask):

        _, pooled_output = self.bert(input_ids= input_id, attention_mask=mask,return_dict=False)
        dropout_output = self.dropout(pooled_output)
        linear_output = self.linear(dropout_output)
        final_layer = self.relu(linear_output)

        return final_layer




def train(model, train_data, val_data, learning_rate, epochs):

    train, val = Dataset(train_data), Dataset(val_data)

    train_dataloader = torch.utils.data.DataLoader(train, batch_size=2, shuffle=True)
    val_dataloader = torch.utils.data.DataLoader(val, batch_size=2)

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    criterion = nn.CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr= learning_rate)

    if use_cuda:

            model = model.cuda()
            criterion = criterion.cuda()
            print('yes--cuda--in--use')
    min_val_loss = 9999
    early_stop = 5

    for epoch_num in range(epochs):

            total_acc_train = 0
            total_loss_train = 0

            for train_input, train_label in tqdm(train_dataloader):

                train_label = train_label.to(device)
                mask = train_input['attention_mask'].to(device)
                input_id = train_input['input_ids'].squeeze(1).to(device)

                output = model(input_id, mask)
                
                batch_loss = criterion(output, train_label.long())
                total_loss_train += batch_loss.item()
                
                acc = (output.argmax(dim=1) == train_label).sum().item()
                total_acc_train += acc

                model.zero_grad()
                batch_loss.backward()
                optimizer.step()
            
            total_acc_val = 0
            total_loss_val = 0

            with torch.no_grad():

                for val_input, val_label in val_dataloader:

                    val_label = val_label.to(device)
                    mask = val_input['attention_mask'].to(device)
                    input_id = val_input['input_ids'].squeeze(1).to(device)

                    output = model(input_id, mask)

                    batch_loss = criterion(output, val_label.long())
                    total_loss_val += batch_loss.item()
                    
                    acc = (output.argmax(dim=1) == val_label).sum().item()
                    total_acc_val += acc
                
            if min_val_loss>total_loss_val:
                torch.save([model.state_dict()], f"models/{epoch_num}.pth")
                min_val_loss = total_loss_val
                early_stop=5
            else:
                early_stop = early_stop-1
                if early_stop<0:
                    print('Early stop triggered!!!')
                    break
            
            print(
                f'Epochs: {epoch_num + 1} | Train Loss: {total_loss_train / len(train_data): .3f} \
                | Train Accuracy: {total_acc_train / len(train_data): .3f} \
                | Val Loss: {total_loss_val / len(val_data): .3f} \
                | Val Accuracy: {total_acc_val / len(val_data): .3f}')
                  
EPOCHS = 10
model = BertClassifier()
LR = 1e-6
              
train(model, df_train, df_val, LR, EPOCHS)

def evaluate(model, test_data):

    test = Dataset(test_data)

    test_dataloader = torch.utils.data.DataLoader(test, batch_size=2)

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")

    if use_cuda:

        model = model.cuda()

    total_acc_test = 0
    actual=[]
    predicted=[]
    with torch.no_grad():

        for test_input, test_label in test_dataloader:

              test_label = test_label.to(device)
              mask = test_input['attention_mask'].to(device)
              input_id = test_input['input_ids'].squeeze(1).to(device)

              output = model(input_id, mask)
              predicted.extend(output.argmax(dim=1).tolist()) 
              actual.extend(test_label.tolist())
              acc = (output.argmax(dim=1) == test_label).sum().item()
              total_acc_test += acc
        precision, recall, fscore, support = score(actual, predicted)

        print('precision: {}'.format(precision))
        print('recall: {}'.format(recall))
        print('fscore: {}'.format(fscore))
        print('support: {}'.format(support))
    
    print(f'Test Accuracy: {total_acc_test / len(test_data): .3f}')
    
evaluate(model, df_test)