data/preprocess_dataset.py

#!/usr/bin/env python3
"""
Pre-tokenize MSC dataset to avoid repeated tokenization during training.
"""

import sys
import os
import torch
from transformers import AutoTokenizer
from tqdm import tqdm

# Ensure absolute imports work when running as module
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data.msc_processor import MSCDataProcessor
from data.shine_dataset import SHINEDataset

def preprocess_dataset(
    msc_root: str,
    tokenizer_name: str,
    output_dir: str,
    max_context_len: int = 512,
    max_qa_len: int = 256
):
    """Pre-tokenize and save dataset to disk"""
    
    print(f"📥 Loading tokenizer: {tokenizer_name}")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    print(f"📥 Processing MSC dataset...")
    processor = MSCDataProcessor(msc_root)
    
    os.makedirs(output_dir, exist_ok=True)
    
    for split in ['train', 'valid']:
        print(f"\n🔄 Preprocessing {split} split...")
        conversations = processor.process_split(split=split)
        
        if not conversations:
            print(f"⚠️ No conversations found for {split} split.")
            continue
        
        #  Use SHINEDataset to handle all text preparation and tokenization logic
        dataset = SHINEDataset(
            conversations=conversations,
            tokenizer=tokenizer,
            max_context_len=max_context_len,
            max_qa_len=max_qa_len
        )
        
        output_file = os.path.join(output_dir, f'{split}_preprocessed.pt')
        preprocessed_data = []
        
        # Iterate and save tokenized tensors
        for idx in tqdm(range(len(dataset)), desc=f"Tokenizing {split}"):
            try:
                item = dataset[idx]
                preprocessed_data.append(item)
            except Exception as e:
                print(f"⚠️ Error processing item {idx}: {e}")
                continue
        
        # Save to disk
        torch.save(preprocessed_data, output_file)
        print(f"✅ Saved {len(preprocessed_data)} conversations to {output_file}")

if __name__ == "__main__":
    preprocess_dataset(
        msc_root="/content/SHINE-LR/data/msc_raw_data/msc",
        tokenizer_name="Qwen/Qwen2.5-3B-Instruct",
        output_dir="/content/SHINE-LR/data/preprocessed",
        max_context_len=512,
        max_qa_len=256
    )