You are a strict worker, your task is to merge the two following code snippets into a single python program, if you can't generate the full code in the first try, stop and complete afterwards, **code 1**: 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")   **code 2**: import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from collections import Counter
import json
from tqdm import tqdm
import math



# Check if CUDA is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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)

# Expert Transformer Model
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)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        self.output_fc = nn.Linear(d_model, output_size)

    def forward(self, x):
        x = self.input_fc(x) + positional_encoding(x.size(1), self.d_model, x.device)
        transformer_output = self.transformer_encoder(x)
        output = self.output_fc(transformer_output)  # Apply output_fc to each time step in the sequence
        return output

# Gating Network
class GatingNetwork(nn.Module):
    def __init__(self, input_feature_dim, num_experts, hidden_dims=None, dropout_rate=0.0):
        super(GatingNetwork, self).__init__()
        layers = []
        last_dim = input_feature_dim
        
        # If hidden layers are specified, create them
        if hidden_dims is not None:
            for hidden_dim in hidden_dims:
                layers.append(nn.Linear(last_dim, hidden_dim))
                layers.append(nn.ReLU())  # You could make this a hyperparameter as well
                if dropout_rate > 0.0:
                    layers.append(nn.Dropout(dropout_rate))
                last_dim = hidden_dim
        
        # Final layer projecting to the number of experts
        layers.append(nn.Linear(last_dim, num_experts))
        self.fc_layers = nn.Sequential(*layers)
        self.softmax = nn.Softmax(dim=1)
    
    def forward(self, x):
        # Assuming x is of shape [batch_size, seq_len, d_model], aggregate across the sequence length
        x = x.mean(dim=1)  # Aggregate feature per instance
        x = self.fc_layers(x)  # Pass through gating network layers
        return self.softmax(x)
    
# Define hyperparameters specific to the transformer
d_model = 128 #128
nhead = 2 #8
dim_feedforward = 96 #256
num_encoder_layers = 2 #2
num_experts = 2 #2

hidden_dims = [256, 128]  # List of hidden layer sizes
dropout_rate = 0.5  # Dropout rate

gating_network = GatingNetwork(
    input_feature_dim=d_model,
    num_experts=num_experts,
    hidden_dims=hidden_dims,
    dropout_rate=dropout_rate,
)

model_name = "MoE_Transformer_X2-Talk"

# Mixture of Experts Model
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  # Store output_size as an instance variable
        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)  # [batch_size, num_experts]

        expert_outputs = [expert(x) for expert in self.experts]  # List of [batch_size, seq_len, output_size]
        stacked_expert_outputs = torch.stack(expert_outputs)  # Shape: [num_experts, batch_size, seq_len, output_size]

        # Expand gating scores
        expanded_gating_scores = gating_scores.unsqueeze(2).unsqueeze(3)  # Shape: [batch_size, num_experts, 1, 1]
        expanded_gating_scores = expanded_gating_scores.expand(-1, -1, x.size(1), self.output_size)
        expanded_gating_scores = expanded_gating_scores.transpose(0, 1)  # Shape: [num_experts, batch_size, seq_len, output_size]

        # Now the shape of expanded_gating_scores matches stacked_expert_outputs, and broadcasting will work
        mixed_output = torch.sum(stacked_expert_outputs * expanded_gating_scores, dim=0)  # Sum weighted expert outputs for each time step
        return mixed_output

class QAJsonlDataset(Dataset):
    def __init__(self, path, seq_len):
        self.seq_len = seq_len
        self.pairs = self.load_data(path)
        # Flatten the pairs completely before passing them to build_vocab
        self.vocab, self.idx2token = self.build_vocab([word for pair in self.pairs for sublist in pair for word in sublist])
        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("question", ""), data.get("answer", "")
                pairs.append((question.split(), answer.split()))
        return pairs

    def tokenize(self, words):
    # Tokenize a sentence and pad if necessary
    # Add <eos> token at the end if there’s room
        tokens = [self.vocab.get(w, self.vocab["<unk>"]) for w in words]
        if len(tokens) < self.seq_len:
            tokens.append(self.vocab["<eos>"])  # Add <eos> token
            tokens.extend([self.vocab["<pad>"]] * (self.seq_len - len(tokens)))  # Pad the rest
        else:
            tokens = tokens[:self.seq_len - 1] + [self.vocab["<eos>"]]
        return tokens

    def build_vocab(self, words):
    # Start with special tokens with fixed indices
        vocab = {"<unk>": 0, "<pad>": 1, "<eos>": 2}
        start_index = len(vocab)

    # Use Counter to count word frequencies in the corpus
        counts = Counter(words)
    
    # Create the vocab dictionary with all words, starting indices after the special tokens
        for word, _ in counts.most_common():
            if word not in vocab:  # Skip special tokens
                vocab[word] = len(vocab)

    # Create the reverse mapping from indices to words
        idx2token = {idx: token for token, idx in vocab.items()}
        return vocab, idx2token
        
    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)

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.125)  # Dropout added for regularization

    def forward(self, x):
        embedded = self.dropout(self.embedding(x))
        return self.moe(embedded)  # Remove positional encoding addition here, as it’s already added in TransformerExpert

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

# Set the path to your jsonl file and define sequence length
path_to_text = 'test.jsonl'  # replace with the path to your jsonl file
seq_len = 32  # sequence length

# Create a dataset and data loader
dataset = QAJsonlDataset(path_to_text, seq_len)

# Save vocabulary to a text file
vocab_file = f"{model_name}_vocab.txt"
with open(vocab_file, "w", encoding="utf-8") as f:
    for token, id in dataset.vocab.items():
        f.write(f"{token}\t{id}\n")

# Model configuration parameters to be saved
model_config = {
    "d_model": d_model,
    "nhead": nhead,
    "dim_feedforward": dim_feedforward,
    "num_encoder_layers": num_encoder_layers,
    "num_experts": num_experts,
    "hidden_dims": hidden_dims,
    "dropout_rate": dropout_rate
}

# Save configuration to a JSON file
config_file = f"{model_name}_config.json"
with open(config_file, "w", encoding="utf-8") as f:
    json.dump(model_config, f, indent=4)

data_loader = DataLoader(dataset, batch_size=24, shuffle=True, collate_fn=collate_fn, pin_memory=True)
    
# Training loop - added gradient clipping to avoid exploding gradients
def train_model(model, criterion, optimizer, num_epochs, data_loader):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        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()
            predictions = model(inputs)
            predictions = predictions.view(-1, predictions.size(-1))
            targets = targets.view(-1)  # Flatten the targets to [batch_size * seq_len]
            loss = criterion(predictions, targets)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Gradient clipping
            optimizer.step()
            total_loss += loss.item()
            progress_bar.set_postfix({"Loss": loss.item()})
        average_loss = total_loss / len(data_loader.dataset)
        print(f"Epoch {epoch+1}, Average Loss: {average_loss}")

def generate_text(model, dataset, seed_text, num_generate, temperature=1.0):
    model.eval()  # Put the model in evaluation mode
    
    # List to store the generated tokens
    generated_tokens = []

    # Initial sequence (prefix) to start the generation process
    input_sequence = [dataset.vocab.get(word, dataset.vocab["<pad>"]) for word in seed_text.split()]  # Convert to token IDs
    current_sequence = torch.tensor(input_sequence, dtype=torch.long).unsqueeze(0)
    current_sequence = current_sequence.to(device)

    # Generate num_generate tokens
    for _ in range(num_generate):
        # Forward pass through the model
        with torch.no_grad():
            output = model(current_sequence)
        
        # Get probabilities, apply temperature scaling, and sample from the distribution
        probabilities = F.softmax(output[:, -1, :] / temperature, dim=-1).detach()
        next_token_idx = torch.multinomial(probabilities, 1).item()

        # Append token to the current sequence and to the generated tokens
        generated_tokens.append(next_token_idx)
        current_sequence = torch.cat((current_sequence, torch.tensor([[next_token_idx]])), 1).to(device)
    
    # Convert tokens to words
    generated_text = " ".join([dataset.idx2token.get(token, "<unk>") for token in generated_tokens])  # Use .get() to provide a default value for missing keys
    return generated_text

# Function to count the number of tokens in the dataset
def count_tokens_in_dataset(dataset):
    return sum([len(pair[0]) + len(pair[1]) for pair in dataset.pairs])

num_tokens = count_tokens_in_dataset(dataset)
print(f"Total number of tokens in the dataset: {num_tokens}")

vocab_size = len(dataset.vocab)  # Assume dataset.vocab is defined in the QAJsonlDataset class

# Instantiate resulting MoE transformer model and move it to device
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)

# Instantiate the MoE transformer model and move it to device
moe_transformer_model = MoETransformerModel(vocab_size, d_model, moe).to(device)

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

# Example usage with your model:
total_params = count_parameters(moe_transformer_model)
print(f"Total trainable parameters: {total_params}")

# Training parameters
num_epochs = 2
learning_rate = 0.0001

# Define Loss Function and Optimizer for MoE model - using Label Smoothing for better generalization
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = torch.optim.AdamW(moe_transformer_model.parameters(), lr=learning_rate, weight_decay=0.01)  # Using AdamW with weight decay

# Train the model with the text data
train_model(moe_transformer_model, criterion, optimizer, num_epochs, data_loader)

# Start a loop for the interactive chat-like text generation
while True:
    try:
        # Get user input
        seed_text = input("Enter seed text (type 'quit' to stop): ")
        
        # Check if user wants to quit the interaction
        if seed_text.lower() == "quit":
            print("Exiting text generation chat.")
            break
        
        # User input is not empty and not “quit”, generate text
        if seed_text.strip():
            num_generate = 20  # Number of words to generate
            temperature = 1.0  # Sampling temperature, higher will increase diversity
            
            # Use the trained model to generate text
            generated_text = generate_text(moe_transformer_model, dataset, seed_text, num_generate, temperature)
            print("Generated Text:", generated_text)
        else:
            print("Seed text cannot be empty.")
    except KeyboardInterrupt:
        # Handle KeyboardInterrupt (Ctrl+C) to gracefully exit
        print("\nExiting text generation chat.")
        break

torch.save(moe_transformer_model.state_dict(), f"{model_name}.pth")