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
from torch.utils.data.dataloader import default_collate

# Import PEFT for LoRA
from peft import get_peft_model, LoraConfig, TaskType

# Set environment variables and warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
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))  # Adjust resolution as needed
                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)
        
        pixel_values = self.processor(image, return_tensors="pt").pixel_values.squeeze()
        labels = self.processor.tokenizer(
            prompt,
            padding="longest", 
            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 function
def train_model(model, train_dataloader, optimizer, device, epochs):
    model.to(device)
    model.train()

    for epoch in range(epochs):
        total_loss = 0.0

        for step, batch in enumerate(train_dataloader):
            if batch is None:
                continue

            pixel_values = batch["pixel_values"].to(device).float()
            labels = batch["labels"].to(device)

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

            # Backward pass and optimization
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            total_loss += loss.item()

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

        avg_loss = total_loss / len(train_dataloader)
        print(f"Epoch {epoch+1}/{epochs}, Average Loss: {avg_loss:.4f}")

def main():
    # Set the device to GPU if available
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    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")
    
    # Add special tokens
    special_tokens = ["
", "
"]
    processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
    model.decoder.resize_token_embeddings(len(processor.tokenizer))

    # Set decoder parameters
    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

    # Enable gradient checkpointing
    model.gradient_checkpointing_enable()
    model.config.use_cache = False

    # Freeze all parameters
    for param in model.parameters():
        param.requires_grad = False

    # Apply LoRA to the model
    peft_config = LoraConfig(
        task_type=TaskType.SEQ_2_SEQ_LM,
        inference_mode=False,
        r=8,  # Rank of LoRA matrices (adjustable)
        lora_alpha=32,
        lora_dropout=0.1,
    )

    model = get_peft_model(model, peft_config)
    print("LoRA model is created and applied.")

    # Print trainable parameters
    model.print_trainable_parameters()

    print("Initializing dataset and dataloader...")
    dataset = DonutDataset(train_annotations, processor)
    train_dataloader = DataLoader(
        dataset,
        batch_size=1,  # Keep batch size small to manage memory
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=0
    )

    print("Starting training...")
    optimizer = optim.AdamW(model.parameters(), lr=1e-3)  # Adjust learning rate if needed
    try:
        train_model(model, train_dataloader, optimizer, device, epochs=3)
    except RuntimeError as e:
        print(f"RuntimeError: {e}")
        print("Training interrupted. Please check the GPU memory usage.")

    print("Saving trained model and processor...")
    # Save the LoRA adapter separately
    model.save_pretrained('trained_model_lora')
    processor.save_pretrained('trained_model_lora')
    print("Model and processor saved successfully.")

if __name__ == "__main__":
    main()

PS C:\Users\VICTUS\Product-Identification> & "C:/Program Files/Python312/python.exe" c:/Users/VICTUS/Product-Identification/c.py
Using device: cuda
Loading training data...
Creating training annotations...
Loading Donut processor and model...
2024-09-15 14:37:17.716275: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-09-15 14:37:22.632271: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\transformers\tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884
  warnings.warn(
Traceback (most recent call last):
  File "c:\Users\VICTUS\Product-Identification\c.py", line 197, in <module>
    main()
  File "c:\Users\VICTUS\Product-Identification\c.py", line 166, in main
    model = get_peft_model(model, peft_config)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\mapping.py", line 183, in get_peft_model
    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\peft_model.py", line 1757, in __init__
    super().__init__(model, peft_config, adapter_name, **kwargs)
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\peft_model.py", line 155, in __init__
    self.base_model = cls(model, {adapter_name: peft_config}, adapter_name)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\tuners\lora\model.py", line 139, in __init__
    super().__init__(model, config, adapter_name)
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\tuners\tuners_utils.py", line 175, in __init__
    self.inject_adapter(self.model, adapter_name)
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\tuners\tuners_utils.py", line 394, in inject_adapter
    peft_config = self._prepare_adapter_config(peft_config, model_config)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\VICTUS\AppData\Roaming\Python\Python312\site-packages\peft\tuners\lora\model.py", line 459, in _prepare_adapter_config
    raise ValueError("Please specify `target_modules` in `peft_config`")
ValueError: Please specify `target_modules` in `peft_config`
PS C:\Users\VICTUS\Product-Identification>