# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="mgZTKjbZ6kfz"
# <div align="center">
# <h1><img width="30" src="https://madewithml.com/static/images/rounded_logo.png">&nbsp;<a href="https://madewithml.com/">Made With ML</a></h1>
# Applied ML · MLOps · Production
# <br>
# Join 30K+ developers in learning how to responsibly <a href="https://madewithml.com/about/">deliver value</a> with ML.
#     <br>
# </div>
#
# <br>
#
# <div align="center">
#     <a target="_blank" href="https://newsletter.madewithml.com"><img src="https://img.shields.io/badge/Subscribe-30K-brightgreen"></a>&nbsp;
#     <a target="_blank" href="https://github.com/GokuMohandas/MadeWithML"><img src="https://img.shields.io/github/stars/GokuMohandas/MadeWithML.svg?style=social&label=Star"></a>&nbsp;
#     <a target="_blank" href="https://www.linkedin.com/in/goku"><img src="https://img.shields.io/badge/style--5eba00.svg?label=LinkedIn&logo=linkedin&style=social"></a>&nbsp;
#     <a target="_blank" href="https://twitter.com/GokuMohandas"><img src="https://img.shields.io/twitter/follow/GokuMohandas.svg?label=Follow&style=social"></a>
#     <br>
#     🔥&nbsp; Among the <a href="https://github.com/topics/deep-learning" target="_blank">top ML</a> repositories on GitHub
# </div>
#
# <br>
# <hr>

# + [markdown] id="eTdCMVl9YAXw"
# # Attention
#
# In this lesson we will learn how to incorporate attention mechanisms to create more context-aware representations.

# + [markdown] id="xuabAj4PYj57"
# <div align="left">
# <a target="_blank" href="https://madewithml.com/courses/foundations/attention/"><img src="https://img.shields.io/badge/📖 Read-blog post-9cf"></a>&nbsp;
# <a href="https://github.com/GokuMohandas/MadeWithML/blob/main/notebooks/14_Attention.ipynb" role="button"><img src="https://img.shields.io/static/v1?label=&amp;message=View%20On%20GitHub&amp;color=586069&amp;logo=github&amp;labelColor=2f363d"></a>&nbsp;
# <a href="https://colab.research.google.com/github/GokuMohandas/MadeWithML/blob/main/notebooks/14_Attention.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
# </div>

# + [markdown] id="m5JJY0tl8StF"
# # Overview

# + [markdown] id="wDYX2d_eySKj"
# In the <a target="_blank" href="https://madewithml.com/courses/foundations/recurrent-neural-networks/">RNN lesson</a>, we were constrained to using the representation at the very end but what if we could give contextual weight to each encoded input ($h_i$) when making our prediction? This is also preferred because it can help mitigate the vanishing gradient issue which stems from processing very long sequences.
#
# Below is attention applied to the outputs from an RNN. In theory, the outputs can come from anywhere where we want to learn how to weight amongst them but since we're working with the context of an RNN from the previous lesson , we'll continue with that.

# + [markdown] id="34Yf1pdURtLn"
# <div align="left">
# <img src="https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/images/foundations/attention/attention.png" width="500">
# </div>
#
# $ \alpha = softmax(W_{attn}h) $
#
# $c_t = \sum_{i=1}^{n} \alpha_{t,i}h_i $
#
#
# *where*:
# * $ h $ = RNN outputs (or any group of outputs you want to attend to) $\in \mathbb{R}^{NXMXH}$ ($N$ is the batch size, $M$ is the max sequence length in the batch, $H$ is the hidden dim)
# * $ \alpha_{t,i} $ = alignment function for output $ y_t $ using input $ h_i $ (we also concatenate other useful representations with $h_i$ here). In our case, this would be the attention value to attribute to each input $h_i$. 
# * $W_{attn}$ = attention weights to learn $\in \mathbb{R}^{HX1}$. We can also apply activations functions, transformations, etc. here
# * $c_t$ = context vector that accounts for the different inputs with attention. We can pass this context vector to downstream processes.

# + [markdown] id="JqxyljU18hvt"
# * **Objective:**  At it's core, attention is about learning how to weigh a group of encoded representations to produce a context-aware representation to use for downstream tasks. This is done by learning a set of attention weights and then using softmax to create attention values that sum to 1.
# * **Advantages:** 
#     * Learn how to account for the appropriate encoded representations regardless of position.
# * **Disadvantages:** 
#     * Another compute step that involves learning weights.
# * **Miscellaneous:** 
#     * Several state-of-the-art approaches extend on basic attention to deliver highly context-aware representations (ex. self-attention). 

# + [markdown] id="V88LZeTqDePU"
# # Set up

# + id="DlRv80v5DfdP"
import numpy as np
import pandas as pd
import random
import torch
import torch.nn as nn

# + id="NoNhL0wTDgLW"
SEED = 1234


# + id="gwVLi0UPDgOd"
def set_seeds(seed=1234):
    """Set seeds for reproducibility."""
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # multi-GPU# Set seeds for reproducibility
set_seeds(seed=SEED)

# + id="HK5Onfz9DiPO"
# Set seeds for reproducibility
set_seeds(seed=SEED)

# + colab={"base_uri": "https://localhost:8080/"} id="4T1rZDhsDiRt" outputId="372e1317-1ed0-46c5-c136-553f1064f31f"
# Set device
cuda = True
device = torch.device('cuda' if (
    torch.cuda.is_available() and cuda) else 'cpu')
torch.set_default_tensor_type('torch.FloatTensor')
if device.type == 'cuda':
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
print (device)

# + [markdown] id="c69z9wpJ56nE"
# ## Load data

# + [markdown] id="2V_nEp5G58M0"
# We will download the [AG News dataset](http://www.di.unipi.it/~gulli/AG_corpus_of_news_articles.html), which consists of 120K text samples from 4 unique classes (`Business`, `Sci/Tech`, `Sports`, `World`)

# + id="y3qKSoEe57na"
import numpy as np
import pandas as pd
import re
import urllib

# + id="BqaHCtToXWAb" colab={"base_uri": "https://localhost:8080/", "height": 204} outputId="f87fb125-c1d6-4250-c863-16bf5397e73b"
# Load data
url = "https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/datasets/news.csv"
df = pd.read_csv(url, header=0) # load
df = df.sample(frac=1).reset_index(drop=True) # shuffle
df.head()

# + [markdown] id="HR_lLiVBDxW1"
# ## Preprocessing

# + [markdown] id="u00AZ0O8DxaY"
# We're going to clean up our input data first by doing operations such as lower text, removing stop (filler) words, filters using regular expressions, etc.

# + id="Jq6a11YsDvEU"
import nltk
from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
import re

# + colab={"base_uri": "https://localhost:8080/"} id="x5IpVaKqDvG1" outputId="2f59b82e-d839-441c-86e7-4243c6db6ba4"
nltk.download('stopwords')
STOPWORDS = stopwords.words('english')
print (STOPWORDS[:5])
porter = PorterStemmer()


# + id="1n4w3KLQDvJk"
def preprocess(text, stopwords=STOPWORDS):
    """Conditional preprocessing on our text unique to our task."""
    # Lower
    text = text.lower()

    # Remove stopwords
    pattern = re.compile(r'\b(' + r'|'.join(stopwords) + r')\b\s*')
    text = pattern.sub('', text)

    # Remove words in paranthesis
    text = re.sub(r'\([^)]*\)', '', text)

    # Spacing and filters
    text = re.sub(r"([-;;.,!?<=>])", r" \1 ", text)
    text = re.sub('[^A-Za-z0-9]+', ' ', text) # remove non alphanumeric chars
    text = re.sub(' +', ' ', text)  # remove multiple spaces
    text = text.strip()

    return text


# + colab={"base_uri": "https://localhost:8080/", "height": 35} id="U9JmpiCmDvMz" outputId="1edc59cb-7f4e-43ea-c008-ded886588ac0"
# Sample
text = "Great week for the NYSE!"
preprocess(text=text)

# + colab={"base_uri": "https://localhost:8080/"} id="8RZ45-N9DvQD" outputId="4f03b82f-bf7e-480b-9dfd-1a27c1bbc0b3"
# Apply to dataframe
preprocessed_df = df.copy()
preprocessed_df.title = preprocessed_df.title.apply(preprocess)
print (f"{df.title.values[0]}\n\n{preprocessed_df.title.values[0]}")

# + [markdown] id="hxo6RKCQ71dl"
# ## Split data

# + id="eS6kCcfY6IHE"
import collections
from sklearn.model_selection import train_test_split

# + id="4qkWnO3uYPnk"
TRAIN_SIZE = 0.7
VAL_SIZE = 0.15
TEST_SIZE = 0.15


# + id="0gbh2TLBXVDu"
def train_val_test_split(X, y, train_size):
    """Split dataset into data splits."""
    X_train, X_, y_train, y_ = train_test_split(X, y, train_size=TRAIN_SIZE, stratify=y)
    X_val, X_test, y_val, y_test = train_test_split(X_, y_, train_size=0.5, stratify=y_)
    return X_train, X_val, X_test, y_train, y_val, y_test


# + id="kqiQd2j_76gP"
# Data
X = preprocessed_df["title"].values
y = preprocessed_df["category"].values

# + colab={"base_uri": "https://localhost:8080/"} id="AwdPM136EBMA" outputId="9abe0d25-ed9d-49cb-a6d3-629b34be0ea4"
# Create data splits
X_train, X_val, X_test, y_train, y_val, y_test = train_val_test_split(
    X=X, y=y, train_size=TRAIN_SIZE)
print (f"X_train: {X_train.shape}, y_train: {y_train.shape}")
print (f"X_val: {X_val.shape}, y_val: {y_val.shape}")
print (f"X_test: {X_test.shape}, y_test: {y_test.shape}")
print (f"Sample point: {X_train[0]} → {y_train[0]}")

# + [markdown] id="kUAn5IOjEF8P"
# ## LabelEncoder

# + [markdown] id="fNTjndI0EKzm"
# Next we'll define a `LabelEncoder` to encode our text labels into unique indices

# + id="OgLNA5yTELB1"
import itertools


# + id="GeoPR77QELJT"
class LabelEncoder(object):
    """Label encoder for tag labels."""
    def __init__(self, class_to_index={}):
        self.class_to_index = class_to_index
        self.index_to_class = {v: k for k, v in self.class_to_index.items()}
        self.classes = list(self.class_to_index.keys())

    def __len__(self):
        return len(self.class_to_index)

    def __str__(self):
        return f"<LabelEncoder(num_classes={len(self)})>"

    def fit(self, y):
        classes = np.unique(y)
        for i, class_ in enumerate(classes):
            self.class_to_index[class_] = i
        self.index_to_class = {v: k for k, v in self.class_to_index.items()}
        self.classes = list(self.class_to_index.keys())
        return self

    def encode(self, y):
        encoded = np.zeros((len(y)), dtype=int)
        for i, item in enumerate(y):
            encoded[i] = self.class_to_index[item]
        return encoded

    def decode(self, y):
        classes = []
        for i, item in enumerate(y):
            classes.append(self.index_to_class[item])
        return classes

    def save(self, fp):
        with open(fp, 'w') as fp:
            contents = {'class_to_index': self.class_to_index}
            json.dump(contents, fp, indent=4, sort_keys=False)

    @classmethod
    def load(cls, fp):
        with open(fp, 'r') as fp:
            kwargs = json.load(fp=fp)
        return cls(**kwargs)


# + colab={"base_uri": "https://localhost:8080/"} id="0MtTxMB5ELL4" outputId="bd92474e-c8d9-4f85-9449-13fd912f5e37"
# Encode
label_encoder = LabelEncoder()
label_encoder.fit(y_train)
NUM_CLASSES = len(label_encoder)
label_encoder.class_to_index

# + colab={"base_uri": "https://localhost:8080/"} id="0IzPFnoDELSr" outputId="a3ee0554-b8c4-455d-8084-9b6fddd6f616"
# Convert labels to tokens
print (f"y_train[0]: {y_train[0]}")
y_train = label_encoder.encode(y_train)
y_val = label_encoder.encode(y_val)
y_test = label_encoder.encode(y_test)
print (f"y_train[0]: {y_train[0]}")

# + colab={"base_uri": "https://localhost:8080/"} id="xjESU2TOEQz_" outputId="1c87bf2d-03f3-40c1-8d1c-06392c9f5019"
# Class weights
counts = np.bincount(y_train)
class_weights = {i: 1.0/count for i, count in enumerate(counts)}
print (f"counts: {counts}\nweights: {class_weights}")

# + [markdown] id="dIfmW7vJ8Jx1"
# ## Tokenizer

# + [markdown] id="IYREXERdETdQ"
# We'll define a `Tokenizer` to convert our text input data into token indices.

# + id="0V6v1yeVGxNg"
import json
from collections import Counter
from more_itertools import take


# + id="ChhTN8WDY4Z5"
class Tokenizer(object):
    def __init__(self, char_level, num_tokens=None, 
                 pad_token='<PAD>', oov_token='<UNK>',
                 token_to_index=None):
        self.char_level = char_level
        self.separator = '' if self.char_level else ' '
        if num_tokens: num_tokens -= 2 # pad + unk tokens
        self.num_tokens = num_tokens
        self.pad_token = pad_token
        self.oov_token = oov_token
        if not token_to_index:
            token_to_index = {pad_token: 0, oov_token: 1}
        self.token_to_index = token_to_index
        self.index_to_token = {v: k for k, v in self.token_to_index.items()}

    def __len__(self):
        return len(self.token_to_index)

    def __str__(self):
        return f"<Tokenizer(num_tokens={len(self)})>"

    def fit_on_texts(self, texts):
        if not self.char_level:
            texts = [text.split(" ") for text in texts]
        all_tokens = [token for text in texts for token in text]
        counts = Counter(all_tokens).most_common(self.num_tokens)
        self.min_token_freq = counts[-1][1]
        for token, count in counts:
            index = len(self)
            self.token_to_index[token] = index
            self.index_to_token[index] = token
        return self

    def texts_to_sequences(self, texts):
        sequences = []
        for text in texts:
            if not self.char_level:
                text = text.split(' ')
            sequence = []
            for token in text:
                sequence.append(self.token_to_index.get(
                    token, self.token_to_index[self.oov_token]))
            sequences.append(np.asarray(sequence))
        return sequences

    def sequences_to_texts(self, sequences):
        texts = []
        for sequence in sequences:
            text = []
            for index in sequence:
                text.append(self.index_to_token.get(index, self.oov_token))
            texts.append(self.separator.join([token for token in text]))
        return texts

    def save(self, fp):
        with open(fp, 'w') as fp:
            contents = {
                'char_level': self.char_level,
                'oov_token': self.oov_token,
                'token_to_index': self.token_to_index
            }
            json.dump(contents, fp, indent=4, sort_keys=False)

    @classmethod
    def load(cls, fp):
        with open(fp, 'r') as fp:
            kwargs = json.load(fp=fp)
        return cls(**kwargs)


# + id="ylpHSbcZYsXa" colab={"base_uri": "https://localhost:8080/"} outputId="5f3c6bef-a08a-42c1-b14f-6af7995cc963"
# Tokenize
tokenizer = Tokenizer(char_level=False, num_tokens=5000)
tokenizer.fit_on_texts(texts=X_train)
VOCAB_SIZE = len(tokenizer)
print (tokenizer)

# + id="bVgBSnFTfE5D" colab={"base_uri": "https://localhost:8080/"} outputId="f97eb536-3e81-42cd-f1b6-0167cb62f809"
# Sample of tokens
print (take(5, tokenizer.token_to_index.items()))
print (f"least freq token's freq: {tokenizer.min_token_freq}") # use this to adjust num_tokens

# + id="WcscM_vL8KvP" colab={"base_uri": "https://localhost:8080/"} outputId="b2eeb74d-d290-4306-fc4d-a1a5bc5e6716"
# Convert texts to sequences of indices
X_train = tokenizer.texts_to_sequences(X_train)
X_val = tokenizer.texts_to_sequences(X_val)
X_test = tokenizer.texts_to_sequences(X_test)
preprocessed_text = tokenizer.sequences_to_texts([X_train[0]])[0]
print ("Text to indices:\n"
    f"  (preprocessed) → {preprocessed_text}\n"
    f"  (tokenized) → {X_train[0]}")


# + [markdown] id="9o91b6kKHjN4"
# ## Padding

# + [markdown] id="bD6CuPxOHlsw"
# We'll need to do 2D padding to our tokenized text.

# + id="dcNEooUHHkwn"
def pad_sequences(sequences, max_seq_len=0):
    """Pad sequences to max length in sequence."""
    max_seq_len = max(max_seq_len, max(len(sequence) for sequence in sequences))
    padded_sequences = np.zeros((len(sequences), max_seq_len))
    for i, sequence in enumerate(sequences):
        padded_sequences[i][:len(sequence)] = sequence
    return padded_sequences


# + colab={"base_uri": "https://localhost:8080/"} id="sfWHamHZHtT-" outputId="1f079464-4a50-4916-bb41-e65affd2a944"
# 2D sequences
padded = pad_sequences(X_train[0:3])
print (padded.shape)
print (padded)


# + [markdown] id="eoWQk0hO9bK2"
# ## Datasets

# + [markdown] id="dSPbGwOkHvfj"
# We're going to create Datasets and DataLoaders to be able to efficiently create batches with our data splits.

# + id="aSdfc0-cb9fc"
class Dataset(torch.utils.data.Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __len__(self):
        return len(self.y)

    def __str__(self):
        return f"<Dataset(N={len(self)})>"

    def __getitem__(self, index):
        X = self.X[index]
        y = self.y[index]
        return [X, len(X), y]

    def collate_fn(self, batch):
        """Processing on a batch."""
        # Get inputs
        batch = np.array(batch, dtype=object)
        X = batch[:, 0]
        seq_lens = batch[:, 1]
        y = np.stack(batch[:, 2], axis=0)

        # Pad inputs
        X = pad_sequences(sequences=X)

        # Cast
        X = torch.LongTensor(X.astype(np.int32))
        seq_lens = torch.LongTensor(seq_lens.astype(np.int32))
        y = torch.LongTensor(y.astype(np.int32))

        return X, seq_lens, y

    def create_dataloader(self, batch_size, shuffle=False, drop_last=False):
        return torch.utils.data.DataLoader(
            dataset=self, batch_size=batch_size, collate_fn=self.collate_fn,
            shuffle=shuffle, drop_last=drop_last, pin_memory=True)


# + colab={"base_uri": "https://localhost:8080/"} id="KYZtzbWMIXY7" outputId="b85875dc-b6a8-41c9-a327-aa42662ac34b"
# Create datasets
train_dataset = Dataset(X=X_train, y=y_train)
val_dataset = Dataset(X=X_val, y=y_val)
test_dataset = Dataset(X=X_test, y=y_test)
print ("Datasets:\n"
    f"  Train dataset:{train_dataset.__str__()}\n"
    f"  Val dataset: {val_dataset.__str__()}\n"
    f"  Test dataset: {test_dataset.__str__()}\n"
    "Sample point:\n"
    f"  X: {train_dataset[0][0]}\n"
    f"  seq_len: {train_dataset[0][1]}\n"
    f"  y: {train_dataset[0][2]}")

# + id="1w6wVKJe9fxk" colab={"base_uri": "https://localhost:8080/"} outputId="9b79b3ec-6a35-4afc-e071-863bd57e7996"
# Create dataloaders
batch_size = 64
train_dataloader = train_dataset.create_dataloader(
    batch_size=batch_size)
val_dataloader = val_dataset.create_dataloader(
    batch_size=batch_size)
test_dataloader = test_dataset.create_dataloader(
    batch_size=batch_size)
batch_X, batch_seq_lens, batch_y = next(iter(train_dataloader))
print ("Sample batch:\n"
    f"  X: {list(batch_X.size())}\n"
    f"  seq_lens: {list(batch_seq_lens.size())}\n"
    f"  y: {list(batch_y.size())}\n"
    "Sample point:\n"
    f"  X: {batch_X[0]}\n"
    f" seq_len: {batch_seq_lens[0]}\n"
    f"  y: {batch_y[0]}")


# + [markdown] id="dB-g9F1TJOPx"
# ## Trainer

# + [markdown] id="dL5FcX0hJR37"
# Let's create the `Trainer` class that we'll use to facilitate training for our experiments.

# + id="k5KWF-DXJN2a"
class Trainer(object):
    def __init__(self, model, device, loss_fn=None, optimizer=None, scheduler=None):

        # Set params
        self.model = model
        self.device = device
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.scheduler = scheduler

    def train_step(self, dataloader):
        """Train step."""
        # Set model to train mode
        self.model.train()
        loss = 0.0

        # Iterate over train batches
        for i, batch in enumerate(dataloader):

            # Step
            batch = [item.to(self.device) for item in batch]  # Set device
            inputs, targets = batch[:-1], batch[-1]
            self.optimizer.zero_grad()  # Reset gradients
            z = self.model(inputs)  # Forward pass
            J = self.loss_fn(z, targets)  # Define loss
            J.backward()  # Backward pass
            self.optimizer.step()  # Update weights

            # Cumulative Metrics
            loss += (J.detach().item() - loss) / (i + 1)

        return loss

    def eval_step(self, dataloader):
        """Validation or test step."""
        # Set model to eval mode
        self.model.eval()
        loss = 0.0
        y_trues, y_probs = [], []

        # Iterate over val batches
        with torch.no_grad():
            for i, batch in enumerate(dataloader):

                # Step
                batch = [item.to(self.device) for item in batch]  # Set device
                inputs, y_true = batch[:-1], batch[-1]
                z = self.model(inputs)  # Forward pass
                J = self.loss_fn(z, y_true).item()

                # Cumulative Metrics
                loss += (J - loss) / (i + 1)

                # Store outputs
                y_prob = torch.sigmoid(z).cpu().numpy()
                y_probs.extend(y_prob)
                y_trues.extend(y_true.cpu().numpy())

        return loss, np.vstack(y_trues), np.vstack(y_probs)

    def predict_step(self, dataloader):
        """Prediction step."""
        # Set model to eval mode
        self.model.eval()
        y_probs = []

        # Iterate over val batches
        with torch.no_grad():
            for i, batch in enumerate(dataloader):

                # Forward pass w/ inputs
                inputs, targets = batch[:-1], batch[-1]
                y_prob = self.model(inputs, apply_softmax=True)

                # Store outputs
                y_probs.extend(y_prob)

        return np.vstack(y_probs)
    
    def train(self, num_epochs, patience, train_dataloader, val_dataloader):
        best_val_loss = np.inf
        for epoch in range(num_epochs):
            # Steps
            train_loss = self.train_step(dataloader=train_dataloader)
            val_loss, _, _ = self.eval_step(dataloader=val_dataloader)
            self.scheduler.step(val_loss)

            # Early stopping
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                best_model = self.model
                _patience = patience  # reset _patience
            else:
                _patience -= 1
            if not _patience:  # 0
                print("Stopping early!")
                break

            # Logging
            print(
                f"Epoch: {epoch+1} | "
                f"train_loss: {train_loss:.5f}, "
                f"val_loss: {val_loss:.5f}, "
                f"lr: {self.optimizer.param_groups[0]['lr']:.2E}, "
                f"_patience: {_patience}"
            )
        return best_model


# + [markdown] id="bUMibnFtuq_i"
# # Attention

# + [markdown] id="hO-KQgMF1jEs"
# Attention applied to the outputs from an RNN. In theory, the outputs can come from anywhere where we want to learn how to weight amongst them but since we're working with the context of an RNN from the previous lesson , we'll continue with that.

# + [markdown] id="KSPxJXxZu4Jo"
# <div align="left">
# <img src="https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/images/foundations/attention/attention.png" width="500">
# </div>
#
# $ \alpha = softmax(W_{attn}h) $
#
# $c_t = \sum_{i=1}^{n} \alpha_{t,i}h_i $
#
#
# *where*:
# * $ h $ = RNN outputs (or any group of outputs you want to attend to) $\in \mathbb{R}^{NXMXH}$ ($N$ is the batch size, $M$ is the max length of each sequence in the batch, $H$ is the hidden dim)
# * $ \alpha_{t,i} $ = alignment function for output $ y_t $ using input $ h_i $ (we also concatenate other useful representations with $h_i$ here). In our case, this would be the attention value to attribute to each input $h_i$. 
# * $W_{attn}$ = attention weights to learn $\in \mathbb{R}^{HX1}$. We can also apply activations functions, transformations, etc. here
# * $c_t$ = context vector that accounts for the different inputs with attention. We can pass this context vector to downstream processes.

# + id="NeGSvWrl73Nz"
import torch.nn.functional as F

# + [markdown] id="AUf0SD2N3IUy"
# The RNN will create an encoded representation for each word in our input resulting in a stacked vector that has dimensions $NXMXH$, where N is the # of samples in the batch, M is the max sequence length in the batch, and H is the number of hidden units in the RNN. 

# + id="MZovWtFy3FuC"
BATCH_SIZE = 64
SEQ_LEN = 8
EMBEDDING_DIM = 100
RNN_HIDDEN_DIM = 128

# + id="dFK6L0nj5H9s"
# Embed
x = torch.rand((BATCH_SIZE, SEQ_LEN, EMBEDDING_DIM))

# + colab={"base_uri": "https://localhost:8080/"} id="PWoF4mU-3FxC" outputId="fcf27520-e380-4c6f-9a70-adf742d85743"
# Encode
rnn = nn.RNN(EMBEDDING_DIM, RNN_HIDDEN_DIM, batch_first=True)
out, h_n = rnn(x) # h_n is the last hidden state
print ("out: ", out.shape)
print ("h_n: ", h_n.shape)

# + colab={"base_uri": "https://localhost:8080/"} id="RYMes2oj3Fzn" outputId="cab8a9cb-7a95-4583-eece-5da74e8d110c"
# Attend
attn = nn.Linear(RNN_HIDDEN_DIM, 1)
e = attn(out)
attn_vals = F.softmax(e.squeeze(2), dim=1)
c = torch.bmm(attn_vals.unsqueeze(1), out).squeeze(1)
print ("e: ", e.shape)
print ("attn_vals: ", attn_vals.shape)
print ("attn_vals[0]: ", attn_vals[0])
print ("sum(attn_vals[0]): ", sum(attn_vals[0]))
print ("c: ", c.shape)

# + colab={"base_uri": "https://localhost:8080/"} id="kdAbPbqM3F20" outputId="da0d39c2-ce21-47b2-ddf0-f7de56945946"
# Predict
fc1 = nn.Linear(RNN_HIDDEN_DIM, NUM_CLASSES)
output = F.softmax(fc1(c), dim=1)
print ("output: ", output.shape)

# + [markdown] id="7U_s9SSY3xmt"
# > In a many-to-many task such as machine translation, our attentional interface will also account for the encoded representation of token in the output as well (via concatenation) so we can know which encoded inputs to attend to based on the encoded output we're focusing on. For more on this, be sure to explore <a target="_blank" href="https://arxiv.org/abs/1409.0473">Bahdanau's attention paper</a>.

# + [markdown] id="pfhjWZRD94hK"
# ## Model

# + [markdown] id="W1mtPyJ725-t"
# Now let's create our RNN based model but with the addition of the attention layer on top of the RNN's outputs.

# + id="qoBC0u8XpbVv"
RNN_HIDDEN_DIM = 128
DROPOUT_P = 0.1
HIDDEN_DIM = 100


# + id="UPP5ROd69mXC"
class RNN(nn.Module):
    def __init__(self, embedding_dim, vocab_size, rnn_hidden_dim,
                 hidden_dim, dropout_p, num_classes, padding_idx=0):
        super(RNN, self).__init__()
        
        # Initialize embeddings
        self.embeddings = nn.Embedding(
            embedding_dim=embedding_dim, num_embeddings=vocab_size,
            padding_idx=padding_idx)
        
        # RNN
        self.rnn = nn.RNN(embedding_dim, rnn_hidden_dim, batch_first=True)

        # Attention
        self.attn = nn.Linear(rnn_hidden_dim, 1)
     
        # FC weights
        self.dropout = nn.Dropout(dropout_p)
        self.fc1 = nn.Linear(rnn_hidden_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, num_classes)

    def forward(self, inputs, apply_softmax=False):
        # Embed
        x_in, seq_lens = inputs
        x_in = self.embeddings(x_in)
            
        # Encode
        out, h_n = self.rnn(x_in)
        
        # Attend
        e = self.attn(out)
        attn_vals = F.softmax(e.squeeze(2), dim=1)
        c = torch.bmm(attn_vals.unsqueeze(1), out).squeeze(1)

        # Predict
        z = self.fc1(c)
        z = self.dropout(z)
        y_pred = self.fc2(z)
        
        if apply_softmax:
            y_pred = F.softmax(y_pred, dim=1)
        return y_pred


# + id="LH_z2FqUn5DR" colab={"base_uri": "https://localhost:8080/"} outputId="4153e4b7-8bfc-4906-a72e-719dde0bdce1"
# Simple RNN cell
model = RNN(
    embedding_dim=EMBEDDING_DIM, vocab_size=VOCAB_SIZE, 
    rnn_hidden_dim=RNN_HIDDEN_DIM, hidden_dim=HIDDEN_DIM, 
    dropout_p=DROPOUT_P, num_classes=NUM_CLASSES)
model = model.to(device) # set device
print (model.named_parameters)

# + [markdown] id="S9ggYO6yHIm2"
# ## Training

# + id="geKOPVzVK6S9"
from torch.optim import Adam

# + id="OfSBvenCiIce"
NUM_LAYERS = 1
LEARNING_RATE = 1e-4
PATIENCE = 10
NUM_EPOCHS = 50

# + id="4ilmlIMUzNIj"
# Define Loss
class_weights_tensor = torch.Tensor(list(class_weights.values())).to(device)
loss_fn = nn.CrossEntropyLoss(weight=class_weights_tensor)

# + id="Ucn3tYq1_sE1"
# Define optimizer & scheduler
optimizer = Adam(model.parameters(), lr=LEARNING_RATE) 
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.1, patience=3)

# + id="X5ac-uJXb-F7"
# Trainer module
trainer = Trainer(
    model=model, device=device, loss_fn=loss_fn, 
    optimizer=optimizer, scheduler=scheduler)

# + id="F-uKcZrABKVZ" colab={"base_uri": "https://localhost:8080/"} outputId="cad1f9ab-330b-44b1-83e4-63f739b0920a"
# Train
best_model = trainer.train(
    NUM_EPOCHS, PATIENCE, train_dataloader, val_dataloader)

# + [markdown] id="axytMjGQmz6T"
# ## Evaluation

# + id="-Ly3zjVAm4Wr"
import json
from sklearn.metrics import precision_recall_fscore_support


# + id="ojOfz9rjm4ZI"
def get_performance(y_true, y_pred, classes):
    """Per-class performance metrics."""
    # Performance
    performance = {"overall": {}, "class": {}}

    # Overall performance
    metrics = precision_recall_fscore_support(y_true, y_pred, average="weighted")
    performance["overall"]["precision"] = metrics[0]
    performance["overall"]["recall"] = metrics[1]
    performance["overall"]["f1"] = metrics[2]
    performance["overall"]["num_samples"] = np.float64(len(y_true))

    # Per-class performance
    metrics = precision_recall_fscore_support(y_true, y_pred, average=None)
    for i in range(len(classes)):
        performance["class"][classes[i]] = {
            "precision": metrics[0][i],
            "recall": metrics[1][i],
            "f1": metrics[2][i],
            "num_samples": np.float64(metrics[3][i]),
        }

    return performance


# + id="tho3e_yQzqZ1"
# Get predictions
test_loss, y_true, y_prob = trainer.eval_step(dataloader=test_dataloader)
y_pred = np.argmax(y_prob, axis=1)

# + id="fylHduxTK-0N" colab={"base_uri": "https://localhost:8080/"} outputId="5813ddb2-b772-42d9-94cc-69dffd4c9eed"
# Determine performance
performance = get_performance(
    y_true=y_test, y_pred=y_pred, classes=label_encoder.classes)
print (json.dumps(performance['overall'], indent=2))

# + id="1r9EvThuTLli"
# Save artifacts
from pathlib import Path
dir = Path("rnn")
dir.mkdir(parents=True, exist_ok=True)
label_encoder.save(fp=Path(dir, "label_encoder.json"))
tokenizer.save(fp=Path(dir, "tokenizer.json"))
torch.save(best_model.state_dict(), Path(dir, "model.pt"))
with open(Path(dir, "performance.json"), "w") as fp:
    json.dump(performance, indent=2, sort_keys=False, fp=fp)


# + [markdown] id="KcMSfbYwQKz2"
# ## Inference

# + id="N1CQ5nqNQM9R"
def get_probability_distribution(y_prob, classes):
    """Create a dict of class probabilities from an array."""
    results = {}
    for i, class_ in enumerate(classes):
        results[class_] = np.float64(y_prob[i])
    sorted_results = {k: v for k, v in sorted(
        results.items(), key=lambda item: item[1], reverse=True)}
    return sorted_results


# + colab={"base_uri": "https://localhost:8080/"} id="zUEYVoHLQNu_" outputId="07665209-2212-4806-d709-2dc959ca5cff"
# Load artifacts
device = torch.device("cpu")
label_encoder = LabelEncoder.load(fp=Path(dir, "label_encoder.json"))
tokenizer = Tokenizer.load(fp=Path(dir, "tokenizer.json"))
model = RNN(
    embedding_dim=EMBEDDING_DIM, vocab_size=VOCAB_SIZE, 
    rnn_hidden_dim=RNN_HIDDEN_DIM, hidden_dim=HIDDEN_DIM, 
    dropout_p=DROPOUT_P, num_classes=NUM_CLASSES)
model.load_state_dict(torch.load(Path(dir, "model.pt"), map_location=device))
model.to(device)

# + id="Vk3Jawc4QNxm"
# Initialize trainer
trainer = Trainer(model=model, device=device)

# + colab={"base_uri": "https://localhost:8080/"} id="n_PXGD8hQgYX" outputId="34648570-6ca6-49f5-c953-7258416255b2"
# Dataloader
text = "The final tennis tournament starts next week."
X = tokenizer.texts_to_sequences([preprocess(text)])
print (tokenizer.sequences_to_texts(X))
y_filler = label_encoder.encode([label_encoder.classes[0]]*len(X))
dataset = Dataset(X=X, y=y_filler)
dataloader = dataset.create_dataloader(batch_size=batch_size)

# + colab={"base_uri": "https://localhost:8080/"} id="fWcn1RNuQgaw" outputId="628daabc-538f-4b46-ddcc-157f845b6c78"
# Inference
y_prob = trainer.predict_step(dataloader)
y_pred = np.argmax(y_prob, axis=1)
label_encoder.decode(y_pred)

# + colab={"base_uri": "https://localhost:8080/"} id="HdlfiMBeQgdO" outputId="28fa5aa7-3527-418d-9a45-30a57fceb0f0"
# Class distributions
prob_dist = get_probability_distribution(y_prob=y_prob[0], classes=label_encoder.classes)
print (json.dumps(prob_dist, indent=2))

# + [markdown] id="fckDDlLrNyqK"
# # Interpretability

# + [markdown] id="YkXEjvrYN0Wi"
# Let's use the attention values to see which encoded tokens were most useful in predicting the appropriate label.

# + id="szYYuvRCNxkG"
import collections
import seaborn as sns


# + id="IyapDiMQNx6a"
class InterpretAttn(nn.Module):
    def __init__(self, embedding_dim, vocab_size, rnn_hidden_dim,
                 hidden_dim, dropout_p, num_classes, padding_idx=0):
        super(InterpretAttn, self).__init__()
        
        # Initialize embeddings
        self.embeddings = nn.Embedding(
            embedding_dim=embedding_dim, num_embeddings=vocab_size,
            padding_idx=padding_idx)
        
        # RNN
        self.rnn = nn.RNN(embedding_dim, rnn_hidden_dim, batch_first=True)

        # Attention
        self.attn = nn.Linear(rnn_hidden_dim, 1)
     
        # FC weights
        self.dropout = nn.Dropout(dropout_p)
        self.fc1 = nn.Linear(rnn_hidden_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, num_classes)

    def forward(self, inputs, apply_softmax=False):
        # Embed
        x_in, seq_lens = inputs
        x_in = self.embeddings(x_in)
            
        # Encode
        out, h_n = self.rnn(x_in)
        
        # Attend
        e = self.attn(out)  # could add optional activation function (ex. tanh)
        attn_vals = F.softmax(e.squeeze(2), dim=1)

        return attn_vals


# + colab={"base_uri": "https://localhost:8080/"} id="zt0cUyubNx9M" outputId="573e75d1-5ff9-4d2b-b9ec-d03a7decf302"
# Initialize model
interpretable_model = InterpretAttn(
    embedding_dim=EMBEDDING_DIM, vocab_size=VOCAB_SIZE, 
    rnn_hidden_dim=RNN_HIDDEN_DIM, hidden_dim=HIDDEN_DIM, 
    dropout_p=DROPOUT_P, num_classes=NUM_CLASSES)
interpretable_model.load_state_dict(torch.load(Path(dir, "model.pt"), map_location=device))
interpretable_model.to(device)

# + id="AfGP5_Uaqi6q"
# Initialize trainer
interpretable_trainer = Trainer(model=interpretable_model, device=device)

# + colab={"base_uri": "https://localhost:8080/"} id="C11Mz8iqqn3l" outputId="f6405c4c-5de5-40fa-df93-e4d3c986ab0e"
# Get attention values
attn_vals  = interpretable_trainer.predict_step(dataloader)
print (attn_vals.shape) # (N, max_seq_len)

# + colab={"base_uri": "https://localhost:8080/", "height": 122} id="2ZOtnF0Jqn6i" outputId="61185ff7-59da-4d87-cfaf-1cea98818cbe"
# Visualize a bi-gram filter's outputs
sns.set(rc={"figure.figsize":(10, 1)})
tokens = tokenizer.sequences_to_texts(X)[0].split(' ')
sns.heatmap(attn_vals, xticklabels=tokens)

# + [markdown] id="1kI48Yh_rQKH"
# The word `tennis` was attended to the most to result in the `Sports` label.

# + [markdown] id="dH1-PtTiXAd2"
# # Types of attention

# + [markdown] id="5WQQaq7AZaAt"
# We'll briefly look at the different types of attention and when to use each them.

# + [markdown] id="VryQtkJyZhH7"
# ## Soft (global) attention

# + [markdown] id="l9FEf9XfZ6Fp"
# Soft attention the type of attention we've implemented so far, where we attend to all encoded inputs when creating our context vector.
#
# - **advantages**: we always have the ability to attend to all inputs in case something we saw much earlier/ see later are crucial for determing the output.
# - **disadvantages**: if our input sequence is very long, this can lead to expensive compute.

# + [markdown] id="oPUxLrMEZkou"
# ## Hard attention

# + [markdown] id="SyUrS6kqbCUT"
# Hard attention is focusing on a specific set of the encoded inputs at each time step.
#
# - **advantages**: we can save a lot of compute on long sequences by only focusing on a local patch each time.
# - **disadvantages**: non-differentiable and so we need to use more complex techniques (variance reduction, reinforcement learning, etc.) to train.

# + [markdown] id="RcWjZT0mcVWZ"
# <div align="left">
# <img src="https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/images/foundations/attention/soft_attention.png" width="700">
# </div>
# <div align="left">
# <small><a href="https://arxiv.org/abs/1502.03044" target="_blank">Show, Attend and Tell: Neural Image Caption Generation with Visual Attention</a></small>
# </div>

# + [markdown] id="0D56P8LIc_66"
# ## Local attention

# + [markdown] id="38RvQIF0dBlg"
# [Local attention](https://arxiv.org/abs/1508.04025) blends the advantages of soft and hard attention. It involves learning an aligned position vector and empirically determining a local window of encoded inputs to attend to.
#
# - **advantages**: apply attention to a local patch of inputs yet remain differentiable.
# - **disadvantages**: need to determine the alignment vector for each output but it's a worthwhile trade off to determine the right window of inputs to attend to in order to avoid attending to all of them.

# + [markdown] id="aK_oNwI7d3KD"
# <div align="left">
# <img src="https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/images/foundations/attention/local_attention.png" width="700">
# </div>
# <div align="left">
# <small><a href="https://arxiv.org/abs/1508.04025" target="_blank">Effective Approaches to Attention-based Neural Machine Translation
# </a></small>
# </div>

# + [markdown] id="2asXNHQbZmyu"
# ## Self-attention

# + [markdown] id="F_VlFfoQeLM0"
# We can also use attention within the encoded input sequence to create a weighted representation that based on the similarity between input pairs. This will allow us to create rich representations of the input sequence that are aware of the relationships between each other. For example, in the image below you can see that when composing the representation of the token "its", this specific attention head will be incorporating signal from the token "Law" (it's learned that "its" is referring to the "Law").
#
# <div align="left">
# <img src="https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/images/foundations/attention/self_attention.png" width="300">
# </div>
# <div align="left">
# <small><a href="https://arxiv.org/abs/1706.03762" target="_blank">Attention Is All You Need</a></small>
# </div>

# + [markdown] id="M5QM3Q2ZeMmw"
# # Transformers

# + [markdown] id="Yw7rVfqVAZJB"
# Transformers are a very popular architecture that leverage and extend the concept of self-attention to create very useful representations of our input data for a downstream task.

# + [markdown] id="5Bfaa_08Ugmj"
# ## Scaled dot-product attention

# + [markdown] id="IyAtst8GUrPh"
# The most popular type of self-attention is scaled dot-product attention from the widely-cited [Attention is all you need](https://arxiv.org/abs/1706.03762) paper. This type of attention involves projecting our encoded input sequences onto three matrices, queries (Q), keys (K) and values (V), whose weights we learn.

# + [markdown] id="f_1jet76X8qV"
# $ inputs \in \mathbb{R}^{NXMXH} $ ($N$ = batch size, $M$ = sequence length, $H$ = hidden dim)
#
# $ Q = XW_q $ where $ W_q \in \mathbb{R}^{HXd_q} $
#
# $ K = XW_k $ where $ W_k \in \mathbb{R}^{HXd_k} $
#
# $ V = XW_v $ where $ W_v \in \mathbb{R}^{HXd_v} $
#
# $ attention (Q, K, V) = softmax( \frac{Q K^{T}}{\sqrt{d_k}} )V \in \mathbb{R}^{MXd_v} $

# + [markdown] id="C0xT7ieqT-zY"
# ## Multi-head attention

# + [markdown] id="L7cwLUkoUBsc"
# Instead of applying self-attention only once across the entire encoded input, we can also separate the input and apply self-attention in parallel (heads) to each input section and concatenate them. This allows the different head to learn unique representations while maintaining the complexity since we split the input into smaller subspaces.

# + [markdown] id="qfLNJqhxX6rE"
# $ MultiHead(Q, K, V) = concat({head}_1, ..., {head}_{h})W_O $ 
#
# * ${head}_i = attention(Q_i, K_i, V_i) $
# * $h$ = # of self-attention heads
# * $W_O \in \mathbb{R}^{hd_vXH} $
# * $H$ = hidden dim. (or dimension of the model $d_{model}$)
#

# + [markdown] id="OCD6cbqgT8Ot"
# ## Positional encoding

# + [markdown] id="zlY6qzamkYoN"
# With self-attention, we aren't able to account for the sequential position of our input tokens. To address this, we can use positional encoding to create a representation of the location of each token with respect to the entire sequence. This can either be learned (with weights) or we can use a fixed function that can better extend to create positional encoding for lengths during inference that were not observed during training.

# + [markdown] id="H_eN-dKerkt5"
# $ PE_{(pos,2i)} = sin({pos}/{10000^{2i/H}}) $
#
# $ PE_{(pos,2i+1)} = cos({pos}/{10000^{2i/H}}) $
#
# where:
#
# * $pos$ = position of the token $(1...M)$
# * $i$ = hidden dim $(1..H)$

# + [markdown] id="t70xLrCZrk4X"
# This effectively allows us to represent each token's relative position using a fixed function for very large sequences. And because we've constrained the positional encodings to have the same dimensions as our encoded inputs, we can simply concatenate them before feeding them into the multi-head attention heads.

# + [markdown] id="gkFSjW6QZcFC"
# ## Architecture

# + [markdown] id="67PTrZhhZaj0"
# And here's how it all fits together! It's an end-to-end architecture that creates these contextual representations and uses an encoder-decoder architecture to predict the outcomes (one-to-one, many-to-one, many-to-many, etc.) Due to the complexity of the architecture, they require massive amounts of data for training without overfitting, however, they can be leveraged as pretrained models to finetune with smaller datasets that are similar to the larger set it was initially trained on.

# + [markdown] id="e9j-97khZ5Pl"
# <div align="left">
# <img src="https://raw.githubusercontent.com/GokuMohandas/MadeWithML/main/images/foundations/attention/transformer.png" width="800">
# </div>
# <div align="left">
# <small><a href="https://arxiv.org/abs/1706.03762" target="_blank">Attention Is All You Need</a></small>
# </div>

# + [markdown] id="akYFH3ru_6Tg"
# > We're not going to the implement the Transformer [from scratch](https://nlp.seas.harvard.edu/2018/04/03/attention.html) but we will use the[ Hugging Face library](https://github.com/huggingface/transformers) to do so in the [baselines](https://madewithml.com/courses/mlops/baselines/#transformers-w-contextual-embeddings) lesson!

StarcoderdataJupyterScriptsDedupFiltered