Your job is to read the following code, analyze it, detect weakness points, detect potential improvemenrs, detect potential optimizations, detect potential simplification, and after all of that you will return a much improved and better version of the entering code with labels, comments, organization and modularity, split your generation into 3 parts so you don't run out of generation tokens, here is the code to work with, follow the instructions strictly, you are not allowed to return anything other that desired code, **raw code**: import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import json
import math
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
from tokenizers import Tokenizer
from torch.optim.lr_scheduler import SequentialLR, StepLR, LinearLR



# ---------- Device Configuration ----------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------- Utility Functions ----------
def positional_encoding(seq_len, d_model, device):
    pos = torch.arange(seq_len, dtype=torch.float, device=device).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)).to(device)
    pe = torch.zeros(seq_len, d_model, device=device)
    pe[:, 0::2] = torch.sin(pos * div_term)
    pe[:, 1::2] = torch.cos(pos * div_term)
    return pe.unsqueeze(0)

# -------- Performance ----------

def evaluate_model(model, data_loader, device):
    model.eval()
    all_preds, all_targets = [], []

    with torch.no_grad():
        for inputs, targets in data_loader:
            inputs, targets = inputs.to(device), targets.to(device)

            outputs = model(inputs)
            predictions = torch.argmax(outputs, dim=-1).view(-1)  # Flatten predicted indices

            all_preds.extend(predictions.cpu().numpy())
            all_targets.extend(targets.view(-1).cpu().numpy())  # Ensure targets are also flattened

    # Calculate precision, recall, and F1 score after ensuring all_preds and all_targets are correctly aligned.
    accuracy = accuracy_score(all_targets, all_preds)
    precision = precision_score(all_targets, all_preds, average='macro', zero_division=0)
    recall = recall_score(all_targets, all_preds, average='macro', zero_division=0)
    f1 = f1_score(all_targets, all_preds, average='macro', zero_division=0)

    print(f"Accuracy: {accuracy:.4f}")
    print(f"Precision: {precision:.4f}")
    print(f"Recall: {recall:.4f}")
    print(f"F1 Score: {f1:.4f}")

    return accuracy ,precision, recall, f1

# Function to plot loss over time
def plot_loss(loss_history):
    plt.figure(figsize=(10, 5))
    plt.plot(loss_history, label='Training Loss')
    plt.xlabel('Batches')
    plt.ylabel('Loss')
    plt.title('Training Loss Over Time')
    plt.legend()
    plt.show()

# ---------- Model Definitions ----------
class TransformerExpert(nn.Module):
    def __init__(self, input_size, d_model, output_size, nhead, dim_feedforward, num_encoder_layers=1):
        super(TransformerExpert, self).__init__()
        self.d_model = d_model
        self.input_fc = nn.Linear(input_size, d_model)
        self.pos_encoder = nn.Parameter(positional_encoding(1, d_model, device), requires_grad=True)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
                                                   dim_feedforward=dim_feedforward, 
                                                   batch_first=True, 
                                                   norm_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        self.output_fc = nn.Linear(d_model, output_size)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x):
        seq_len = x.shape[1]
        pos_encoder = positional_encoding(seq_len, self.d_model, device)
        x = self.norm(self.input_fc(x)) + pos_encoder
        transformer_output = self.transformer_encoder(x)
        output = self.output_fc(transformer_output)
        return output


class GatingNetwork(nn.Module):
    def __init__(self, input_feature_dim, num_experts, hidden_dims=[256], dropout_rate=0.2):
        super(GatingNetwork, self).__init__()
        layers = []
        last_dim = input_feature_dim

        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(last_dim, hidden_dim),
                nn.GELU(),
                nn.Dropout(dropout_rate),
            ])
            last_dim = hidden_dim

        layers.append(nn.Linear(last_dim, num_experts))
        self.fc_layers = nn.Sequential(*layers)
        self.softmax = nn.Softmax(dim=1)
    
    def forward(self, x):
        x = x.mean(dim=1)  # To ensure gating is based on overall features across the sequence
        x = self.fc_layers(x)
        return self.softmax(x)

class MixtureOfTransformerExperts(nn.Module):
    def __init__(self, input_size, d_model, output_size, nhead, dim_feedforward, num_experts, num_encoder_layers=1):
        super(MixtureOfTransformerExperts, self).__init__()
        self.num_experts = num_experts
        self.output_size = output_size
        self.experts = nn.ModuleList([TransformerExpert(input_size, d_model, output_size, nhead, dim_feedforward, num_encoder_layers) for _ in range(num_experts)])
        self.gating_network = GatingNetwork(d_model, num_experts)

    def forward(self, x):
        gating_scores = self.gating_network(x)

        expert_outputs = [expert(x) for expert in self.experts]
        stacked_expert_outputs = torch.stack(expert_outputs)

        expanded_gating_scores = gating_scores.unsqueeze(2).unsqueeze(3)
        expanded_gating_scores = expanded_gating_scores.expand(-1, -1, x.size(1), self.output_size)
        expanded_gating_scores = expanded_gating_scores.transpose(0, 1)

        mixed_output = torch.sum(stacked_expert_outputs * expanded_gating_scores, dim=0)
        return mixed_output

class MoETransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, moe):
        super(MoETransformerModel, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model)
        self.moe = moe
        self.dropout = nn.Dropout(p=0.1)

    def forward(self, x):
        embedded = self.dropout(self.embedding(x))
        return self.moe(embedded)

# ---------- Dataset Definitions ----------
class QAJsonlDataset(Dataset):
    def __init__(self, path, seq_len, tokenizer_path):
        # Load the trained tokenizer
        self.tokenizer = Tokenizer.from_file(tokenizer_path)
        
        self.seq_len = seq_len
        self.pairs = self.load_data(path)
        # Using BPE, so no need for manual vocab or idx2token.
        # Tokenization will now happen using self.tokenizer.
        self.tokenized_pairs = [(self.tokenize(q), self.tokenize(a)) for q, a in self.pairs]

    def load_data(self, path):
        pairs = []
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                data = json.loads(line.strip())
                question, answer = data.get("user", ""), data.get("content", "")
                pairs.append((question, answer))  # Store questions and answers as raw strings
        return pairs

    def tokenize(self, text):
        # Tokenizing using the BPE tokenizer
        encoded = self.tokenizer.encode(text)
        tokens = encoded.ids
        # Padding/truncation
        if len(tokens) < self.seq_len:
            # Padding
            tokens += [self.tokenizer.token_to_id("<pad>")] * (self.seq_len - len(tokens))
        else:
            # Truncation
            tokens = tokens[:self.seq_len - 1] + [self.tokenizer.token_to_id("<eos>")]
        return tokens
        
    def __len__(self):
        return len(self.tokenized_pairs)

    def __getitem__(self, idx):
        tokenized_question, tokenized_answer = self.tokenized_pairs[idx]
        return torch.tensor(tokenized_question, dtype=torch.long), torch.tensor(tokenized_answer, dtype=torch.long)

def collate_fn(batch):
    questions, answers = zip(*batch)
    questions = pad_sequence(questions, batch_first=True, padding_value=0)
    answers = pad_sequence(answers, batch_first=True, padding_value=0)
    return questions, answers

# ---------- Training and Inference Functions ----------

def train_model(model, criterion, optimizer, num_epochs, data_loader, label_smoothing=0.1):
    criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)
    model.train()
    loss_history = []  # Initialize a list to keep track of losses
    for epoch in range(num_epochs):
        total_loss = 0
        total_items = 0  # Keep track of total items processed
        progress_bar = tqdm(enumerate(data_loader), total=len(data_loader), desc=f"Epoch {epoch+1}", leave=False)
        for i, (inputs, targets) in progress_bar:
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            
            # Predict
            predictions = model(inputs)
            predictions = predictions.view(-1, predictions.size(-1))  # Make sure predictions are the right shape
            targets = targets.view(-1)  # Flatten targets to match prediction shape if necessary
            
            # Calculate loss
            loss = criterion(predictions, targets)
            loss.backward()
            
            # Gradient clipping for stabilization
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            # scheduler.step()
            
            # Update total loss and the number of items
            total_loss += loss.item() * inputs.size(0)  # Multiply loss by batch size
            total_items += inputs.size(0)
            loss_history.append(loss.item())
            
            progress_bar.set_postfix({"Loss": loss.item()})
        
        average_loss = total_loss / total_items  # Correctly compute average loss
        print(f"Epoch {epoch+1}, Average Loss: {average_loss:.6f}")
    
    return loss_history

class WarmupLR(torch.optim.lr_scheduler._LRScheduler):
    def __init__(self, optimizer, warmup_steps, scheduler_step_lr):
        self.warmup_steps = warmup_steps
        self.scheduler_step_lr = scheduler_step_lr  # The subsequent scheduler
        super(WarmupLR, self).__init__(optimizer)
    
    def get_lr(self):
        if self._step_count <= self.warmup_steps:
            warmup_factor = float(self._step_count) / float(max(1, self.warmup_steps))
            for base_lr in self.base_lrs:
                yield base_lr * warmup_factor
        else:
            self.scheduler_step_lr.step()  # Update the subsequent scheduler
            for param_group in self.optimizer.param_groups:
                yield param_group['lr']

class GERU(nn.Module):
    def __init__(self, in_features):
        super(GERU, self).__init__()
        self.alpha = nn.Parameter(torch.rand(in_features))
        
    def forward(self, x):
        return torch.max(x, torch.zeros_like(x)) + self.alpha * torch.min(x, torch.zeros_like(x))

def generate_text(model, tokenizer, seed_text, num_generate, temperature=1.0):
    model.eval()
    
    generated_tokens = []
    
    # Encode the seed text using the tokenizer
    encoded_input = tokenizer.encode(seed_text)
    input_ids = torch.tensor(encoded_input.ids, dtype=torch.long).unsqueeze(0).to(device)

    # Generate num_generate tokens
    with torch.no_grad():
        for _ in range(num_generate):
            output = model(input_ids)
            
            # Get the last logits and apply temperature
            logits = output[:, -1, :] / temperature
            probabilities = F.softmax(logits, dim=-1)
            next_token_id = torch.argmax(probabilities, dim=-1).item()

            # Append generated token ID and prepare the new input_ids
            generated_tokens.append(next_token_id)
            input_ids = torch.cat([input_ids, torch.tensor([[next_token_id]], dtype=torch.long).to(device)], dim=1)

    # Decode the generated token IDs back to text
    generated_text = tokenizer.decode(generated_tokens)
    return generated_text

def count_tokens_in_dataset(dataset):
    return sum([len(pair[0]) + len(pair[1]) for pair in dataset.pairs])

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# ---------- Hyperparameters and Model Instantiation ----------
# Transformer :
d_model = 128
nhead = 16
dim_feedforward = 256
num_encoder_layers = 1
num_experts = 2

# Training Parameters
batch_size = 64  # Adjustable batch size
optimizer_type = "AdamW"  # Could be “SGD”, “RMSprop”, etc.
learning_rate = 200
weight_decay = 0.01  # For L2 regularization
num_epochs = 500

# Dataset :
path_to_dataset = "C:/Users/L14/Documents/Projets/Easy-MoE/Easy-MoE/data/Real_talk.jsonl"
tokenizer_path = "tokenizer/BPE_tokenizer(Real-Talk).json"
seq_len = 64
dataset = QAJsonlDataset(path_to_dataset, seq_len, tokenizer_path)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, pin_memory=True)
    
num_tokens = count_tokens_in_dataset(dataset)
print(f"Total number of tokens in the dataset: {num_tokens}")

# Load the tokenizer
tokenizer = Tokenizer.from_file(tokenizer_path)

# Determine the vocabulary size
vocab_size = tokenizer.get_vocab_size()

moe = MixtureOfTransformerExperts(
    input_size=d_model,
    d_model=d_model,
    output_size=vocab_size,
    nhead=nhead,
    dim_feedforward=dim_feedforward,
    num_experts=num_experts,
    num_encoder_layers=num_encoder_layers
).to(device)

moe_transformer_model = MoETransformerModel(vocab_size, d_model, moe).to(device)

# Count of total parameters :
total_params = count_parameters(moe_transformer_model)
print(f"Total trainable parameters: {total_params}")

# ---------- Training ----------

# Adjusting optimizer setup to include weight decay and allow switching between types
if optimizer_type == "AdamW":
    optimizer = torch.optim.AdamW(moe_transformer_model.parameters(), lr=learning_rate, weight_decay=weight_decay)
elif optimizer_type == "SGD":
    optimizer = torch.optim.SGD(moe_transformer_model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay)
elif optimizer_type == "Adam":
    optimizer = torch.optim.Adam(moe_transformer_model.parameters(), lr=learning_rate, weight_decay=weight_decay)

# Setup optimizers just like before
warmup_epochs = 1
scheduler1 = LinearLR(optimizer, start_factor=1e-5, total_iters=warmup_epochs)
scheduler2 = StepLR(optimizer, step_size=10, gamma=0.9)

scheduler = SequentialLR(optimizer, schedulers=[scheduler1, scheduler2], milestones=[warmup_epochs])

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# Train the model
loss_history = train_model(moe_transformer_model, criterion, optimizer, num_epochs, data_loader)

# Evaluating the model
plot_loss(loss_history)
train_accuracy = evaluate_model(moe_transformer_model, data_loader, device)

# ---------- Inference ----------
def interactive_text_generation(model, dataset, max_length=32, temperature=1.0):
    while True:
        try:
            # Get user input
            seed_text = input("Enter seed text (type 'quit' to exit and save the model): ").strip()
            
            # Check if the user wants to quit the interaction
            if seed_text.lower() == 'quit':
                print("Exiting text generation mode.")
                break

            # Generate text based on the seed text
            if seed_text:
                generated_text = generate_text(model, dataset, seed_text, max_length, temperature)  # Modify max_length/temperature as needed
                print("Generated Text:", generated_text)
            else:
                print("Seed text cannot be empty. Please enter some text.")
                
        except Exception as e:
            # Handle potential errors gracefully
            print(f"An error occurred: {e}. Try again.")

interactive_text_generation(moe_transformer_model, tokenizer)

# ---------- Save Trained Model ----------

def save_model_with_config(model, config, save_dir, model_name):
    """
    Saves the model weights, configuration, and performance metrics.

    Parameters:
    - model: the PyTorch model to save.
    - config: a dictionary with the model's configuration.
    - metrics: a dictionary with the model's performance metrics.
    - save_dir: the root directory to save the model and its info.
    - model_name: the name of the model, used to create a subdirectory.
    """
    model_path = os.path.join(save_dir, model_name)
    os.makedirs(model_path, exist_ok=True)

    # Save model weigths
    torch.save(model.state_dict(), os.path.join(model_path, '.pth'))

    # Save configuration
    with open(os.path.join(model_path, 'config.json'), 'w') as config_file:
        json.dump(config, config_file, indent=4)

    # Save metrics
    #with open(os.path.join(model_path, 'metrics.json'), 'w') as metrics_file:
    #    json.dump(metrics, metrics_file, indent=4)

    print(f"Model, configuration, and metrics saved in {model_path}")

config = {
    'd_model': d_model,'nhead': nhead,'dim_feedforward': dim_feedforward,'num_encoder_layers': num_encoder_layers,
    'num_experts': num_experts,'seq_len': seq_len,'batch_size': batch_size,'learning_rate': learning_rate,
    'weight_decay': weight_decay,'num_epochs': num_epochs,
}

save_model_with_config(moe_transformer_model, config, "Trained_models", "Transformer-Alpha-v04")