<jupyter_start><jupyter_text>Imports<jupyter_code>try:
    from typing import List, Tuple
    import tensorflow as tf
    import tensorflow_text as tf_text
    import pandas as pd
    import seaborn as sns
    import numpy as np
    import tqdm
    import matplotlib.pyplot as plt
except ModuleNotFoundError:
    %pip install -q tensorflow_text
    %pip install -q tensorflow
    %pip install -q pandas
    %pip -q install -q seaborn
    %pip -q install -q numpy
    %pip install -q tqdm
    %pip install -q matplotlib
    import tensorflow as tf
    import tensorflow_text as tf_text
    import pandas as pd
    import seaborn as sns
    import numpy as np
    import tqdm
    import matplotlib.pyplot as plt
import sys
f"Python version: {sys.version}"
tf.__version__
tf_text.__version__
tf.random.set_seed(0)
device: str
if 'google' in sys.modules:
    device = 'colab'
else:
    device = 'locally'<jupyter_output><empty_output><jupyter_text>Data loading<jupyter_code>if device == 'colab':  # If notebook is ran on colab
    try:
        from google.colab import drive
        drive.mount('/content/drive')
        df: pd.DataFrame = pd.read_csv('/content/drive/MyDrive/final_project/wikipedia_articles.csv')
    except ModuleNotFoundError:
        pass
else:  # If notebook is ran on my laptop
  df: pd.DataFrame = pd.read_csv('C:/yoni/final_project/model/wiki_data/good.csv')
df.shape
df: pd.Series = df['text']
data_list: List[str] = df.to_list()
DATA_SIZE = len(data_list)
print(f"There are {DATA_SIZE} data points")
string_lengths: List[int] = [len(data_point) for data_point in data_list]
max_string_len = max(string_lengths)
print(f"The length of the longest text IN CHARACTERS is: {max_string_len}")
min_string_len = min(string_lengths)
print(f"The length of the shortest text IN CHARACTERS is: {min_string_len}")<jupyter_output>There are 30279 data points
The length of the longest text IN CHARACTERS is: 141803
The length of the shortest text IN CHARACTERS is: 816<jupyter_text>Tokenizing  Creating the vocabulary<jupyter_code># tensor_list: List = [tf.convert_to_tensor(data_point) for data_point in data_list]
# print(type(tensor_list[0]))
# %%time
bert_tokenizer_params: dict = dict(lower_case=True)
reserved_tokens: List[str] = ["[PAD]", "[UNK]", "[START]", "[END]", "[MASK]"]
VOCAB_SIZE: int = 8192
#
# bert_vocab_args: dict = dict(
#     # The target vocabulary size
#     vocab_size = VOCAB_SIZE,
#     # Reserved tokens that must be included in the vocabulary
#     reserved_tokens=reserved_tokens,
#     # Arguments for `tf_text.BertTokenizer`
#     bert_tokenizer_params=bert_tokenizer_params,
#     # Arguments for `wordpiece_vocab.wordpiece_tokenizer_learner_lib.learn`
#     learn_params={},
# )
# data_set: tf.data.Dataset = tf.data.Dataset.from_tensor_slices(tensor_list)
# #I already ran this code and saved the file to C:/yoni/final_project/model/look_up_table.txt
# vocab: List[str] = bert_vocab_from_dataset.bert_vocab_from_dataset(
#     data_set,
#     **bert_vocab_args,)
# len(vocab), vocab[:10]
# with open('C:/yoni/final_project/model/look_up_table.txt', 'w') as f:
#     for token in vocab:
#         f.write(token + ' ')
if device == 'azure':  # If notebook is ran on azure
    path = "replace_that.txt"
elif device == 'colab':  # If notebook is ran on colab
    path = '/content/drive/MyDrive/final_project/look_up_table.txt'
else:  # If notebook is ran on my laptop
    path = 'C:/yoni/final_project/model/look_up_table.txt'
with open(path, 'r') as f:
    vocab: List[str] = f.read().split()
print(type(vocab[0]))
print(vocab[:15])
print(len(vocab))
tensor_vocab: List[tf.Tensor] = [tf.convert_to_tensor(token_key, dtype=tf.string) for token_key in vocab]  # dtype = tf.String
print(type(tensor_vocab[0]))
lookup_table = tf.lookup.StaticVocabularyTable(
    tf.lookup.KeyValueTensorInitializer(
        keys=tensor_vocab,
        key_dtype=tf.string,
        values=tf.range(tf.size(vocab, out_type=tf.int64), dtype=tf.int64),
        value_dtype=tf.int64),
    num_oov_buckets=1
)<jupyter_output><empty_output><jupyter_text>Creating the tokenizer<jupyter_code>tokenizer = tf_text.BertTokenizer(lookup_table, **bert_tokenizer_params)<jupyter_output><empty_output><jupyter_text>Tokenizing the data<jupyter_code>START: int = tf.argmax(tf.constant(reserved_tokens) == "[START]")  # The value of the start token
END: int = tf.argmax(tf.constant(reserved_tokens) == "[END]")  # The value of the end token
starts = tf.cast(tf.Variable([START]), dtype = tf.int32)  # Tensor of shape [1] and dtype int
ends = tf.cast(tf.Variable([END]), dtype = tf.int32)  # Tensor of shape [1] and dtype int
def tokenize(text: str) -> tf.Tensor:
    """Converts string to tensor"""
    ragged = tokenizer.tokenize(text)[0, :]
    eager: tf.Tensor = ragged.to_tensor(default_value=0, shape=[None, 1])  # 0 is the value of the padding token
    sqeezed: tf.Tensor = tf.squeeze(eager, axis=1)
    typed: tf.Tensor = tf.cast(sqeezed, tf.int32)
    edited: tf.Tensor = tf.concat([starts, typed, ends], axis=0)
    return edited
tokenized_data: List[tf.Tensor] = [tokenize(data_point) for data_point in tqdm.tqdm(data_list)]  # tqdm is a progress bar
print(len(tokenized_data))
print(tokenized_data[0].shape)
print(tokenized_data[0][:10])
lengths_tokenized: List[int] = [text.shape[0] for text in tokenized_data]
print(max(lengths_tokenized))
print(min(lengths_tokenized))
sns.displot(lengths_tokenized);<jupyter_output>25317
180<jupyter_text>This is a plot of the distribution of the lengths of the tokenized data (AKA how many tokens in each data point of the data set).We can see that there is very few data points with length over 16384 (2^14).A model that can handle long texts is very expensive (both in emory and runtime).Therefore I will filter out all texts longer than 16384 tokens.Note that 16384 is also a lot.  Filter too long texts<jupyter_code>MAX_LENGTH = 16384
def shorter_than_max_len(tokenized_text: tf.Tensor) -> bool:
    """Returns True if the tokenized text is shorter than 16385 tokens and False otherwise"""
    l = tokenized_text.shape[0] #length of the text in tokens
    if l <= MAX_LENGTH:
        return True
    return False
filtered_data: List[tf.Tensor] = list(filter(shorter_than_max_len, tokenized_data))
DATA_SIZE: int = len(filtered_data)
print(DATA_SIZE)<jupyter_output>30195<jupyter_text>We only removed 32 data points.<jupyter_code>print(filtered_data[0].shape)
def pad(tensor: tf.Tensor) -> tf.Tensor:
    """Pads the tensor to the length of the longest text in the data set"""
    padded: tf.Tensor = tf.pad(tensor=tensor, paddings=[[0, MAX_LENGTH - tensor.shape[0]]], mode='CONSTANT', constant_values=0)
    # 0 is the padding token
    return padded
padded_data: List[tf.Tensor] = [pad(text) for text in filtered_data]
padded_data[0]<jupyter_output><empty_output><jupyter_text>Train test val split<jupyter_code>import random
random.seed(0)
random.shuffle(padded_data)
train_size: int = int(DATA_SIZE * 0.8)
val_test_size: int = int(DATA_SIZE * 0.1)  # Both validation and test get 10% of the data

train_tokenized: List[tf.Tensor] = padded_data[:train_size]
val_tokenized: List[tf.Tensor] = padded_data[train_size:val_test_size]
test_tokenized: List[tf.Tensor] = padded_data[(train_size + val_test_size):val_test_size]
BATCH_SIZE: int = 64
def list_to_dataset(tokenized_list: List[tf.Tensor]) -> tf.data.Dataset:
    """Converts a list of tokenized texts after all preprocessing to a tf.data.Dataset"""
    dataset: tf.data.Dataset = tf.data.Dataset.from_tensor_slices(tokenized_list)
    batched: tf.data.Dataset = dataset.batch(BATCH_SIZE)
    return batched

train_set: tf.data.Dataset = list_to_dataset(train_tokenized)
val_set: tf.data.Dataset = list_to_dataset(val_tokenized)
test_set: tf.data.Dataset = list_to_dataset(test_tokenized)<jupyter_output><empty_output><jupyter_text>Clear memory<jupyter_code>del train_tokenized, test_tokenized, val_tokenized
del padded_data, filtered_data, tokenized_data, data_list, df, lengths_tokenized<jupyter_output><empty_output><jupyter_text>ModelIf you run it inside the directory final_project/model replace the entire section with:'from transformer import Transformer'<jupyter_code>from transformer import Transformer<jupyter_output><empty_output><jupyter_text>Positional encoding The formula for calculating the positional encoding is as follows:$$\Large{PE_{(pos, 2i)} = \sin(pos / 10000^{2i / d_{model}})} $$$$\Large{PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i / d_{model}})} $$where $d_{model}$ is the model dimension, $pos$ is the position and $i$ is the index of the embedding.<NAME>., <NAME>., <NAME>., <NAME>., <NAME>., <NAME>., ... & <NAME>. (2017). Attention is all you need. Advances in neural information processing systems, 30.<jupyter_code>def get_angles(pos: int, i: int, d_model: int) -> tf.Tensor:
    """Returns the angle in radians for a given position and timestamp"""
    angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
    return pos * angle_rates
def positional_encoding(position: int, d_model: int) -> tf.Tensor:
    """Returns the positional encoding for a given position and timestamp"""
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                            np.arange(d_model)[np.newaxis, :],
                            d_model)

    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

    pos_encode = angle_rads[np.newaxis, ...]

    return tf.cast(pos_encode, dtype=tf.float32)
n, d = 2048, 512
pos_encoding = positional_encoding(n, d)
print(pos_encoding.shape)
pos_encoding = pos_encoding[0]<jupyter_output>(1, 2048, 512)<jupyter_text>Masking Mask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. The mask indicates where pad value 0 is present: it outputs a 1 at those locations, and a 0 otherwise.<jupyter_code>def create_padding_mask(seq: tf.Tensor) -> tf.Tensor:
    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)

    # add extra dimensions to add the padding
    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)<jupyter_output><empty_output><jupyter_text>Just to check if it works<jupyter_code>x = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])
create_padding_mask(x)<jupyter_output><empty_output><jupyter_text>The look-ahead mask is used to mask the future tokens in a sequence. In other words, the mask indicates which entries should not be used.This means that to predict the third token, only the first and second token will be used. Similarly to predict the fourth token, only the first, second and the third tokens will be used and so on.<jupyter_code>def create_look_ahead_mask(size: int) -> tf.Tensor:
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)
exm = tf.random.uniform((1, 3))
print(exm)
temp = create_look_ahead_mask(exm.shape[1])
print(temp)
def create_masks(inp: tf.Tensor, tar: tf.Tensor, set_size: int):
        # Encoder padding mask
        enc_padding_mask = create_padding_mask(inp)

        # Used in the 2nd attention block in the decoder.
        # This padding mask is used to mask the encoder outputs.
        dec_padding_mask = create_padding_mask(inp)

        # Used in the 1st attention block in the decoder.
        # It is used to pad and mask future tokens in the input received by
        # the decoder.
        look_ahead_mask = create_look_ahead_mask(set_size)
        dec_target_padding_mask = create_padding_mask(tar)
        look_ahead_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)

        return enc_padding_mask, look_ahead_mask, dec_padding_mask<jupyter_output><empty_output><jupyter_text>Layers and blocks<jupyter_code>def scaled_dot_product_attention(q: tf.Tensor, k: tf.Tensor, v: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
    """Returns the attention weights.
    Args:
      q: query shape == (..., inp_len, depth)
      k: key shape == (..., inp_len, depth)
      v: value shape == (..., inp_len, depth_v)
      mask: Float tensor with shape broadcast-able
            to (..., inp_len, inp_len). Defaults to None.

    Returns:
      output, attention_weights
    """

    matmul_qk: tf.Tensor = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

    # scale matmul_qk
    dk: tf.Tensor = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits: tf.Tensor = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None:
        # noinspection PyTypeChecker
        scaled_attention_logits += (mask * -1e9)

    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output
class MultiHeadAttention(tf.keras.layers.Layer):
    """U can use the built-in tf.keras.layers.multihead_attention but is caused a bug for me"""
    def __init__(self, d_model: int, num_heads: int):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model

        assert d_model % self.num_heads == 0

        self.depth = d_model // self.num_heads

        self.wq = tf.keras.layers.Dense(d_model)
        self.wk = tf.keras.layers.Dense(d_model)
        self.wv = tf.keras.layers.Dense(d_model)

        self.dense = tf.keras.layers.Dense(d_model)

    def split_heads(self, x: tf.Tensor, batch_size) -> tf.Tensor:
        """Split the last dimension into (num_heads, depth).
        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
        """
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, v: tf.Tensor, k: tf.Tensor, q: tf.Tensor, mask: tf.Tensor) -> Tuple[tf.Tensor]:
        batch_size = tf.shape(q)[0]

        q = self.wq(q)  # (batch_size, seq_len, d_model)
        k = self.wk(k)  # (batch_size, seq_len, d_model)
        v = self.wv(v)  # (batch_size, seq_len, d_model)

        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

        # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
        # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
        scaled_attention = scaled_dot_product_attention(
            q, k, v, mask)

        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

        concat_attention = tf.reshape(scaled_attention,
                                      (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

        return output
def point_wise_feed_forward_network(d_model: int, dff: int) -> tf.keras.Model:
    return tf.keras.Sequential([
        tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
        tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
    ])
class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model: int, num_heads: int, dff: int, rate: float = 0.1):
        super(EncoderLayer, self).__init__()

        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = point_wise_feed_forward_network(d_model, dff)

        self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def call(self, x: tf.Tensor, training, mask: tf.Tensor) -> tf.Tensor:

        attn_output = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layer_norm1(x + attn_output)  # (batch_size, input_seq_len, d_model)

        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layer_norm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)

        return out2
        
class DecoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model: int, num_heads: int, dff: int, rate: float = 0.1):
        super(DecoderLayer, self).__init__()

        self.mha1 = MultiHeadAttention(d_model, num_heads)
        self.mha2 = MultiHeadAttention(d_model, num_heads)

        self.ffn = point_wise_feed_forward_network(d_model, dff)

        self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
        self.dropout3 = tf.keras.layers.Dropout(rate)

    def call(self, x: tf.Tensor, enc_output: tf.Tensor, training,
             look_ahead_mask: tf.Tensor, padding_mask: tf.Tensor) -> tf.Tensor:
        # enc_output.shape == (batch_size, input_seq_len, d_model)

        attn1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, set_size, d_model)
        attn1 = self.dropout1(attn1, training=training)
        out1 = self.layer_norm1(attn1 + x)

        attn2 = self.mha2(
            enc_output, enc_output, out1, padding_mask)  # (batch_size, set_size, d_model)
        attn2 = self.dropout2(attn2, training=training)
        out2 = self.layer_norm2(attn2 + out1)  # (batch_size, set_size, d_model)

        ffn_output = self.ffn(out2)  # (batch_size, set_size, d_model)
        ffn_output = self.dropout3(ffn_output, training=training)
        out3 = self.layer_norm3(ffn_output + out2)  # (batch_size, set_size, d_model)

        return out3
        
VOCAB_SIZE = 8192


class Encoder(tf.keras.layers.Layer):
    def __init__(self, num_layers: int, d_model: int, num_heads: int, dff: int, input_vocab_size: int,
                 maximum_position_encoding: int, rate=0.1):
        super(Encoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
        self.pos_encoding = positional_encoding(maximum_position_encoding,
                                                self.d_model)

        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate)
                           for _ in range(num_layers)]

        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self, x: tf.Tensor, training, mask: tf.Tensor) -> tf.Tensor:

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.enc_layers[i](x, training, mask)

        return x  # (batch_size, input_seq_len, d_model)
        
class Decoder(tf.keras.layers.Layer):
    def __init__(self, num_layers: int, d_model: int, num_heads: int, dff: int, vocab_size: int,
                 maximum_position_encoding: int, rate: float = 0.1):
        super(Decoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

        self.embedding = tf.keras.layers.Embedding(vocab_size, d_model)
        self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)

        self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate)
                           for _ in range(num_layers)]
        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self, x: tf.Tensor, enc_output: tf.Tensor, training: bool,
             look_ahead_mask: tf.Tensor, padding_mask: tf.Tensor) -> tf.Tensor:

        seq_len = tf.shape(x)[1]

        x = self.embedding(x)  # (batch_size, set_size, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        x += self.pos_encoding[:, :seq_len, :]

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x = self.dec_layers[i](x, enc_output, training,
                                                   look_ahead_mask, padding_mask)

        # x.shape == (batch_size, set_size, d_model)
        return x<jupyter_output><empty_output><jupyter_text>The full model<jupyter_code>class Transformer(tf.keras.Model):
    """The base architecture of my models in this project."""
    def __init__(self, set_size: int, batch_size: int, num_layers: int, d_model: int, num_heads: int, dff: int,
                 vocab_size: int, max_len: int, rate: float = 0.1):
        super().__init__()  # Initializes the base classes of this class.
        self.set_size: int = set_size
        self.batch_size: int = batch_size
        self.encoder = Encoder(num_layers, d_model, num_heads, dff,
                               vocab_size, max_len, rate)

        self.decoder = Decoder(num_layers, d_model, num_heads, dff,
                               vocab_size, max_len, rate)

        self.final_layer = tf.keras.layers.Dense(vocab_size)
        # The call function is compiled only when the model is created.        
        self.call_func = tf.function(func=self.call_not_graph,
                                     input_signature=[tf.TensorSpec(shape=[self.batch_size, None], dtype=tf.int32),
                                                      tf.TensorSpec(shape=[self.batch_size, self.set_size], dtype=tf.int32),
                                                      bool])

    def call(self, inputs: List[tf.Tensor], training: bool) -> tf.Tensor:
        # The method that gets called once the model is called
        # Keras models prefer if you pass all your inputs in the first argument
        inp, tar = inputs
        return self.call_func(inp, tar, training)

    def call_not_graph(self, inp: tf.Tensor, tar: tf.Tensor, training: bool) -> tf.Tensor:
        enc_padding_mask, look_ahead_mask, dec_padding_mask = create_masks(inp, tar, self.set_size)

        enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)

        # dec_output.shape == (batch_size, set_size, d_model)
        dec_output = self.decoder(tar, enc_output, training, look_ahead_mask, dec_padding_mask)

        final_output = self.final_layer(dec_output)
        return final_output
        
sample_transformer = Transformer(
    set_size=4, batch_size=BATCH_SIZE, num_layers=2, d_model=512, num_heads=8, dff=2048,
    vocab_size=VOCAB_SIZE,
    max_len=MAX_LENGTH)

temp_input = tf.random.uniform((64, 38), dtype=tf.int64, minval=0, maxval=200)
temp_target = tf.random.uniform((64, 4), dtype=tf.int64, minval=0, maxval=200)

fn_out = sample_transformer([temp_input, temp_target], training=False)

fn_out.shape  # (batch_size, tar_seq_len, target_vocab_size)
sample_transformer.summary()
del sample_transformer<jupyter_output><empty_output><jupyter_text>Training the model  Hyper-parameters<jupyter_code>set_size: int = 512
num_layers: int = 4
d_model: int = 128
dff: int = 512
num_heads: int = 8
learning_rate: float = 0.01
TAKE_TO_ACCOUNT: int = 512
dropout_rate: float = 0.1<jupyter_output><empty_output><jupyter_text>Optimizer<jupyter_code>optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)<jupyter_output><empty_output><jupyter_text>Loss Function<jupyter_code>loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')
@tf.function(input_signature=[
    tf.TensorSpec(shape=[BATCH_SIZE, set_size], dtype=tf.int32), 
    tf.TensorSpec(shape=[BATCH_SIZE, set_size], dtype=tf.int32)])
def loss_function(real: tf.Tensor, pred: tf.Tensor) -> tf.Tensor:
    """TODO: change to embedding distance
    returns a scalar tensor representing the loss"""
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask

    return tf.reduce_sum(loss_)/tf.reduce_sum(mask)<jupyter_output><empty_output><jupyter_text>Create the model<jupyter_code>model = Transformer(
    batch_size=BATCH_SIZE,
    set_size=set_size,
    num_layers=num_layers,
    d_model=d_model,
    num_heads=num_heads,
    dff=dff,
    vocab_size=VOCAB_SIZE,
    max_len=MAX_LENGTH,
    rate=dropout_rate)<jupyter_output><empty_output><jupyter_text>Train step function Since I can't use model.compile and model.fit and I dont want this function to be compiled every time I run it, I use tf.function to pre-compile the function into tf.graph<jupyter_code>@tf.function(input_signature=[tf.TensorSpec(shape=[BATCH_SIZE, None], dtype=tf.int32), tf.TensorSpec(shape=[BATCH_SIZE, set_size], dtype=tf.int32)])
def train_step(inp: tf.Tensor, outp: tf.Tensor) -> tf.Tensor:
    with tf.GradientTape() as tape:
        pred: tf.Tensor = model([inp, outp], training=True)
        loss_val: tf.Tensor = loss_function(real=outp, pred=pred)
    grads: tf.Tensor = tape.gradient(loss_val, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    return loss_val<jupyter_output><empty_output><jupyter_text>The training loop itself<jupyter_code>EPOCHS: int = 1
NUM_SETS: int = (MAX_LENGTH // set_size) - 1 # Because we dont predict the first set
per_epoch_train_loss: List[float] = []
for epoch in range(EPOCHS):
    per_batch_train_loss: List[float] = []
    for batch_num, batch in tqdm.tqdm(enumerate(train_set)):  # tqdm is a progress bar
        per_generation_loss: List[float] = []
        for i in range(NUM_SETS):
            # The input is of size set_size-TAKE_TO_ACCOUNT
            already_predicted: int = i*(set_size+1)
            start_from: int = max(0, already_predicted - TAKE_TO_ACCOUNT)
            inp: tf.Tensor = batch[:, start_from:(i+1)*set_size]
            outp: tf.Tensor = batch[:, (i+1)*set_size:(i+2)*set_size]
            loss_val: tf.Tensor = train_step(inp, outp)
            float_loss: float = float(loss_val)
            print(float_loss)
            per_generation_loss.append(float_loss)
        per_batch_train_loss.append(sum(per_generation_loss) / NUM_SETS)
    per_epoch_train_loss.append(sum(per_batch_train_loss) / len(per_generation_loss))<jupyter_output>0it [00:00, ?it/s]
{'fraction_non_alphanumeric': 0.11387595455392066, 'fraction_numerical': 0.014080834419817471, 'mean_word_length': 3.6430300933932895, 'pattern_counts': {'":': 0, '<': 92, '<?xml version=': 0, '>': 111, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '22', 'dir': 'jupyter-structured-clean-dedup', 'id': '118791', 'license': 'Apache-2.0', 'n_tokens_mistral': 9597, 'n_tokens_neox': 9239, 'n_words': 2161, 'path': 'final_project.ipynb', 'repo_name': 'yonikremer/final_project', 'size': '26876'}