Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.

TASK : we have end to end pipeline , (generalised , robustness any scale of dataset like 2gb to 2000GB) thing all possible 
Make pipeline robustness, scalable, generalised , 


 pre_processing , loading , batching , etc features 
train_loader, test_loader, eval_loader


import os
import io
import requests
from PIL import Image
import torch
from torchvision import transforms
from typing import Union, Tuple, Optional

class ImageLoader:
    def __init__(self, resize: Optional[Tuple[int, int]] = None,
                 mean: Optional[Tuple[float, float, float]] = (0.485, 0.456, 0.406),
                 std: Optional[Tuple[float, float, float]] = (0.229, 0.224, 0.225)):
        """
        Initialize the ImageLoader with optional resizing and normalization parameters.

        Args:
            resize (tuple, optional): Tuple containing the target (width, height). Defaults to None.
            mean (tuple, optional): Tuple containing mean values for normalization. Defaults to (0.485, 0.456, 0.406).
            std (tuple, optional): Tuple containing standard deviation values for normalization. Defaults to (0.229, 0.224, 0.225).
        """
        self.resize = resize
        self.mean = mean
        self.std = std

    def load_image(self, source: Union[str, os.PathLike]) -> torch.Tensor:
        """
        Load an image from a local path or an HTTP URL, preprocess it (optional resizing and normalization), and return a tensor.

        Args:
            source (str or PathLike): File path or URL from which to load the image.

        Returns:
            torch.Tensor: Preprocessed image as a tensor.
        """
        try:
            # Load image from URL or local path
            if self._is_url(source):
                image = self._load_from_url(source)
            else:
                image = self._load_from_local(source)

            # Apply transformations: resize (if specified), convert to tensor, and normalize
            transform_pipeline = self._build_transform_pipeline()
            image_tensor = transform_pipeline(image)

            return image_tensor

        except Exception as e:
            raise RuntimeError(f"Error loading image from {source}: {str(e)}")

    def _load_from_local(self, path: Union[str, os.PathLike]) -> Image.Image:
        """
        Load image from a local file system path.

        Args:
            path (str or PathLike): Local file path.

        Returns:
            Image.Image: Loaded image.
        """
        if not os.path.exists(path):
            raise FileNotFoundError(f"File not found: {path}")

        try:
            image = Image.open(path).convert('RGB')
            return image
        except Exception as e:
            raise RuntimeError(f"Failed to load image from path: {path} - {str(e)}")

    def _load_from_url(self, url: str) -> Image.Image:
        """
        Load image from an HTTP URL.

        Args:
            url (str): URL of the image to be loaded.

        Returns:
            Image.Image: Loaded image.
        """
        try:
            response = requests.get(url, timeout=10)
            response.raise_for_status()  # Raise HTTPError for bad responses
            image = Image.open(io.BytesIO(response.content)).convert('RGB')
            return image
        except requests.RequestException as req_err:
            raise RuntimeError(f"Failed to load image from URL: {url} - {str(req_err)}")
        except Exception as e:
            raise RuntimeError(f"Error processing image from URL: {url} - {str(e)}")

    def _build_transform_pipeline(self):
        """
        Build the transformation pipeline that includes optional resizing, conversion to tensor, and normalization.

        Returns:
            torchvision.transforms.Compose: Transformation pipeline.
        """
        transformations = []
        if self.resize:
            transformations.append(transforms.Resize(self.resize))

        transformations.extend([
            transforms.ToTensor(),
            transforms.Normalize(mean=self.mean, std=self.std)
        ])

        return transforms.Compose(transformations)

    @staticmethod
    def _is_url(path: str) -> bool:
        """
        Check if the provided path is a URL.

        Args:
            path (str): String to check.

        Returns:
            bool: True if the path is a URL, False otherwise.
        """
        return path.startswith('http://') or path.startswith('https://')


# Example Usage
if __name__ == "__main__":
    # Initialize the loader with optional resize settings
    loader = ImageLoader(resize=(224, 224))

    # Load an image from a local path or URL
    # local_image_path = r'E:\LLMS\Fine-tuning\cat.png'
    url_image = 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/cb/The_Prime_Minister%2C_Shri_Narendra_Modi_meeting_Mr._Elon_Musk_in_New_York%2C_USA_on_June_20%2C_2023_%282%29_%28cropped%29_%28b%29.jpg/330px-The_Prime_Minister%2C_Shri_Narendra_Modi_meeting_Mr._Elon_Musk_in_New_York%2C_USA_on_June_20%2C_2023_%282%29_%28cropped%29_%28b%29.jpg'
    try:
        # # Load from local
        # local_tensor = loader.load_image(local_image_path)
        # print(f"Loaded image from local path with shape: {local_tensor.shape}")

        # Load from URL
        url_tensor = loader.load_image(url_image)
        print(f"Loaded image from URL with shape: {url_tensor.shape}")
    except RuntimeError as e:
        print(e)

from torch.utils.data.dataloader import (
    _DatasetKind,
    DataLoader,
    default_collate,
    default_convert,
    get_worker_info,
)
from torch.utils.data.datapipes._decorator import (
    argument_validation,
    functional_datapipe,
    guaranteed_datapipes_determinism,
    non_deterministic,
    runtime_validation,
    runtime_validation_disabled,
)
from torch.utils.data.datapipes.datapipe import (
    DataChunk,
    DFIterDataPipe,
    IterDataPipe,
    MapDataPipe,
)
from torch.utils.data.dataset import (
    ChainDataset,
    ConcatDataset,
    Dataset,
    IterableDataset,
    random_split,
    StackDataset,
    Subset,
    TensorDataset,
)
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data.sampler import (
    BatchSampler,
    RandomSampler,
    Sampler,
    SequentialSampler,
    SubsetRandomSampler,
    WeightedRandomSampler,
)

import torch
from torch.utils.data import DataLoader, Dataset, random_split
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import TensorDataset
from typing import Tuple, Union, Optional, List
import numpy as np

# Define the custom dataset class
class CustomImageDataset(Dataset):
    """
    Custom Dataset class to handle input tensors and associated labels.
    This class extends PyTorch's Dataset class and provides basic functionality to
    retrieve samples and labels.
    """

    def __init__(self, inputs: torch.Tensor, labels: torch.Tensor):
        """
        Initializes the CustomImageDataset with input tensors and labels.

        Args:
            inputs (torch.Tensor): Tensor of shape (N, C, H, W) where N is the number of samples.
            labels (torch.Tensor): Tensor of corresponding labels.
        """
        assert inputs.size(0) == labels.size(0), "Inputs and labels must have the same number of samples."
        self.inputs = inputs
        self.labels = labels

    def __len__(self) -> int:
        """
        Returns the total number of samples in the dataset.
        """
        return self.inputs.size(0)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Retrieves the input tensor and corresponding label at the specified index.

        Args:
            idx (int): Index of the sample to retrieve.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: A tuple containing the input tensor and label.
        """
        return self.inputs[idx], self.labels[idx]


# Preprocessing function to scale or normalize the data if needed
def preprocess_data(inputs: torch.Tensor, labels: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Preprocesses the input data and labels. This function can be customized to include
    transformations like normalization, augmentations, etc.

    Args:
        inputs (torch.Tensor): Input tensor of shape (N, C, H, W).
        labels (torch.Tensor): Labels tensor.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Preprocessed inputs and labels.
    """
    # Example: Normalize inputs to [0, 1] (assuming the images are in range [0, 255])
    inputs = inputs.float() / 255.0
    return inputs, labels


# Function to create data loaders for training, validation, and testing
def create_data_loaders(
    dataset: Dataset,
    batch_size: int,
    val_split: float = 0.2,
    test_split: float = 0.1,
    shuffle: bool = True,
    num_workers: int = 4,
    distributed: bool = False
) -> Tuple[DataLoader, DataLoader, DataLoader]:
    """
    Creates DataLoaders for training, validation, and testing.

    Args:
        dataset (Dataset): The full dataset to be split.
        batch_size (int): The batch size for DataLoader.
        val_split (float): The proportion of the dataset to use for validation.
        test_split (float): The proportion of the dataset to use for testing.
        shuffle (bool): Whether to shuffle the dataset before splitting.
        num_workers (int): The number of worker threads for data loading.
        distributed (bool): Whether to use DistributedSampler for distributed training.

    Returns:
        Tuple[DataLoader, DataLoader, DataLoader]: DataLoaders for training, validation, and testing.
    """
    dataset_size = len(dataset)
    indices = list(range(dataset_size))
    
    test_size = int(test_split * dataset_size)
    val_size = int(val_split * dataset_size)
    train_size = dataset_size - val_size - test_size

    if shuffle:
        np.random.seed(42)  # For reproducibility
        np.random.shuffle(indices)

    train_indices, val_indices, test_indices = indices[:train_size], indices[train_size:train_size + val_size], indices[train_size + val_size:]

    train_sampler = SubsetRandomSampler(train_indices) if not distributed else DistributedSampler(Subset(dataset, train_indices))
    val_sampler = SubsetRandomSampler(val_indices)
    test_sampler = SubsetRandomSampler(test_indices)

    train_loader = DataLoader(dataset, batch_size=batch_size, sampler=train_sampler, num_workers=num_workers)
    val_loader = DataLoader(dataset, batch_size=batch_size, sampler=val_sampler, num_workers=num_workers)
    test_loader = DataLoader(dataset, batch_size=batch_size, sampler=test_sampler, num_workers=num_workers)

    return train_loader, val_loader, test_loader


# Example of a full pipeline
def data_pipeline(
    input_tensor: torch.Tensor,
    labels_tensor: torch.Tensor,
    batch_size: int = 32,
    val_split: float = 0.2,
    test_split: float = 0.1,
    shuffle: bool = True,
    num_workers: int = 4,
    distributed: bool = False
) -> Tuple[DataLoader, DataLoader, DataLoader]:
    """
    Full data pipeline that preprocesses the data and creates train, validation, and test DataLoaders.

    Args:
        input_tensor (torch.Tensor): Input tensor of shape (N, C, H, W).
        labels_tensor (torch.Tensor): Corresponding labels tensor.
        batch_size (int): Batch size for DataLoaders.
        val_split (float): Proportion of dataset to use for validation.
        test_split (float): Proportion of dataset to use for testing.
        shuffle (bool): Whether to shuffle the dataset before splitting.
        num_workers (int): Number of workers for DataLoader.
        distributed (bool): Whether to use DistributedSampler for distributed training.

    Returns:
        Tuple[DataLoader, DataLoader, DataLoader]: DataLoaders for training, validation, and testing.
    """
    # Preprocess the data
    inputs, labels = preprocess_data(input_tensor, labels_tensor)

    # Create a dataset from the preprocessed data
    dataset = CustomImageDataset(inputs, labels)

    # Create data loaders for train, validation, and test sets
    train_loader, val_loader, test_loader = create_data_loaders(
        dataset,
        batch_size=batch_size,
        val_split=val_split,
        test_split=test_split,
        shuffle=shuffle,
        num_workers=num_workers,
        distributed=distributed
    )

    return train_loader, val_loader, test_loader


# Example usage
if __name__ == '__main__':
    # Example input tensor (N, C, H, W) and labels tensor (N,)
    N, C, H, W = 1000, 3, 64, 64  # Example dimensions (1000 samples, 3 channels, 64x64 images)
    inputs = torch.randint(0, 256, (N, C, H, W), dtype=torch.uint8)
    labels = torch.randint(0, 10, (N,), dtype=torch.int64)

    # Create the data pipeline
    train_loader, val_loader, test_loader = data_pipeline(
        input_tensor=inputs,
        labels_tensor=labels,
        batch_size=32,
        val_split=0.2,
        test_split=0.1,
        shuffle=True,
        num_workers=4
    )

    # Example iteration over the train_loader
    for batch_idx, (data, target) in enumerate(train_loader):
        print(f"Batch {batch_idx}: data shape = {data.shape}, target shape = {target.shape}")