import os
import warnings
import pandas as pd
import time
from PIL import Image
import requests
from io import BytesIO
import torch
from torch.utils.data import DataLoader
from torch import nn, optim
from transformers import DonutProcessor, VisionEncoderDecoderModel
import pytorch_lightning as pl
from torch.utils.data.dataloader import default_collate

# Set environment variables and warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
warnings.filterwarnings("ignore", category=UserWarning)

# Function to create annotations
def create_train_annotations(df):
    annotations = []
    for idx, row in df.iterrows():
        image_url = row['image_link']
        entity_value = row.get('entity_value', '')  # Only 'entity_value' used here
        annotations.append({
            'image_url': image_url,
            'text': f"
{entity_value}
"
        })
    return annotations

# Dataset class for handling the data
class DonutDataset(torch.utils.data.Dataset):
    def __init__(self, data_list, processor, max_length=512, ignore_id=-100):
        self.data_list = data_list
        self.processor = processor
        self.max_length = max_length
        self.ignore_id = ignore_id

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

    def __getitem__(self, idx):
        sample = self.data_list[idx]
        image_url = sample['image_url']
        prompt = sample['text']

        # Retry logic for downloading the image
        max_retries = 4
        num_retries = 0
        success = False
        while num_retries < max_retries and not success:
            try:
                response = requests.get(image_url, timeout=10)
                response.raise_for_status()
                image = Image.open(BytesIO(response.content)).convert("RGB")
                image = image.resize((32, 32))  # Change resolution here
                success = True
            except Exception as e:
                num_retries += 1
                print(f"Error loading image from {image_url} on attempt {num_retries}: {e}")
                if num_retries >= max_retries:
                    print(f"Skipping image {image_url} after {max_retries} failed attempts.")
                    return None
                else:
                    time.sleep(1)  # Optional delay between retries
        
        pixel_values = self.processor(image, return_tensors="pt").pixel_values.squeeze()
        labels = self.processor.tokenizer(
            prompt,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        ).input_ids.squeeze()

        labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_id

        return {"pixel_values": pixel_values, "labels": labels}

# Custom collate function to filter None samples
def collate_fn(batch):
    batch = [sample for sample in batch if sample is not None]
    if len(batch) == 0:
        return None  # Skip this batch
    else:
        return default_collate(batch)

# Training model class
class DonutModel(pl.LightningModule):
    def __init__(self, processor, model, lr=1e-5):
        super().__init__()
        self.processor = processor
        self.model = model
        self.lr = lr

    def on_train_start(self):
        self.model.train()
        for param in self.model.parameters():
            param.requires_grad = True
        num_train = sum(1 for m in self.model.modules() if m.training)
        num_eval = sum(1 for m in self.model.modules() if not m.training)
        print(f"Modules in train mode: {num_train}, eval mode: {num_eval}")

    def training_step(self, batch, batch_idx):
        if batch is None:
            return None
        outputs = self.model(pixel_values=batch["pixel_values"], labels=batch["labels"])
        loss = outputs.loss
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

def train_model(model, train_dataloader, processor, loss_fn, optimizer, epochs, accumulation_steps):
    model.to('cuda')
    model.train()

    for epoch in range(epochs):
        optimizer.zero_grad()
        for step, batch in enumerate(train_dataloader):
            if batch is None:
                continue

            pixel_values = batch.get("pixel_values")
            labels = batch.get("labels")
            
            if pixel_values is None or labels is None:
                raise ValueError("Expected 'pixel_values' and 'labels' in batch but found None.")
            
            pixel_values = pixel_values.to('cuda').float()  # Ensure float32
            labels = labels.to('cuda')

            # Forward pass with autocast
            with autocast():
                outputs = model(pixel_values=pixel_values, labels=labels)
                loss = outputs.loss

            # Scale loss and backward pass
            loss = loss / accumulation_steps
            scaler.scale(loss).backward()

            # Update weights
            if (step + 1) % accumulation_steps == 0:
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()

            # Print progress
            if step % 10 == 0:
                print(f"Epoch {epoch+1}/{epochs}, Step {step}, Loss: {loss.item()}")

        # Clear GPU cache
        torch.cuda.empty_cache()
        
def save_checkpoint(model, optimizer, epoch, path='checkpoint.pth'):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, path)

def load_checkpoint(model, optimizer, path='checkpoint.pth'):
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    return checkpoint['epoch']

# Main function
def main():
    import os
    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    
    train_csv = 'shortened_dataset.csv'
    print("Loading training data...")
    train_df = pd.read_csv(train_csv)

    print("Creating training annotations...")
    train_annotations = create_train_annotations(train_df)

    print("Loading Donut processor and model...")
    processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base")
    model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base")
    
    for name, param in model.named_parameters():
        if not param.requires_grad:
            param.requires_grad = True

    model.train()

    special_tokens = ["
", "
"]
    processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
    model.decoder.resize_token_embeddings(len(processor.tokenizer))
    model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids("
")
    model.config.eos_token_id = processor.tokenizer.convert_tokens_to_ids("
")
    model.config.pad_token_id = processor.tokenizer.pad_token_id
    model.gradient_checkpointing_enable()
    model.config.use_cache = False
    model.config.gradient_checkpointing = True  # Enable gradient checkpointing

    print("Initializing dataset and dataloader...")
    dataset = DonutDataset(train_annotations, processor)
    train_dataloader = DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=collate_fn, num_workers=2)

    print("Starting training...")
    try:
        train_model(model, train_dataloader, processor, nn.CrossEntropyLoss(), optim.AdamW(model.parameters(), lr=1e-4), epochs=3, accumulation_steps=4)
    except RuntimeError as e:
        print(f"RuntimeError: {e}")
        torch.cuda.empty_cache()
        print("Training interrupted. Please check the GPU memory usage.")

    print("Saving trained model and processor...")
    model.save_pretrained('trained_model')
    processor.save_pretrained('trained_model')
    print("Model and processor saved successfully.")

if __name__ == "__main__":
    main()

RuntimeError: CUDA error: out of memory
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.