Speed and power: fine-tuning distilled transformers models

%%capture
!pip uninstall -y kaggle && pip install kaggle 
# Needed to get the latest version of the Kaggle CLI

from getpass import getpass
import os

# We'll use the Kaggle-CLI to download the dataset
# To create an authentication token on Kaggle check https://www.kaggle.com/docs/api
# You'll also have to accept the competition rules here: 
# https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge/rules

os.environ["KAGGLE_USERNAME"] = getpass(prompt='Kaggle username: ')
os.environ["KAGGLE_KEY"] = getpass(prompt='Token: ')

!kaggle --version
!kaggle competitions download -c jigsaw-toxic-comment-classification-challenge
!unzip jigsaw-toxic-comment-classification-challenge.zip && unzip train.csv.zip

import pandas as pd
from sklearn.model_selection import train_test_split

dataset = pd.read_csv("train.csv")
texts = list(dataset["comment_text"])
label_names = dataset.drop(["id", "comment_text"], axis=1).columns
labels = dataset[label_names].values

train_texts, test_texts, train_labels, test_labels = train_test_split(
    texts, labels, test_size=0.2, random_state=42
)

sample_idx = 23
print(f'Sample: "{train_texts[sample_idx]}"')
print(f"Labels: {pd.Series(train_labels[sample_idx], label_names).to_dict()}")

!pip install -q transformers > /dev/null

import transformers
print(f"Transformers package version: {transformers.__version__}")

from transformers import AutoConfig, AutoTokenizer, AutoModel

MAX_LENGTH = 200
BATCH_SIZE = 16
LEARNING_RATE = 1e-05

MODEL_NAME = "google/bert_uncased_L-2_H-128_A-2"

config = AutoConfig.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) 
transformer_model = AutoModel.from_pretrained(MODEL_NAME)

import numpy as np
import pandas as pd
from sklearn import metrics
import transformers
import torch
from torch.utils.data import Dataset, DataLoader

device = 'cuda' if torch.cuda.is_available() else 'cpu'

class TransformersDataset(Dataset):
    """Encode texts and return them with their corresponding labels."""
    def __init__(self, texts, labels, tokenizer, max_length):
        """Tokenize the corpus of texts and store them with the labels,
        as PyTorch tensors."""
        self.encoded_inputs = tokenizer(texts, truncation=True, padding=True, 
            max_length=max_length, return_tensors="pt")
        self.labels = torch.tensor(labels, dtype=torch.float)

    def __len__(self) -> int:
        """In PyTorch datasets have to override the length method."""
        return len(self.labels)

    def __getitem__(self, index: int) -> dict:
        """This method defines how to feed the data during model training."""
        inputs_ands_labels = dict()
        for key in ["input_ids", "attention_mask", "token_type_ids"]:
            inputs_ands_labels[key] = self.encoded_inputs[key][index]
        inputs_ands_labels['labels'] = self.labels[index]
        return inputs_ands_labels

train_dataset = TransformersDataset(train_texts, train_labels, tokenizer, MAX_LENGTH)
test_dataset = TransformersDataset(test_texts, test_labels, tokenizer, MAX_LENGTH)

training_loader = DataLoader(train_dataset, BATCH_SIZE, shuffle=True, num_workers=0)
testing_loader = DataLoader(test_dataset, BATCH_SIZE, shuffle=False, num_workers=0)

from torch.nn.functional import binary_cross_entropy_with_logits

class MultiLabelTransformer(torch.nn.Module):
    """PyTorch module used to fine-tune a transformer model. 
    
    The structure of this class is meant to be compatible with the Trainer from the 
    Transformers library: https://huggingface.co/transformers/main_classes/trainer.html
    """
    def __init__(self, transformer_model, config, n_labels):
        super().__init__()
        self.transformer_model = transformer_model
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        self.classifier = torch.nn.Linear(config.hidden_size, n_labels)
    
    def forward(self, input_ids, attention_mask=None, token_type_ids=None, labels=None):
        """During fine-tuning, attention masks, token types ids have to be passed as
        argument, with the labels. During inference, the model simply accepts a sequence 
        of tokens ids and produces a set of binary logit scores.
        """
        pooled_output = self.transformer_model(
            input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids
        ).pooler_output
        pooled_output = self.dropout(pooled_output)
        logits  = self.classifier(pooled_output)
        
        if labels is None:
            return logits
        else:
            return binary_cross_entropy_with_logits(logits, labels), logits

model = MultiLabelTransformer(transformer_model, config, train_labels.shape[1])
model.to(device) # Move our model from memory to the GPU's memory

optimizer = torch.optim.Adam(params=model.parameters(), lr=LEARNING_RATE)

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    evaluation_strategy="epoch",
    logging_first_step=True,
    logging_steps=10000
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset
)

trainer.train()

# Uncomment the lines below if you want to visualize metrics with TensorBoard
# %load_ext tensorboard
# %tensorboard --logdir logs

torch.save(model, "./model.pt")

import pandas as pd
from time import time

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) 
loaded_model = torch.load("./model.pt")
loaded_model.eval()  # Run to finish loading the model

def score_text(text, tokenizer=tokenizer, model=loaded_model):
    """Tokenize and score a single text."""
    inputs = tokenizer.encode_plus(
        text, max_length=MAX_LENGTH, padding='max_length', 
        return_token_type_ids=False, return_attention_mask=False
    )

    token_ids = torch.tensor(inputs['input_ids'], dtype=torch.long).reshape(1, -1)
    token_ids = token_ids.to(device)

    logits = model(token_ids)
    scores = torch.sigmoid(logits[0])
    return scores.cpu().detach().numpy()

text = """I am a nice Wikipedia user, I mean no harm, 
I will not insult anybody or be offensive anyhow."""

t0 = time()
scores = score_text(text)
latency = time()-t0

scores = pd.Series(scores, label_names, name="scores")
print(scores.to_frame())
print(f"\nLatency: {latency:.3f} seconds")

%%capture
!pip install pytorch-lightning pytorch-lightning-bolts

import pytorch_lightning as pl
from pytorch_lightning.metrics import Accuracy, AveragePrecision
from torch.nn.functional import binary_cross_entropy_with_logits

class LightningTransformer(pl.LightningModule):
    def __init__(self, transformer_model, config, n_labels):
        super().__init__()
        self.transformer_model = transformer_model
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        self.n_labels = n_labels
        self.classifier = torch.nn.Linear(config.hidden_size, n_labels)
    
    def forward(self, batch):
        input_ids = batch['input_ids']
        attention_mask = batch['attention_mask']
        token_type_ids = batch['token_type_ids']

        pooled_output = self.transformer_model(
            input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids
        ).pooler_output
        pooled_output = self.dropout(pooled_output)
        logits  = self.classifier(pooled_output)
        return logits

    def training_step(self, batch, batch_idx):
        labels = batch["labels"]
        outputs = self.forward(batch)
        loss = binary_cross_entropy_with_logits(outputs, labels)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        labels = batch["labels"]
        outputs = self.forward(batch)
        loss = binary_cross_entropy_with_logits(outputs, labels)
        self.log('val_loss', loss)
        
        scores = torch.sigmoid(outputs)

        preds = scores.round()
        matches = torch.tensor(preds == labels, dtype=torch.float)
        self.log('binary_accuracy', matches.mean())
        self.log('multilabel_accuracy', matches.min(axis=1).values.mean())
        return loss

    def test_step(self, batch, batch_idx):
        return self.validation_step(batch, batch_idx)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=LEARNING_RATE)
        return optimizer

n_labels = train_labels.shape[1]
lightning_transformer = LightningTransformer(transformer_model, config, n_labels)
trainer = pl.Trainer(gpus=1, max_epochs=3, progress_bar_refresh_rate=20)

trainer.fit(lightning_transformer, training_loader, testing_loader)