main.py

# coding: utf-8

# # Keyword-Extraction using BERT

# Use BERT Token Classification Model to extract keyword tokens from a sentence.

# ## Prepare Dataset for BERT.
# 
# Convert key-text recognition dataset to BIO format dataset.



import os
import argparse
from nltk.tokenize import sent_tokenize
import abbreviation
from abbreviation import limits

parser = argparse.ArgumentParser(description='BERT Keyword Extraction Model')

parser.add_argument('--data', type=str, default='maui-semeval2010-train',
                    help='location of the data corpus')
parser.add_argument('--epochs', type=int, default=4,
                    help='upper epoch limit')
parser.add_argument('--batch_size', type=int, default=32, metavar='N',
                    help='batch size')
parser.add_argument('--seq_len', type=int, default=75, metavar='N',
                    help='sequence length')
parser.add_argument('--lr', type=float, default=3e-5,
                    help='initial learning rate')
parser.add_argument('--save', type=str, default='model.pt',
                    help='path to save the final model')
args = parser.parse_args()


train_path = args.data

txt = sorted([f for f in os.listdir(train_path) if not f.endswith("-justTitle.txt") and not f.endswith(".key") and not f.endswith("-CrowdCountskey")])
key = sorted([f for f in os.listdir(train_path) if f.endswith(".key")])


filekey = dict()
for i, k in enumerate(txt):
    filekey[key[i]] = k


def convert(key):
    sentences = ""
    for line in open(train_path + "/" + filekey[key], 'r'):
        sentences += (" " + line.rstrip())
    tokens = sent_tokenize(sentences)
    key_file = open(train_path + "/" + str(key),'r')
    keys = [line.strip() for line in key_file]
    key_sent = []
    labels = []

    for token in tokens:
        
        for word in token.split():
            
            if word.lower() in limits:
                token = token.replace(word, limits[word.lower()])
             
        punctuations = '''!()-[]{};:\,<>'"./?@#$%^&*_~'''

 

        no_punct = ""
        for char in token:
            if char not in punctuations:
                no_punct = no_punct + char
        

        
        token=no_punct
       
        z = ['O'] * len(token.split())
        for k in keys:

            if k in token.lower():

                if len(k.split())==1:
                    try:
                        
                        z[token.lower().split().index(k.lower().split()[0])] = 'B'

                    except ValueError:
                        continue
                        
                    
                elif len(k.split())>1:

                    try:
                        h = k.split()
                        def index(arr, num):
                            a = []
                            for i, x in enumerate(arr):
                                if x == num:
                                    a.append(i)
                            return(a)

                     
                            
                            
                        lis1 = index(token.lower().split(),k.lower().split()[0])
                        lis2 =  index(token.lower().split(),k.lower().split()[1])
                        x = []
                        for i in lis1:
                            for j in lis2:
                                if abs(i - j) == 1:
                                    x.append(i)
                        num = x[0]  
                        z[num] = 'B'
                        
                        for j in range(1, len(h)):
                            z[num + j] = 'I'
                            
                    except:
                        continue
                     
            elif k.lower().split()[0] in token.lower():
                
                    try:
                        h = k.split()
#                         print(token)
#                         print(len(token.split()))
                        def index(arr, num):
                            a = []
                            for i, x in enumerate(arr):
                                if x == num:
                                    a.append(i)
                            return(a)

                     
                            
                            
                        lis1 = index(token.lower().split(),k.lower().split()[0])
                        lis2 =  index(token.lower().split(),k.lower().split()[1])
                        x = []
                        for i in lis1:
                            for j in lis2:
                                if abs(i - j) == 1:
                                    x.append(i)
                                    
                        num = x[0]  
                        z[num] = 'B'
                        p= len(h)-1
                        y=0
                        for j in range(1,len(token.split())):
                            if (k.lower().split()[j] in token.lower().split()):
#                                 print('rahul')
                                print(j)
                                z[num + j] = 'I'
                                print(k.lower().split()[j])
                                y+=1
                                if y==p:
                                  break
                            else:
                                continue
                            
                    except:
                        continue      
        
        for m, n in enumerate(z):
            if z[m] == 'I' and z[m-1] == 'O':
                z[m] = 'O'

        if set(z) != {'O'}:
            labels.append(z) 
            key_sent.append(token)
    return key_sent, labels


sentences_ = []
labels_ = []
for key, value in filekey.items():
    s, l = convert(key)
    sentences_.append(s)
    labels_.append(l)
sentences = [item for sublist in sentences_ for item in sublist]
labels = [item for sublist in labels_ for item in sublist]
# print(len(sentences), len(labels))


# ## Model

import pandas as pd
import numpy as np
from tqdm import tqdm, trange
import torch
from torch.optim import Adam
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from torch.nn.utils.rnn import pad_sequence
from keras.preprocessing.sequence import pad_sequences
from sklearn.model_selection import train_test_split
from pytorch_pretrained_bert import BertTokenizer, BertConfig
from pytorch_pretrained_bert import BertForTokenClassification, BertAdam


MAX_LEN = args.seq_len
bs = args.batch_size


tag2idx = {'B': 0, 'I': 1, 'O': 2}
tags_vals = ['B', 'I', 'O']


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


tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)


tokenized_texts = [tokenizer.tokenize(sent) for sent in sentences]


input_ids = pad_sequences([tokenizer.convert_tokens_to_ids(txt) for txt in tokenized_texts],
                          maxlen=MAX_LEN, dtype="long", truncating="post", padding="post")
tags = pad_sequences([[tag2idx.get(l) for l in lab] for lab in labels],
                     maxlen=MAX_LEN, value=tag2idx["O"], padding="post",
                     dtype="long", truncating="post")


attention_masks = [[float(i>0) for i in ii] for ii in input_ids]


tr_inputs, val_inputs, tr_tags, val_tags = train_test_split(input_ids, tags, 
                                                            random_state=2018, test_size=0.1)
tr_masks, val_masks, _, _ = train_test_split(attention_masks, input_ids,
                                             random_state=2018, test_size=0.1)


tr_inputs = torch.tensor(tr_inputs)
val_inputs = torch.tensor(val_inputs)
tr_tags = torch.tensor(tr_tags)
val_tags = torch.tensor(val_tags)
tr_masks = torch.tensor(tr_masks)
val_masks = torch.tensor(val_masks)


train_data = TensorDataset(tr_inputs, tr_masks, tr_tags)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=bs)

valid_data = TensorDataset(val_inputs, val_masks, val_tags)
valid_sampler = SequentialSampler(valid_data)
valid_dataloader = DataLoader(valid_data, sampler=valid_sampler, batch_size=bs)


model = BertForTokenClassification.from_pretrained("bert-base-uncased", num_labels=len(tag2idx))


model = model.cuda()


FULL_FINETUNING = True
if FULL_FINETUNING:
    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'gamma', 'beta']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
         'weight_decay_rate': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
         'weight_decay_rate': 0.0}
    ]
else:
    param_optimizer = list(model.classifier.named_parameters()) 
    optimizer_grouped_parameters = [{"params": [p for n, p in param_optimizer]}]
optimizer = Adam(optimizer_grouped_parameters, lr=args.lr)


from seqeval.metrics import f1_score

def flat_accuracy(preds, labels):
    pred_flat = np.argmax(preds, axis=2).flatten()
    labels_flat = labels.flatten()
    return np.sum(pred_flat == labels_flat) / len(labels_flat)


epochs = args.epochs
max_grad_norm = 1.0

for _ in trange(epochs, desc="Epoch"):
    # TRAIN loop
    model.train()
    tr_loss = 0
    nb_tr_examples, nb_tr_steps = 0, 0
    for step, batch in enumerate(train_dataloader):
        # add batch to gpu
        batch = tuple(t.to(device) for t in batch)
        b_input_ids, b_input_mask, b_labels = batch
        # forward pass
        loss = model(b_input_ids, token_type_ids=None,
                     attention_mask=b_input_mask, labels=b_labels)
        # backward pass
        loss.backward()
        # track train loss
        tr_loss += loss.item()
        nb_tr_examples += b_input_ids.size(0)
        nb_tr_steps += 1
        # gradient clipping
        torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=max_grad_norm)
        # update parameters
        optimizer.step()
        model.zero_grad()
    # print train loss per epoch
    print("Train loss: {}".format(tr_loss/nb_tr_steps))
    # VALIDATION on validation set
    model.eval()
    eval_loss, eval_accuracy = 0, 0
    nb_eval_steps, nb_eval_examples = 0, 0
    predictions , true_labels = [], []
    for batch in valid_dataloader:
        batch = tuple(t.to(device) for t in batch)
        b_input_ids, b_input_mask, b_labels = batch
        
        with torch.no_grad():
            tmp_eval_loss = model(b_input_ids, token_type_ids=None,
                                  attention_mask=b_input_mask, labels=b_labels)
            logits = model(b_input_ids, token_type_ids=None,
                           attention_mask=b_input_mask)
        logits = logits.detach().cpu().numpy()
        label_ids = b_labels.to('cpu').numpy()
        predictions.extend([list(p) for p in np.argmax(logits, axis=2)])
        true_labels.append(label_ids)
        
        tmp_eval_accuracy = flat_accuracy(logits, label_ids)
        
        eval_loss += tmp_eval_loss.mean().item()
        eval_accuracy += tmp_eval_accuracy
        
        nb_eval_examples += b_input_ids.size(0)
        nb_eval_steps += 1
    eval_loss = eval_loss/nb_eval_steps
    print("Validation loss: {}".format(eval_loss))
    print("Validation Accuracy: {}".format(eval_accuracy/nb_eval_steps))
    pred_tags = [tags_vals[p_i] for p in predictions for p_i in p]
    valid_tags = [tags_vals[l_ii] for l in true_labels for l_i in l for l_ii in l_i]
    print("F1-Score: {}".format(f1_score(pred_tags, valid_tags)))

torch.save(model, args.save)
    
model.eval()
predictions = []
true_labels = []
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for batch in valid_dataloader:
    batch = tuple(t.to(device) for t in batch)
    b_input_ids, b_input_mask, b_labels = batch

    with torch.no_grad():
        tmp_eval_loss = model(b_input_ids, token_type_ids=None,
                              attention_mask=b_input_mask, labels=b_labels)
        logits = model(b_input_ids, token_type_ids=None,
                       attention_mask=b_input_mask)
        
    logits = logits.detach().cpu().numpy()
    predictions.extend([list(p) for p in np.argmax(logits, axis=2)])

    label_ids = b_labels.to('cpu').numpy()
    true_labels.append(label_ids)
    tmp_eval_accuracy = flat_accuracy(logits, label_ids)

    eval_loss += tmp_eval_loss.mean().item()
    eval_accuracy += tmp_eval_accuracy

    nb_eval_examples += b_input_ids.size(0)
    nb_eval_steps += 1

pred_tags = [[tags_vals[p_i] for p_i in p] for p in predictions]
valid_tags = [[tags_vals[l_ii] for l_ii in l_i] for l in true_labels for l_i in l ]
print("Validation loss: {}".format(eval_loss/nb_eval_steps))
print("Validation Accuracy: {}".format(eval_accuracy/nb_eval_steps))
print("Validation F1-Score: {}".format(f1_score(pred_tags, valid_tags)))