scripts/train.py

#!/usr/bin/env python3
"""
Training entry point for SHINE-LR-v3 (Optimized for Pre-tokenized Data)
Bypasses heavy JSON parsing and tokenization during startup.
"""

import sys
import os
import yaml
import torch
from torch.utils.data import Dataset, DataLoader

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.shine_model import SHINEModel
from trainers.shine_trainer import SHINETrainer


class PreprocessedDataset(Dataset):
    """
    Lightweight dataset to load pre-tokenized tensors directly from disk.
    Eliminates CPU bottleneck during training.
    """
    def __init__(self, file_path: str):
        print(f"📥 Loading preprocessed data: {file_path}")
        # map_location='cpu' ensures fast loading regardless of GPU state
        self.data = torch.load(file_path, map_location='cpu', weights_only=False)
        print(f"✅ Loaded {len(self.data)} samples.")

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

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


def simple_collate_fn(batch):
    """Pass-through collate function for pre-tokenized data"""
    return batch


def load_config(config_path: str) -> dict:
    """Load configuration from YAML file"""
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)


def main():
    config_path = "/content/SHINE-LR/configs/default.yaml"
    config = load_config(config_path)

    print("=" * 60)
    print(" SHINE-LR-v3 Training (Optimized)")
    print("=" * 60)

    # 1. Check and Load Preprocessed Data
    preprocessed_dir = config['data'].get('preprocessed_dir', "/content/SHINE-LR/data/preprocessed")
    train_file = os.path.join(preprocessed_dir, 'train_preprocessed.pt')
    val_file = os.path.join(preprocessed_dir, 'valid_preprocessed.pt')

    if not os.path.exists(train_file) or not os.path.exists(val_file):
        raise FileNotFoundError(
            f"❌ Preprocessed files not found in {preprocessed_dir}. "
            "Please run `python -m data.preprocess_dataset` first."
        )

    train_dataset = PreprocessedDataset(train_file)
    val_dataset = PreprocessedDataset(val_file)

    # 2. Create DataLoaders (Optimized for 2-core CPU environment)
    print("\n🔄 Initializing DataLoaders...")
    train_loader = DataLoader(
        train_dataset,
        batch_size=config['training']['batch_size'],
        shuffle=True,
        collate_fn=simple_collate_fn,
        num_workers=0,          # 0 is optimal for 2-core CPUs (avoids context switching)
        pin_memory=True,        # Fast transfer to GPU
        drop_last=True          # Drops the last incomplete batch for stable gradient accumulation
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config['training']['val_batch_size'],
        shuffle=False,
        collate_fn=simple_collate_fn,
        num_workers=0,
        pin_memory=True,
        drop_last=True
    )

    # 3. Initialize Model
    print("\n🏗️ Initializing SHINE model...")
    model = SHINEModel(
        base_llm_name=config['model']['base_llm'],
        context_encoder_name=config['model']['context_encoder'],
        num_layers=config['model']['num_layers'],
        rank=config['model']['rank'],
        hidden_dim=config['model']['hidden_dim'],
        param_state_dim=config['model']['param_state_dim'],
        hyper_hidden_dim=config['model']['hyper_hidden_dim'],
        persona_embed_dim=config['model']['persona_embed_dim'],
        decay_gamma=config['model']['decay_gamma'],
        device=config['device']
    )

    # 4. Initialize Trainer
    trainer_config = {
        **config['training'],
        **config['loss'],
        'checkpoint_dir': config['checkpoint_dir'],
        'device': config['device']
    }

    trainer = SHINETrainer(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        config=trainer_config
    )

    # 5. Start Training
    trainer.train()


if __name__ == "__main__":
    main()