data_module.py

from datasets import load_dataset
from transformers import AutoTokenizer
from dataclasses import dataclass
import os, random, numpy as np, torch
from tqdm import tqdm

@dataclass
class TokenizationConfig:
    dataset_name: str = "roneneldan/TinyStories"
    text_column: str = "text"                       # TinyStories uses "text" column
    tokenizer_name: str = "openai-community/gpt2"   # Use pretrained GPT-2 Tokenizer
    block_size: int = 128                          
    num_proc: int = 4                               # Parallel CPU processes
    output_dir: str = './tokenized_tinystories'
    seed: int = 42

def set_seed(seed: int):
    """Ensure deterministic preprocessing."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    print(f"Random seed set to {seed}")

def get_tokenizer(cfg: TokenizationConfig):
    tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_name, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer

def tokenize_and_group(cfg: TokenizationConfig):
    set_seed(cfg.seed)
    print(f"Loading dataset: {cfg.dataset_name}")
    raw = load_dataset(cfg.dataset_name)

    # Ensure validation split exists
    if "validation" not in raw.keys():
        print("No validation split found, creating one manually (1% of train).")
        raw = raw["train"].train_test_split(test_size=0.01, seed=cfg.seed)
        raw["validation"] = raw["test"]
        del raw["test"]

    print(f"Loading tokenizer: {cfg.tokenizer_name}")
    tokenizer = get_tokenizer(cfg)

    # Step 1: Tokenize
    def tokenize_function(examples):
        return tokenizer(examples[cfg.text_column], return_attention_mask=False)
    
    print(f"Tokenizing dataset (batched)")
    tokenized = raw.map(
        tokenize_function,
        batched=True,
        num_proc=cfg.num_proc,
        remove_columns=[cfg.text_column]
    )

    # Step 2: group into fixed size blocks
    def group_texts(examples):
        all_ids = sum(examples['input_ids'], [])
        total_length = (len(all_ids)//cfg.block_size) * cfg.block_size
        result = {
            'input_ids': [all_ids[i: i+cfg.block_size] for i in range(0, total_length, cfg.block_size)]
        }
        result['labels'] = result['input_ids'].copy()
        return result
    
    print(f"Group into {cfg.block_size}-token blocks...")
    grouped = tokenized.map(
        group_texts,
        batched=True,
        num_proc=cfg.num_proc
    )

   # Step 3: save both splits
    os.makedirs(cfg.output_dir, exist_ok=True)
    print(f"Saving tokenized dataset to {cfg.output_dir}")

    for split in ["train", "validation"]:
        split_path = os.path.join(cfg.output_dir, split)
        os.makedirs(split_path, exist_ok=True)
        print(f"  → Saving {split} split ({len(grouped[split])} samples)")
        grouped[split].save_to_disk(split_path)

    # Save tokenizer too
    tokenizer.save_pretrained(cfg.output_dir)

    print("Done! Tokenized dataset (train + validation) saved.")
    print(f"Each sample length: {cfg.block_size}")
    return grouped

if __name__ == '__main__':
    cfg = TokenizationConfig()
    tokenize_and_group(cfg)