# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
"""
# 0. Setup
"""
## Transformer
!pip install transformers --quiet
!pip install datasets transformers[sentencepiece] --quiet
!pip install "transformers[sentencepiece]" --quiet
"""
# 1. Basic Exploratory Analysis
"""
# import library
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import seaborn as sns 
import time
import datetime
df = pd.read_csv('/kaggle/input/jigsaw-toxic-comment-classification-challenge/train.csv.zip')
test_csv = pd.read_csv('/kaggle/input/jigsaw-toxic-comment-classification-challenge/test.csv.zip')
test_csv_labels = pd.read_csv('/kaggle/input/jigsaw-toxic-comment-classification-challenge/test_labels.csv.zip')
print(df.columns)
print(df.shape)
target_col= df.columns[2:]
feature_col= df.columns[1:2]
df.head()
for col in target_col:
    print(f"The unique value for {col} are {df[col].unique()}")
for col in target_col:
    print(f"The unique value for {col} are {test_csv_labels[col].unique()}")
print(target_col)
print(feature_col)
df.dtypes
categories = ['toxic','severe_toxic','obscene','threat','insult','identity_hate']
sns.barplot(x=categories,y=df[categories].sum())
plt.show()
# Create "y" variable as binary target variable for at least one identification of class
condition = ((df["toxic"]==1) |
             (df["severe_toxic"]==1) |
             (df["obscene"]==1) |
             (df["threat"]==1) |
             (df["insult"]==1) |
             (df["identity_hate"]==1)
            )

df.loc[condition, "y"]=1
df.loc[~condition, "y"]=0
df["y"] = df["y"].astype(int)
agg_df = df["y"]\
    .value_counts(normalize=True)\
    .rename('Proportion')\
    .reset_index()
    
agg_df.columns = ["y", "Proportion"]
splot = sns.barplot(x="y", y='Proportion',data=agg_df)
plt.ylim(0,1)
plt.title("Proportion of Non-Toxic/Toxic Post")
for p in splot.patches:
    splot.annotate(format(p.get_height(), '.2f'), 
                   (p.get_x() + p.get_width() / 2., p.get_height()), 
                   ha = 'center', va = 'center', 
                   xytext = (0, 9), 
                   textcoords = 'offset points')
# Create variable to check on length
df['len'] = df["comment_text"].str.split().str.len()
print(f"Average length is {df['len'].mean()} while max length is {df['len'].max()}.")
df['len'].hist(bins=20)
plt.show()
# Split the distribution chart by y = 1 or 0
sns.histplot(data=df, x="len", hue="y", bins=50, alpha=0.2)
# Cumulative Distribution
sns.histplot(data=df, x="len", hue="y",
             log_scale=True, element="step", fill=False,
             cumulative=True, stat="density", common_norm=False)
df = df.rename(columns={"id": "idx"})
df.head()
"""
# 2. Pre-processing & Modelling
"""
# import pyTorch
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
# import other parts of transformers as well as tqdm
from transformers import (AutoTokenizer, AutoModel, 
                          AutoModelForSequenceClassification, 
                          DataCollatorWithPadding, AdamW, get_scheduler,
                          get_linear_schedule_with_warmup,
                          )
# import pyarrow (can only import this with GPU on kaggle notebook)
import pyarrow as pa
from tqdm.auto import tqdm
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
import datasets
import random
from sklearn.metrics import classification_report
# Setting up seed value
seed_value = 42
random.seed(seed_value)
np.random.seed(seed_value)
torch.manual_seed(seed_value)
torch.cuda.manual_seed_all(seed_value)
"""
## Train, Validation and Test Split
"""
# train & validation, test split
train_val_df, test_df = train_test_split(df[["idx", "comment_text"] + categories], test_size=0.2, random_state=seed_value)
# train, validation split
train_df, val_df, = train_test_split(train_val_df[["idx", "comment_text"] + categories], test_size=0.25, random_state=seed_value)
print(f"Size of train, validation and test are {len(train_df)}, {len(val_df)}, {len(test_df)} respectively.")
print(f"Proportion of train, validation and test are {round(len(train_df)/len(df),2)}, {round(len(val_df)/len(df),2)}, {round(len(test_df)/len(df),2)} respectively.")
train_df.reset_index(inplace=True)
train_df.drop("index", axis=1, inplace=True)

val_df.reset_index(inplace=True)
val_df.drop("index", axis=1, inplace=True)

test_df.reset_index(inplace=True)
test_df.drop("index", axis=1, inplace=True)
train_df.head()
"""
## Tokenizing Text in Batches
"""
# instantiate tokenizer, model using "bert-base-uncased"
# checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
# tokenize and encode sequences in the training set
train_tokens = tokenizer.batch_encode_plus(train_df["comment_text"].tolist(),
                                           max_length = 200,
                                           pad_to_max_length=True,
                                           truncation=True,
                                           return_token_type_ids=False
                                           )

# tokenize and encode sequences in the validation set
val_tokens = tokenizer.batch_encode_plus(val_df["comment_text"].tolist(),
                                         max_length = 200,
                                         pad_to_max_length=True,
                                         truncation=True,
                                         return_token_type_ids=False
                                         )

# tokenize and encode sequences in the test set
test_tokens = tokenizer.batch_encode_plus(test_df["comment_text"].tolist(),
                                          max_length = 200,
                                          pad_to_max_length=True,
                                          truncation=True,
                                          return_token_type_ids=False
                                          )
## convert lists to tensors
train_seq = torch.tensor(train_tokens['input_ids'])
train_mask = torch.tensor(train_tokens['attention_mask'])
# change from a list of 1 & 0 to array then to tensor
train_y = torch.tensor(np.array(train_df[categories]))  

val_seq = torch.tensor(val_tokens['input_ids'])
val_mask = torch.tensor(val_tokens['attention_mask'])
# change from a list of 1 & 0 to array then to tensor
val_y = torch.tensor(np.array(val_df[categories]))

test_seq = torch.tensor(test_tokens['input_ids'])
test_mask = torch.tensor(test_tokens['attention_mask'])
# change from a list of 1 & 0 to array then to tensor
test_y = torch.tensor(np.array(test_df[categories]))
train_y
# wrap tensors
train_data = TensorDataset(train_seq, train_mask, train_y)
# sampler for sampling the data during training
train_sampler = RandomSampler(train_data)

# wrap tensors
val_data = TensorDataset(val_seq, val_mask, val_y)
# sampler for sampling the data during training
val_sampler = SequentialSampler(val_data)

# wrap tensors
test_data = TensorDataset(test_seq, test_mask, test_y)
test_sampler = SequentialSampler(test_data)
# define a batch size
batch_size = 32

# dataLoader for train set
train_dataloader = DataLoader(train_data, 
                              sampler=train_sampler, 
                              batch_size=batch_size,
                              # collate_fn=data_collator
                              )

# dataLoader for validation set
val_dataloader = DataLoader(val_data, 
                            sampler = val_sampler, 
                            batch_size=batch_size,
                            # collate_fn=data_collator
                            )
"""
No dynamic padding. Above accomodate for fix padding size of 200
"""
for step,batch in enumerate(train_dataloader):
    break
print(batch[0])
print(batch[1])
print(batch[2])
"""
## Setting up DistilBert Model (Smaller version of Bert)
"""
# instantiate model using bert-base-uncased"
# Will be training for the whole network instead of just finetuning the weight of the head layer
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels = 6)

# In the event that we only need to finetune the head layer and not the whole network,
# just need to uncomment the below code

# for param in model.bert.parameters():
#     param.requires_grad = False

# for name, param in model.named_parameters():                
#     if param.requires_grad:
#         print(name)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
"""
Also need to set up the following before training:

1. Optimizer
2. Scheduler (if want to go add a little decay in the learning rate)
3. Get the right loss function/criterion
4. Specify the right measurement for accuracy (since this is a multi-label problem)
"""
# Set up optimizer
# Will be training for the whole network instead of just finetuning the weight of the head layer
LEARN_RATE = 3e-5
optimizer = AdamW(model.parameters(),
                  lr = LEARN_RATE, # args.learning_rate - default is 5e-5, our notebook had 2e-5
                  eps = 1e-8 # args.adam_epsilon  - default is 1e-8.
                  )

# In the event that we only need to finetune the head layer and not the whole network,
# just need to uncomment the below code

# optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()),
#                   lr = LEARN_RATE, # args.learning_rate - default is 5e-5, our notebook had 2e-5
#                   eps = 1e-8 # args.adam_epsilon  - default is 1e-8.
#                   )

# optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = 1e-5)
# Number of training epochs. The BERT authors recommend between 2 and 4. 
# We chose to run for 3
epochs = 3

# Total number of training steps is [number of batches] x [number of epochs]. 
# (Note that this is not the same as the number of training samples).
total_steps = len(train_dataloader) * epochs

# Create the learning rate scheduler.
scheduler = get_linear_schedule_with_warmup(optimizer, 
                                            num_warmup_steps = 0, # default value
                                            num_training_steps = total_steps
                                            )
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
# Specify loss function
criterion = nn.BCEWithLogitsLoss()

# https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
# This criterion combines LogSoftmax and NLLLoss in one single class.
# More relevant for multiclass (because of the softmax)

# https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html
# More relevant for multilabel. For more explanation, can refer to the link below.
# https://discuss.pytorch.org/t/using-bcewithlogisloss-for-multi-label-classification/67011/2
# Specify measurement of accuracy
def accuracy_thresh(y_pred, y_true, thresh:float=0.4, sigmoid:bool=True):
    "Compute accuracy when `y_pred` and `y_true` are the same size."
    if sigmoid: 
        y_pred = y_pred.sigmoid()
    # return ((y_pred>thresh)==y_true.byte()).float().mean().item()
    return np.mean(((y_pred>thresh).float()==y_true.float()).float().cpu().numpy(), axis=1).sum()
def format_time(elapsed):
    '''
    Takes a time in seconds and returns a string hh:mm:ss
    '''
    # Round to the nearest second.
    elapsed_rounded = int(round((elapsed)))
    
    # Format as hh:mm:ss
    return str(datetime.timedelta(seconds=elapsed_rounded))
# We'll store a number of quantities such as training and validation loss, 
# validation accuracy, and timings.
training_stats = []

# Measure the total training time for the whole run.
total_t0 = time.time()

# For each epoch...
for epoch_i in range(0, epochs):
    # ========================================
    #               Training
    # ========================================
    # Perform one full pass over the training set.
    print("")
    print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
    print('Training...')

    # Measure how long the training epoch takes.
    t0 = time.time()
    # Reset the total loss & accuracy for this epoch.
    total_train_loss = 0
    total_train_accuracy = 0

    # Put the model into training mode. Don't be mislead--the call to 
    # `train` just changes the *mode*, it doesn't *perform* the training.
    # `dropout` and `batchnorm` layers behave differently during training
    # vs. test (source: https://stackoverflow.com/questions/51433378/what-does-model-train-do-in-pytorch)
    model.train()

    # For each batch of training data...
    for step, batch in enumerate(train_dataloader):
        # Progress update every 40 batches.
        if step % 40 == 0 and not step == 0:
            # Calculate elapsed time in minutes.
            elapsed = format_time(time.time() - t0)
            # Report progress.
            print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_dataloader), elapsed))

        # Unpack this training batch from our dataloader. 
        # As we unpack the batch, we'll also copy each tensor to the GPU using the `to` method.
        # `batch` contains three pytorch tensors:
        #   [0]: input ids 
        #   [1]: attention masks
        #   [2]: labels 
        b_input_ids = batch[0].to(device)
        b_input_mask = batch[1].to(device)
        b_labels = batch[2].to(device)

        # Always clear any previously calculated gradients before performing a backward pass. 
        # PyTorch doesn't do this automatically because accumulating the gradients is "convenient while training RNNs". 
        # (source: https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch)
        model.zero_grad()        

        # Perform a forward pass (evaluate the model on this training batch).
        # The documentation for this `model` function is here: 
        # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
        # It returns different numbers of parameters depending on what arguments are given and what flags are set. 
        # For our useage here, it returns the loss (because we provided labels) and the "logits"--the model outputs prior to activation.
        output = model(b_input_ids, 
                       # token_type_ids=None, 
                       attention_mask=b_input_mask, 
                       #labels=b_labels
                       )
        logits = output.logits
        loss = criterion(logits, b_labels.float())

        # Accumulate the training loss over all of the batches so that we can calculate the average loss at the end. 
        # `loss` is a Tensor containing a single value; the `.item()` function just returns the Python value from the tensor.
        total_train_loss += loss.item()
        
        # Calculate the accuracy for this batch of test sentences, and accumulate it over all batches.
        total_train_accuracy += accuracy_thresh(logits,b_labels.float())

        # Perform a backward pass to calculate the gradients.
        loss.backward()

        # Clip the norm of the gradients to 1.0.
        # This is to help prevent the "exploding gradients" problem.
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        # Update parameters and take a step using the computed gradient.
        # The optimizer dictates the "update rule"--how the parameters are modified based on their gradients, the learning rate, etc.
        optimizer.step()

        # Update the learning rate.
        scheduler.step()

    # Report the final accuracy for this validation run.
    # avg_train_accuracy = total_train_accuracy / len(train_dataloader)
    # print("  Accuracy: {0:.5f}".format(avg_train_accuracy))
    train_accuracy = total_train_accuracy / len(train_df)
    print("  Accuracy: {0:.5f}".format(train_accuracy))
    
    # Calculate the average loss over all of the batches.
    avg_train_loss = total_train_loss / len(train_dataloader)            
    
    # Measure how long this epoch took.
    training_time = format_time(time.time() - t0)

    print("")
    print("  Average training loss: {0:.2f}".format(avg_train_loss))
    print("  Training epcoh took: {:}".format(training_time))
        
    # ========================================
    #               Validation
    # ========================================
    # After the completion of each training epoch, measure our performance on our validation set.

    print("")
    print("Running Validation...")

    t0 = time.time()

    # Put the model in evaluation mode--the dropout layers behave differently during evaluation.
    model.eval()

    # Tracking variables 
    total_eval_accuracy = 0
    total_eval_loss = 0
    nb_eval_steps = 0

    # Evaluate data for one epoch
    for batch in val_dataloader:
        
        # Unpack this training batch from our dataloader. 
        # As we unpack the batch, we'll also copy each tensor to the GPU using the `to` method.
        # `batch` contains three pytorch tensors:
        #   [0]: input ids 
        #   [1]: attention masks
        #   [2]: labels 
        b_input_ids = batch[0].to(device)
        b_input_mask = batch[1].to(device)
        b_labels = batch[2].to(device)
        
        # Tell pytorch not to bother with constructing the compute graph during the forward pass, since this is only needed for backprop (training).
        with torch.no_grad():        

            # Forward pass, calculate logit predictions.
            # token_type_ids is the same as the "segment ids", which differentiates sentence 1 and 2 in 2-sentence tasks.
            # The documentation for this `model` function is here: 
            # https://huggingface.co/transformers/v2.2.0/model_doc/bert.html#transformers.BertForSequenceClassification
            # Get the "logits" output by the model. The "logits" are the output values prior to applying an activation function like the softmax.
            output = model(b_input_ids, 
                           # token_type_ids=None, 
                           attention_mask=b_input_mask, 
                           # labels=b_labels
                           )
            
            logits = output.logits
            loss = criterion(logits, b_labels.float())
            
        # Accumulate the validation loss.
        total_eval_loss += loss.item()
        
        # Calculate the accuracy for this batch of test sentences, and accumulate it over all batches.
        total_eval_accuracy += accuracy_thresh(logits,b_labels.float())

        # Move logits and labels to CPU
        logits = logits.detach().cpu().numpy() # maybe irrelevant over here
        label_ids = b_labels.to('cpu').numpy() # maybe irrelevant over here

    # Report the final accuracy for this validation run.
    # avg_val_accuracy = total_eval_accuracy / len(val_dataloader)
    # print("  Accuracy: {0:.5f}".format(avg_val_accuracy))
    val_accuracy = total_eval_accuracy / len(val_df)
    print("  Accuracy: {0:.5f}".format(val_accuracy))

    # Calculate the average loss over all of the batches.
    avg_val_loss = total_eval_loss / len(val_dataloader)
    
    # Measure how long the validation run took.
    validation_time = format_time(time.time() - t0)
    
    print("  Validation Loss: {0:.2f}".format(avg_val_loss))
    print("  Validation took: {:}".format(validation_time))

    # Record all statistics from this epoch.
    training_stats.append(
        {
            'epoch': epoch_i + 1,
            'Training Loss': avg_train_loss,
            'Valid. Loss': avg_val_loss,
            # 'Training Accur': avg_train_accuracy,
            # 'Valid. Accur.': avg_val_accuracy,
            'Training Accur': train_accuracy,
            'Valid. Accur.': val_accuracy,
            'Training Time': training_time,
            'Validation Time': validation_time
        }
    )

print("")
print("Training complete!")

print("Total training took {:} (h:mm:ss)".format(format_time(time.time()-total_t0)))
# Display floats with two decimal places.
pd.set_option('precision', 5)
# Create a DataFrame from our training statistics.
df_stats = pd.DataFrame(data=training_stats)
# Use the 'epoch' as the row index.
df_stats = df_stats.set_index('epoch')
# A hack to force the column headers to wrap.
#df = df.style.set_table_styles([dict(selector="th",props=[('max-width', '70px')])])
# Display the table.
df_stats
# Plot the learning curve.
plt.plot(df_stats['Training Loss'], 'b-o', label="Training")
plt.plot(df_stats['Valid. Loss'], 'g-o', label="Validation")

# Label the plot.
plt.title("Training & Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.xticks([1, 2, 3, 4])

plt.show()
# Plot the accuracy curve.
plt.plot(df_stats['Training Accur'], 'b-o', label="Training")
plt.plot(df_stats['Valid. Accur.'], 'g-o', label="Validation")

# Label the plot.
plt.title("Training & Validation Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.xticks([1, 2, 3, 4])

plt.show()
"""
## Get Test Metrics Result
"""
# get predictions for test data
thresh = 0.4
with torch.no_grad():
    outputs = model(test_seq[0:1000].to(device), test_mask[0:1000].to(device))
    pred_probs = torch.sigmoid(outputs.logits)
    pred_probs = pred_probs.cpu().detach().numpy()
pred_probs
y_pred = (pred_probs > thresh).astype(int)
y_pred
y_true = np.array(test_df[categories])[0:1000]
from sklearn.metrics import hamming_loss, accuracy_score
def hamming_score(y_true, y_pred, normalize=True, sample_weight=None):
    '''
    Compute the Hamming score (a.k.a. label-based accuracy) for the multi-label case
    http://stackoverflow.com/q/32239577/395857
    
    Take in np.array for y_true and y_pred. E.g.
    y_true = np.array([[0,1,0],
                       [0,1,1],
                       [1,0,1],
                       [0,0,1]])

    y_pred = np.array([[0,1,1],
                       [0,1,1],
                       [0,1,0],
                       [0,0,0]])
    '''
    acc_list = []
    for i in range(y_true.shape[0]):
        set_true = set( np.where(y_true[i])[0] )
        set_pred = set( np.where(y_pred[i])[0] )
        #print('\nset_true: {0}'.format(set_true))
        #print('set_pred: {0}'.format(set_pred))
        tmp_a = None
        if len(set_true) == 0 and len(set_pred) == 0:
            tmp_a = 1
        else:
            tmp_a = len(set_true.intersection(set_pred))/\
                    float( len(set_true.union(set_pred)) )
        #print('tmp_a: {0}'.format(tmp_a))
        acc_list.append(tmp_a)
    return np.mean(acc_list)
print("accuracy_score:", accuracy_score(y_true, y_pred))
print("Hamming_score:", hamming_score(y_true, y_pred))
print("Hamming_loss:", hamming_loss(y_true, y_pred))
"""
# 3. Saving the Model
"""
PATH = "./toxic_distilBERT_multilabel"
torch.save(model.state_dict(), PATH)
"""
# 4. Submission
"""
test_csv.head()
len(test_csv)
# tokenize and encode sequences in the actual test set
sub_tokens = tokenizer.batch_encode_plus(test_csv["comment_text"].tolist(),
                                         max_length = 200,
                                         pad_to_max_length=True,
                                         truncation=True,
                                         return_token_type_ids=False
                                         )
sub_seq = torch.tensor(sub_tokens['input_ids'])
sub_mask = torch.tensor(sub_tokens['attention_mask'])
sub_data = TensorDataset(sub_seq, sub_mask)
# dataLoader for validation set
sub_dataloader = DataLoader(sub_data, 
                            batch_size=batch_size)
# Measure how long the evaluation going to takes.
t0 = time.time()

for step, batch in enumerate(sub_dataloader):
    # Progress update every 40 batches.
    if step % 40 == 0 and not step == 0:
        # Calculate elapsed time in minutes.
        elapsed = format_time(time.time() - t0)
        # Report progress.
        print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(sub_dataloader), elapsed))
    b_input_ids = batch[0].to(device)
    b_input_mask = batch[1].to(device)
    with torch.no_grad():
        outputs = model(b_input_ids, b_input_mask)
        pred_probs = torch.sigmoid(outputs.logits)
        if step == 0:
            predictions = pred_probs.cpu().detach().numpy()
        else:
            predictions = np.append(predictions, pred_probs.cpu().detach().numpy(), axis=0)
predictions_df = pd.DataFrame(predictions, columns = categories)
len(predictions_df)
submission = pd.concat([test_csv["id"], predictions_df], axis=1)
submission.head()
submission.to_csv('submission.csv', index=False, header=True)
"""
# 5. Test on run time for CPU
"""
# instantiate model using "distilbert-base-uncased"
checkpoint = "distilbert-base-uncased"
PATH = "./toxic_distilBERT_multilabel"
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels = 6)
model.load_state_dict(torch.load(PATH, map_location=torch.device('cpu')))
device = torch.device("cpu")
model.to(device)
# tokenize and encode sequences in the actual test set
sub_tokens = tokenizer.batch_encode_plus(test_csv["comment_text"].tolist(),
                                         max_length = 200,
                                         pad_to_max_length=True,
                                         truncation=True,
                                         return_token_type_ids=False
                                         )
sub_seq = torch.tensor(sub_tokens['input_ids'])
sub_mask = torch.tensor(sub_tokens['attention_mask'])
sub_data = TensorDataset(sub_seq, sub_mask)
# dataLoader for validation set
sub_dataloader = DataLoader(sub_data, 
                            batch_size=batch_size)
# Measure how long the evaluation going to takes.
t0 = time.time()

for step, batch in enumerate(sub_dataloader):
    # Progress update every 40 batches.
    if step % 40 == 0 and not step == 0:
        # Calculate elapsed time in minutes.
        elapsed = format_time(time.time() - t0)
        # Report progress.
        print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(sub_dataloader), elapsed))
    b_input_ids = batch[0].to(device)
    b_input_mask = batch[1].to(device)
    with torch.no_grad():
        outputs = model(b_input_ids, b_input_mask)
        pred_probs = torch.sigmoid(outputs.logits)
        if step == 0:
            predictions = pred_probs.cpu().detach().numpy()
        else:
            predictions = np.append(predictions, pred_probs.cpu().detach().numpy(), axis=0)