Notes

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Embedding, LayerNormalization, Add

# # Simulated tokenized data (replace with your actual data)
# input_texts = ["What is a cat?", "Tell me about dogs."]
# output_texts = ["A cat is a domesticated carnivorous mammal.", "Dogs are loyal companions."]


# Tokenize input and output texts
tokenizer = tf.keras.layers.TextVectorization(max_tokens=1000, output_sequence_length=20)
tokenizer.adapt(df['question'] + df['answer'])
train_input_ids = tokenizer(df['question'])
train_output_ids = tokenizer(df['answer'])

# Constants
vocab_size = len(tokenizer.get_vocabulary())
max_seq_length = 20
embedding_dim = 256
num_transformer_layers = 4

# Custom attention mask layer
class AttentionMaskLayer(tf.keras.layers.Layer):
def __init__(self, **kwargs):
super(AttentionMaskLayer, self).__init__(**kwargs)

def call(self, inputs):
input_ids, attention_mask = inputs
return tf.math.logical_not(tf.math.equal(attention_mask,0))

# Input layer
inputs = Input(shape=(max_seq_length,), dtype=tf.int32)
attention_mask = tf.cast(inputs != 0, dtype=tf.int32) # Assuming 0 is the padding token index

# Embedding layer
embedding = Embedding(input_dim=vocab_size, output_dim=embedding_dim)(inputs)

# Transformer layers
for _ in range(num_transformer_layers):
attention_mask_layer = AttentionMaskLayer()([inputs, attention_mask])
attention = tf.keras.layers.Attention(use_scale=True)([embedding, embedding], mask=[attention_mask_layer, attention_mask_layer])
add_attention = Add()([embedding, attention])
normalized_attention = LayerNormalization()(add_attention)
feed_forward = Dense(units=embedding_dim, activation='relu')(normalized_attention)
add_feed_forward = Add()([normalized_attention, feed_forward])
normalized_feed_forward = LayerNormalization()(add_feed_forward)
embedding = normalized_feed_forward

# Output layer
outputs = Dense(units=vocab_size, activation='softmax')(embedding)

# Build the model
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# Compile the model
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
model.compile(optimizer=optimizer, loss=loss_fn)

# Train the model
model.fit(train_input_ids, train_output_ids, epochs=500)