Please complete this code and return only the missing parts without regenerating the hole code again, **code**: import tensorflow as tf

class Embedding(tf.keras.layers.Layer):
  def __init__(self, vocab_size, embedding_dim):
    super(Embedding, self).__init__()
    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)

  def call(self, inputs):
    return self.embedding(inputs)

class PositionalEncoding(tf.keras.layers.Layer):
  def __init__(self, embedding_dim, max_len=5000):
    super(PositionalEncoding, self).__init__()
    self.embedding_dim = embedding_dim
    self.pos_encoding = tf.keras.layers.Embedding(max_len, embedding_dim)

  def call(self, inputs):
    positions = tf.range(tf.shape(inputs)[1])[:, None]
    return inputs + self.pos_encoding(positions)

class DecoderLayer(tf.keras.layers.Layer):
  def __init__(self, embedding_dim, heads, ff_dim):
    super(DecoderLayer, self).__init__()
    self.self_attn = tf.keras.layers.MultiHeadAttention(heads, embedding_dim)
    self.ffn = tf.keras.Sequential([
      tf.keras.layers.Dense(ff_dim, activation="relu"),
      tf.keras.layers.Dense(embedding_dim)
    ])
    self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

  def call(self, inputs, training=False):
    attn_output, _ = self.self_attn(inputs, inputs, inputs, training=training)
    out1 = self.layer_norm1(inputs + attn_output)
    ffn_output = self.ffn(out1)
    out2 = self.layer_norm2(out1 + ffn_output)
    return out2

class Decoder(tf.keras.layers.Layer):
  def __init__(self, vocab_size, embedding_dim, num_layers, heads, ff_dim):
    super(Decoder, self).__init__()

    self.embedding = Embedding(vocab_size, embedding_dim)
    self.pos_encoding = PositionalEncoding(embedding_dim)
    self.decoder_layers = [DecoderLayer(embedding_dim, heads, ff_dim) for _ in range(num_layers)]
    self.final_layer = tf.keras.layers.Dense(vocab_size)

  def call(self, inputs, training=False):
    positions = tf.range(tf.shape(inputs)[1])[:, None]
    embedded_inputs = self.embedding(inputs)
    encoded_inputs = embedded_inputs + self.pos_encoding(positions)

    for layer in self.decoder_layers:
      encoded_inputs = layer(encoded_inputs, training=training)

    decoder_output = self.final_layer(encoded_inputs)
    return decoder_output
  
def load_data(self, path):
    pairs = []
    with open(path, "rb") as f:  # Open in binary mode
        for line in f:
            data = json.loads(line.decode("utf-8"))  # Decode binary data to JSON
            question, answer = self.decode_binary_data(data["text"], data["text"])  # Decode binary fields
            pairs.append((question.split(), answer.split()))
    return pairs

def decode_binary_data(self, question_data, answer_data):
    # Implement your binary decoding logic here
    # For example, if it's base64 encoded:
    question = base64.b64decode(question_data).decode("utf-8")
    answer = base64.b64decode(answer_data).decode("utf-8")
    return question, answer

def tokenize(self, integer_sequence):
    tokens = integer_sequence[:self.seq_len]
    if len(tokens) < self.seq_len:
       # Pad with appropriate binary token
        padding_token = self.vocab["<pad>"]  # Consider a binary padding token if needed
        tokens.extend([padding_token] * (self.seq_len - len(tokens)))
       # Add <eos> token only if appropriate for binary representation
    return tokens

def build_vocab(self, binary_tokens):
    # Create a vocabulary for unique binary tokens if needed
    return


# Define model parameters
vocab_size = 10000  # Adjust based on your vocabulary size
embedding_dim = 256
num_layers = 2  # Number of decoder layers
heads = 4  # Number of attention heads
ff_dim = 1024  # Feed-forward dimension

# Create decoder-only transformer model
model = Decoder(vocab_size, embedding_dim, num_layers, heads, ff_dim)

# Dataloading here...
dataset = QAJsonlDataset(path_to_binary_jsonl_file, seq_len=512)  # Adjust seq_len as needed
train_loader = DataLoader(dataset, batch_size=32)  # Adjust batch_size as needed

# Define optimizer and loss function (consider using sparse_categorical_crossentropy for large vocab)
optimizer = tf.keras.optimizers.Adam()
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Compile the model
model.compile(optimizer=optimizer, loss=loss_fn, metrics=['accuracy'])

# Train the model on your prepared training data
history = model.fit(
    x=train_inputs,
    y=train_outputs,
    epochs=10,  # Adjust number of epochs
    validation_data=(val_inputs, val_outputs)
)