Running this code: import tensorflow as tf
from keras.layers import TextVectorization, Embedding, LayerNormalization, Dropout, Dense, MultiHeadAttention, GlobalAveragePooling1D, Input
from keras.models import Sequential, Model
from keras import Sequential

# Standard imports for saving tokenizer configuration
import json

# Define some constants
MAX_TOKENS = 250  # Size of the vocabulary
EMBED_DIM = 64    # Embedding size
NUM_HEADS = 4     # Number of attention heads
FF_DIM = 64       # Hidden layer size in feed forward network
MAX_LEN = 40      # Max length of sequence
BATCH_SIZE = 64   # Training batch size

# Prepare your dataset text file
text_file_path = "dataset.txt"

# Define the text vectorization layer with padding
vectorize_layer = TextVectorization(
    max_tokens=MAX_TOKENS,
    output_mode="int",
    output_sequence_length=MAX_LEN,
    pad_to_max_tokens=True
)

# Load the text data and adapt the vectorization layer
with open(text_file_path, "r", encoding="utf-8") as f:
    lines = f.readlines()
    vectorize_layer.adapt(lines)

# Token and position embedding as a custom layer
class TokenAndPositionEmbedding(tf.keras.layers.Layer):
    def init(self, maxlen, vocab_size, embed_dim):
        super().init()
        self.token_emb = Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = Embedding(input_dim=maxlen, output_dim=embed_dim)

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

# Define the Transformer block as a function
def transformer_encoder(inputs, embed_dim, num_heads, ff_dim, rate=0.1):
    # Multi-headed self-attention
    attention_output = MultiHeadAttention(
        num_heads=num_heads, key_dim=embed_dim
    )(inputs, inputs)
    attention_output = Dropout(rate)(attention_output)
    attention_output = LayerNormalization(epsilon=1e-6)(inputs + attention_output)

    # Feed-forward layer
    ff_output = Sequential(
        [
            Dense(ff_dim, activation="relu"),
            Dense(embed_dim),
        ]
    )(attention_output)
    ff_output = Dropout(rate)(ff_output)
    ff_output = LayerNormalization(epsilon=1e-6)(attention_output + ff_output)
    return ff_output

# Create the Transformer model
def create_transformer_model(maxlen, vocab_size, embed_dim, num_heads, ff_dim):
    inputs = Input(shape=(maxlen,))
    embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)(inputs)
    transformer_block_output = transformer_encoder(embedding_layer, embed_dim, num_heads, ff_dim)
    outputs = GlobalAveragePooling1D()(transformer_block_output)
    outputs = Dropout(0.1)(outputs)
    outputs = Dense(vocab_size, activation="softmax")(outputs)  # Assuming next-token prediction
    model = Model(inputs=inputs, outputs=outputs)
    return model

transformer_model = create_transformer_model(MAX_LEN, MAX_TOKENS, EMBED_DIM, NUM_HEADS, FF_DIM)
transformer_model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")

# Prepare the dataset for training with next-token prediction
text_dataset = tf.data.TextLineDataset(text_file_path)
text_vectorized = vectorize_layer(text_dataset)
text_dataset = text_vectorized.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
text_labels = text_vectorized.map(lambda x: x[:, 1:])  # Shifted by one position for labels
text_inputs = text_vectorized.map(lambda x: x[:, :-1])  # Original sequence without the last token as input
dataset = tf.data.Dataset.zip((text_inputs, text_labels)).shuffle(1024).batch(BATCH_SIZE)

# Train the model
transformer_model.fit(dataset, epochs=10)

# Save the model
transformer_model.save("STM-Experiment")

# Save the tokenizer to a json file
tokenizer_json = vectorize_layer.get_vocabulary()
with open("tokenizer.json", "w", encoding="utf-8") as f:
    json.dump(tokenizer_json, f, ensure_ascii=False)

print("Model and tokenizer have been saved.")        leaves me with this error: ValueError: Exception encountered when calling layer 'text_vectorization' (type TextVectorization).

Attempt to convert a value (<TextLineDatasetV2 element_spec=TensorSpec(shape=(), dtype=tf.string, name=None)>) with an unsupported type (<class 'tensorflow.python.data.ops.readers.TextLineDatasetV2'>) to a Tensor.

Call arguments received by layer 'text_vectorization' (type TextVectorization):
  • inputs=<TextLineDatasetV2 element_spec=TensorSpec(shape=(), dtype=tf.string, name=None)>