After running this python 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):
    # Reshape the input tensor
    inputs = tf.expand_dims(inputs, axis=2)

    # 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)
    
    # Reshape the output tensor
    ff_output = tf.squeeze(ff_output, axis=2)
    
    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
def prepare_dataset(text_data, vectorize_layer, batch_size):
    text_dataset = text_data.map(vectorize_layer).batch(batch_size)
    text_labels = text_dataset.map(lambda x: x[:, 1:])  # Shifted by one position for labels
    text_inputs = text_dataset.map(lambda x: x[:, :-1])  # Original sequence without the last token as input
    dataset = tf.data.Dataset.zip((text_inputs, text_labels))
    dataset = dataset.shuffle(buffer_size=1024).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return dataset

text_data = tf.data.TextLineDataset(text_file_path)
dataset = prepare_dataset(text_data, vectorize_layer, 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.")    i get this error: Traceback (most recent call last):
  File "c:\Users\Dell-PC\Desktop\Projets\Simple-But-Functional\main.py", line 101, in <module>
    transformer_model.fit(dataset, epochs=10)
  File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\src\utils\traceback_utils.py", line 70, in error_handler  
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\Dell-PC\AppData\Local\Temp\__autograph_generated_fileyw_jno0t.py", line 15, in tf__train_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
ValueError: in user code:

    File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\src\engine\training.py", line 1338, in train_function  *
        return step_function(self, iterator)
    File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\src\engine\training.py", line 1322, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\src\engine\training.py", line 1303, in run_step  **
        outputs = model.train_step(data)
    File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\src\engine\training.py", line 1080, in train_step
        y_pred = self(x, training=True)
    File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\keras\src\utils\traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None

    ValueError: Exception encountered when calling layer 'query' (type EinsumDense).

    Shape must be rank 4 but is rank 5
         for 0th input and equation: abcd,def->abcef for '{{node model/multi_head_attention/query/einsum/Einsum}} = Einsum[N=2, T=DT_FLOAT, equation="abcd,def->abcef"](model/tf.expand_dims/ExpandDims, model/multi_head_attention/query/einsum/Einsum/ReadVariableOp)' with input shapes: [?,?,1,?,64], [64,4,64].

    Call arguments received by layer 'query' (type EinsumDense):
      • inputs=tf.Tensor(shape=(None, None, 1, None, 64), dtype=float32)