We get this error: Starting text tokenization and embedding...
An error occurred while loading the dataset: len() of unsized object
Traceback (most recent call last):
  File "C:\Users\L14\Documents\Projets\Model_Creator(NEW)\model_creator\lib\site-packages\jax\_src\array.py", line 254, in __len__
    return self.shape[0]
IndexError: tuple index out of range

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "c:\Users\L14\Documents\Projets\Model_Creator(NEW)\App.py", line 242, in load_dataset
    batch_tensor = torch.tensor(batch_array)
  File "C:\Users\L14\Documents\Projets\Model_Creator(NEW)\model_creator\lib\site-packages\jax\_src\array.py", line 256, in __len__
    raise TypeError("len() of unsized object") from err  # same as numpy error
TypeError: len() of unsized object    ; When we run this 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
from tkinter import scrolledtext

# 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

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

def load_dataset(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = [json.loads(line) for i, line in enumerate(f) if i < 5]
            keys = set()
            for line in lines:
                keys.update(line.keys())

        input_key, output_key = get_input_output_keys(keys)

        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])

        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")

            max_length = max(len(tokenizer.encode(text, add_special_tokens=True)) for text in inputs)

            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:
                    batch_tokenized = [tokenizer.encode(text, add_special_tokens=True, max_length=max_length, truncation=True) for text in batch]
                    attention_masks = [[1] * len(tokenized_text) + [0] * (max_length - len(tokenized_text)) for tokenized_text in batch_tokenized]

                    # Pad the lists in batch_tokenized to the maximum length
                    padded_batch_tokenized = [tokenized_text + [tokenizer.pad_token_id] * (max_length - len(tokenized_text)) for tokenized_text in batch_tokenized]

                    if padded_batch_tokenized and attention_masks:
                        batch_array = jnp.array(padded_batch_tokenized, dtype=jnp.int32)
                        attention_mask_array = jnp.array(attention_masks, dtype=jnp.int32)

                        batch_tensor = torch.tensor(batch_array)
                        mask_tensor = torch.tensor(attention_mask_array)

                        embedded_batch = embedding_model(input_ids=batch_tensor, attention_mask=mask_tensor).pooler_output
                        batched_inputs.append(jnp.array(embedded_batch.detach().numpy()))

            inputs = jnp.array(batched_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)
        self.output_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()

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

        # Text area to display messages
        self.output_text = scrolledtext.ScrolledText(self.output_frame, wrap=tk.WORD, height=5)
        self.output_text.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:
                self.show_error("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":
                self.show_message("Info", "Training 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":
                self.show_message("Info", "Training 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()

            self.show_message("Success", "Training finished successfully.")

        except ValueError as e:
            self.show_error("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
            self.show_error("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)

            self.show_message("Success", 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)
        self.show_message("Error", message)
        print(f"{title}: {message}")

    def show_message(self, title, message):
        """Displays a message box and appends the message to the output text area."""
        messagebox.showinfo(title, message)
        self.output_text.config(state=tk.NORMAL)  # Enable text editing
        self.output_text.insert(tk.END, f"{title}: {message}\n")
        self.output_text.config(state=tk.DISABLED)  # Disable text editing


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