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 : make generalised below api robustness , 
1)all dataset will may not have  train , test ,eval 
2) some may one split train or two train , test depend on type of datsets 

drawback back

1)PS E:\LLMS\Fine-tuning> & C:/Users/hemanthk.LAP53-FJS.000/AppData/Local/Programs/Python/Python310/python.exe e:/LLMS/Fine-tuning/2024notes/LLMs/preprocessing.py
C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\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(
2024-10-28 11:42:19,524 - INFO - __main__ - Tokenizer 'gpt2' loaded successfully.
2024-10-28 11:42:19,524 - INFO - __main__ - Pad token not found. Setting pad_token to eos_token.
2024-10-28 11:42:19,524 - INFO - __main__ - Loading dataset 'fka/awesome-chatgpt-prompts' from Hugging Face.
2024-10-28 11:42:28,944 - INFO - __main__ - Dataset 'fka/awesome-chatgpt-prompts' loaded successfully.
2024-10-28 11:42:28,944 - INFO - __main__ - Loaded dataset is a DatasetDict with splits.
2024-10-28 11:42:28,944 - INFO - __main__ - All required columns are present in the dataset.
Applying prompt template to train set: 100%|████████████████████████████████| 170/170 [00:00<00:00, 9726.65 examples/s]
Tokenizing train set: 100%|█████████████████████████████████████████████████| 170/170 [00:00<00:00, 1518.60 examples/s]
2024-10-28 11:42:29,152 - INFO - __main__ - Dataset processing complete.
2024-10-28 11:42:29,168 - INFO - __main__ - Total chunks created: 5
2024-10-28 11:42:29,168 - ERROR - __main__ - Error creating DataLoaders: 'test'
2024-10-28 11:42:29,168 - ERROR - __main__ - Failed to create DataLoaders: 'test'
PS E:\LLMS\Fine-tuning> 


import logging
import os
import random
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
from torch.utils.data import BatchSampler, DataLoader, Dataset
from transformers import PreTrainedTokenizer, AutoTokenizer
from datasets import Dataset as HFDataset, DatasetDict, load_from_disk, load_dataset

# ------------------------------
# Configuration Classes
# ------------------------------


@dataclass
class DataConfig:
    """Configuration for data loading."""
    dataset_name: Optional[str] = None  # Name of the dataset to load from Hugging Face
    dataset_config_name: Optional[str] = None  # Specific configuration name for the dataset
    folder_path: Optional[str] = None  # Path to local dataset folder
    input_columns: List[str] = None
    target_column: str = ""
    custom_prompt: str = ""
    max_length: int = 512
    batch_size: int = 32


@dataclass
class DatasetConfig:
    """Configuration for dataset processing."""
    eval_ratio: float = 0.1
    test_ratio: float = 0.1


# ------------------------------
# Logging Setup
# ------------------------------

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)


# ------------------------------
# Prompt Template
# ------------------------------


class PromptTemplate:
    """Handles formatting of prompts based on input variables."""

    def __init__(self, template: str, input_variables: List[str]) -> None:
        self.template = template
        self.input_variables = input_variables

    def format(self, **kwargs) -> str:
        formatted_kwargs = {k: kwargs.get(k, "") for k in self.input_variables}
        try:
            return self.template.format(**formatted_kwargs)
        except KeyError as e:
            missing_key = e.args[0]
            logger.error(f"Missing key in prompt formatting: {missing_key}")
            raise ValueError(f"Missing key in prompt formatting: {missing_key}") from e


# ------------------------------
# Dataset Loader
# ------------------------------


class DatasetLoader:
    """Loads datasets from a specified source."""

    def __init__(self, config: DataConfig) -> None:
        self.config = config

    def load(self) -> HFDataset:
        """Loads the dataset using either `load_dataset` or `load_from_disk`."""
        try:
            if self.config.dataset_name:
                logger.info(
                    f"Loading dataset '{self.config.dataset_name}' from Hugging Face."
                )
                dataset = load_dataset(
                    self.config.dataset_name,
                    self.config.dataset_config_name,
                )
                logger.info(f"Dataset '{self.config.dataset_name}' loaded successfully.")
            elif self.config.folder_path:
                if not os.path.exists(self.config.folder_path):
                    logger.error(
                        f"Folder path does not exist: {self.config.folder_path}"
                    )
                    raise FileNotFoundError(
                        f"Folder path does not exist: {self.config.folder_path}"
                    )
                logger.info(
                    f"Loading dataset from local path: {self.config.folder_path}"
                )
                dataset = load_from_disk(self.config.folder_path)
                logger.info(
                    f"Successfully loaded dataset from {self.config.folder_path}"
                )
            else:
                logger.error(
                    "Either `dataset_name` or `folder_path` must be provided in DataConfig."
                )
                raise ValueError(
                    "Either `dataset_name` or `folder_path` must be provided in DataConfig."
                )
            return dataset
        except Exception as e:
            logger.error(f"Failed to load dataset: {e}")
            raise


# ------------------------------
# Dataset Processor
# ------------------------------


class DatasetProcessor:
    """Processes the dataset: loading, splitting, validating, tokenizing."""

    def __init__(
        self,
        dataloader_config: DataConfig,
        dataset_config: DatasetConfig,
        tokenizer: PreTrainedTokenizer,
        prompt_template: PromptTemplate,
    ) -> None:
        self.config = dataset_config
        self.dataloader_config = dataloader_config
        self.tokenizer = tokenizer
        self.prompt_template = prompt_template

    def load_and_split_dataset(self) -> DatasetDict:
        """Loads and splits the dataset into train, test, and evaluation sets."""
        try:
            loader = DatasetLoader(config=self.dataloader_config)
            dataset = loader.load()

            if isinstance(dataset, DatasetDict):
                logger.info("Loaded dataset is a DatasetDict with splits.")
                # If dataset already has splits, we might want to handle them differently
                # For simplicity, we'll proceed to split if not already split
            else:
                logger.info("Splitting the dataset into train, test, and eval splits.")
                split_total = self.config.eval_ratio + self.config.test_ratio
                if split_total >= 1.0:
                    raise ValueError(
                        "Sum of eval_ratio and test_ratio must be less than 1."
                    )

                test_size_ratio = self.config.test_ratio / split_total
                eval_size_ratio = self.config.eval_ratio / split_total
                train_size_ratio = 1 - split_total

                split_dataset = dataset.train_test_split(
                    test_size=test_size_ratio + eval_size_ratio,
                    shuffle=True,
                    seed=42,
                )

                split_test_eval = split_dataset["test"].train_test_split(
                    test_size=eval_size_ratio / (test_size_ratio + eval_size_ratio),
                    shuffle=True,
                    seed=42,
                )

                dataset_dict = DatasetDict(
                    {
                        "train": split_dataset["train"],
                        "test": split_test_eval["train"],
                        "eval": split_test_eval["test"],
                    }
                )

                logger.info("Dataset successfully split into train, test, and eval.")
                dataset = dataset_dict

            return dataset

        except Exception as e:
            logger.error(f"Error loading and splitting dataset: {e}")
            raise

    def validate_columns(self, dataset: HFDataset) -> None:
        """Validates that all required columns are present in the dataset."""
        required_columns = self.dataloader_config.input_columns + [
            self.dataloader_config.target_column
        ]
        missing_columns = [
            col for col in required_columns if col not in dataset.column_names
        ]
        if missing_columns:
            logger.error(f"Missing columns in dataset: {missing_columns}")
            raise ValueError(f"Missing columns in dataset: {missing_columns}")
        logger.info("All required columns are present in the dataset.")

    def apply_prompt_template(self, batch: Dict[str, List]) -> Dict[str, List]:
        """Applies the prompt template to the batch of data."""
        prompts = []
        targets = []
        for i in range(len(batch[self.dataloader_config.input_columns[0]])):
            input_dict = {
                col: batch[col][i] for col in self.dataloader_config.input_columns
            }
            try:
                prompt = self.prompt_template.format(**input_dict)
                target = batch[self.dataloader_config.target_column][i]
                prompts.append(prompt)
                targets.append(target)
            except Exception as e:
                logger.error(f"Error applying prompt template: {e}")
                raise

        return {"prompt": prompts, "target": targets}

    def tokenize_and_add_labels(self, batch: Dict[str, List]) -> Dict[str, List]:
        """Tokenizes the prompts and targets, and prepares labels for model training."""
        try:
            input_ids = []
            attention_masks = []
            labels = []

            for prompt, target in zip(batch["prompt"], batch["target"]):
                encoded_prompt = self.tokenizer.encode(
                    self.tokenizer.bos_token + prompt,
                    add_special_tokens=False,
                    truncation=True,
                    max_length=self.dataloader_config.max_length // 2,
                )
                encoded_target = self.tokenizer.encode(
                    target + self.tokenizer.eos_token,
                    add_special_tokens=False,
                    truncation=True,
                    max_length=self.dataloader_config.max_length
                    - len(encoded_prompt),
                )

                combined = encoded_prompt + encoded_target
                padding_length = (
                    self.dataloader_config.max_length - len(combined)
                )

                input_id = combined + [self.tokenizer.pad_token_id] * padding_length
                attention_mask = [1] * len(combined) + [0] * padding_length
                label = (
                    [-100] * len(encoded_prompt)
                    + encoded_target
                    + [-100] * padding_length
                )

                input_ids.append(input_id)
                attention_masks.append(attention_mask)
                labels.append(label)

            return {
                "input_ids": input_ids,
                "attention_mask": attention_masks,
                "labels": labels,
            }

        except Exception as e:
            logger.error(f"Error during tokenization and label preparation: {e}")
            raise

    def process_dataset(self) -> DatasetDict:
        """Processes the dataset by applying templates and tokenization."""
        try:
            dataset_dict = self.load_and_split_dataset()

            for split in dataset_dict:
                self.validate_columns(dataset_dict[split])
                dataset_dict[split] = (
                    dataset_dict[split]
                    .map(
                        self.apply_prompt_template,
                        batched=True,
                        batch_size=self.dataloader_config.batch_size,
                        remove_columns=self.dataloader_config.input_columns
                        + [self.dataloader_config.target_column],
                        desc=f"Applying prompt template to {split} set",
                    )
                    .map(
                        self.tokenize_and_add_labels,
                        batched=True,
                        batch_size=self.dataloader_config.batch_size,
                        remove_columns=["prompt", "target"],
                        desc=f"Tokenizing {split} set",
                    )
                )

            logger.info("Dataset processing complete.")
            return dataset_dict

        except Exception as e:
            logger.error(f"Error processing dataset: {e}")
            raise


# ------------------------------
# Custom Batch Samplers
# ------------------------------


class LengthBasedBatchSampler(BatchSampler):
    """Batches samples based on sequence length to optimize padding."""

    def __init__(
        self,
        data_source: Dataset,
        batch_size: int,
        drop_last: bool = False,
        shuffle: bool = True,
    ) -> None:
        lengths = [len(sample["input_ids"]) for sample in data_source]
        self.lengths = lengths
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.shuffle = shuffle

    def __iter__(self) -> Any:
        # Sort by lengths for efficient padding
        sorted_indices = sorted(range(len(self.lengths)), key=lambda x: self.lengths[x])

        if self.shuffle:
            sorted_indices = random.sample(sorted_indices, len(sorted_indices))

        batches = [
            sorted_indices[i : i + self.batch_size]
            for i in range(0, len(sorted_indices), self.batch_size)
        ]

        if self.drop_last and len(batches[-1]) < self.batch_size:
            batches = batches[:-1]

        for batch in batches:
            yield batch

    def __len__(self) -> int:
        if self.drop_last:
            return len(self.lengths) // self.batch_size
        return (len(self.lengths) + self.batch_size - 1) // self.batch_size


class DistributedLengthBasedBatchSampler(BatchSampler):
    """Distributed sampler for length-based batching."""

    def __init__(
        self,
        data_source: Dataset,
        batch_size: int,
        num_replicas: int,
        rank: int,
        shuffle: bool = True,
        seed: int = 0,
    ) -> None:
        self.batch_sampler = LengthBasedBatchSampler(
            data_source,
            batch_size=batch_size,
            drop_last=True,
            shuffle=shuffle,
        )
        self.num_replicas = num_replicas
        self.rank = rank
        self.seed = seed

    def __iter__(self) -> Any:
        random.seed(self.seed)
        batches = list(self.batch_sampler)
        num_batches = len(batches)
        per_replica = num_batches // self.num_replicas
        remainder = num_batches % self.num_replicas

        if self.rank < remainder:
            per_replica += 1
            start = self.rank * per_replica
        else:
            start = self.rank * per_replica + remainder

        end = start + per_replica
        selected_batches = batches[start:end]

        for batch in selected_batches:
            yield batch

    def __len__(self) -> int:
        return len(self.batch_sampler) // self.num_replicas


# ------------------------------
# DataLoader Creation
# ------------------------------


def create_dataloaders(
    dataset_dict: DatasetDict,
    batch_size: int,
    tokenizer: PreTrainedTokenizer,
    num_replicas: int = 1,
    rank: int = 0,
) -> Tuple[DataLoader, Optional[DataLoader], Optional[DataLoader]]:
    """
    Creates DataLoaders for train, test, and eval datasets.

    Args:
        dataset_dict: The processed DatasetDict containing train, test, and eval splits.
        batch_size: The size of each batch.
        tokenizer: The tokenizer used for encoding.
        num_replicas: Number of distributed replicas.
        rank: Rank of the current replica.

    Returns:
        A tuple containing DataLoaders for train, test, and eval.
    """
    try:
        train_dataset = ConcatDataset(dataset_dict["train"])
        test_dataset = ConcatDataset(dataset_dict["test"])
        eval_dataset = ConcatDataset(dataset_dict["eval"])

        train_sampler = (
            DistributedLengthBasedBatchSampler(
                train_dataset,
                batch_size=batch_size,
                num_replicas=num_replicas,
                rank=rank,
                shuffle=True,
                seed=42,
            )
            if num_replicas > 1
            else LengthBasedBatchSampler(
                train_dataset,
                batch_size=batch_size,
                drop_last=True,
                shuffle=True,
            )
        )

        test_sampler = LengthBasedBatchSampler(
            test_dataset,
            batch_size=batch_size,
            drop_last=False,
            shuffle=False,
        )

        eval_sampler = LengthBasedBatchSampler(
            eval_dataset,
            batch_size=batch_size,
            drop_last=False,
            shuffle=False,
        )

        train_loader = DataLoader(
            train_dataset,
            batch_sampler=train_sampler,
            collate_fn=default_collate,
        )

        test_loader = DataLoader(
            test_dataset,
            batch_sampler=test_sampler,
            collate_fn=default_collate,
        )

        eval_loader = DataLoader(
            eval_dataset,
            batch_sampler=eval_sampler,
            collate_fn=default_collate,
        )

        logger.info("DataLoaders successfully created.")
        return train_loader, test_loader, eval_loader

    except Exception as e:
        logger.error(f"Error creating DataLoaders: {e}")
        raise


def default_collate(batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
    """
    Collate function to convert a list of samples into batched tensors.

    Args:
        batch: List of samples.

    Returns:
        A dictionary with batched tensors.
    """
    batched = {}
    for key in batch[0]:
        batched[key] = torch.stack([torch.tensor(sample[key]) for sample in batch])
    return batched


# ------------------------------
# ConcatDataset Class
# ------------------------------


class ConcatDataset(Dataset):
    """Concatenates multiple sequences into chunks of a specified size."""

    def __init__(self, dataset: HFDataset, chunk_size: int = 4096) -> None:
        self.dataset = dataset
        self.chunk_size = chunk_size
        self.samples = self._create_chunks()

    def _create_chunks(self) -> List[Dict[str, torch.Tensor]]:
        samples = []
        buffer: Dict[str, List[int]] = {
            "input_ids": [],
            "attention_mask": [],
            "labels": [],
        }

        for sample in self.dataset:
            for key in buffer:
                buffer[key].extend(sample[key])

            while len(buffer["input_ids"]) >= self.chunk_size:
                chunk = {k: buffer[k][: self.chunk_size] for k in buffer}
                samples.append(chunk)
                buffer = {k: buffer[k][self.chunk_size :] for k in buffer}

        # Handle remaining samples
        if any(buffer[k] for k in buffer):
            padding_length = self.chunk_size - len(buffer["input_ids"])
            for key in buffer:
                buffer[key].extend([0] * padding_length)  # Assuming padding with 0
            samples.append(buffer)

        logger.info(f"Total chunks created: {len(samples)}")
        return samples

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        return self.samples[idx]

    def __len__(self) -> int:
        return len(self.samples)


# ------------------------------
# Tokenizer Capture Logger
# ------------------------------


class CaptureLogger:
    """Context manager to capture logger output."""

    def __init__(self, logger: logging.Logger) -> None:
        self.logger = logger
        self.records: List[str] = []

    def __enter__(self) -> "CaptureLogger":
        self.handler = logging.StreamHandler(self._record)
        self.logger.addHandler(self.handler)
        return self

    def _record(self, record: logging.LogRecord) -> None:
        self.records.append(self.logger.format(record))

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.logger.removeHandler(self.handler)

    @property
    def out(self) -> str:
        return "\n".join(self.records)


# ------------------------------
# Tokenization Function
# ------------------------------


def tokenize_function(
    examples: Dict[str, Any],
    input_column_names: List[str],
    target_column_name: str,
    tokenizer: PreTrainedTokenizer,
    tok_logger: logging.Logger,
    max_length: int = 512,
) -> Dict[str, Any]:
    """
    Tokenizes the input and target columns in the dataset.

    Args:
        examples: A dictionary containing the data batch.
        input_column_names: A list of column names to be used as input.
        target_column_name: The target column name for labels.
        tokenizer: The tokenizer to be used for tokenization.
        tok_logger: Logger to capture tokenizer warnings or output.
        max_length: Maximum length for tokenized sequences.

    Returns:
        A dictionary with tokenized input and labels.
    """
    with CaptureLogger(tok_logger) as cl:
        try:
            input_texts = [
                " ".join(str(examples[col][i]) for col in input_column_names)
                for i in range(len(examples[input_column_names[0]]))
            ]
            target_texts = [
                str(examples[target_column_name][i])
                for i in range(len(examples[target_column_name]))
            ]

            model_inputs = tokenizer(
                input_texts,
                add_special_tokens=True,
                padding=False,
                truncation=False,
                max_length=max_length // 2,
                return_attention_mask=True,
            )
            labels = tokenizer(
                target_texts,
                add_special_tokens=False,
                padding=False,
                truncation=True,
                max_length=max_length - (max_length // 2),
            )

            combined_input_ids = []
            combined_attention_masks = []
            combined_labels = []

            for input_id, attention_mask, label in zip(
                model_inputs["input_ids"],
                model_inputs["attention_mask"],
                labels["input_ids"],
            ):
                input_id = input_id[: max_length // 2]
                label = label[: max_length - len(input_id) - 1]  # Reserving space for EOS

                input_ids = input_id + [tokenizer.eos_token_id] + label
                attention_mask_combined = attention_mask + [1] * len(label)
                labels_combined = [-100] * (len(input_id) + 1) + label

                padding_length = max_length - len(input_ids)
                if padding_length > 0:
                    input_ids += [tokenizer.pad_token_id] * padding_length
                    attention_mask_combined += [0] * padding_length
                    labels_combined += [-100] * padding_length
                else:
                    input_ids = input_ids[:max_length]
                    attention_mask_combined = attention_mask_combined[:max_length]
                    labels_combined = labels_combined[:max_length]

                combined_input_ids.append(input_ids)
                combined_attention_masks.append(attention_mask_combined)
                combined_labels.append(labels_combined)

            return {
                "input_ids": combined_input_ids,
                "attention_mask": combined_attention_masks,
                "labels": combined_labels,
            }

        except Exception as e:
            logger.error(f"Error in tokenization function: {e}")
            raise

    if "Token indices sequence length is longer than the" in cl.out:
        tok_logger.warning("Long input sequences have been truncated or chunked.")

    return {}


# ------------------------------
# Example Usage
# ------------------------------


def main() -> None:
    # Example configurations
    data_config = DataConfig(
        dataset_name="fka/awesome-chatgpt-prompts",  # Specify the dataset name from Hugging Face
        input_columns=["act", "prompt"],
        target_column="prompt",
        custom_prompt="Translate the following text: {prompt} and {act}",
        max_length=100,
        batch_size=16,
    )

    dataset_config = DatasetConfig(
        eval_ratio=0.1,
        test_ratio=0.1,
    )

    # Initialize tokenizer
    tokenizer_name = "gpt2"
    try:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        logger.info(f"Tokenizer '{tokenizer_name}' loaded successfully.")
    except Exception as e:
        logger.error(f"Failed to load tokenizer '{tokenizer_name}': {e}")
        raise

    # Ensure the tokenizer has a pad token
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        logger.info(
            "Pad token not found. Setting pad_token to eos_token."
        )

    # Initialize prompt template
    prompt_template = PromptTemplate(
        template=data_config.custom_prompt,
        input_variables=data_config.input_columns,
    )

    # Process the dataset
    processor = DatasetProcessor(
        dataloader_config=data_config,
        dataset_config=dataset_config,
        tokenizer=tokenizer,
        prompt_template=prompt_template,
    )

    try:
        processed_dataset = processor.process_dataset()
    except Exception as e:
        logger.error(f"Dataset processing failed: {e}")
        return

    # Create DataLoaders
    try:
        train_loader, test_loader, eval_loader = create_dataloaders(
            processed_dataset,
            batch_size=data_config.batch_size,
            tokenizer=tokenizer,
        )
    except Exception as e:
        logger.error(f"Failed to create DataLoaders: {e}")
        return

    # Example iteration
    try:
        for batch in train_loader:
            input_ids = batch["input_ids"]
            attention_mask = batch["attention_mask"]
            labels = batch["labels"]
            logger.info(f"Batch input_ids shape: {input_ids.shape}")
            logger.info(f"Batch attention_mask shape: {attention_mask.shape}")
            logger.info(f"Batch labels shape: {labels.shape}")
            break  # Remove this break to iterate through the entire dataset
    except Exception as e:
        logger.error(f"Error during DataLoader iteration: {e}")
        raise


if __name__ == "__main__":
    main()