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 : from datasets module we have generalised pipeline preprocessing , user_input : str , path of folder , with columns name and custopm prompt 

tokenizer from transformers module 

loading 
split( may be train or custom name ("sum_train" de dont user input )
split into 80:10:10 train , test, eval

pre_processing tokenization

make sure robustness , genmeralised , make size and shaping 

dataloader

for batchs, labels in train_loader:
    print(batch) Batch_size,seq_length


loading 
import logging
import copy
from typing import Dict, List
from datasets import Dataset, DatasetDict
from transformers import PreTrainedTokenizer



logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)



class PromptTemplate:
    def __init__(self, template: str, input_variables: List[str]):
        self.template = template
        self.input_variables = input_variables

    def format(self, **kwargs) -> str:
        return self.template.format(**{k: kwargs.get(k, '') for k in self.input_variables})

class DatasetProcessor:
    def __init__(
        self,
        dataloader_config:  DataConfig,
        dataset_config: DatasetConfig,
        tokenizer: PreTrainedTokenizer,
        prompt_template: PromptTemplate,
    ):
        self.config = dataset_config
        self.datasetloader_config=dataloader_config
        self.tokenizer = tokenizer
        self.prompt_template = prompt_template
        

    def load_and_split_dataset(self) -> DatasetDict:
        try:
            dataset = DatasetLoader(config=self.datasetloader_config).load()
            return dataset.train_test_split(
                test_size=self.config.eval_ratio + self.config.test_ratio,
                shuffle=True,
                seed=42,
            )
        except Exception as e:
            logger.error(f"Error loading dataset: {e}")
            raise

    def validate_columns(self, dataset: Dataset):
        all_columns = self.config.input_columns + [self.config.target_column]
        missing_columns = [col for col in all_columns if col not in dataset.column_names]
        if missing_columns:
            raise ValueError(f"Missing columns in dataset: {missing_columns}")

    def apply_prompt_template(self, batch: Dict[str, List]) -> Dict[str, List]:
        prompts = []
        targets = []
        for i in range(len(batch[self.config.input_columns[0]])):
            input_dict = {col: batch[col][i] for col in self.config.input_columns}
            prompts.append(self.prompt_template.format(**input_dict))
            targets.append(batch[self.config.target_column][i])
        return {"prompt": prompts, "target": targets}

    def tokenize_and_add_labels(self, batch: Dict[str, List]) -> Dict[str, List]:
        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,
        }

    def process_dataset(self) -> DatasetDict:
        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=dataset_dict[split].column_names,
                    )
                    .map(
                        self.tokenize_and_add_labels,
                        batched=True,
                        batch_size=self.config.batch_size,
                        remove_columns=["prompt", "target"],
                    )
                )
            return dataset_dict
        except Exception as e:
            logger.error(f"Error processing dataset: {e}")
            raise

def tokenize_function(
    examples: Dict[str, Any],
    input_column_names: List[str],
    target_column_name: str,
    tokenizer: Any,
    tok_logger: logging.Logger,
    max_length: int=512
) -> Dict[str, Any]:
    """
    Tokenize 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:
        # Concatenate input columns to form the input text for each example in the batch
        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)
        labels = tokenizer(target_texts, add_special_tokens=False, padding=False, truncation=False)

        for i in range(len(input_texts)):
            sample_input_ids = model_inputs["input_ids"][i]
            label_input_ids = labels["input_ids"][i] + [tokenizer.eos_token_id]

            combined_input_ids = sample_input_ids + label_input_ids
            combined_labels = [-100] * len(sample_input_ids) + label_input_ids

            padding_length = max(0, max_length - len(combined_input_ids))

            # Pad input_ids, attention_mask, and labels
            padded_input_ids = [tokenizer.pad_token_id] * padding_length + combined_input_ids
            padded_attention_mask = [0] * padding_length + [1] * len(combined_input_ids)
            padded_labels = [-100] * padding_length + combined_labels

            # Truncate to ensure no overflow beyond max_length
            model_inputs["input_ids"][i] = torch.tensor(padded_input_ids[:max_length])
            model_inputs["attention_mask"][i] = torch.tensor(padded_attention_mask[:max_length])
            labels["input_ids"][i] = torch.tensor(padded_labels[:max_length])

        model_inputs["labels"] = labels["input_ids"]

    if "Token indices sequence length is longer than the" in cl.out:
        tok_logger.warning(
            "Please ignore the warning above - this long input will be chunked "
            "into smaller bits before being passed to the model."
        )

    return model_inputs

from tqdm import tqdm
from itertools import chain

from torch.utils.data import Dataset


class ConcatDataset(Dataset):
    def __init__(self, dataset, chunk_size=4096):
        self.dataset = dataset
        self.chunk_size = chunk_size

        self.samples = []

        buffer = {
            "input_ids": [],
            "attention_mask": [],
            "labels": [],
            }

        for sample in tqdm(self.dataset, desc="Preprocessing dataset", dynamic_ncols=True):
            buffer = {k: v + sample[k] for k,v in buffer.items()}

            while len(next(iter(buffer.values()))) > self.chunk_size:
                self.samples.append({k: v[:self.chunk_size] for k,v in buffer.items()})
                buffer = {k: v[self.chunk_size:] for k,v in buffer.items()}

    def __getitem__(self, idx):
        return self.samples[idx]

    def __len__(self):
        return len(self.samples)

ideA 2
import random
from itertools import islice

import numpy as np
import torch


class LengthBasedBatchSampler(torch.utils.data.BatchSampler):
    def __init__(self, data_source, batch_size: int, drop_last: bool, shuffle: bool=True) -> None:
        if isinstance(next(iter(data_source)), dict):
            first_key = next(iter(next(iter(data_source)).keys()))
            self.lengths = [len(d[first_key]) for d in data_source]
        else:
            self.lengths = [len(d) for d in data_source]
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.shuffle = shuffle

    def __iter__(self):
        ids = np.argsort(self.lengths, kind='mergesort')
        if self.drop_last:
            ids = ids[:len(ids) // self.batch_size * self.batch_size]

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

        if self.shuffle:
            random.shuffle(batches)

        for b in batches:
            yield b

    def __len__(self):
        if self.drop_last:
            return len(self.lengths) // self.batch_size
        else:
            return len(self.lengths) // self.batch_size + (len(self.lengths) % self.batch_size > 0)


class DistributedLengthBasedBatchSampler(torch.utils.data.BatchSampler):
    def __init__(self, data_source, batch_size: int, num_replicas: int, rank: int, shuffle: bool = True, seed: int = 0) -> None:
        random.seed(seed)
        self.batch_sampler = LengthBasedBatchSampler(
            data_source, batch_size=batch_size, drop_last=True, shuffle=shuffle
            )
        self.num_replicas = num_replicas
        self.rank = rank

    def __iter__(self):
        max_length = len(self.batch_sampler) // self.num_replicas * self.num_replicas
        return islice(self.batch_sampler, self.rank, max_length, self.num_replicas)

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