help me in coding a full implementation of a decoder transformer and two lists: inputs and targets to get the model trained on them, i have already coded a decoder transformer but for another purpose, here is the code for that: import torch
import torch.nn.functional as F
from torch.optim.lr_scheduler import StepLR  # Example scheduler
import torch.nn as nn
import torch.optim as optim
import json
from tokenizers import Tokenizer
from torch.utils.data import Dataset
import random
import math


class PositionalEncoding(nn.Module):
    def __init__(self, embedding_dim, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.embedding_dim = embedding_dim
        
        # Creating positional encoding matrix
        pe = torch.zeros(max_len, embedding_dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embedding_dim, 2).float() * (-math.log(10000.0) / embedding_dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        
        # Registering pe as a buffer since it’s not a model parameter
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return x

class DecoderLayer(nn.Module):
    def __init__(self, embedding_dim, heads, ff_dim):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(embedding_dim, heads)
        
        # Feedforward network
        self.ffn = nn.Sequential(
            nn.Linear(embedding_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embedding_dim),
        )
        
        self.layer_norm1 = nn.LayerNorm(embedding_dim)
        self.layer_norm2 = nn.LayerNorm(embedding_dim)
    
    def forward(self, src):
        src2 = self.layer_norm1(src)
        attn_output, _ = self.self_attn(src2, src2, src2)
        src = src + attn_output
        src2 = self.layer_norm2(src)
        src = src + self.ffn(src2)
        return src

class Decoder(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_layers, heads, ff_dim):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.pos_encoding = PositionalEncoding(embedding_dim)
        self.decoder_layers = nn.ModuleList([DecoderLayer(embedding_dim, heads, ff_dim) for _ in range(num_layers)])
        self.final_layer = nn.Linear(embedding_dim, vocab_size)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.pos_encoding(x)
        
        for layer in self.decoder_layers:
            x = layer(x)
        
        output = self.final_layer(x)
        return output
  
class QAJsonlDataset(Dataset):
    def __init__(self, path, tokenizer_path, seq_len=512):
        super().__init__()
        # Load the trained tokenizer
        self.tokenizer = Tokenizer.from_file(tokenizer_path)
        self.seq_len = seq_len

        self.pairs = self.load_data(path)

    def load_data(self, path):
        pairs = []
        with open(path, "r", encoding="utf-8") as f:  # Assume UTF-8 encoded JSON
            for line in f:
                data = json.loads(line)
                # Assuming ‘question’ and ‘answer’ are the keys
                # Tokenize directly here and truncate/pad as necessary
                question = self.tokenize(data['user'])
                answer = self.tokenize(data['content'])
                pairs.append((question, answer))
        return pairs

    def tokenize(self, text):
        # Encode the text, truncating or padding to seq_len as required
        output = self.tokenizer.encode(text).ids
        # Truncate if needed
        output = output[:self.seq_len]
        # Pad if needed
        if len(output) < self.seq_len:
            output += [self.tokenizer.token_to_id('<pad>')] * (self.seq_len - len(output))
        return output

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx):
        question, answer = self.pairs[idx]
        return torch.tensor(question, dtype=torch.long), torch.tensor(answer, dtype=torch.long)

class CustomDataLoader:
    def __init__(self, dataset, batch_size=32):
        self.dataset = dataset
        self.batch_size = batch_size

    def len(self):
        return len(self.dataset) // self.batch_size

    def __getitem__(self, idx):
        batch = self.dataset[idx * self.batch_size:(idx + 1) * self.batch_size]
        inputs, targets = zip(*batch)
        inputs_padded = torch.nn.utils.rnn.pad_sequence(inputs, batch_first=True, padding_value=0)
        targets_padded = torch.nn.utils.rnn.pad_sequence(targets, batch_first=True, padding_value=0)
        return inputs_padded, targets_padded


# Paths for tokenizer and dataset
path_to_data = "data/Real_talk.jsonl"
tokenizer_path = "Tokenizer-Max.json"
tokenizer = Tokenizer.from_file(tokenizer_path)

# Define model parameters
vocab_size = tokenizer.get_vocab_size()
embedding_dim = 24
num_layers = 2  # Number of decoder layers
heads = 4  # Number of attention heads
ff_dim = 96  # Feed-forward dimension

# Initialize dataset
dataset = QAJsonlDataset(path_to_data, tokenizer_path)

# Initialize decoder model with correct vocab size
model = Decoder(vocab_size, embedding_dim, num_layers, heads, ff_dim)

# Shuffle the dataset
random.shuffle(dataset.pairs)

train_size = int(0.9 * len(dataset))  # 80% of data for training, adjust as necessary
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32)

# Define optimizer
optimizer = optim.Adam(model.parameters())
loss_fn = nn.CrossEntropyLoss()  # PyTorch version of SparseCategoricalCrossentropy with from_logits=True

def print_model_param_count(model):
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Total Trainable Parameters: {total_params}")

# Assuming your model’s variable name is model
print_model_param_count(model)

# Define an optimizer and a learning rate scheduler
optimizer = optim.Adam(model.parameters(), lr=0.005)
scheduler = StepLR(optimizer, step_size=10, gamma=0.7)  # Example scheduler, adjust according to needs

def train(model, train_loader, optimizer, scheduler, num_epochs):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = F.cross_entropy(outputs.view(-1, vocab_size), targets.view(-1))
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()

            if batch_idx % 50 == 0:  # Adjust logging frequency as needed
                print(f'Epoch {epoch} [{batch_idx * len(inputs)}/{len(train_loader.dataset)}] Loss: {loss.item()}')

        # Step the scheduler
        scheduler.step()
        
        # Call evaluation function here (assuming you have one)
        val_accuracy = evaluate(model, val_loader)
        
        print(f'Epoch {epoch} Val Accuracy {val_accuracy}')

def evaluate(model, val_loader):
    model.eval()
    total_accuracy = 0
    with torch.no_grad():
        for inputs, targets in val_loader:
            outputs = model(inputs)
            predictions = torch.argmax(outputs, dim=2)
            correct_predictions = (predictions == targets).float()
            mask = targets != tokenizer.token_to_id("<pad>")
            correct_predictions = correct_predictions * mask
            total_accuracy += correct_predictions.sum().item() / mask.sum().item()

    avg_accuracy = total_accuracy / len(val_loader)
    model.train()  # Set the model back to training mode
    return avg_accuracy

# Replace the existing training loop call with:
train(model, train_loader, optimizer, scheduler, num_epochs=80)

def generate_text(model, initial_text, tokenizer, seq_len=512, temperature=1.0):
    model.eval()  # Ensures the model is in evaluation mode
    
    # Tokenize the initial text
    tokens = tokenizer.encode(initial_text).ids
    
    # Generate text
    with torch.no_grad():
        for _ in range(seq_len - len(tokens)):
            input_tensor = torch.tensor([tokens], dtype=torch.long)
            output = model(input_tensor)
            logits = output[0, -1, :] / temperature
            probabilities = torch.softmax(logits, dim=-1)
            next_token = torch.argmax(probabilities).item()
            tokens.append(next_token)
            
            # Assume either a special token for EOS or check if max length reached
            if next_token == tokenizer.token_to_id('<eos>'):
                break

    # Decode the tokens back to text
    generated_text = tokenizer.decode(tokens)
    return generated_text

# During usage, you pass the loaded tokenizer directly
generated_text = generate_text(model, "what is earth ?", tokenizer, seq_len=64)
print(generated_text)