ValueError: Caught ValueError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "<ipython-input-7-91ae6a087dd0>", line 82, in __getitem__
    tokens = self.tokenizer(text, **self.tokenization_params)
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 3055, in __call__
    encodings = self._call_one(text=text, text_pair=text_pair, **all_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 3114, in _call_one
    raise ValueError(
ValueError: text input must be of type `str` (single example), `List[str]` (batch or single pretokenized example) or `List[List[str]]` (batch of pretokenized examples).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/worker.py", line 309, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 52, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "<ipython-input-7-91ae6a087dd0>", line 88, in __getitem__
    raise ValueError(f"Failed to tokenize text at index {idx}: {e}")
ValueError: Failed to tokenize text at index 25663: text input must be of type `str` (single example), `List[str]` (batch or single pretokenized example) or `List[List[str]]` (batch of pretokenized examples).
import os
import random
from io import BytesIO
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
from datasets import Dataset as HFDataset
from datasets import DatasetDict, load_dataset
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from transformers import PreTrainedTokenizerBase


class ImageTextDataset(Dataset):
    """
    A custom Dataset class that handles image transformations and text tokenization.
    """

    def __init__(
        self,
        hf_dataset: HFDataset,
        image_column: Optional[str],
        text_column: Optional[str],
        image_transforms: Optional[Callable[[Image.Image], torch.Tensor]],
        tokenizer: Optional[PreTrainedTokenizerBase],
        tokenization_params: Dict[str, Any],
        seed: Optional[int] = None,
    ) -> None:
        """
        Initialize the dataset.

        Args:
            hf_dataset (HFDataset): HuggingFace Dataset.
            image_column (Optional[str]): Column name for images.
            text_column (Optional[str]): Column name for text.
            image_transforms (Optional[Callable[[Image.Image], torch.Tensor]]): Image transformations.
            tokenizer (Optional[PreTrainedTokenizerBase]): Tokenizer for text data.
            tokenization_params (Dict[str, Any]): Parameters for tokenizer.
            seed (Optional[int]): Random seed.
        """
        self.dataset = hf_dataset
        self.image_column = image_column
        self.text_column = text_column
        self.image_transforms = image_transforms
        self.tokenizer = tokenizer
        self.tokenization_params = tokenization_params
        self.seed = seed

    def __len__(self) -> int:
        """Returns the length of the dataset."""
        return len(self.dataset)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """
        Retrieves an item from the dataset at the specified index.

        Args:
            idx (int): Index of the item.

        Returns:
            Dict[str, Any]: A dictionary containing processed image and/or text data.
        """
        example = self.dataset[idx]
        data: Dict[str, Any] = {}

        if self.image_column is not None:
            try:
                image_data = example[self.image_column]
                image = self.load_image(image_data)
                if self.image_transforms:
                    image = self.image_transforms(image)
                data['image'] = image
            except Exception as e:
                raise ValueError(f"Failed to process image at index {idx}: {e}")

        if self.text_column is not None:
            try:
                text = example[self.text_column]
                if self.tokenizer:
                    tokens = self.tokenizer(text, **self.tokenization_params)
                    data['input_ids'] = torch.tensor(tokens['input_ids'])
                    data['attention_mask'] = torch.tensor(tokens['attention_mask'])
                else:
                    data['text'] = text
            except Exception as e:
                raise ValueError(f"Failed to tokenize text at index {idx}: {e}")

        if not data:
            raise ValueError(f"No data found at index {idx}.")

        return data

    def load_image(self, image_data: Any) -> Image.Image:
        """
        Loads the image from a local path, URL, or raw data.

        Args:
            image_data (Any): Image path, URL, or image data.

        Returns:
            Image.Image: Loaded image.
        """
        try:
            if isinstance(image_data, Image.Image):
                image = image_data.convert('RGB')
            elif isinstance(image_data, bytes):
                image = Image.open(BytesIO(image_data)).convert('RGB')
            elif isinstance(image_data, str):
                if os.path.isfile(image_data):
                    image = Image.open(image_data).convert('RGB')
                elif image_data.startswith('http'):
                    import requests

                    response = requests.get(image_data)
                    image = Image.open(BytesIO(response.content)).convert('RGB')
                else:
                    raise ValueError(f"Invalid image path or URL: {image_data}")
            else:
                raise ValueError(f"Unsupported image data type: {type(image_data)}")

            return image

        except Exception as e:
            raise ValueError(f"Failed to load image from {image_data}: {e}")


class DataPipeline:
    """
    A flexible data pipeline for handling image and text datasets using PyTorch, torchvision,
    and Hugging Face's datasets and transformers libraries.
    """

    def __init__(
        self,
        dataset_name: str,
        data_split: Optional[Tuple[float, float, float]] = (0.8, 0.1, 0.1),
        batch_size: int = 32,
        shuffle: bool = True,
        num_workers: int = 0,
        seed: Optional[int] = None,
        image_transforms: Optional[Callable[[Image.Image], torch.Tensor]] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        tokenization_params: Optional[Dict[str, Any]] = None,
        image_column: Optional[str] = None,
        text_column: Optional[str] = None,
        **load_dataset_kwargs: Any,
    ) -> None:
        """
        Initialize the data pipeline with the given configurations.

        Args:
            dataset_name (str): Name of the dataset to load using Hugging Face's load_dataset.
            data_split (Optional[Tuple[float, float, float]]): Split ratios for train, val, and test sets.
            batch_size (int): Batch size for the DataLoader.
            shuffle (bool): Whether to shuffle the data.
            num_workers (int): Number of workers for DataLoader.
            seed (Optional[int]): Random seed for reproducibility.
            image_transforms (Optional[Callable[[Image.Image], torch.Tensor]]): Image transformations.
            tokenizer (Optional[PreTrainedTokenizerBase]): Tokenizer for text data.
            tokenization_params (Optional[Dict[str, Any]]): Parameters for the tokenizer.
            image_column (Optional[str]): Column name for images.
            text_column (Optional[str]): Column name for text.
            **load_dataset_kwargs (Any): Additional keyword arguments for load_dataset function.
        """
        self.dataset_name = dataset_name
        self.data_split = data_split
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.num_workers = num_workers
        self.seed = seed
        self.image_transforms = image_transforms
        self.tokenizer = tokenizer
        self.tokenization_params = tokenization_params or {}
        self.load_dataset_kwargs = load_dataset_kwargs
        self.image_column = image_column
        self.text_column = text_column

        if not self.image_column and not self.text_column:
            raise ValueError("At least one of image_column or text_column must be provided.")

        # Set seeds for reproducibility
        self.set_seed()

        # Load the dataset
        self.load_dataset()

    def set_seed(self) -> None:
        """Sets the random seed for reproducibility."""
        if self.seed is not None:
            random.seed(self.seed)
            np.random.seed(self.seed)
            torch.manual_seed(self.seed)
            torch.cuda.manual_seed_all(self.seed)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False

    def load_dataset(self) -> None:
        """
        Loads the dataset using Hugging Face's load_dataset function and splits it.
        """
        try:
            self.dataset = load_dataset(self.dataset_name, 'pairwise')
        except Exception as e:
            raise ValueError(f"Failed to load dataset '{self.dataset_name}': {e}")

        # Check if dataset is already split
        if isinstance(self.dataset, DatasetDict):
            self.train_dataset = self.dataset.get('train')
            self.val_dataset = self.dataset.get('validation')
            self.test_dataset = self.dataset.get('test')

            # If validation or test sets are missing, create splits
            if self.val_dataset is None or self.test_dataset is None:
                combined_dataset = self.train_dataset
                self.train_dataset, self.val_dataset, self.test_dataset = self.split_dataset(
                    combined_dataset
                )
        else:
            # Dataset is not split, split it
            self.train_dataset, self.val_dataset, self.test_dataset = self.split_dataset(
                self.dataset
            )

        # Preprocess datasets
        self.preprocess_datasets()

    def split_dataset(self, dataset: HFDataset) -> Tuple[HFDataset, HFDataset, HFDataset]:
        """
        Splits the dataset into train, validation, and test sets.

        Args:
            dataset (HFDataset): The dataset to split.

        Returns:
            Tuple[HFDataset, HFDataset, HFDataset]: Split train, validation, and test datasets.
        """
        try:
            train_size = self.data_split[0]
            val_size = self.data_split[1] / (1 - train_size)
            test_size = self.data_split[2] / (1 - train_size)

            splits = dataset.train_test_split(
                test_size=1 - train_size, shuffle=self.shuffle, seed=self.seed
            )
            train_dataset = splits['train']
            temp_dataset = splits['test']

            val_test_splits = temp_dataset.train_test_split(
                test_size=test_size / (val_size + test_size),
                shuffle=self.shuffle,
                seed=self.seed,
            )
            val_dataset = val_test_splits['train']
            test_dataset = val_test_splits['test']
        except Exception as e:
            raise ValueError(f"Failed to split dataset: {e}")

        return train_dataset, val_dataset, test_dataset

    def preprocess_datasets(self) -> None:
        """Preprocesses the datasets by applying image transformations and text tokenization."""
        self.train_dataset = ImageTextDataset(
            hf_dataset=self.train_dataset,
            image_column=self.image_column,
            text_column=self.text_column,
            image_transforms=self.image_transforms,
            tokenizer=self.tokenizer,
            tokenization_params=self.tokenization_params,
            seed=self.seed,
        )

        self.val_dataset = ImageTextDataset(
            hf_dataset=self.val_dataset,
            image_column=self.image_column,
            text_column=self.text_column,
            image_transforms=self.image_transforms,
            tokenizer=self.tokenizer,
            tokenization_params=self.tokenization_params,
            seed=self.seed,
        )

        self.test_dataset = ImageTextDataset(
            hf_dataset=self.test_dataset,
            image_column=self.image_column,
            text_column=self.text_column,
            image_transforms=self.image_transforms,
            tokenizer=self.tokenizer,
            tokenization_params=self.tokenization_params,
            seed=self.seed,
        )

    def get_dataloaders(self) -> Tuple[DataLoader, DataLoader, DataLoader]:
        """
        Create DataLoaders for train, validation, and test datasets.

        Returns:
            Tuple[DataLoader, DataLoader, DataLoader]: DataLoaders for train, validation, and test datasets.
        """
        train_loader = DataLoader(
            dataset=self.train_dataset,
            batch_size=self.batch_size,
            shuffle=self.shuffle,
            num_workers=self.num_workers,
        )

        val_loader = DataLoader(
            dataset=self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
        )

        test_loader = DataLoader(
            dataset=self.test_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
        )

        return train_loader, val_loader, test_loader


# Example usage
if __name__ == "__main__":
    from transformers import AutoTokenizer

    # Define image transformations
    image_transforms = transforms.Compose(
        [
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            ),
        ]
    )

    # Initialize tokenizer
    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

    # Initialize data pipeline
    pipeline = DataPipeline(
        dataset_name="lmms-lab/llava-critic-113k",
        # load_dataset_kwargs={'name': 'pairwise'},
        image_column='image',
        text_column='conversations',
        image_transforms=image_transforms,
        tokenizer=tokenizer,
        tokenization_params={
            'padding': 'max_length',
            'truncation': True,
            'max_length': 128,
            'return_tensors': 'pt',
        },
        batch_size=16,
        shuffle=True,
        num_workers=4,
        seed=42,
    )

    # Get DataLoaders
    train_loader, val_loader, test_loader = pipeline.get_dataloaders()

    # Iterate over the train_loader
    for batch in train_loader:
        images = batch.get('image')
        input_ids = batch.get('input_ids')
        attention_mask = batch.get('attention_mask')
        # Perform training operations...
        break  # Remove this break to iterate over the entire dataset