Keras Transformers

# !pip install tensorflow_cpu==2.17.1
# !pip install matplotlib==3.9.2

import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import numpy as np
from keras.models import Model
from keras.layers import Input, LSTM, Dense, Embedding, Dropout
from tensorflow.keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from tensorflow.keras import backend as K
from keras.layers import Layer
import warnings
warnings.simplefilter('ignore', FutureWarning)


# Sample parallel sentences (English -> Spanish)
input_texts = [
    "Hello.", "How are you?", "I am learning machine translation.", "What is your name?", "I love programming."
]
target_texts = [
    "Hola.", "¿Cómo estás?", "Estoy aprendiendo traducción automática.", "¿Cuál es tu nombre?", "Me encanta programar."
]

target_texts = ["startseq " + x + " endseq" for x in target_texts]


# Tokenization - covert words into numerical sequences
input_tokenizer = Tokenizer()
input_tokenizer.fit_on_texts(input_texts)
input_sequences = input_tokenizer.texts_to_sequences(input_texts)

output_tokenizer = Tokenizer()
output_tokenizer.fit_on_texts(target_texts)
output_sequences = output_tokenizer.texts_to_sequences(target_texts)

input_vocab_size = len(input_tokenizer.word_index) + 1
output_vocab_size = len(output_tokenizer.word_index) + 1

# add padding to make sure that all sequences have the same length
max_input_length = max([len(seq) for seq in input_sequences])
max_output_length = max([len(seq) for seq in output_sequences])
input_sequences = pad_sequences(input_sequences, maxlen=max_input_length, padding='post')
output_sequences = pad_sequences(output_sequences, maxlen=max_output_length, padding='post')


#############################################################
# Self-Attention Class
#############################################################
from tensorflow.keras.layers import Layer
from tensorflow.keras import backend as K

class SelfAttention(Layer):
    def __init__(self, **kwargs):
        super(SelfAttention, self).__init__(**kwargs)

    def build(self, input_shape):
        # input_shape is a list: [q_shape, k_shape, v_shape]
        feature_dim = input_shape[0][-1]
        initializer_type = 'glorot_uniform'
        
        # trainable weight matrices, with shape (feature_dim, feature_dim) 
        # which means they transform input into Q, K, and V. 
        self.Wq = self.add_weight(
            shape=(feature_dim, feature_dim),
            initializer=initializer_type,
            trainable=True,
            name='Wq'
        )

        self.Wk = self.add_weight(
            shape=(feature_dim, feature_dim),
            initializer=initializer_type,
            trainable=True,
            name='Wk'
        )

        self.Wv = self.add_weight(
            shape=(feature_dim, feature_dim),
            initializer=initializer_type,
            trainable=True,
            name='Wv'
        )

        super(SelfAttention, self).build(input_shape)

    def call(self, inputs):
        # Expect list: [query, key, value]
        q, k, v = inputs
        
        # Compute Q, K, V by multiplying inputs with weigh matrices
        q = K.dot(q, self.Wq)
        k = K.dot(k, self.Wk)
        v = K.dot(v, self.Wv)

        # Scaled dot-product attention (resulting in a (batch_size, seq_len, seq_len) matrix)
        # axes=[2, 2] means it multiplies along the last dimension of both tensors.
        # The result is a similarity score showing how much each query matches each key.
        # “How related is this query vector to this key vector?”
        scores = K.batch_dot(q, k, axes=[2, 2])
        # gets the size of the last dimension of k, usually the key dimension, often called d_k
        # cast converts it to a floating-point number 
        # so math like square root and division works correctly
        dk = K.cast(K.shape(k)[-1], dtype=K.floatx())
        # scaling keeps the numbers from getting too large
        scores = scores / K.sqrt(dk)
        
        # normalize attention scores 
        # softmax(..., axis=-1) converts those scores into numbers between 0 and 1.
        # Along the last axis, they add up to 1.
        # So now each query has a distribution like:
        # 0.7 attention on token 1
        # 0.2 attention on token 2
        # 0.1 attention on token 3
        attention_weights = K.softmax(scores, axis=-1)
        
        # This line uses the attention weights to take a weighted combination of the value vectors
        # Tokens with higher attention contribute more to the output
        output = K.batch_dot(attention_weights, v)

        return output


# Encoder / Decoder
# Input: defines model inputs
# Embedding: turns token IDs into dense vectors
# LSTM: recurrent layer for sequence modeling
# Concatenate: joins tensors together
# Dense: final classifier layer
# Model: builds the full Keras model
from tensorflow.keras.layers import AdditiveAttention, Concatenate, Dense, Embedding, Input, LSTM
from tensorflow.keras.models import Model

# Encoder
# Each element is usually an integer token ID.
encoder_inputs = Input(shape=(max_input_length,))
# This maps each input token ID to a 256-dimensional vector.
encoder_embedding = Embedding(input_vocab_size, 256)(encoder_inputs)
# encoder_outputs = all intermediate encoder hidden states
# state_h, state_c = the final encoder state, used to initialize the decoder
# return_sequences=True: return the hidden state at every time step
# return_state=True: also return the final hidden and cell states
encoder_lstm = LSTM(256, return_sequences=True, return_state=True)
encoder_outputs, state_h, state_c = encoder_lstm(encoder_embedding)
encoder_states = [state_h, state_c]
 
# Decoder
# Why max_output_length - 1? Because in training, the decoder usually gets the target sequence shifted right.
decoder_inputs = Input(shape=(max_output_length - 1,))
# This turns decoder token IDs into 256-dimensional vectors.
decoder_embedding = Embedding(output_vocab_size, 256)(decoder_inputs)
# decoder_outputs: shape (batch_size, max_output_length - 1, 256)
decoder_lstm = LSTM(256, return_sequences=True, return_state=True)
decoder_outputs, _, _ = decoder_lstm(decoder_embedding, initial_state=encoder_states)
 
# Attention: decoder attends to encoder outputs
self_attention = SelfAttention()
attention_output = self_attention(
    [decoder_outputs, encoder_outputs, encoder_outputs]
)

 
# Combine decoder outputs with attention context along the last dimension.
decoder_concat = Concatenate(axis=-1)([decoder_outputs, attention_output])
 
# Final Dense layer, This maps each decoder time step to a probability distribution over the output vocabulary.
decoder_dense = Dense(output_vocab_size, activation='softmax')
decoder_outputs = decoder_dense(decoder_concat)
 
# Full Model
model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
 
# Summary
model.summary()


# train the model 
history_glorot_adam = model.fit([input_sequences, decoder_input_data], decoder_output_data, epochs=100, batch_size=16)

# Plotting training loss
import matplotlib.pyplot as plt
plt.plot(history_glorot_adam.history['loss'])
plt.title('Training Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.show()