transformer-template.py

import tensorflow as tf


class TransformerEncoder(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()

    def call(self, hidden_state):
        return hidden_state


if __name__ == "__main__":
    hidden_size = 256
    num_heads = 2
    intermediate_size = 768
    activation_function = "relu"

    input_node = tf.keras.Input((None, hidden_size), dtype=tf.float32)
    output_node = TransformerEncoder()(input_node)
    model = tf.keras.Model(input_node, output_node, name="model")

    tf.keras.callbacks.TensorBoard(log_dir="./logs").set_model(model)