Please fix my code, especially load_dataset function, add a section to the GUI to give feedback messages on action that did or did not go successfully with red and green text, here is the code: import jax
import jax.numpy as jnp
import jax.random as random
from flax import linen as nn
from flax.training import train_state
import optax
import matplotlib.pyplot as plt
from tqdm import tqdm
import tkinter as tk
from tkinter import filedialog, ttk
import json
from tkinter import messagebox
from transformers import AutoTokenizer, AutoModel
import traceback
import torch

# Define standard transformer components
class MultiHeadSelfAttention(nn.Module):
    d_model: int
    num_heads: int

    def setup(self):
        assert self.d_model % self.num_heads == 0
        self.depth = self.d_model // self.num_heads

        self.wq = nn.Dense(self.d_model)
        self.wk = nn.Dense(self.d_model)
        self.wv = nn.Dense(self.d_model)
        self.dense = nn.Dense(self.d_model)

    def split_heads(self, x, batch_size):
        x = x.reshape((batch_size, -1, self.num_heads, self.depth))
        return x.transpose((0, 2, 1, 3))

    def __call__(self, x):
        batch_size = x.shape[0]
        q = self.split_heads(self.wq(x), batch_size)
        k = self.split_heads(self.wk(x), batch_size)
        v = self.split_heads(self.wv(x), batch_size)

        matmul_qk = jnp.matmul(q, k.transpose((0, 1, 3, 2)))

        dk = jnp.float32(k.shape[-1])
        scaled_attention_logits = matmul_qk / jnp.sqrt(dk)
        attention_weights = nn.softmax(scaled_attention_logits, axis=-1)

        output = jnp.matmul(attention_weights, v)
        output = output.transpose((0, 2, 1, 3)).reshape((batch_size, -1, self.d_model))
        return self.dense(output)

class StandardTransformerBlock(nn.Module):
    d_model: int
    num_heads: int
    dff: int

    def setup(self):
        self.mha = MultiHeadSelfAttention(self.d_model, self.num_heads)
        self.ffn = nn.Sequential([
            nn.Dense(self.dff),
            nn.relu,
            nn.Dense(self.d_model)
        ])

        self.layernorm1 = nn.LayerNorm()
        self.layernorm2 = nn.LayerNorm()

    def __call__(self, x):
        attn_output = self.mha(x)
        out1 = self.layernorm1(x + attn_output)
        ffn_output = self.ffn(out1)
        return self.layernorm2(out1 + ffn_output)

class StandardTransformer(nn.Module):
    d_model: int
    num_heads: int
    num_layers: int
    dff: int
    num_classes: int

    def setup(self):
        self.layers = [StandardTransformerBlock(self.d_model, self.num_heads, self.dff) for _ in range(self.num_layers)]
        self.output_layer = nn.Dense(self.num_classes)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        x = x.mean(axis=1)  # Reduce the sequence dimension
        return self.output_layer(x)

# Define the SaganTransformer components
class ScaledDotProductAttention(nn.Module):
    d_k: int
    d_model: int  # Add d_model to the constructor

    def setup(self):
        self.query_dense = nn.Dense(self.d_k)
        self.key_dense = nn.Dense(self.d_k)
        self.value_dense = nn.Dense(self.d_k)
        self.output_dense = nn.Dense(self.d_model)  # Change to project back to d_model

    def __call__(self, query, key, value):
        q = self.query_dense(query)
        k = self.key_dense(key)
        v = self.value_dense(value)

        scores = jnp.matmul(q, k.transpose((0, 2, 1))) / jnp.sqrt(self.d_k)
        weights = nn.softmax(scores, axis=-1)
        output = jnp.matmul(weights, v)

        return self.output_dense(output)  # Ensure the output dimension is d_model

class SaganTransformerBlock(nn.Module):
    d_model: int
    n_head: int

    def setup(self):
        self.attention = ScaledDotProductAttention(self.d_model // self.n_head, self.d_model)  # Pass d_model
        self.dense = nn.Dense(self.d_model)
        self.layernorm1 = nn.LayerNorm()
        self.layernorm2 = nn.LayerNorm()
        self.feed_forward = nn.Sequential([nn.Dense(self.d_model * 4), nn.relu, nn.Dense(self.d_model)])

    def __call__(self, x):
        attended = self.attention(x, x, x)
        x = self.layernorm1(x + attended)
        feed_forward_output = self.feed_forward(x)
        x = self.layernorm2(x + feed_forward_output)
        return x

class SaganTransformer(nn.Module):
    d_model: int
    n_head: int
    n_layer: int
    num_classes: int

    def setup(self):
        self.layers = [SaganTransformerBlock(self.d_model, self.n_head) for _ in range(self.n_layer)]
        self.output_layer = nn.Dense(self.num_classes)

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        x = x.mean(axis=1)  # Reduce the sequence dimension
        return self.output_layer(x)

# Utility functions for training
def train_step(state, batch):
    def loss_fn(params):
        logits = state.apply_fn({'params': params}, batch['inputs'])
        loss = jnp.mean(optax.softmax_cross_entropy(logits, batch['labels']))
        return loss, logits

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, logits), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    accuracy = jnp.mean(jnp.argmax(logits, -1) == jnp.argmax(batch['labels'], -1))
    return state, loss, accuracy

def create_train_state(rng, model, learning_rate, n_layer, num_classes):
    # Initialize the model (this returns the module state, including params)
    init_state = model.init(rng, jnp.ones([1, 10, model.d_model])) 

    # Extract the parameters from the init_state dictionary
    params = init_state['params']

    tx = optax.adam(learning_rate)
    return train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

def train_and_evaluate(train_data, test_data, model_cls, d_model, num_heads, num_layers, num_classes, epochs, learning_rate=0.001):
    rng = jax.random.PRNGKey(0)
    
    if model_cls == StandardTransformer:
        model = model_cls(d_model=d_model, num_heads=num_heads, num_layers=num_layers, dff=d_model, num_classes=num_classes)
    else:
        model = model_cls(d_model=d_model, n_head=num_heads, n_layer=num_layers, num_classes=num_classes)

    state = create_train_state(rng, model, learning_rate, num_layers, num_classes)

    train_loss_history = []
    train_accuracy_history = []

    # Use tqdm to track the progress of epochs
    for epoch in tqdm(range(epochs)):
        epoch_loss = 0
        epoch_accuracy = 0
        for batch in train_data:
            inputs, labels = batch
            batch = {'inputs': inputs, 'labels': labels}
            state, loss, accuracy = train_step(state, batch)
            epoch_loss += loss
            epoch_accuracy += accuracy
        train_loss_history.append(epoch_loss / len(train_data))
        train_accuracy_history.append(epoch_accuracy / len(train_data))

    return train_loss_history, train_accuracy_history, state

# Simulated recurrence task dataset
def generate_data(num_samples, sequence_length, d_model, num_classes):
    x = jax.random.normal(random.PRNGKey(0), (num_samples, sequence_length, d_model))
    y = jax.random.randint(random.PRNGKey(1), (num_samples,), 0, num_classes)
    y_one_hot = jax.nn.one_hot(y, num_classes)
    return x, y_one_hot

def abbreviate_number(num):
    for unit in ['','K','M','B']:
        if abs(num) < 1000.0:
            return f"{num:.1f}{unit}"
        num /= 1000.0
    return f"{num:.1f}T"

# Main comparison
def main():
    num_samples = 2000
    sequence_length = 8
    d_model = 32
    num_heads = 4
    num_layers = 8
    epochs = 3
    num_classes = 10

    x, y = generate_data(num_samples, sequence_length, d_model, num_classes)
    train_data = [(x[i:i+32], y[i:i+32]) for i in range(0, num_samples, 32)]
    test_data = train_data  # For simplicity, using the same data for test

    # Train and evaluate both models
    sagan_loss, sagan_acc, sagan_state = train_and_evaluate(
        train_data,
        test_data,
        SaganTransformer,
        d_model,
        num_heads,
        num_layers,
        num_classes,
        epochs
    )
    standard_loss, standard_acc, standard_state = train_and_evaluate(
        train_data,
        test_data,
        StandardTransformer,
        d_model,
        num_heads,  # Use 'num_heads' here
        num_layers,
        num_classes,
        epochs
    )

    # Print number of parameters
    def count_params(params):
        return sum(p.size for p in jax.tree_leaves(params))
    
    sagan_p=count_params(sagan_state.params)
    trans_p=count_params(standard_state.params)

    # Plot loss over epochs
    plt.figure(figsize=(12, 6))
    plt.plot(sagan_loss, label=f'Sagan Transformer ({abbreviate_number(sagan_p)} params)')  # Add parameter count to label
    plt.plot(standard_loss, label=f'Standard Transformer ({abbreviate_number(trans_p)} params)')  # Add parameter count to label
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Loss over Epochs')
    plt.legend()
    plt.show()

    print(f"Sagan Transformer Parameters: {count_params(sagan_state.params)}")
    print(f"Standard Transformer Parameters: {count_params(standard_state.params)}")

# --- Data Loading and Key Selection Functions ---

def load_dataset(file_path):
    """Loads and analyzes the dataset from a JSONL file."""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            # Read the first few lines to analyze keys
            lines = [json.loads(line) for i, line in enumerate(f) if i < 5]
            keys = set()
            for line in lines:
                keys.update(line.keys())

        # Identify input and output keys based on user selection
        input_key, output_key = get_input_output_keys(keys)

        # Create data lists
        inputs = []
        outputs = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                data = json.loads(line)
                inputs.append(data[input_key])
                outputs.append(data[output_key])

        # Tokenize and embed text if necessary
        if isinstance(inputs[0], str):
            print("Starting text tokenization and embedding...")
            tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
            embedding_model = AutoModel.from_pretrained("bert-base-uncased")

            # Find the maximum sequence length
            max_length = max(len(tokenizer.encode(text, add_special_tokens=True)) for text in inputs)

            # Batch the input_ids and pad them to the maximum length
            batch_size = 16
            batched_inputs = []
            batched_attention_masks = []
            for i in range(0, len(inputs), batch_size):
                batch = inputs[i:i + batch_size]
                mask = outputs[i:i + batch_size]  # Assuming mask is available
                if batch and mask:
                    # Tokenize each batch
                    batch_tokenized = [tokenizer.encode(text, add_special_tokens=True, max_length=max_length, truncation=True) for text in batch]
                    # Pad to the maximum length
                    batch_tokenized = [tokenized_text + [tokenizer.pad_token_id] * (max_length - len(tokenized_text)) for tokenized_text in batch_tokenized]

                    # Convert to JAX array and append to the batch list
                    batched_inputs.append(jnp.array(batch_tokenized))
                    # Create attention mask for the batch
                    attention_masks = [[1] * len(tokenized_text) + [0] * (max_length - len(tokenized_text)) for tokenized_text in batch_tokenized]
                    batched_attention_masks.append(jnp.array(attention_masks))

                else:
                    print("Empty batch or mask, skipping tokenization.")

            if batched_inputs and batched_attention_masks:
                # Embed the batched input_ids
                inputs = []
                for batch, mask in zip(batched_inputs, batched_attention_masks):
                    # Convert JAX arrays to torch tensors
                    batch_tensor = torch.tensor(batch)
                    mask_tensor = torch.tensor(mask)

                    # Embed the batch using the embedding model
                    embedded_batch = embedding_model(input_ids=batch_tensor, attention_mask=mask_tensor).pooler_output
                    # Convert back to JAX array
                    inputs.append(jnp.array(embedded_batch.detach().numpy()))

            # Convert to JAX arrays
            inputs = jnp.array(inputs)
            outputs = jnp.array(outputs)

            print("Dataset loading and processing complete.")
            return inputs, outputs

    except FileNotFoundError:
        messagebox.showerror("Error", f"File not found: {file_path}")
        print(f"File not found: {file_path}")
        return None, None
    except json.JSONDecodeError:
        messagebox.showerror("Error", "Invalid JSONL file format.")
        print("Invalid JSONL file format.")
        return None, None
    except KeyError as e:
        messagebox.showerror("Error", f"Key '{e}' not found in the dataset.")
        print(f"Key '{e}' not found in the dataset.")
        return None, None
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while loading the dataset: {e}")
        print(f"An error occurred while loading the dataset: {e}")
        traceback.print_exc()  # Print the full traceback to the terminal
        return None, None



def get_input_output_keys(keys):
    """Allows the user to choose input and output keys from the dataset."""
    # Create a separate window for key selection
    key_selection_window = tk.Toplevel()  # No need to pass self.master here
    key_selection_window.title("Select Input and Output Keys")

    # Create dropdown menus for input and output keys
    input_key_var = tk.StringVar(key_selection_window)
    input_key_var.set(list(keys)[0])  # Default value
    input_key_menu = ttk.Combobox(key_selection_window, textvariable=input_key_var, values=list(keys))
    input_key_menu.pack()

    output_key_var = tk.StringVar(key_selection_window)
    output_key_var.set(list(keys)[1])  # Default value
    output_key_menu = ttk.Combobox(key_selection_window, textvariable=output_key_var, values=list(keys))
    output_key_menu.pack()

    # Create a button to confirm the selection
    def confirm_selection():
        input_key = input_key_var.get()
        output_key = output_key_var.get()
        key_selection_window.destroy()
        return input_key, output_key

    confirm_button = tk.Button(key_selection_window, text="Confirm", command=confirm_selection)
    confirm_button.pack()

    # Wait for the user to confirm the selection
    key_selection_window.wait_window()
    return confirm_selection()

# --- Model Saving Function ---

def save_model(state):
    """Saves the trained model to a file."""
    file_path = filedialog.asksaveasfilename(defaultextension=".ckpt", filetypes=(("Checkpoint files", "*.ckpt"), ("all files", "*.*")))
    if file_path:
        # Save model state to file (replace with your actual saving logic)
        # Example using JAX's save function (adjust based on your specific saving method)
        jax.experimental.save(file_path, state.params)

        messagebox.showinfo("Model Saved", f"Model saved successfully to: {file_path}")


# --- GUI Implementation ---

class TransformerGUI:
    def __init__(self, master):
        self.master = master
        master.title("Transformer Training GUI")

        # Create frames
        self.dataset_frame = tk.Frame(master)
        self.model_frame = tk.Frame(master)
        self.training_frame = tk.Frame(master)

        # Dataset frame
        self.dataset_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.dataset_frame, text="Dataset Settings").pack()

        # Dataset File
        self.dataset_file_label = tk.Label(self.dataset_frame, text="Dataset File:")
        self.dataset_file_label.pack()

        self.dataset_file_entry = tk.Entry(self.dataset_frame, width=50)
        self.dataset_file_entry.pack()

        self.dataset_file_button = tk.Button(self.dataset_frame, text="Browse", command=self.browse_dataset_file)
        self.dataset_file_button.pack()

        # Preview Dataset
        self.preview_button = tk.Button(self.dataset_frame, text="Preview Dataset", command=self.preview_dataset)
        self.preview_button.pack()

        # Model frame
        self.model_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.model_frame, text="Model Settings").pack()

        # Model Type
        self.model_type_label = tk.Label(self.model_frame, text="Model Type:")
        self.model_type_label.pack()

        self.model_type_var = tk.StringVar(self.model_frame)
        self.model_type_var.set("Standard Transformer")  # Default value

        self.model_type_options = ["Standard Transformer", "Sagan Transformer"]
        self.model_type_menu = ttk.Combobox(self.model_frame, textvariable=self.model_type_var, values=self.model_type_options)
        self.model_type_menu.pack()

        # D Model
        self.d_model_label = tk.Label(self.model_frame, text="D Model:")
        self.d_model_label.pack()

        self.d_model_entry = tk.Entry(self.model_frame, width=10)
        self.d_model_entry.insert(0, "32")
        self.d_model_entry.pack()

        # Num Heads
        self.num_heads_label = tk.Label(self.model_frame, text="Num Heads:")
        self.num_heads_label.pack()

        self.num_heads_entry = tk.Entry(self.model_frame, width=10)
        self.num_heads_entry.insert(0, "4")
        self.num_heads_entry.pack()

        # Num Layers
        self.num_layers_label = tk.Label(self.model_frame, text="Num Layers:")
        self.num_layers_label.pack()

        self.num_layers_entry = tk.Entry(self.model_frame, width=10)
        self.num_layers_entry.insert(0, "8")
        self.num_layers_entry.pack()

        # Num Classes
        self.num_classes_label = tk.Label(self.model_frame, text="Num Classes:")
        self.num_classes_label.pack()

        self.num_classes_entry = tk.Entry(self.model_frame, width=10)
        self.num_classes_entry.insert(0, "10")
        self.num_classes_entry.pack()

        # Training frame
        self.training_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        tk.Label(self.training_frame, text="Training Settings").pack()

        # Epochs
        self.epochs_label = tk.Label(self.training_frame, text="Epochs:")
        self.epochs_label.pack()

        self.epochs_entry = tk.Entry(self.training_frame, width=10)
        self.epochs_entry.insert(0, "3")
        self.epochs_entry.pack()

        # Learning Rate
        self.learning_rate_label = tk.Label(self.training_frame, text="Learning Rate:")
        self.learning_rate_label.pack()

        self.learning_rate_entry = tk.Entry(self.training_frame, width=10)
        self.learning_rate_entry.insert(0, "0.001")
        self.learning_rate_entry.pack()

        # Start Training Button
        self.start_training_button = tk.Button(self.training_frame, text="Start Training", command=self.start_training)
        self.start_training_button.pack()

        # Save Model Button
        self.save_model_button = tk.Button(self.training_frame, text="Save Model", command=self.save_model)
        self.save_model_button.pack()

    def browse_dataset_file(self):
        """Opens a file dialog to select the dataset file."""
        file_path = filedialog.askopenfilename(initialdir="/", title="Select Dataset File",
                                              filetypes=(("JSONL files", "*.jsonl"), ("all files", "*.*")))
        self.dataset_file_entry.delete(0, tk.END)
        self.dataset_file_entry.insert(0, file_path)

    def preview_dataset(self):
        """Opens a new window to preview the selected dataset file."""
        file_path = self.dataset_file_entry.get()
        if not file_path:
            self.show_error("Error", "Please select a dataset file.")
            return

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                # Read the first few lines to analyze keys
                lines = [json.loads(line) for i, line in enumerate(f) if i < 5]
                keys = set()
                for line in lines:
                    keys.update(line.keys())

            # Create a new window for the preview
            preview_window = tk.Toplevel(self.master)
            preview_window.title("Dataset Preview")

            # Create a table to display the data
            preview_table = ttk.Treeview(preview_window, columns=list(keys), show="headings")
            for key in keys:
                preview_table.heading(key, text=key)
            preview_table.pack()

            # Add the data to the table
            for i, line in enumerate(lines):
                values = [line[key] for key in keys]
                preview_table.insert("", "end", values=values)

        except FileNotFoundError:
            self.show_error("Error", "File not found: " + file_path)
        except json.JSONDecodeError:
            self.show_error("Error", "Invalid JSONL file format.")
        except Exception as e:
            self.show_error("Error", f"An error occurred while previewing the dataset: {e}")

    def start_training(self):
        """Starts the training process."""
        try:
            # Get model parameters
            model_type = self.model_type_var.get()
            d_model = int(self.d_model_entry.get())
            num_heads = int(self.num_heads_entry.get())
            num_layers = int(self.num_layers_entry.get())
            num_classes = int(self.num_classes_entry.get())
            epochs = int(self.epochs_entry.get())
            learning_rate = float(self.learning_rate_entry.get())

            # Load dataset
            file_path = self.dataset_file_entry.get()
            if not file_path:
                messagebox.showerror("Error", "Please select a dataset file.")
                return

            train_data, test_data = load_dataset(file_path)
            if train_data is None or test_data is None:
                return  # Dataset loading failed

            # Start training
            if model_type == "Standard Transformer":
                train_loss_history, train_accuracy_history, state = train_and_evaluate(
                    train_data, test_data, StandardTransformer, d_model, num_heads, num_layers,
                    num_classes, epochs, learning_rate
                )
            elif model_type == "Sagan Transformer":
                train_loss_history, train_accuracy_history, state = train_and_evaluate(
                    train_data, test_data, SaganTransformer, d_model, num_heads, num_layers,
                    num_classes, epochs, learning_rate
                )

            # Save the trained model
            self.save_model(state)  # Make sure to pass 'state' to save_model

            # Plot training progress
            plt.figure(figsize=(12, 6))
            plt.plot(train_loss_history, label="Training Loss")
            plt.xlabel('Epochs')
            plt.ylabel('Loss')
            plt.title('Training Loss over Epochs')
            plt.legend()
            plt.show()

            messagebox.showinfo("Training Complete", "Training finished successfully.")

        except ValueError as e:
            messagebox.showerror("Error", f"Invalid input for parameters: {e}")
            print(f"Invalid input for parameters: {e}")
            traceback.print_exc()  # Print the full traceback to the terminal
        except Exception as e:  # Catch all exceptions
            messagebox.showerror("Error", f"An error occurred during training: {e}")
            print(f"An error occurred during training: {e}")
            traceback.print_exc()  # Print the full traceback to the terminal


    def save_model(self, state):
        """Saves the trained model to a file."""
        file_path = filedialog.asksaveasfilename(defaultextension=".ckpt", filetypes=(("Checkpoint files", "*.ckpt"), ("all files", "*.*")))
        if file_path:
            # Save model state to file (replace with your actual saving logic)
            # Example using JAX's save function (adjust based on your specific saving method)
            jax.experimental.save(file_path, state.params)

            messagebox.showinfo("Model Saved", f"Model saved successfully to: {file_path}")

    def show_error(self, title, message):
        """Displays an error message box and logs the error to the terminal."""
        messagebox.showerror(title, message)
        print(f"{title}: {message}")


# Create the main window and start the GUI
root = tk.Tk()
app = TransformerGUI(root)
root.mainloop()