Upload 9 files

0ec1bb9 verified over 1 year ago

12.2 kB

import os
import argparse
import json
import numpy as np
from tqdm import tqdm
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split
from torch.cuda.amp import GradScaler, autocast
from transformers import GPT2TokenizerFast, AdamW, get_linear_schedule_with_warmup
from datasets import load_dataset
from transformers import logging as hf_logging

# Suppress unnecessary logging
hf_logging.set_verbosity_error()

# Import your transformer model
# from your_transformer_module import create_transformer_model

class Config:
    """Centralized configuration for training"""
    def __init__(self):
        # Model hyperparameters
        self.vocab_size = 50257  # GPT-2 vocab size
        self.d_model = 512
        self.nhead = 8
        self.num_layers = 6
        self.dim_feedforward = 2048
        self.dropout = 0.1
        
        # Training hyperparameters
        self.batch_size = 32
        self.num_epochs = 3
        self.learning_rate = 5e-5
        self.weight_decay = 0.01
        self.warmup_steps = 0.1  # Percentage of total steps
        self.max_seq_length = 512
        self.gradient_accumulation_steps = 1
        self.max_grad_norm = 1.0
        self.seed = 42
        
        # Paths
        self.output_dir = "./checkpoints"
        self.model_save_prefix = "reasoning_model"
        
        # Device configuration
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.fp16 = torch.cuda.is_available()

    def save(self, path):
        """Save configuration to file"""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, 'w') as f:
            json.dump(self.__dict__, f, indent=2)
    
    @classmethod
    def from_file(cls, path):
        """Load configuration from file"""
        config = cls()
        with open(path, 'r') as f:
            config.__dict__.update(json.load(f))
        return config

def load_and_preprocess_data(config):
    """Load and preprocess the dataset"""
    # Load dataset
    dataset = load_dataset("ag2428/reasoningDataV4", split="train")
    
    # Initialize tokenizer
    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token  # Set padding token
    
    # Tokenization function
    def tokenize_function(examples):
        # Combine instruction and answer with a separator
        texts = [f"{inst}\n{ans}" for inst, ans in zip(examples["instruction"], examples["answer"])]
        
        # Tokenize
        tokenized = tokenizer(
            texts,
            max_length=config.max_seq_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt"
        )
        
        # Create labels (shifted input_ids for language modeling)
        tokenized["labels"] = tokenized["input_ids"].clone()
        return tokenized
    
    # Tokenize dataset
    tokenized_datasets = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=dataset.column_names,
        desc="Tokenizing dataset"
    )
    
    # Split into train and validation sets
    train_val = tokenized_datasets.train_test_split(test_size=0.1, seed=config.seed)
    train_dataset = train_val["train"]
    val_dataset = train_val["test"]
    
    # Convert to PyTorch format
    train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
    val_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config.batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )
    
    return train_loader, val_loader, tokenizer

def train_epoch(model, train_loader, optimizer, scheduler, scaler, config, epoch):
    """Train for one epoch"""
    model.train()
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}")
    
    for step, batch in enumerate(progress_bar):
        # Move batch to device
        input_ids = batch['input_ids'].to(config.device)
        attention_mask = batch['attention_mask'].to(config.device)
        labels = batch['labels'].to(config.device)
        
        # Forward pass with mixed precision
        with autocast(enabled=config.fp16):
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )
            loss = outputs.loss / config.gradient_accumulation_steps
        
        # Backward pass and optimize
        if config.fp16:
            scaler.scale(loss).backward()
        else:
            loss.backward()
        
        # Gradient clipping and optimization step
        if (step + 1) % config.gradient_accumulation_steps == 0:
            if config.fp16:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
                scaler.step(optimizer)
                scaler.update()
            else:
                torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
                optimizer.step()
            
            scheduler.step()
            optimizer.zero_grad()
        
        total_loss += loss.item() * config.gradient_accumulation_steps
        
        # Update progress bar
        progress_bar.set_postfix({
            'loss': f"{total_loss / (step + 1):.4f}",
            'lr': f"{scheduler.get_last_lr()[0]:.2e}"
        })
    
    return total_loss / len(train_loader)

def evaluate(model, val_loader, config):
    """Evaluate the model on the validation set"""
    model.eval()
    total_loss = 0
    
    with torch.no_grad():
        for batch in tqdm(val_loader, desc="Evaluating"):
            input_ids = batch['input_ids'].to(config.device)
            attention_mask = batch['attention_mask'].to(config.device)
            labels = batch['labels'].to(config.device)
            
            with autocast(enabled=config.fp16):
                outputs = model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    labels=labels
                )
                loss = outputs.loss
            
            total_loss += loss.item()
    
    return total_loss / len(val_loader)

def save_checkpoint(model, optimizer, scheduler, epoch, config, is_best=False):
    """Save model checkpoint"""
    os.makedirs(config.output_dir, exist_ok=True)
    
    # Prepare checkpoint
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'config': config.__dict__,
    }
    
    # Save checkpoint
    if is_best:
        filename = os.path.join(config.output_dir, f"{config.model_save_prefix}_best.pt")
    else:
        filename = os.path.join(config.output_dir, f"{config.model_save_prefix}_epoch_{epoch}.pt")
    
    torch.save(checkpoint, filename)
    print(f"Checkpoint saved to {filename}")

def main():
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Train a reasoning model")
    parser.add_argument('--config', type=str, default=None, help="Path to config file")
    parser.add_argument('--output_dir', type=str, default=None, help="Output directory for checkpoints")
    parser.add_argument('--batch_size', type=int, default=None, help="Batch size")
    parser.add_argument('--num_epochs', type=int, default=None, help="Number of epochs")
    parser.add_argument('--learning_rate', type=float, default=None, help="Learning rate")
    parser.add_argument('--fp16', action='store_true', help="Use mixed precision training")
    args = parser.parse_args()
    
    # Initialize config
    if args.config:
        config = Config.from_file(args.config)
    else:
        config = Config()
    
    # Override config with command line arguments
    if args.output_dir:
        config.output_dir = args.output_dir
    if args.batch_size:
        config.batch_size = args.batch_size
    if args.num_epochs:
        config.num_epochs = args.num_epochs
    if args.learning_rate:
        config.learning_rate = args.learning_rate
    if args.fp16:
        config.fp16 = True
    
    # Set random seed for reproducibility
    torch.manual_seed(config.seed)
    np.random.seed(config.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(config.seed)
    
    # Create output directory
    os.makedirs(config.output_dir, exist_ok=True)
    
    # Save config
    config.save(os.path.join(config.output_dir, "config.json"))
    
    # Load data
    print("Loading and preprocessing data...")
    train_loader, val_loader, tokenizer = load_and_preprocess_data(config)
    
    # Initialize model
    print("Initializing model...")
    # TODO: Replace with your model initialization
    # model = create_transformer_model(
    #     vocab_size=config.vocab_size,
    #     d_model=config.d_model,
    #     nhead=config.nhead,
    #     num_layers=config.num_layers,
    #     dim_feedforward=config.dim_feedforward,
    #     dropout=config.dropout,
    #     max_seq_length=config.max_seq_length
    # )
    
    # For now, let's use a placeholder that will raise an error
    class PlaceholderModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
        
        def forward(self, input_ids, attention_mask, labels=None):
            # This is a placeholder that will raise an error
            # Replace with your actual model implementation
            raise NotImplementedError(
                "Please implement your transformer model and replace this placeholder. "
                "See the TODO comment in the code for more details."
            )
    
    model = PlaceholderModel()
    model = model.to(config.device)
    
    # Initialize optimizer and scheduler
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {
            'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            'weight_decay': config.weight_decay,
        },
        {
            'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            'weight_decay': 0.0,
        }
    ]
    
    optimizer = AdamW(optimizer_grouped_parameters, lr=config.learning_rate)
    
    # Calculate total training steps
    total_steps = len(train_loader) * config.num_epochs // config.gradient_accumulation_steps
    warmup_steps = int(total_steps * config.warmup_steps)
    
    # Initialize learning rate scheduler
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps
    )
    
    # Initialize gradient scaler for mixed precision training
    scaler = GradScaler(enabled=config.fp16)
    
    # Training loop
    print("Starting training...")
    best_val_loss = float('inf')
    
    for epoch in range(config.num_epochs):
        # Train for one epoch
        train_loss = train_epoch(model, train_loader, optimizer, scheduler, scaler, config, epoch)
        
        # Evaluate on validation set
        val_loss = evaluate(model, val_loader, config)
        
        print(f"Epoch {epoch + 1}/{config.num_epochs}:")
        print(f"  Train loss: {train_loss:.4f}")
        print(f"  Val loss: {val_loss:.4f}")
        
        # Save checkpoint
        save_checkpoint(model, optimizer, scheduler, epoch, config)
        
        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            save_checkpoint(model, optimizer, scheduler, epoch, config, is_best=True)
    
    print("Training complete!")

if __name__ == "__main__":
    main()

	import os
	import argparse
	import json
	import numpy as np
	from tqdm import tqdm
	import torch
	import torch.nn as nn
	from torch.utils.data import DataLoader, random_split
	from torch.cuda.amp import GradScaler, autocast
	from transformers import GPT2TokenizerFast, AdamW, get_linear_schedule_with_warmup
	from datasets import load_dataset
	from transformers import logging as hf_logging

	# Suppress unnecessary logging
	hf_logging.set_verbosity_error()

	# Import your transformer model
	# from your_transformer_module import create_transformer_model

	class Config:
	"""Centralized configuration for training"""
	def __init__(self):
	# Model hyperparameters
	self.vocab_size = 50257 # GPT-2 vocab size
	self.d_model = 512
	self.nhead = 8
	self.num_layers = 6
	self.dim_feedforward = 2048
	self.dropout = 0.1

	# Training hyperparameters
	self.batch_size = 32
	self.num_epochs = 3
	self.learning_rate = 5e-5
	self.weight_decay = 0.01
	self.warmup_steps = 0.1 # Percentage of total steps
	self.max_seq_length = 512
	self.gradient_accumulation_steps = 1
	self.max_grad_norm = 1.0
	self.seed = 42

	# Paths
	self.output_dir = "./checkpoints"
	self.model_save_prefix = "reasoning_model"

	# Device configuration
	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	self.fp16 = torch.cuda.is_available()

	def save(self, path):
	"""Save configuration to file"""
	os.makedirs(os.path.dirname(path), exist_ok=True)
	with open(path, 'w') as f:
	json.dump(self.__dict__, f, indent=2)

	@classmethod
	def from_file(cls, path):
	"""Load configuration from file"""
	config = cls()
	with open(path, 'r') as f:
	config.__dict__.update(json.load(f))
	return config

	def load_and_preprocess_data(config):
	"""Load and preprocess the dataset"""
	# Load dataset
	dataset = load_dataset("ag2428/reasoningDataV4", split="train")

	# Initialize tokenizer
	tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
	tokenizer.pad_token = tokenizer.eos_token # Set padding token

	# Tokenization function
	def tokenize_function(examples):
	# Combine instruction and answer with a separator
	texts = [f"{inst}\n{ans}" for inst, ans in zip(examples["instruction"], examples["answer"])]

	# Tokenize
	tokenized = tokenizer(
	texts,
	max_length=config.max_seq_length,
	truncation=True,
	padding="max_length",
	return_tensors="pt"
	)

	# Create labels (shifted input_ids for language modeling)
	tokenized["labels"] = tokenized["input_ids"].clone()
	return tokenized

	# Tokenize dataset
	tokenized_datasets = dataset.map(
	tokenize_function,
	batched=True,
	remove_columns=dataset.column_names,
	desc="Tokenizing dataset"
	)

	# Split into train and validation sets
	train_val = tokenized_datasets.train_test_split(test_size=0.1, seed=config.seed)
	train_dataset = train_val["train"]
	val_dataset = train_val["test"]

	# Convert to PyTorch format
	train_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
	val_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

	# Create data loaders
	train_loader = DataLoader(
	train_dataset,
	batch_size=config.batch_size,
	shuffle=True,
	num_workers=4,
	pin_memory=True
	)

	val_loader = DataLoader(
	val_dataset,
	batch_size=config.batch_size,
	shuffle=False,
	num_workers=4,
	pin_memory=True
	)

	return train_loader, val_loader, tokenizer

	def train_epoch(model, train_loader, optimizer, scheduler, scaler, config, epoch):
	"""Train for one epoch"""
	model.train()
	total_loss = 0
	progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}")

	for step, batch in enumerate(progress_bar):
	# Move batch to device
	input_ids = batch['input_ids'].to(config.device)
	attention_mask = batch['attention_mask'].to(config.device)
	labels = batch['labels'].to(config.device)

	# Forward pass with mixed precision
	with autocast(enabled=config.fp16):
	outputs = model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	labels=labels
	)
	loss = outputs.loss / config.gradient_accumulation_steps

	# Backward pass and optimize
	if config.fp16:
	scaler.scale(loss).backward()
	else:
	loss.backward()

	# Gradient clipping and optimization step
	if (step + 1) % config.gradient_accumulation_steps == 0:
	if config.fp16:
	scaler.unscale_(optimizer)
	torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
	scaler.step(optimizer)
	scaler.update()
	else:
	torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
	optimizer.step()

	scheduler.step()
	optimizer.zero_grad()

	total_loss += loss.item() * config.gradient_accumulation_steps

	# Update progress bar
	progress_bar.set_postfix({
	'loss': f"{total_loss / (step + 1):.4f}",
	'lr': f"{scheduler.get_last_lr()[0]:.2e}"
	})

	return total_loss / len(train_loader)

	def evaluate(model, val_loader, config):
	"""Evaluate the model on the validation set"""
	model.eval()
	total_loss = 0

	with torch.no_grad():
	for batch in tqdm(val_loader, desc="Evaluating"):
	input_ids = batch['input_ids'].to(config.device)
	attention_mask = batch['attention_mask'].to(config.device)
	labels = batch['labels'].to(config.device)

	with autocast(enabled=config.fp16):
	outputs = model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	labels=labels
	)
	loss = outputs.loss

	total_loss += loss.item()

	return total_loss / len(val_loader)

	def save_checkpoint(model, optimizer, scheduler, epoch, config, is_best=False):
	"""Save model checkpoint"""
	os.makedirs(config.output_dir, exist_ok=True)

	# Prepare checkpoint
	checkpoint = {
	'epoch': epoch,
	'model_state_dict': model.state_dict(),
	'optimizer_state_dict': optimizer.state_dict(),
	'scheduler_state_dict': scheduler.state_dict(),
	'config': config.__dict__,
	}

	# Save checkpoint
	if is_best:
	filename = os.path.join(config.output_dir, f"{config.model_save_prefix}_best.pt")
	else:
	filename = os.path.join(config.output_dir, f"{config.model_save_prefix}_epoch_{epoch}.pt")

	torch.save(checkpoint, filename)
	print(f"Checkpoint saved to {filename}")

	def main():
	# Parse command line arguments
	parser = argparse.ArgumentParser(description="Train a reasoning model")
	parser.add_argument('--config', type=str, default=None, help="Path to config file")
	parser.add_argument('--output_dir', type=str, default=None, help="Output directory for checkpoints")
	parser.add_argument('--batch_size', type=int, default=None, help="Batch size")
	parser.add_argument('--num_epochs', type=int, default=None, help="Number of epochs")
	parser.add_argument('--learning_rate', type=float, default=None, help="Learning rate")
	parser.add_argument('--fp16', action='store_true', help="Use mixed precision training")
	args = parser.parse_args()

	# Initialize config
	if args.config:
	config = Config.from_file(args.config)
	else:
	config = Config()

	# Override config with command line arguments
	if args.output_dir:
	config.output_dir = args.output_dir
	if args.batch_size:
	config.batch_size = args.batch_size
	if args.num_epochs:
	config.num_epochs = args.num_epochs
	if args.learning_rate:
	config.learning_rate = args.learning_rate
	if args.fp16:
	config.fp16 = True

	# Set random seed for reproducibility
	torch.manual_seed(config.seed)
	np.random.seed(config.seed)
	if torch.cuda.is_available():
	torch.cuda.manual_seed_all(config.seed)

	# Create output directory
	os.makedirs(config.output_dir, exist_ok=True)

	# Save config
	config.save(os.path.join(config.output_dir, "config.json"))

	# Load data
	print("Loading and preprocessing data...")
	train_loader, val_loader, tokenizer = load_and_preprocess_data(config)

	# Initialize model
	print("Initializing model...")
	# TODO: Replace with your model initialization
	# model = create_transformer_model(
	# vocab_size=config.vocab_size,
	# d_model=config.d_model,
	# nhead=config.nhead,
	# num_layers=config.num_layers,
	# dim_feedforward=config.dim_feedforward,
	# dropout=config.dropout,
	# max_seq_length=config.max_seq_length
	# )

	# For now, let's use a placeholder that will raise an error
	class PlaceholderModel(nn.Module):
	def __init__(self):
	super().__init__()
	self.loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

	def forward(self, input_ids, attention_mask, labels=None):
	# This is a placeholder that will raise an error
	# Replace with your actual model implementation
	raise NotImplementedError(
	"Please implement your transformer model and replace this placeholder. "
	"See the TODO comment in the code for more details."
	)

	model = PlaceholderModel()
	model = model.to(config.device)

	# Initialize optimizer and scheduler
	no_decay = ['bias', 'LayerNorm.weight']
	optimizer_grouped_parameters = [
	{
	'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
	'weight_decay': config.weight_decay,
	},
	{
	'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
	'weight_decay': 0.0,
	}
	]

	optimizer = AdamW(optimizer_grouped_parameters, lr=config.learning_rate)

	# Calculate total training steps
	total_steps = len(train_loader) * config.num_epochs // config.gradient_accumulation_steps
	warmup_steps = int(total_steps * config.warmup_steps)

	# Initialize learning rate scheduler
	scheduler = get_linear_schedule_with_warmup(
	optimizer,
	num_warmup_steps=warmup_steps,
	num_training_steps=total_steps
	)

	# Initialize gradient scaler for mixed precision training
	scaler = GradScaler(enabled=config.fp16)

	# Training loop
	print("Starting training...")
	best_val_loss = float('inf')

	for epoch in range(config.num_epochs):
	# Train for one epoch
	train_loss = train_epoch(model, train_loader, optimizer, scheduler, scaler, config, epoch)

	# Evaluate on validation set
	val_loss = evaluate(model, val_loader, config)

	print(f"Epoch {epoch + 1}/{config.num_epochs}:")
	print(f" Train loss: {train_loss:.4f}")
	print(f" Val loss: {val_loss:.4f}")

	# Save checkpoint
	save_checkpoint(model, optimizer, scheduler, epoch, config)

	# Save best model
	if val_loss < best_val_loss:
	best_val_loss = val_loss
	save_checkpoint(model, optimizer, scheduler, epoch, config, is_best=True)

	print("Training complete!")

	if __name__ == "__main__":
	main()