data_utils.py

"""
Data loading and preprocessing utilities for training and evaluation.
"""
import json
import os
from typing import List, Dict, Optional, Tuple, Union
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
# Use PIL transforms instead of torchvision to avoid version conflicts
try:
    from PIL import ImageEnhance
    HAS_PIL_ENHANCE = True
except ImportError:
    HAS_PIL_ENHANCE = False

# Avoid torchvision import to prevent Pillow compatibility issues
# Use PIL-only transforms instead
HAS_TORCHVISION = False

# Create minimal transform replacements using PIL only
class Compose:
    def __init__(self, transforms_list):
        self.transforms = transforms_list
    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

class RandomHorizontalFlip:
    def __init__(self, p=0.5):
        self.p = p
    def __call__(self, img):
        import random
        if random.random() < self.p:
            return img.transpose(Image.FLIP_LEFT_RIGHT)
        return img

class ColorJitter:
    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        self.brightness = brightness
        self.contrast = contrast
        self.saturation = saturation
        self.hue = hue
    def __call__(self, img):
        import random
        if HAS_PIL_ENHANCE:
            if self.brightness > 0:
                enhancer = ImageEnhance.Brightness(img)
                img = enhancer.enhance(1 + random.uniform(-self.brightness, self.brightness))
            if self.contrast > 0:
                enhancer = ImageEnhance.Contrast(img)
                img = enhancer.enhance(1 + random.uniform(-self.contrast, self.contrast))
            if self.saturation > 0:
                enhancer = ImageEnhance.Color(img)
                img = enhancer.enhance(1 + random.uniform(-self.saturation, self.saturation))
        return img

class RandomAffine:
    def __init__(self, degrees=0, translate=None):
        self.degrees = degrees
        self.translate = translate
    def __call__(self, img):
        import random
        angle = random.uniform(-self.degrees, self.degrees)
        if self.translate:
            max_dx = self.translate[0] * img.width
            max_dy = self.translate[1] * img.height
            dx = random.uniform(-max_dx, max_dx)
            dy = random.uniform(-max_dy, max_dy)
        else:
            dx, dy = 0, 0
        return img.rotate(angle, translate=(dx, dy))

# Create transforms module-like object
class TransformsModule:
    Compose = Compose
    RandomHorizontalFlip = RandomHorizontalFlip
    ColorJitter = ColorJitter
    RandomAffine = RandomAffine

transforms = TransformsModule()

from processing_paligemma import PaliGemmaProcessor
import logging

logger = logging.getLogger(__name__)


class PaliGemmaDataset(Dataset):
    """
    Dataset for PaliGemma training and evaluation.
    Supports various data formats (JSON, JSONL, CSV).
    """
    
    def __init__(
        self,
        data_path: str,
        processor: PaliGemmaProcessor,
        image_column: str = "image",
        text_column: str = "text",
        prompt_template: str = "{text}",
        max_samples: Optional[int] = None,
        image_augmentation: bool = False,
        image_size: int = 224,
    ):
        """
        Initialize dataset.
        
        Args:
            data_path: Path to data file or directory
            processor: PaliGemmaProcessor instance
            image_column: Column name for image paths
            text_column: Column name for text
            prompt_template: Template for formatting prompts
            max_samples: Maximum number of samples to load
            image_augmentation: Whether to apply data augmentation
            image_size: Target image size
        """
        self.processor = processor
        self.data_path = data_path
        self.image_column = image_column
        self.text_column = text_column
        self.prompt_template = prompt_template
        self.image_augmentation = image_augmentation
        self.image_size = image_size
        self.max_length = None  # Will be set if needed for consistent padding
        
        # Load data
        self.data = self._load_data(data_path)
        
        if max_samples:
            self.data = self.data[:max_samples]
        
        # Setup image augmentation
        if image_augmentation:
            self.transform = transforms.Compose([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
                transforms.RandomAffine(degrees=5, translate=(0.05, 0.05)),
            ])
        else:
            self.transform = None
        
        logger.info(f"Loaded {len(self.data)} samples from {data_path}")
    
    def _load_data(self, data_path: str) -> List[Dict]:
        """Load data from various formats"""
        if data_path.endswith('.jsonl'):
            data = []
            with open(data_path, 'r') as f:
                for line in f:
                    data.append(json.loads(line))
            return data
        elif data_path.endswith('.json'):
            with open(data_path, 'r') as f:
                return json.load(f)
        elif data_path.endswith('.csv'):
            import pandas as pd
            df = pd.read_csv(data_path)
            return df.to_dict('records')
        else:
            raise ValueError(f"Unsupported file format: {data_path}")
    
    def __len__(self) -> int:
        return len(self.data)
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Get a single sample"""
        item = self.data[idx]
        
        # Load image
        image_path = item[self.image_column]
        if not os.path.isabs(image_path):
            # Handle relative paths - try multiple locations
            base_dir = os.path.dirname(self.data_path)
            # Try relative to data file first
            candidate_path = os.path.join(base_dir, image_path)
            if not os.path.exists(candidate_path):
                # Try relative to dataset folder (parent of data file)
                candidate_path = os.path.join(os.path.dirname(base_dir), image_path)
            if not os.path.exists(candidate_path):
                # Try just the filename in the same directory as data file
                candidate_path = os.path.join(base_dir, os.path.basename(image_path))
            image_path = candidate_path
        
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image not found: {image_path} (from {item[self.image_column]})")
        
        image = Image.open(image_path).convert('RGB')
        
        # Apply augmentation if training
        if self.transform:
            image = self.transform(image)
        
        # Get text
        text = item[self.text_column]
        prompt = self.prompt_template.format(text=text)
        
        # Process with processor
        # Set max_length to ensure consistent padding - CRITICAL for batching
        if self.max_length:
            max_length = self.max_length
        else:
            # Try to get from tokenizer, but use safe default
            model_max = getattr(self.processor.tokenizer, 'model_max_length', None)
            if model_max and model_max < 1e30 and model_max < 10000:
                max_length = int(model_max)
            else:
                max_length = 1024  # Safe default for memory
        
        # Ensure max_length is always set
        if max_length is None or max_length <= 0:
            max_length = 1024
        
        model_inputs = self.processor(
            text=[prompt],
            images=[image],
            padding="max_length",
            truncation=True,
            max_length=max_length,  # Always provide max_length
        )
        
        # Verify padding worked - all sequences should be same length
        input_ids = model_inputs["input_ids"].squeeze(0)
        if input_ids.shape[0] != max_length:
            # Force pad if needed (safety check)
            import torch.nn.functional as F
            current_len = input_ids.shape[0]
            if current_len < max_length:
                pad_len = max_length - current_len
                pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0
                input_ids = F.pad(input_ids, (0, pad_len), value=pad_token_id)
                model_inputs["input_ids"] = input_ids.unsqueeze(0)
                # Also pad attention mask
                attention_mask = model_inputs["attention_mask"].squeeze(0)
                attention_mask = F.pad(attention_mask, (0, pad_len), value=0)
                model_inputs["attention_mask"] = attention_mask.unsqueeze(0)
        
        # Create labels (shifted for language modeling)
        input_ids = model_inputs["input_ids"].squeeze(0)
        attention_mask = model_inputs["attention_mask"].squeeze(0)
        labels = input_ids.clone()
        
        # Mask out image tokens and padding tokens in labels
        image_token_id = self.processor.image_token_id
        pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0
        
        labels[labels == image_token_id] = -100
        labels[labels == pad_token_id] = -100
        
        # Also mask based on attention mask (where attention is 0, label should be -100)
        labels[attention_mask == 0] = -100
        
        return {
            "input_ids": input_ids,
            "attention_mask": model_inputs["attention_mask"].squeeze(0),
            "pixel_values": model_inputs["pixel_values"].squeeze(0),
            "labels": labels,
        }


def create_dataloader(
    dataset: Dataset,
    batch_size: int = 4,
    shuffle: bool = True,
    num_workers: int = 4,
    pin_memory: bool = True,
    collate_fn: Optional[callable] = None,
) -> DataLoader:
    """
    Create a DataLoader with optimal settings.
    
    Args:
        dataset: PyTorch dataset
        batch_size: Batch size
        shuffle: Whether to shuffle
        num_workers: Number of worker processes
        pin_memory: Whether to pin memory
        collate_fn: Custom collate function
    
    Returns:
        DataLoader instance
    """
    if collate_fn is None:
        collate_fn = default_collate_fn
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        collate_fn=collate_fn,
    )


def default_collate_fn(batch: List[Dict]) -> Dict[str, torch.Tensor]:
    """Default collate function for batching with proper padding"""
    import torch.nn.functional as F
    
    # Get max length in batch
    max_len = max(item["input_ids"].shape[0] for item in batch)
    
    # Pad all sequences to same length
    padded_input_ids = []
    padded_attention_mask = []
    padded_labels = []
    
    for item in batch:
        seq_len = item["input_ids"].shape[0]
        
        # Pad input_ids (use 0 as default pad_token_id)
        if seq_len < max_len:
            pad_len = max_len - seq_len
            padding = F.pad(item["input_ids"], (0, pad_len), value=0)
        else:
            padding = item["input_ids"]
        padded_input_ids.append(padding)
        
        # Pad attention_mask
        if seq_len < max_len:
            padding_mask = F.pad(item["attention_mask"], (0, pad_len), value=0)
        else:
            padding_mask = item["attention_mask"]
        padded_attention_mask.append(padding_mask)
        
        # Pad labels (use -100 for padding tokens)
        if seq_len < max_len:
            padding_labels = F.pad(item["labels"], (0, pad_len), value=-100)
        else:
            padding_labels = item["labels"]
        padded_labels.append(padding_labels)
    
    return {
        "input_ids": torch.stack(padded_input_ids),
        "attention_mask": torch.stack(padded_attention_mask),
        "pixel_values": torch.stack([item["pixel_values"] for item in batch]),
        "labels": torch.stack(padded_labels),
    }


class VQADataset(PaliGemmaDataset):
    """Dataset for Visual Question Answering tasks"""
    
    def __init__(self, *args, question_column: str = "question", **kwargs):
        super().__init__(*args, **kwargs)
        self.question_column = question_column
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        item = self.data[idx]
        
        # Load image
        image_path = item[self.image_column]
        image = Image.open(image_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
        
        # Format as question-answer pair
        question = item[self.question_column]
        answer = item.get("answer", "")
        prompt = f"Question: {question}\nAnswer:"
        
        model_inputs = self.processor(
            text=[prompt],
            images=[image],
            padding="max_length",
            truncation=True,
        )
        
        input_ids = model_inputs["input_ids"].squeeze(0)
        labels = input_ids.clone()
        # Only compute loss on answer tokens
        # This is simplified - in practice, you'd mask question tokens
        labels[labels == self.processor.tokenizer.pad_token_id] = -100
        
        return {
            "input_ids": input_ids,
            "attention_mask": model_inputs["attention_mask"].squeeze(0),
            "pixel_values": model_inputs["pixel_values"].squeeze(0),
            "labels": labels,
            "answer": answer,
        }