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 : 

below drawback :
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
  File "e:\LLMS\Fine-tuning\2024notes\LLMs\preprocessing.py", line 123
    test_size=split_size=test_size + eval_size,
                        ^
SyntaxError: invalid syntax
PS E:\LLMS\Fine-tuning>

2) incorporate  this api aslo load_dataset( download from huggingface datasets


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."""
    folder_path: str
    input_columns: List[str]
    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(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


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

class DatasetLoader:
    """Loads datasets from a specified folder path."""

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

    def load(self) -> HFDataset:
        if not os.path.exists(self.folder_path):
            logger.error(f"Folder path does not exist: {self.folder_path}")
            raise FileNotFoundError(f"Folder path does not exist: {self.folder_path}")

        try:
            dataset = load_from_disk(self.folder_path)
            logger.info(f"Successfully loaded dataset from {self.folder_path}")
            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()

            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 = self.config.test_ratio / (split_total)
            eval_size = self.config.eval_ratio / (split_total)
            train_size = 1 - split_total

            split_dataset = dataset.train_test_split(
                test_size=split_size=test_size + eval_size,
                shuffle=True,
                seed=42,
            )

            split_test_eval = split_dataset["test"].train_test_split(
                test_size=eval_size / (test_size + eval_size),
                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.")
            return dataset_dict

        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.config.max_length // 2,
                )
                encoded_target = self.tokenizer.encode(
                    target + self.tokenizer.eos_token,
                    add_special_tokens=False,
                    truncation=True,
                    max_length=self.config.max_length - len(encoded_prompt),
                )

                combined = encoded_prompt + encoded_target
                padding_length = self.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.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.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:
            random.shuffle(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():
    # Example configurations
    data_config = DataConfig(
        folder_path="fka/awesome-chatgpt-prompts",
        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

    tokenizer.pad_token = tokenizer.eos_token  # Ensuring pad token is defined

    # 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()