import torch
import torch.nn as nn
import torch.optim as optim
from torch.amp import autocast, GradScaler
from transformers import AutoTokenizer, PreTrainedModel, DataCollatorForLanguageModeling, PretrainedConfig
from datasets import load_dataset
from torch.utils.data import DataLoader
from typing import List
import math
from tqdm import tqdm
import gc

# Configuration for the advanced model
class AdvancedLLMConfig(PretrainedConfig):
    model_type = "advanced_llm"

    def __init__(self, vocab_size=30000, d_model=768, nhead=12, dim_feedforward=3072,
                 max_position_embeddings=2048, memory_size=2000, num_layers=12, eos_token_id=50256, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.nhead = nhead
        self.dim_feedforward = dim_feedforward
        self.max_position_embeddings = max_position_embeddings
        self.memory_size = memory_size
        self.num_layers = num_layers
        self.eos_token_id = eos_token_id

# Positional encoding implementation
class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model, dtype=torch.half)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:x.size(1)].transpose(0, 1)

# Encoder implementation
class MultiResolutionEncoder(nn.Module):
    def __init__(self, d_model: int, nhead: int, dim_feedforward: int, num_layers: int):
        super().__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, batch_first=True, dtype=torch.half)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Ensure input is 3D: [batch_size, seq_len, d_model]
        assert x.dim() == 3, f"Expected 3D input tensor but got {x.dim()}D tensor"
        return self.encoder(x)


class AdaptiveAttentionDecoder(nn.Module):
    def __init__(self, d_model: int, nhead: int, dim_feedforward: int):
        super().__init__()
        self.multi_head_attention = nn.MultiheadAttention(d_model, nhead, batch_first=True, dtype=torch.half)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, dim_feedforward, dtype=torch.half),
            nn.ReLU(),
            nn.Linear(dim_feedforward, d_model, dtype=torch.half)
        )
        self.layer_norm1 = nn.LayerNorm(d_model, dtype=torch.half)
        self.layer_norm2 = nn.LayerNorm(d_model, dtype=torch.half)

    def forward(self, x: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        # Ensure x and memory are correctly shaped
        assert x.dim() == 3 and memory.dim() == 3, f"Expected 3D input tensors but got {x.dim()}D and {memory.dim()}D tensors"
        
        attended, _ = self.multi_head_attention(x, memory, memory)
        x = self.layer_norm1(x + attended)
        x = self.layer_norm2(x + self.feed_forward(x))
        return x

    
# Model definition
class AdvancedLLM(PreTrainedModel):
    config_class = AdvancedLLMConfig

    def __init__(self, config):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.d_model, dtype=torch.half)
        self.positional_encoding = PositionalEncoding(config.d_model, config.max_position_embeddings)
        self.encoder = MultiResolutionEncoder(config.d_model, config.nhead, config.dim_feedforward, config.num_layers)
        self.decoder = AdaptiveAttentionDecoder(config.d_model, config.nhead, config.dim_feedforward)
        self.output_layer = nn.Linear(config.d_model, config.vocab_size, dtype=torch.half)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids.long())
        x = self.positional_encoding(x.half())
        encoded = self.encoder(x)

        # Make sure to use encoded correctly as the memory input
        decoded = self.decoder(encoded, encoded)  # Use encoded as both inputs
        return self.output_layer(decoded)

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, max_length: int, temperature: float = 1.0) -> List[int]:
        self.eval()
        
        generated = input_ids.tolist()
        
        for _ in range(max_length - len(generated)):
            inputs = torch.tensor(generated).unsqueeze(0).to(input_ids.device)
            print(f"Input shape before generating: {input_ids.shape}")
            if input_ids.dim() == 4:
                batch_size, num_heads, sequence_length, embedding_dim = input_ids.shape
                input_ids = input_ids.view(batch_size, sequence_length, embedding_dim)
                print(f"Reshaped input to: {input_ids.shape}")
            # Forward pass should be using the proper shape
            outputs = self(inputs)  # Ensure inputs are correctly shaped
            next_token_logits = outputs[0, -1, :] / temperature
            
            probabilities = torch.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probabilities, num_samples=1).item()
            generated.append(next_token)
            
            if next_token == self.config.eos_token_id:
                break
        
        return generated


def train_model(model, train_dataloader, num_epochs=3, learning_rate=5e-5,
                device='cuda', accumulation_steps=8, max_grad_norm=1.0):
    model.to(device)
    model.train()

    scaler = GradScaler()  # Initialize GradScaler for mixed precision
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate)
    
    # Ensure the model parameters are in FP32 for the optimizer
    for param in model.parameters():
        param.data = param.data.float()

    criterion = nn.CrossEntropyLoss()  # Define loss function

    for epoch in range(num_epochs):
        total_loss = 0
        optimizer.zero_grad()  # Reset gradients at the beginning of each epoch

        progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch + 1}/{num_epochs}")

        for i, batch in enumerate(progress_bar):
            input_ids = batch['input_ids'].to(device).long()  # Move to device
            labels = batch['labels'].to(device)  # Move to device

            with autocast('cuda'):  # Enable mixed precision
                outputs = model(input_ids)  # Forward pass
                outputs = outputs.view(-1, outputs.size(-1))
                labels = labels.view(-1)
                loss = criterion(outputs, labels)  # Compute loss

            scaler.scale(loss).backward()  # Scale loss and perform backpropagation

            if (i + 1) % accumulation_steps == 0:  # Check if it's time to step the optimizer
                # Check for any inf or NaN values in gradients
                for param in model.parameters():
                    if param.grad is not None and (torch.isinf(param.grad).any() or torch.isnan(param.grad).any()):
                        print("Inf or NaN gradients detected. Skipping optimizer step.")
                        scaler.update()
                        return model  # Skip this step

                scaler.unscale_(optimizer)  # Unscale gradients for clipping
                nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)  # Gradient clipping
                scaler.step(optimizer)  # Update parameters
                scaler.update()  # Update the scaler
                optimizer.zero_grad()  # Zero gradients for the next accumulation

            total_loss += loss.item()  # Accumulate loss
            progress_bar.set_postfix({"Loss": loss.item()})

        avg_loss = total_loss / len(train_dataloader)  # Compute average loss for the epoch
        print(f"Epoch {epoch + 1}/{num_epochs}, Average Loss: {avg_loss:.4f}")

        # Save model checkpoint after each epoch
        torch.save(model.state_dict(), f'checkpoint_epoch_{epoch + 1}.pt')
        gc.collect()  # Clear GPU memory

    return model  # Return the trained model


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    dataset = load_dataset("wikipedia", "20220301.en", split="train[:500]")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    def tokenize_function(examples):
        return tokenizer(examples["text"], truncation=True, max_length=512, padding="max_length")

    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
    tokenized_dataset = tokenized_dataset.map(lambda examples: {'labels': examples['input_ids']})

    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    train_dataloader = DataLoader(tokenized_dataset, batch_size=5, collate_fn=data_collator, shuffle=True)

    config = AdvancedLLMConfig(
        vocab_size=tokenizer.vocab_size,
        d_model=384,
        nhead=12,
        dim_feedforward=2304,
        max_position_embeddings=2048,
        memory_size=2000,
        num_layers=12,
        eos_token_id=tokenizer.eos_token_id
    )

    model = AdvancedLLM(config)
    
    model_size = sum(p.numel() for p in model.parameters()) * 2 / 1024 / 1024  
    print(f"Model size: {model_size:.2f} MB")

    # After training the model
    trained_model = train_model(model, train_dataloader, num_epochs=3, learning_rate=5e-5, device=device, accumulation_steps=8)

    # Move the model to CPU for saving
    trained_model = trained_model.to('cpu')
    trained_model.save_pretrained("./advanced_llm_model")
    tokenizer.save_pretrained("./advanced_llm_model")
    print("Model training complete and saved in Hugging Face format.")

    # Prepare input for generation
    input_text = "The history of artificial intelligence"
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to('cpu') 

    # Generate text
    generated_ids = trained_model.generate(input_ids, max_length=50)
    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    print("Generated Text:", generated_text)


if __name__ == "__main__":
    main()

please correct this and all othere possible errors and write updated codes:
Traceback (most recent call last):
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/Decoder/ai.py", line 238, in <module>
    main()
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/Decoder/ai.py", line 232, in main
    generated_ids = trained_model.generate(input_ids, max_length=50)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/Decoder/ai.py", line 113, in generate
    outputs = self(inputs)  # Ensure inputs are correctly shaped
              ^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/Decoder/ai.py", line 93, in forward
    encoded = self.encoder(x)
              ^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ramachandra/Documents/AI_Testing_LLM/LLM/Decoder/ai.py", line 52, in forward
    assert x.dim() == 3, f"Expected 3D input tensor but got {x.dim()}D tensor"
           ^^^^^^^^^^^^
AssertionError: Expected 3D input tensor but got 4D tensor