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 : Generalised robustness tokenization  Pipeline , (we rewrite using transformers module ) avoid using tiketoken
user_inputs
1)df.dataframe multi columns and rows
 2)tokenizer(transformers module) 
special tokens eos, bos, pos , unk etc..

max_length


return: input_ids, masked_ids,

sample or reference code :
import pandas as pd
import torch
from transformers import AutoTokenizer
from typing import Dict, List, Tuple, Union
import numpy as np
from dataclasses import dataclass
from torch.nn.utils.rnn import pad_sequence
import math
@dataclass
class TokenizerConfig:
    model_name: str
    max_length: int
    padding_side: str = "right"
    truncation_side: str = "right"
    special_tokens: Dict[str, str] = None
    
class AdvancedWorldModelPreprocessor:
    def __init__(
        self,
        config: TokenizerConfig
    ):
        self.config = config
        self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
        self._setup_special_tokens()
        
    def _setup_special_tokens(self):
        """Setup special tokens with advanced configurations"""
        special_tokens_dict = {
            "bos_token": "
",
            "eos_token": "
",
            "pad_token": "
",
            "unk_token": "
",
            "sep_token": "
",
            "mask_token": "
",
            # Add custom tokens for world model
            "task_token": "
",
            "context_token": "
",
            "response_token": "
",
            "system_token": "
",
        }
        
        # Add custom special tokens if provided
        if self.config.special_tokens:
            special_tokens_dict.update(self.config.special_tokens)
            
        # Add special tokens to tokenizer
        self.tokenizer.add_special_tokens({"additional_special_tokens": list(special_tokens_dict.values())})
        
        # Store special token IDs
        self.special_token_ids = {
            name: self.tokenizer.convert_tokens_to_ids(token)
            for name, token in special_tokens_dict.items()
        }

    def process_dataframe(
        self,
        df: pd.DataFrame,
        text_columns: List[str],
        context_column: str = None,
        response_column: str = None
    ) -> Dict[str, torch.Tensor]:
        """
        Process multiple columns from DataFrame with advanced formatting
        """
        # Initialize lists for collecting processed data
        all_input_ids = []
        all_attention_masks = []
        all_position_ids = []
        all_segment_ids = []
        
        for idx in range(len(df)):
            # Collect text from all columns
            texts = [df[col].iloc[idx] for col in text_columns]
            
            # Special handling for context and response
            context = df[context_column].iloc[idx] if context_column else None
            response = df[response_column].iloc[idx] if response_column else None
            
            # Process single example
            processed = self._process_single_example(
                texts=texts,
                context=context,
                response=response
            )
            
            all_input_ids.append(processed["input_ids"])
            all_attention_masks.append(processed["attention_mask"])
            all_position_ids.append(processed["position_ids"])
            all_segment_ids.append(processed["segment_ids"])

        # Pad sequences
        padded = self._pad_sequences(
            input_ids=all_input_ids,
            attention_masks=all_attention_masks,
            position_ids=all_position_ids,
            segment_ids=all_segment_ids
        )
        
        return padded

    def _process_single_example(
        self,
        texts: List[str],
        context: str = None,
        response: str = None
    ) -> Dict[str, torch.Tensor]:
        """
        Process single example with advanced token handling
        """
        # Initialize tokens
        tokens = [self.special_token_ids["bos_token"]]
        segment_ids = [0]
        
        # Add context if provided
        if context:
            tokens.extend([
                self.special_token_ids["context_token"],
                *self.tokenizer.encode(context, add_special_tokens=False)
            ])
            segment_ids.extend([1] * (len(tokens) - len(segment_ids)))
        
        # Add main texts
        for i, text in enumerate(texts):
            tokens.extend([
                self.special_token_ids["sep_token"],
                *self.tokenizer.encode(text, add_special_tokens=False)
            ])
            segment_ids.extend([i + 2] * (len(tokens) - len(segment_ids)))
        
        # Add response if provided
        if response:
            tokens.extend([
                self.special_token_ids["response_token"],
                *self.tokenizer.encode(response, add_special_tokens=False)
            ])
            segment_ids.extend([len(texts) + 2] * (len(tokens) - len(segment_ids)))
        
        # Add EOS token
        tokens.append(self.special_token_ids["eos_token"])
        segment_ids.append(segment_ids[-1])
        
        # Truncate if necessary
        if len(tokens) > self.config.max_length:
            tokens = tokens[:self.config.max_length-1] + [self.special_token_ids["eos_token"]]
            segment_ids = segment_ids[:self.config.max_length]
        
        # Create attention mask and position IDs
        attention_mask = [1] * len(tokens)
        position_ids = list(range(len(tokens)))
        
        return {
            "input_ids": torch.tensor(tokens),
            "attention_mask": torch.tensor(attention_mask),
            "position_ids": torch.tensor(position_ids),
            "segment_ids": torch.tensor(segment_ids)
        }

    def _pad_sequences(
        self,
        input_ids: List[torch.Tensor],
        attention_masks: List[torch.Tensor],
        position_ids: List[torch.Tensor],
        segment_ids: List[torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Pad sequences with advanced handling
        """
        # Pad sequences
        padded_input_ids = pad_sequence(
            input_ids,
            batch_first=True,
            padding_value=self.special_token_ids["pad_token"]
        )
        padded_attention_masks = pad_sequence(
            attention_masks,
            batch_first=True,
            padding_value=0
        )
        padded_position_ids = pad_sequence(
            position_ids,
            batch_first=True,
            padding_value=0
        )
        padded_segment_ids = pad_sequence(
            segment_ids,
            batch_first=True,
            padding_value=0
        )
        
        return {
            "input_ids": padded_input_ids,
            "attention_mask": padded_attention_masks,
            "position_ids": padded_position_ids,
            "segment_ids": padded_segment_ids
        }

# Example usage:
if __name__ == "__main__":
    # Create sample DataFrame
    df = pd.DataFrame({
        "instruction": ["Solve this math problem", "Translate this text"],
        "input": ["What is 2+2?", "Hello world"],
        "context": ["Math context", "Translation context"],
        "response": ["The answer is 4", "Hola mundo"]
    })
    
    # Initialize preprocessor
    config = TokenizerConfig(
        model_name="gpt2",
        max_length=512,
        special_tokens={"custom_token": "
"}
    )
    
    preprocessor = AdvancedWorldModelPreprocessor(config)
    
    # Process DataFrame
    processed_data = preprocessor.process_dataframe(
        df=df,
        text_columns=["instruction", "input"],
        context_column="context",
        response_column="response"
    )
    
    # Print shapes and sample outputs
    print("Input IDs shape:", processed_data["input_ids"].shape)
    print("Attention Mask shape:", processed_data["attention_mask"].shape)
    print("Position IDs shape:", processed_data["position_ids"].shape)
    print("Segment IDs shape:", processed_data["segment_ids"].shape)