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

# + [markdown] colab_type="text" id="view-in-github"
# <a href="https://colab.research.google.com/github/zerotodeeplearning/ztdl-masterclasses/blob/master/solutions_do_not_open/Attention_with_Keras_solution.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] colab_type="text" id="2bwH96hViwS7"
# ## Learn with us: www.zerotodeeplearning.com
#
# Copyright © 2021: Zero to Deep Learning ® Catalit LLC.

# + colab={} colab_type="code" id="bFidPKNdkVPg"
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# + [markdown] colab_type="text" id="DvoukA2tkGV4"
# # Attention with Keras

# + [markdown] id="jeAWZNIdxFI3"
# This exercise follows:
# https://keras.io/examples/nlp/text_classification_with_transformer/

# + id="usTEmKM7kdVl"
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.datasets import imdb
import tensorflow as tf
from tensorflow.keras.layers import Layer, Embedding, MultiHeadAttention, Dense, GlobalAveragePooling1D, Dropout, Dense, LayerNormalization, Input
from tensorflow.keras.models import Sequential, Model
import matplotlib.pyplot as plt
from tensorflow.keras.utils import plot_model
# -

# ## Load the IMDB dataset and its word index

# + colab={"base_uri": "https://localhost:8080/"} id="QfN9FD8mkqp-" outputId="65b818d8-c761-4dbc-cce4-c20100cefcd5"
vocab_size = 20000
maxlen = 200

(X_train, y_train), (X_test, y_test) = imdb.load_data(num_words=vocab_size)
X_train = pad_sequences(X_train, maxlen=maxlen)
X_test = pad_sequences(X_test, maxlen=maxlen)

word_index = imdb.get_word_index()
inverted_word_index = dict((i+3, word) for (word, i) in word_index.items())
inverted_word_index[0] = ''
inverted_word_index[1] = '<start>'
inverted_word_index[2] = '<oov>'
# -

# Let's check the data shape

# + colab={"base_uri": "https://localhost:8080/"} id="v5tkLay9ksg6" outputId="88c45658-b4a4-4d27-8b48-6865fc96140b"
X_train.shape

# + colab={"base_uri": "https://localhost:8080/"} id="SuoO5CgAlLIz" outputId="166ecfcc-5d83-4c7f-f3d1-8ea1ecc4c03e"
X_train[0]
# -

# Let's check a couple of sentences for consistency

# + colab={"base_uri": "https://localhost:8080/", "height": 154} id="hrdoWLJhlM2c" outputId="25085b4b-f948-46fc-b2dc-86e6177a5268"
" ".join(inverted_word_index[i] for i in X_train[0])

# + colab={"base_uri": "https://localhost:8080/", "height": 154} id="NM03LbYWlj00" outputId="99e7673a-1719-4e83-ca01-e76c3efaaf2f"
" ".join(inverted_word_index[i] for i in X_train[1])


# -

# ## Token and Position Embedding

# + id="V9thqY9anz8g"
class TokenAndPositionEmbedding(Layer):
  def __init__(self, maxlen, vocab_size, embed_dim, **kwargs):
    super(TokenAndPositionEmbedding, self).__init__(**kwargs)
    self.token_emb = Embedding(input_dim=vocab_size, output_dim=embed_dim)
    self.pos_emb = Embedding(input_dim=maxlen, output_dim=embed_dim)

  def call(self, x):
    maxlen = tf.shape(x)[-1]
    positions = tf.range(start=0, limit=maxlen, delta=1)
    positions = self.pos_emb(positions)
    x = self.token_emb(x)
    return x + positions


# -

# Let's display a few sentences:

# + id="Amo9V3vpoRF_"
example_tpe = Sequential([TokenAndPositionEmbedding(maxlen, vocab_size, 32)])

# + id="rn1rrjXyrOOX"
n_reviews = 5

# + id="bZYDWHBOoW8n"
embedded_sentences = example_tpe(X_train[:n_reviews])

# + colab={"base_uri": "https://localhost:8080/", "height": 734} id="2HTHE3BCq_m_" outputId="e9c859d0-d614-4e40-b56f-c6b2b39d56de"
plt.figure(figsize=(10, 10))
for i in range(n_reviews):
  plt.subplot(n_reviews, 1, i+1)
  plt.imshow(embedded_sentences.numpy()[i].transpose())
  plt.xlabel("word in sentence -->")
  plt.ylabel("<-- embedding dim")
  plt.title(f"movie review {i}")

plt.tight_layout();


# -

# ## Transformer Block

# + id="iivbLxlSolvb"
class TransformerBlock(Layer):
  def __init__(self, embed_dim, n_att_heads, n_dense_nodes, rate=0.1, **kwargs):
    super(TransformerBlock, self).__init__(**kwargs)
    self.att = MultiHeadAttention(num_heads=n_att_heads, key_dim=embed_dim)
    self.ffn = Sequential([
        Dense(n_dense_nodes, activation="relu"),
        Dense(embed_dim)]
    )
    self.layernorm2 = LayerNormalization(epsilon=1e-6)
    self.layernorm1 = LayerNormalization(epsilon=1e-6)
    self.dropout1 = Dropout(rate)
    self.dropout2 = Dropout(rate)

  def call(self, inputs, training):
    attn_output = self.att(inputs, inputs)
    attn_output = self.dropout1(attn_output, training=training)
    out1 = self.layernorm1(inputs + attn_output)
    ffn_output = self.ffn(out1)
    ffn_output = self.dropout2(ffn_output, training=training)
    return self.layernorm2(out1 + ffn_output)


# -

# ## Exercise 1:

# Using either the Sequential or the Functional API in Keras build a transformer classification model with the following architecture:
#
# ```
#     TokenAndPositionEmbedding(...
#     TransformerBlock(...
#     GlobalAveragePooling1D(...
#     Dropout(...
#     Dense(...
#     Dropout(...
#     Dense(2, activation="softmax")
# ````
#
# Once the model is built, print out the summary.
#
# You will need to decide a few hyperparameters including:
#
# - Embedding size
# - Number of attention heads
# - Size of the dense hidden layer inside the transformer block
# - Size of the other dense layers
# - Dropout rate

# + id="o1DMTFYip7Ct" tags=["solution", "empty"]
embed_dim = 32  # Embedding size for each token
n_att_heads = 2  # Number of attention heads
n_dense_nodes = 32  # Hidden layer size in feed forward network inside transformer


model = Sequential([
  TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim, input_shape=(maxlen,)),
  TransformerBlock(embed_dim, n_att_heads, n_dense_nodes),
  GlobalAveragePooling1D(),
  Dropout(0.1),
  Dense(20, activation="relu"),
  Dropout(0.1),
  Dense(2, activation="softmax")
])

model.summary()

# + colab={"base_uri": "https://localhost:8080/"} id="9aHj4LoItnR6" outputId="00167ed1-e662-4180-f7bb-07b78dc8de21" tags=["solution"]
model.summary()
# -

# ## Exercise 2
#
# Compile, train, and evaluate the model. Pay attention to the loss function. We defined the output layer as a `Dense(2, activation="softmax")` so you will need to choose the loss accordingly.

# + colab={"base_uri": "https://localhost:8080/"} id="1kjWZocLqGdf" outputId="4ab87a62-7854-4303-b78e-1758bd8f0114" tags=["solution", "empty"]
model.compile("adam", "sparse_categorical_crossentropy", metrics=["accuracy"])

h = model.fit(X_train, y_train, batch_size=32, epochs=2, validation_split=0.1)

# + colab={"base_uri": "https://localhost:8080/"} id="yShtCMbDqQeI" outputId="3a71f85c-4237-4bf5-ad5c-fbc73bdca018" tags=["solution"]
model.evaluate(X_test, y_test)

StarcoderdataJupyterScriptsDedupFiltered