Transformers TF

import tensorflow as tf
from tensorflow import keras
import numpy as np

_course_titles = tf.constant([
    "Regional leaders agree on a public health plan",
    "The home team wins the basketball final",
    "Markets rise after a strong business report",
    "Researchers release an efficient intelligence model",
    "Diplomats begin a new round of peace talks",
    "The tennis champion returns after an injury",
    "Retail sales improve with consumer confidence",
    "A software update improves privacy",
])
_course_descriptions = tf.constant([
    "World policy and international cooperation",
    "Sports results and player development",
    "Business revenue and economic indicators",
    "Science and technology research",
    "World diplomacy and humanitarian aid",
    "Sports training and competition",
    "Business investment and employment",
    "Science security and technology",
])
_course_labels = tf.constant([0, 1, 2, 3, 0, 1, 2, 3], dtype=tf.int64)
_course_records = {"title": _course_titles, "description": _course_descriptions, "label": _course_labels}
ds_train = tf.data.Dataset.from_tensor_slices(_course_records).take(6)
ds_test = tf.data.Dataset.from_tensor_slices(_course_records).skip(6)

def extract_text(x):
    return x['title'] + ' ' + x['description']

def tupelize(x):
    return extract_text(x), x['label']

class TokenAndPositionEmbedding(keras.layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super(TokenAndPositionEmbedding, self).__init__()
        self.token_emb = keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = keras.layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
        self.maxlen = maxlen

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

class TransformerBlock(keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim, name='attn')
        self.ffn = keras.Sequential(
            [keras.layers.Dense(ff_dim, activation="relu"), keras.layers.Dense(embed_dim),]
        )
        self.layernorm1 = keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = keras.layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = keras.layers.Dropout(rate)
        self.dropout2 = keras.layers.Dropout(rate)

    def call(self, inputs, training=None):
        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)

embed_dim = 32  # Embedding size for each token
num_heads = 2  # Number of attention heads
ff_dim = 32  # Hidden layer size in feed forward network inside transformer
maxlen = 256
vocab_size = 20000

model = keras.models.Sequential([
    keras.Input(shape=(), dtype=tf.string),
    keras.layers.TextVectorization(max_tokens=vocab_size, output_sequence_length=maxlen),
    TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim),
    TransformerBlock(embed_dim, num_heads, ff_dim),
    keras.layers.GlobalAveragePooling1D(),
    keras.layers.Dropout(0.1),
    keras.layers.Dense(20, activation="relu"),
    keras.layers.Dropout(0.1),
    keras.layers.Dense(4, activation="softmax")
])

model.summary()

print('Training tokenizer')
model.layers[0].adapt(ds_train.map(extract_text))
model.compile(loss='sparse_categorical_crossentropy',metrics=['acc'], optimizer='adam')
model.fit(ds_train.map(tupelize).batch(128),validation_data=ds_test.map(tupelize).batch(128))

vectorizer = keras.layers.TextVectorization(max_tokens=2000, output_sequence_length=32)
vectorizer.adapt(ds_train.map(extract_text))

vectorizer(['I love transformers'])

bert = keras.Sequential([
    keras.layers.Embedding(2000, 32),
    keras.layers.GlobalAveragePooling1D(),
], name="course_local_encoder")

z = bert(vectorizer(["I love transformers"]))
print(f"Local encoder output shape -> {z.shape}")

inp = keras.Input(shape=(), dtype=tf.string)
x = vectorizer(inp)
x = bert(x)
x = keras.layers.Dropout(0.1)(x)
out = keras.layers.Dense(4, activation="softmax")(x)
model = keras.models.Model(inp, out)
bert.trainable = False
model.summary()

model.compile(loss='sparse_categorical_crossentropy',metrics=['acc'], optimizer='adam')
model.fit(ds_train.map(tupelize).batch(128),validation_data=ds_test.map(tupelize).batch(128))

bert.trainable = True
epochs = 2
model.compile(loss="sparse_categorical_crossentropy", metrics=["acc"], optimizer=keras.optimizers.Adam(learning_rate=3e-4))
model.fit(ds_train.map(tupelize).batch(4), validation_data=ds_test.map(tupelize).batch(4), epochs=epochs)

tokenizer = vectorizer
MAX_SEQ_LEN = 32
PAD_INDEX = 0
UNK_INDEX = 1

tokenizer(["TensorFlow is a great framework for NLP"])

tokenizer(["Hello, there"])

def process_fn(x):
    text = tf.expand_dims(extract_text(x), axis=0)
    return vectorizer(text)[0], x["label"]

model = keras.Sequential([
    keras.layers.Embedding(2000, 32),
    keras.layers.GlobalAveragePooling1D(),
    keras.layers.Dense(4, activation="softmax"),
], name="course_local_sequence_classifier")

model.summary()

model.layers[0].trainable = False
model.summary()

model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["acc"])
model.fit(ds_train.map(process_fn).batch(4), validation_data=ds_test.map(process_fn).batch(4), epochs=2)