# -*- coding: utf-8 -*-
"""IMDb Review Classification using CNN Pytorch.ipynb

Automatically generated by Colaboratory.

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

import torch
import torch.nn as nn
from torchtext.datasets import IMDB
from torchtext.data.utils import get_tokenizer
from torchtext.data.functional import to_map_style_dataset
from torch.utils.data.dataset import random_split
from torch.utils.data import DataLoader
from torchtext.vocab import build_vocab_from_iterator
import numpy as np
import string
import re

"""#### Download Dataset"""

train_iter, test_iter = IMDB()

"""#### Dataset overview"""

len(train_iter), len(test_iter)

next(train_iter)

"""#### Build vocab"""

tokenizer = get_tokenizer("basic_english")

def yield_tokens(data_iter):
    for _, text in data_iter:
        yield tokenizer(text)


vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<unk>"])
vocab.set_default_index(vocab["<unk>"])

print(f"Number of vocab: {len(vocab)}")
print(f"Idx for unknown word: {vocab.get_default_index()}")

vocab["<pad>"]

"""#### Reload iterable dataset and build dataloader"""

EMBEDDING_SIZE = 300
MAX_SEQ_LEN = 128
VOCAB_SIZE = len(vocab)
BATCH_SIZE = 32

def encode(text):
    tokenized_text = tokenizer(text)
    vectorized_text = vocab(tokenized_text)
    sliced_text = vectorized_text[:MAX_SEQ_LEN]
    pad_size = MAX_SEQ_LEN - len(sliced_text)
    sliced_text += [vocab.get_default_index()] * pad_size
    return sliced_text

label_pipeline = lambda label: 1 if label == "pos" else 0

def collate_batch(batch):
    label_list, text_list = [], []

    for _label, _text in batch:
        label_list.append(label_pipeline(_label))
        text_list.append(encode(_text))

    text_tensor = torch.tensor(text_list, dtype=torch.int64)
    label_tensor = torch.tensor(label_list, dtype=torch.long)

    return label_tensor, text_tensor

train_iter, test_iter = IMDB()

train_dataset = to_map_style_dataset(train_iter)
test_dataset = to_map_style_dataset(test_iter)

num_train = int(len(train_dataset) * 0.95)
num_val = len(train_dataset) - num_train

split_train, split_val = random_split(train_dataset, [num_train, num_val])

train_dataloader = DataLoader(split_train, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
val_dataloader = DataLoader(split_val, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
test_dataloder = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)

"""#### Build Model"""

class CNNText(torch.nn.Module):
    def __init__(self, vocab_size, emb_size):
        super(CNNText, self).__init__()

        self.emb = nn.Embedding(num_embeddings=vocab_size,
                                embedding_dim=emb_size,
                                padding_idx=0)
        self.conv1 = nn.Conv1d(in_channels=emb_size,
                               out_channels=128,
                               kernel_size=3)
        self.conv2 = nn.Conv1d(in_channels=128,
                               out_channels=256,
                               kernel_size=6)
        self.globalmaxpooling = nn.AdaptiveMaxPool1d(1)
        self.classifier = nn.Linear(256, 2)
        self.dropout = nn.Dropout(0.5)
        self.relu = nn.ReLU()

    def forward(self, x):
        # original: bs, seq_len

        # bs, seq_len, emb_dim
        x = self.emb(x)

        # bs, emb_dim, seq_len
        x = x.permute(0, 2, 1)
        x = self.conv1(x)
        x = self.relu(x)

        # bs, emb_dim, seq_len
        x = self.conv2(x)
        x = self.relu(x)

        # output shape: bs, emb_dim, 1 
        x = self.globalmaxpooling(x)
        x = self.dropout(x)

        # output shape: bs, emb_dim
        x = x.squeeze(dim=2)
        x = self.classifier(x)

        return x

model = CNNText(VOCAB_SIZE, EMBEDDING_SIZE)

"""#### Training"""

EPOCHS = 3
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.CrossEntropyLoss()

def train(dataloader):
    model.train()
    training_count = 0
    training_correct = 0

    for idx, (label, text) in enumerate(dataloader):
        optimizer.zero_grad()

        predictions = model(text)
        loss = criterion(predictions, label)

        loss.backward()
        optimizer.step()
        
        training_correct += (predictions.argmax(1) == label).sum().item()
        training_count += label.size(0)

        if (idx % 50 == 0) and (idx > 0):
            print(f"Epoch: {epoch} on {idx} batches with {training_correct/training_count:.3f} accuracy")

def evaluate(dataloader):
    model.eval()
    total_count = 0
    total_correct = 0

    with torch.no_grad():
        for idx, (label, text) in enumerate(dataloader):
            predictions = model(text)
            total_correct += (predictions.argmax(1) == label).sum().item()
            total_count += label.size(0)

    return total_correct/total_count

for epoch in range(1, EPOCHS+1):
    train(train_dataloader)
    val_acc = evaluate(val_dataloader)

    print(f"Epoch: {epoch} with {val_acc:.3f} validation accuracy")

"""#### Testing"""

test_acc = evaluate(test_dataloder)
print(f"Test acc: {test_acc:.3f}")

"""#### Inferencing"""

test_sentence = "Terrible movie. Bad Bad I hate this so much."
test_input = torch.tensor(encode(test_sentence), dtype=torch.int64).unsqueeze(dim=0)

with torch.no_grad():
    predictions = model(test_input).argmax(1)

print(f"Text: {test_sentence}")
print(f"Predictions: {predictions} ({'Positive' if predictions == 1 else 'Negative'})")

test_sentence = "BEST movie. I love this so much."
test_input = torch.tensor(encode(test_sentence), dtype=torch.int64).unsqueeze(dim=0)

with torch.no_grad():
    predictions = model(test_input).argmax(1)

print(f"Text: {test_sentence}")
print(f"Predictions: {predictions} ({'Positive' if predictions == 1 else 'Negative'})")