# -*- coding: utf-8 -*-
"""AG News Classification using DistilBERT Pytorch.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1JJhh5uAWznDkC_cQZmv4QTd4mik6DXt9

### 0. Install and load library
"""

#Restart kernel after installation: Runtime -> Restart runtime

#!pip install -U transformers sentencepiece datasets

from datasets import load_dataset
from datasets import DatasetDict
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification, pipeline
import torch
from tqdm import tqdm

"""### 1. Load dataset"""

datasets = load_dataset("ag_news")

"""### 2. Split dataset"""

datasets.keys() #dict_keys(['train', 'test'])

# split original train set into train/val with 8/2 ratio
# gather together to process splits at once
train_valid = datasets['train'].train_test_split(test_size=.2)
train_valid_test_datasets = DatasetDict({
    'train': train_valid['train'],
    'valid': train_valid['test'],
    'test': datasets['test']
})

# ['World', 'Sports', 'Business', 'Sci/Tech']
dataset_labels = datasets['train'].features['label'].names
id2label = {id: label for id, label in enumerate(dataset_labels)}
label2id = {label: id for id, label in id2label.items()}

"""### 3. Load model and tokenizer with correct configs """

model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=len(dataset_labels), label2id=label2id)
tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased")

"""### 4. Tokenize and convert dataset to DataLoader"""

encoded_datasets = train_valid_test_datasets.map(lambda examples: tokenizer(examples['text'], truncation=True, padding='max_length'), batched=True)
encoded_datasets = encoded_datasets.map(lambda examples: {'labels': examples['label']}, batched=True)
encoded_datasets.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'], device='cuda')

train_dataloader = torch.utils.data.DataLoader(encoded_datasets['train'], batch_size=8)
val_dataloader = torch.utils.data.DataLoader(encoded_datasets['valid'], batch_size=8)
test_dataloader = torch.utils.data.DataLoader(encoded_datasets['test'], batch_size=8)

"""### 5. Prepare training arguments"""

device = 'cuda' if torch.cuda.is_available() else 'cpu'
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9, verbose=True) # unnecessary, for fun only
epochs = 3

model.to(device)

"""### 6. Training"""

print('*'*50)

for epoch in range(epochs):
    model.train()
     
    training_loss = 0.0 
    correct_pred = {label: 0 for label in dataset_labels}
    total_pred = {label: 0 for label in dataset_labels}

    for i, batch in enumerate(train_dataloader):
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        training_loss += loss.item()
        if i%2000 == 1999:
            print(f"\n[Epoch-{epoch+1} {i+1}] {training_loss/2000}")
            training_loss = 0.0
    
    scheduler.step()
    
    # validation
    model.eval()
    with torch.no_grad():
        for batch in val_dataloader:
            batch = {k: v.to(device) for k,v in batch.items()}
            outputs = model(**batch)
            logits = outputs.logits
            _, predictions = torch.max(logits, 1)
            labels = batch['labels']

            # collect info for each class
            for label, prediction in zip(labels, predictions):
                if label == prediction:
                    correct_pred[id2label[label.item()]] += 1
                total_pred[id2label[label.item()]] += 1
    
    for label, correct_count in correct_pred.items():
        acc = 100 * float(correct_count) / total_pred[label]
        print(f"\nAccuracy for class {label} is {acc:.3f}")

"""### 7. Predicts on Test set"""

device = 'cpu'
model.to(device)

correct_pred = {label: 0 for label in dataset_labels}
total_pred = {label: 0 for label in dataset_labels}

model.eval()
with torch.no_grad():
    for batch in tqdm(test_dataloader, position=0, leave=True):
        batch = {k: v.to(device) for k,v in batch.items()}
        outputs = model(**batch)
        logits = outputs.logits
        _, predictions = torch.max(logits, 1)
        labels = batch['labels']

        # collect info for each class
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[id2label[label.item()]] += 1
            total_pred[id2label[label.item()]] += 1

for label, correct_count in correct_pred.items():
    acc = 100 * float(correct_count) / total_pred[label]
    print(f"\nAccuracy for class {label} is {acc:.3f}")

"""### 8. Inference using pipeline"""

ag_news_pipeline = pipeline('text-classification', model=model, tokenizer=tokenizer)

"""
datasets['test'][0]

>>> {
    'label': 2,
    'text': "Fears for T N pension after talks Unions representing workers at Turner   Newall say they are 'disappointed' after talks with stricken parent firm Federal Mogul."
    }

"""

test_input = datasets['test'][0]['text']
outputs = ag_news_pipeline(test_input)
print(outputs)

"""
>>> [{'label': 'LABEL_2', 'score': 0.9857332110404968}]
"""

