report any faults in the definiton of the transformer architecture of the gating network or the position encoding or any hardcoded parameters of the model that can't be modified or represented by a variable like the already modifieable hyperparameters, **code**: 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 collections import Counter
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


# ---------- 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):
        x = self.norm(self.input_fc(x)) + self.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=[6144, 3072, 1536, 768], 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):
    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.multinomial(probabilities, num_samples=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 = 384
nhead = 12
dim_feedforward = 1536
num_encoder_layers = 12
num_experts = 2

# Training Parameters
batch_size = 384  # Adjustable batch size
warmup_steps = 2000  # Warmup steps for learning rate
optimizer_type = "AdamW"  # Could be “SGD”, “RMSprop”, etc.
learning_rate = 1e-1
weight_decay = 1e-5  # For L2 regularization
num_epochs = 10

# Dataset :
path_to_dataset = "C:/Users/L14/Documents/Projets/Easy-MoE/Easy-MoE/data/Real_talk.jsonl"
tokenizer_path = "BPE_tokenizer(Real-Talk).json"
seq_len = 256
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)

# Instantiate your main scheduler (StepLR)
step_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.95)

# Wrap it with WarmupLR
scheduler = WarmupLR(optimizer, warmup_steps=warmup_steps, scheduler_step_lr=step_lr_scheduler)

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 ----------
torch.save(moe_transformer_model.state_dict(), "Transformer-Alpha-v04.pth")