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 copy.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 12:03:45,216 - INFO - __main__ - Tokenizer 'gpt2' loaded successfully.
2024-10-28 12:03:45,217 - INFO - __main__ - Pad token not found. Setting pad_token to eos_token.
2024-10-28 12:03:45,217 - INFO - __main__ - Loading dataset 'fka/awesome-chatgpt-prompts' from Hugging Face.
2024-10-28 12:03:55,344 - INFO - __main__ - Dataset 'fka/awesome-chatgpt-prompts' loaded successfully.
2024-10-28 12:03:55,345 - INFO - __main__ - Loaded dataset is a DatasetDict with splits.
2024-10-28 12:03:55,345 - INFO - __main__ - Available splits: dict_keys(['train'])
2024-10-28 12:03:55,346 - INFO - __main__ - Processing split: train
2024-10-28 12:03:55,346 - INFO - __main__ - All required columns are present in the dataset.
2024-10-28 12:03:55,427 - INFO - __main__ - Dataset processing complete.
2024-10-28 12:03:55,428 - INFO - __main__ - Dataset processing completed successfully.
2024-10-28 12:03:55,428 - INFO - __main__ - Creating DataLoader for split: train
2024-10-28 12:03:55,451 - INFO - __main__ - Total chunks created: 5
2024-10-28 12:03:55,451 - INFO - __main__ - Train DataLoader successfully created.
2024-10-28 12:03:55,451 - INFO - __main__ - DataLoaders created successfully.0
Traceback (most recent call last):
  File "e:\LLMS\Fine-tuning\2024notes\LLMs\preprocessing copy.py", line 761, in <module>
    main()
  File "e:\LLMS\Fine-tuning\2024notes\LLMs\preprocessing copy.py", line 748, in main
    next(iter(loaders["train"]))
  File "C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\dataloader.py", line 631, in __next__
    data = self._next_data()
  File "C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\dataloader.py", line 674, in _next_data
    index = self._next_index()  # may raise StopIteration
  File "C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\dataloader.py", line 621, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
StopIteration
PS E:\LLMS\Fine-tuning> 
import logging
import os
import random
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

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] = field(default_factory=list)
    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 available splits."""
        try:
            loader = DatasetLoader(config=self.dataloader_config)
            dataset = loader.load()

            if isinstance(dataset, DatasetDict):
                logger.info("Loaded dataset is a DatasetDict with splits.")
                available_splits = dataset.keys()
                logger.info(f"Available splits: {available_splits}")
                # No further splitting; use existing splits
            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()

            available_splits = dataset_dict.keys()
            for split in available_splits:
                logger.info(f"Processing split: {split}")
                split_dataset = dataset_dict[split]
                self.validate_columns(split_dataset)
                split_dataset = (
                    split_dataset
                    .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",
                    )
                )
                dataset_dict[split] = split_dataset

            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:
        self.lengths = [len(sample["input_ids"]) for sample in data_source]
        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,
    splits:List[str],
    batch_size: int,
    tokenizer: PreTrainedTokenizer,
    num_replicas: int = 1,
    rank: int = 0,
) -> Dict[str, Optional[DataLoader]]:
    """
    Creates DataLoaders for available dataset splits.

    Args:
        dataset_dict: The processed DatasetDict containing available 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 dictionary containing DataLoaders for available splits. Missing splits will have None.
    """
    try:
        loaders: Dict[str, Optional[DataLoader]] = {}
        for split in splits:
            if split in dataset_dict:
                logger.info(f"Creating DataLoader for split: {split}")
                dataset = ConcatDataset(dataset_dict[split])
                if split == "train":
                    sampler = (
                        DistributedLengthBasedBatchSampler(
                            dataset,
                            batch_size=batch_size,
                            num_replicas=num_replicas,
                            rank=rank,
                            shuffle=True,
                            seed=42,
                        )
                        if num_replicas > 1
                        else LengthBasedBatchSampler(
                            dataset,
                            batch_size=batch_size,
                            drop_last=True,
                            shuffle=True,
                        )
                    )
                else:
                    sampler = LengthBasedBatchSampler(
                        dataset,
                        batch_size=batch_size,
                        drop_last=False,
                        shuffle=False,
                    )

                loaders[split] = DataLoader(
                    dataset,
                    batch_sampler=sampler,
                    collate_fn=default_collate,
                )
                logger.info(f"{split.capitalize()} DataLoader successfully created.")
            else:
                logger.warning(f"Split '{split}' not found in the dataset. Skipping.")
                loaders[split] = None

        return loaders

    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]:
        try:
            batched[key] = torch.stack([torch.tensor(sample[key]) for sample in batch])
        except Exception as e:
            logger.error(f"Error collating batch for key '{key}': {e}")
            raise
    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: List[Dict[str, torch.Tensor]] = self._create_chunks()

    def _create_chunks(self) -> List[Dict[str, torch.Tensor]]:
        samples: List[Dict[str, List[int]]] = []
        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.formatter = logging.Formatter("%(message)s")
        self.handler.setFormatter(self.formatter)
        self.logger.addHandler(self.handler)
        return self

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

    def format(self, record: logging.LogRecord) -> str:
        return self.formatter.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()
        logger.info("Dataset processing completed successfully.")
    except Exception as e:
        logger.error(f"Dataset processing failed: {e}")
        return

    # Create DataLoaders
    try:
        loaders = create_dataloaders(
            
            processed_dataset,
            batch_size=data_config.batch_size,
            tokenizer=tokenizer,
            splits=["train"],
        )
        logger.info(f"DataLoaders created successfully.{len(loaders['train'])}")
    except Exception as e:
        logger.error(f"Failed to create DataLoaders: {e}")
        return

    next(iter(loaders["train"]))
    # logger.info(f"Iterating over the {split} DataLoader:")
    for batch in loaders["train"]:
        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
        

if __name__ == "__main__":
    main()