Deep Learning with Python - Chapter 11

# Download the dataset
!curl -O https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
!tar -xf aclImdb_v1.tar.gz

# Remove the directory we don't need
!rm -r aclImdb/train/unsup

# Take a look at what the data looks like
!cat aclImdb/train/pos/4077_10.txt

# Prepare the validation set by setting apart 20% f the training text files into
# a new directory
import os, pathlib, shutil, random

base_dir = pathlib.Path("aclImdb")
val_dir = base_dir / "val"
train_dir = base_dir / "train"
for category in ("neg", "pos"):
  os.makedirs(val_dir / category)
  files = os.listdir(train_dir / category)
  # Shuffle the list of training files using a seed, to ensure
  # we always get the same validation set
  random.Random(1337).shuffle(files)
  # Take 20% of the training files to use for validation
  num_val_samples = int(0.2 * len(files))
  val_files = files[-num_val_samples:]
  for fname in val_files:
    # Move the files
    shutil.move(train_dir / category / fname, val_dir / category / fname)

from tensorflow import keras
batch_size = 32
"""
Creating Dataset objects for training, validation, and testing
"""
train_ds = keras.utils.text_dataset_from_directory(
 "aclImdb/train", batch_size=batch_size
)
val_ds = keras.utils.text_dataset_from_directory(
 "aclImdb/val", batch_size=batch_size
)
test_ds = keras.utils.text_dataset_from_directory(
 "aclImdb/test", batch_size=batch_size
)

"""
Displaying the shapes and dtypes of the first batch
"""
for inputs, targets in train_ds:
  print("inputs.shape:", inputs.shape)
  print("inputs.dtype:", inputs.dtype)
  print("targets.shape:", targets.shape)
  print("targets.dtype:", targets.dtype)
  print("inputs[0]:", inputs[0])
  print("targets[0]:", targets[0])
  break

from tensorflow.keras.layers import TextVectorization
"""
Preprocessing our dataset with a TextVectorization layer

Limit the vocabulary to the 20,000 most frequent words. Otherwise, we'd be
indexing every word in the training data. In general, 20,000 is the right
vocabulary for text classification.
"""
text_vectorization = TextVectorization(
 max_tokens=20000,
 # Encode the output tokens as mult-hot binary vectors
 output_mode="multi_hot",
)
# Prepare a dataset that only yeidls raw text inputs
# (no labels)
text_only_train_ds = train_ds.map(lambda x, y: x)
# Use the dataset to index the dataset vocabulary via the
# adapt() method
text_vectorization.adapt(text_only_train_ds)
"""
Prepare processed versions of our training, validation,
and text dataset. Make sure to specify num_parallel_calls to leverage
multiple CPU cores.
"""
binary_1gram_train_ds = train_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
binary_1gram_val_ds = val_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
binary_1gram_test_ds = test_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)

# Inspecting the output of our binary unigram dataset
for inputs, targets in binary_1gram_train_ds:
  print("inputs.shape:", inputs.shape)
  print("inputs.dtype:", inputs.dtype)
  print("targets.shape:", targets.shape)
  print("targets.dtype:", targets.dtype)
  print("inputs[0]:", inputs[0])
  print("targets[0]:", targets[0])
  break

"""
Our model-building utility
"""
from tensorflow import keras
from tensorflow.keras import layers
def get_model(max_tokens=20000, hidden_dim=16):
  inputs = keras.Input(shape=(max_tokens,))
  x = layers.Dense(hidden_dim, activation="relu")(inputs)
  x = layers.Dropout(0.5)(x)
  outputs = layers.Dense(1, activation="sigmoid")(x)
  model = keras.Model(inputs, outputs)
  model.compile(optimizer="rmsprop",
  loss="binary_crossentropy",
  metrics=["accuracy"])
  return model

"""
Training and testing the binary unigram model
"""
model = get_model()
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("binary_1gram.keras",save_best_only=True)
]
"""
We call cache() on the datasets to cache them in memory:
this way we will only do the preprocessing once, during the first epoch and
we'll reuse the prepocessed texts for the following epochs. Thhis can be
done if the data is small enough to fit in memory.
"""
model.fit(binary_1gram_train_ds.cache(),
validation_data=binary_1gram_val_ds.cache(),
 epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("binary_1gram.keras")
print(f"Test acc: {model.evaluate(binary_1gram_test_ds)[1]:.3f}")

"""
Configuring the TextVectorization layer to return bigrams
"""
text_vectorization = TextVectorization(
    ngrams=2,
    max_tokens=20000,
    output_mode="multi_hot",
)

"""
Training and testing the binary bigram model
"""
text_vectorization.adapt(text_only_train_ds)
binary_2gram_train_ds = train_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
binary_2gram_val_ds = val_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
binary_2gram_test_ds = test_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
model = get_model()
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("binary_2gram.keras",
 save_best_only=True)
]
model.fit(binary_2gram_train_ds.cache(),
 validation_data=binary_2gram_val_ds.cache(),
 epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("binary_2gram.keras")
print(f"Test acc: {model.evaluate(binary_2gram_test_ds)[1]:.3f}")

"""
Configurting the TextVectorization layer to return token counts
"""
text_vectorization = TextVectorization(
 ngrams=2,
 max_tokens=20000,
 output_mode="count"
)

"""
Congiguring TextVectorization to return TF-IDF weighted outputs
"""
text_vectorization = TextVectorization(
 ngrams=2,
 max_tokens=20000,
 output_mode="tf_idf",
)

"""
Training and testing the TF-IDF bigram model
"""
# The adapt() call will learn the TF-IDF weights in addition to
# the vocabulary.
text_vectorization.adapt(text_only_train_ds)
tfidf_2gram_train_ds = train_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
tfidf_2gram_val_ds = val_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
tfidf_2gram_test_ds = test_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
model = get_model()
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("tfidf_2gram.keras",
 save_best_only=True)
]
model.fit(tfidf_2gram_train_ds.cache(),
 validation_data=tfidf_2gram_val_ds.cache(),
 epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("tfidf_2gram.keras")
print(f"Test acc: {model.evaluate(tfidf_2gram_test_ds)[1]:.3f}")

"""
Preparing Integer Datasets
"""
from tensorflow.keras import layers
max_length = 600
max_tokens = 20000
text_vectorization = layers.TextVectorization(
 max_tokens=max_tokens,
 output_mode="int",
 # In order to keep a manageable input size, we'll truncate the inputs
 # after the first 600 words
 # This is a reasonable choice, since teh average review length is 233
 # words, and only 5% of reviews are longer than 600 words.
 output_sequence_length=max_length,
)
text_vectorization.adapt(text_only_train_ds)
int_train_ds = train_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
int_val_ds = val_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)
int_test_ds = test_ds.map(
 lambda x, y: (text_vectorization(x), y),
 num_parallel_calls=4)

"""
The simplest way to convert the integer sequences to vector sequences is to one-
hot encode the integers. On top of these one-hot vectros, we add a simple
bidirectional LSTM

A sequence model built on one-hot encoded vector sequences
"""
import tensorflow as tf
# One input is a sequence of integers
inputs = keras.Input(shape=(None,), dtype="int64")
# Encode the integers into binary 20,000-dimensional vectors
embedded = tf.one_hot(inputs, depth=max_tokens)
# Add a bidrectional LSTM
x = layers.Bidirectional(layers.LSTM(32))(embedded)
x = layers.Dropout(0.5)(x)
# Add a classification layer
outputs = layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="rmsprop",
 loss="binary_crossentropy",
 metrics=["accuracy"])
model.summary()

"""
Training a first basic sequence model
"""
callbacks = [
 keras.callbacks.ModelCheckpoint("one_hot_bidir_lstm.keras",
 save_best_only=True)
]
model.fit(int_train_ds, validation_data=int_val_ds, epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("one_hot_bidir_lstm.keras")
print(f"Test acc: {model.evaluate(int_test_ds)[1]:.3f}")
"""
Observations:
The model trains very slowly, espcially when compared to the light-weight model
of the previous section. This is because our inputs are quite large. Second,
the model only gets 87% test accuracy - clealy, using ohe to turn words into
vectors wasn;t a great idea. A better way: *word encodings*.
"""

"""
Model that uses an Embedding layer trained from scratch
"""
inputs = keras.Input(shape=(None,), dtype="int64")
embedded = layers.Embedding(input_dim=max_tokens, output_dim=256)(inputs)
x = layers.Bidirectional(layers.LSTM(32))(embedded)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="rmsprop",
 loss="binary_crossentropy",
 metrics=["accuracy"])
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("embeddings_bidir_gru.keras",
 save_best_only=True)
]
model.fit(int_train_ds, validation_data=int_val_ds, epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("embeddings_bidir_gru.keras")
print(f"Test acc: {model.evaluate(int_test_ds)[1]:.3f}")

"""
Using an Embedidng Model with masking enabled
"""
inputs = keras.Input(shape=(None,), dtype="int64")
embedded = layers.Embedding(
 input_dim=max_tokens, output_dim=256, mask_zero=True)(inputs)
x = layers.Bidirectional(layers.LSTM(32))(embedded)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="rmsprop",
 loss="binary_crossentropy",
 metrics=["accuracy"])
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("embeddings_bidir_gru_with_masking.keras",
 save_best_only=True)
]
model.fit(int_train_ds, validation_data=int_val_ds, epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("embeddings_bidir_gru_with_masking.keras")
print(f"Test acc: {model.evaluate(int_test_ds)[1]:.3f}")

# Download the GloVe word embeddings precomputed on the 2014 English Wikipedia
# Dataset
!wget http://nlp.stanford.edu/data/glove.6B.zip
!unzip -q glove.6B.zip

"""
Parsing the GloVe word-embeddings file

Parse the unzipped .txt file to build an index that maps words
as strings to their vector representation
"""
import numpy as np
path_to_glove_file = "glove.6B.100d.txt"
embeddings_index = {}
with open(path_to_glove_file) as f:
  for line in f:
    word, coefs = line.split(maxsplit=1)
    coefs = np.fromstring(coefs, "f", sep=" ")
    embeddings_index[word] = coefs
print(f"Found {len(embeddings_index)} word vectors.")

"""
Preparing the GloVe word-embeddings matrix

- Build an embedding matrix that you can load into an
Embedding layer
- It must be a matirx of shape (max_words, embedding_dim), where
each entry i contains the embedding_dim-dimensional vector for the
word of index i in the reference word index
"""
embedding_dim = 100
# Retrieve the vocabulary indexed by our previous TextVectorization layer
vocabulary = text_vectorization.get_vocabulary()
# Use it to create a mapping from words to their index in the vocabulary
word_index = dict(zip(vocabulary, range(len(vocabulary))))
# Prepare a matrix that we'll fill with the GloVe vectors
embedding_matrix = np.zeros((max_tokens, embedding_dim))
for word, i in word_index.items():
  if i < max_tokens:
    embedding_vector = embeddings_index.get(word)
  # Fill entry i in the matrix with the word vector for index i
  # Words not found in the embedding index will be all zeros
  if embedding_vector is not None:
    embedding_matrix[i] = embedding_vector

"""
Use a Constant initializer to load the pretrained embeddings in an
Embedding layer. Freeze the pretrained representation during training - set
trainable=False
"""
embedding_layer = layers.Embedding(
  max_tokens,
  embedding_dim,
  embeddings_initializer=keras.initializers.Constant(embedding_matrix),
  trainable=False,
  mask_zero=True,
)

"""
Model that uses a pretrained Embedding layer
"""
inputs = keras.Input(shape=(None,), dtype="int64")
embedded = embedding_layer(inputs)
x = layers.Bidirectional(layers.LSTM(32))(embedded)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="rmsprop",
 loss="binary_crossentropy",
 metrics=["accuracy"])
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("glove_embeddings_sequence_model.keras",
 save_best_only=True)
]
model.fit(int_train_ds, validation_data=int_val_ds, epochs=10,
 callbacks=callbacks)
model = keras.models.load_model("glove_embeddings_sequence_model.keras")
print(f"Test acc: {model.evaluate(int_test_ds)[1]:.3f}")

"""
Transformer encoder implemented as a subclassed Layer
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
class TransformerEncoder(layers.Layer):
  def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
    super().__init__(**kwargs)
    # Size of the input token vectors
    self.embed_dim = embed_dim
    # Size of the inner dense layer
    self.dense_dim = dense_dim
    # Number of attention heads
    self.num_heads = num_heads
    self.attention = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
    self.dense_proj = keras.Sequential([
        layers.Dense(dense_dim, activation="relu"),
        layers.Dense(embed_dim),
      ]
    )
    self.layernorm_1 = layers.LayerNormalization()
    self.layernorm_2 = layers.LayerNormalization()

  def call(self, inputs, mask=None):
    """
    Computation goes in call()
    """
    if mask is not None:
      # The mask that will be generated by the Embedding layer will
      # be 2D, but the attention layer expects to be 3D or 4D,
      # so we expand its rank
      mask = mask[:, tf.newaxis, :]
    attention_output = self.attention(inputs, inputs, attention_mask=mask)
    proj_input = self.layernorm_1(inputs + attention_output)
    proj_output = self.dense_proj(proj_input)
    return self.layernorm_2(proj_input + proj_output)
  def get_config(self):
    """
    Implement serialization so we can save the model

    When you write custom layers, make sure to implement the get_config method:
    this enables the layer to be reinstated from its config dict, which is
    useful during model saving and loading. This method should return a Python
    dict that contains the values of the constructor arguments used to create
    the layer.
    """
    config = super().get_config()
    config.update({
        "embed_dim": self.embed_dim,
        "num_heads": self.num_heads,
        "dense_dim": self.dense_dim,
    })
    return config

"""
Using the Transformer encoder for text classification
"""
vocab_size = 20000
embed_dim = 256
num_heads = 2
dense_dim = 32
inputs = keras.Input(shape=(None,), dtype="int64")
x = layers.Embedding(vocab_size, embed_dim)(inputs)
x = TransformerEncoder(embed_dim, dense_dim, num_heads)(x)
"""
Since TransformerEncoder returns full sequences,
we need to reduce each sequence to a single vector for clasification
via a global pooling layer.
"""
x = layers.GlobalMaxPooling1D()(x)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="rmsprop",
 loss="binary_crossentropy",
 metrics=["accuracy"]
)
model.summary()
"""
Training and evaluating the Transformer encoded based model
"""
callbacks = [
 keras.callbacks.ModelCheckpoint("transformer_encoder.keras",
 save_best_only=True)
]
model.fit(int_train_ds, validation_data=int_val_ds, epochs=20,
callbacks=callbacks)
model = keras.models.load_model(
 "transformer_encoder.keras",
 custom_objects={"TransformerEncoder": TransformerEncoder}) # provide the custom
 # transformerEncoder class to the model-loading process
print(f"Test acc: {model.evaluate(int_test_ds)[1]:.3f}")

"""
Implementing positional embedding as a subclassed layer
"""

class PositionalEmbedding(layers.Layer):
  """
  A downside of position embeddings is that the sequence lengths
  needs to be known in advance
  """
  def __init__(self, sequence_length, input_dim, output_dim, **kwargs):
    super().__init__(**kwargs)
    # Prepare an Embedding layer for the token indices
    self.token_embeddings = layers.Embedding(input_dim=input_dim, output_dim=output_dim)
    # Add another one for the token positions
    self.position_embeddings = layers.Embedding(
    input_dim=sequence_length, output_dim=output_dim)
    self.sequence_length = sequence_length
    self.input_dim = input_dim
    self.output_dim = output_dim
  def call(self,inputs):
    length = tf.shape(inputs)[-1]
    positions = tf.range(start=0, limit=length, delta=1)
    embedded_tokens = self.token_embeddings(inputs)
    embedded_positions = self.position_embeddings(positions)
    # Add both embedding vectors together
    return embedded_tokens + embedded_positions
  def compute_mask(self, inputs, mask=None):
    """
    Like the Embedding Layer, this layer should be able to generate a
    mask so we can ignore padding 0s in the inputs. The compute_mask
    method will be called automatically by the framework, and the mask
    will get propagated to the next layer
    """
    return tf.math.not_equal(inputs, 0)
  def get_config(self):
    """
    Implement serialization so that we can save the model
    """
    config = super().get_config()
    config.update({
      "output_dim": self.output_dim,
      "sequence_length": self.sequence_length,
      "input_dim": self.input_dim,
    })
    return config

"""
Combining the Transformer encoder with positional embedding

This model gives an 88.3% test accuracy, which demonstrates the value
of word order information for text classification, but it is still one
notch below the bag-of-words approach.
"""
vocab_size = 20000
sequence_length = 600
embed_dim = 256
num_heads = 2
dense_dim = 32
inputs = keras.Input(shape=(None,), dtype="int64")
x = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(inputs)
x = TransformerEncoder(embed_dim, dense_dim, num_heads)(x)
x = layers.GlobalMaxPooling1D()(x)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(1, activation="sigmoid")(x)
model = keras.Model(inputs, outputs)
model.compile(optimizer="rmsprop",
 loss="binary_crossentropy",
 metrics=["accuracy"])
model.summary()
callbacks = [
 keras.callbacks.ModelCheckpoint("full_transformer_encoder.keras",
 save_best_only=True)
]
model.fit(int_train_ds, validation_data=int_val_ds, epochs=20,
callbacks=callbacks)
model = keras.models.load_model(
 "full_transformer_encoder.keras",
 custom_objects={"TransformerEncoder": TransformerEncoder,
 "PositionalEmbedding": PositionalEmbedding})
print(f"Test acc: {model.evaluate(int_test_ds)[1]:.3f}")

## Download the English-to-Spanish Dataset
!wget http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip
!unzip -q spa-eng.zip

# Parsing the text files
text_file = "spa-eng/spa.txt"
with open(text_file) as f:
  lines = f.read().split("\n")[:-1]
text_pairs = []

for line in lines: # Iterate over the lines in the file
  # Each line contains an English phrase and its Spanish translation, tab-separated
  english, spanish = line.split("\t")
  # We prepend "[start]" and append "[end]" to the Spanish template
  spanish = "[start] " + spanish + " [end]"
  text_pairs.append((english, spanish))

# What our text_pairs look like
import random
print(random.choice(text_pairs))

"""
Shuffle the pairs and split them into training, validation, and test sets
"""
import random
random.shuffle(text_pairs)
num_val_samples = int(0.15 * len(text_pairs))
num_train_samples = len(text_pairs) - 2 * num_val_samples
train_pairs = text_pairs[:num_train_samples]
val_pairs = text_pairs[num_train_samples:num_train_samples + num_val_samples]
test_pairs = text_pairs[num_train_samples + num_val_samples:]

"""
Vectorizing the English and Spanish text pairs
"""
import tensorflow as tf
import string
import re

"""
Prepare a custom string standardization fucntion for the Spanish
TextVectorization layer: it preserves "[" and "]" but strips the upside down
question mark (as well as all other characters from strings.punctuation)
"""
strip_chars = string.punctuation + "¿"
strip_chars = strip_chars.replace("[", "")
strip_chars = strip_chars.replace("]", "")
def custom_standardization(input_string):
 lowercase = tf.strings.lower(input_string)
 return tf.strings.regex_replace(lowercase, f"[{re.escape(strip_chars)}]", "")

"""
To keep thimngs simple, we will only look at the top 15,000 words in each
language, and we'll restrict sentences to 20 words
"""
vocab_size = 15000
sequence_length = 20
# The English layer
source_vectorization = layers.TextVectorization(
 max_tokens=vocab_size,
 output_mode="int",
 output_sequence_length=sequence_length,
)
# The Spanish Layer
target_vectorization = layers.TextVectorization(
 max_tokens=vocab_size,
 output_mode="int",
 # Generate Spanish sentences that have one extra token
 # Since we'll need to offset the sentence by one step during training
 output_sequence_length=sequence_length + 1,
 standardize=custom_standardization,
)
train_english_texts = [pair[0] for pair in train_pairs]
train_spanish_texts = [pair[1] for pair in train_pairs]
# Learn the vocabulary of each language
source_vectorization.adapt(train_english_texts)
target_vectorization.adapt(train_spanish_texts)

"""
Preparing datasets for the translation task


"""
batch_size = 64
def format_dataset(eng, spa):
  eng = source_vectorization(eng)
  spa = target_vectorization(spa)
  return ({
  "english": eng,
  # The input Spanish doesn't include the last token to keep inputs and
  # targets at the same length
  "spanish": spa[:, :-1],
  # The target Spanish sentence is one step ahead.
  # Both are still the same length (20 words)
  }, spa[:, 1:])
def make_dataset(pairs):
  eng_texts, spa_texts = zip(*pairs)
  eng_texts = list(eng_texts)
  spa_texts = list(spa_texts)
  dataset = tf.data.Dataset.from_tensor_slices((eng_texts, spa_texts))
  dataset = dataset.batch(batch_size)
  dataset = dataset.map(format_dataset, num_parallel_calls=4)
  # Use in memory cacheing to speed up preprocessing
  return dataset.shuffle(2048).prefetch(16).cache()
train_ds = make_dataset(train_pairs)
val_ds = make_dataset(val_pairs)

"""
What out dataset outputs look like
"""
for inputs, targets in train_ds.take(1):
  print(f"inputs['english'].shape: {inputs['english'].shape}")
  print(f"inputs['spanish'].shape: {inputs['spanish'].shape}")
  print(f"targets.shape: {targets.shape}")

"""
GRU Based Encoder
"""
from tensorflow import keras
from tensorflow.keras import layers
embed_dim = 256
latent_dim = 1024
# The English source sentence goes here
source = keras.Input(shape=(None,), dtype="int64", name="english")
# Masking is critical in this setup
x = layers.Embedding(vocab_size, embed_dim, mask_zero=True)(source)

encoded_source = layers.Bidirectional(
 layers.GRU(latent_dim), merge_mode="sum")(x) # Our encoded source sentence is the last output of bidiectional GRU

"""
GRU-based decoder and the end-to-end model
"""
# The spanish target sentence goes here
past_target = keras.Input(shape=(None,), dtype="int64", name="spanish")
# Don't forget masking
x = layers.Embedding(vocab_size, embed_dim, mask_zero=True)(past_target)
decoder_gru = layers.GRU(latent_dim, return_sequences=True)
# The encoded source sentence serves as the initial state of the decoder GRU
x = decoder_gru(x, initial_state=encoded_source)
x = layers.Dropout(0.5)(x)
# Predicts the next token
target_next_step = layers.Dense(vocab_size, activation="softmax")(x)
# End-to-end model: maps the source sentence and the target sentence
# to the target senetence one step in the future
seq2seq_rnn = keras.Model([source, past_target], target_next_step)

"""
Training our recurrent sequence-to-sequence model
"""
seq2seq_rnn.compile(
 optimizer="rmsprop",
 loss="sparse_categorical_crossentropy",
 metrics=["accuracy"])
seq2seq_rnn.fit(train_ds, epochs=15, validation_data=val_ds)

"""
The Transformer Decoder
"""
class TransformerDecoder(layers.Layer):
  def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
    super().__init__(**kwargs)
    self.embed_dim = embed_dim
    self.dense_dim = dense_dim
    self.num_heads = num_heads
    self.attention_1 = layers.MultiHeadAttention(
    num_heads=num_heads, key_dim=embed_dim)
    self.attention_2 = layers.MultiHeadAttention(
    num_heads=num_heads, key_dim=embed_dim)
    self.dense_proj = keras.Sequential(
      [layers.Dense(dense_dim, activation="relu"),
      layers.Dense(embed_dim),]
    )
    self.layernorm_1 = layers.LayerNormalization()
    self.layernorm_2 = layers.LayerNormalization()
    self.layernorm_3 = layers.LayerNormalization()
    # This attribute ensures that the layer will propagate its input mask
    # to its outputs; masking in Keras is explicitly opt-in.
    self.supports_masking = True
  def get_config(self):
    config = super().get_config()
    config.update({
      "embed_dim": self.embed_dim,
      "num_heads": self.num_heads,
      "dense_dim": self.dense_dim,
    })
    return config