train.py

import os
import math
import json
import yaml
import pickle
import argparse
import logging
from pathlib import Path
from collections import Counter

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset, random_split
from torch.optim import AdamW
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

import matplotlib.pyplot as plt
import tqdm

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace

from model import LM_GPT

from utils import save_training_results

os.makedirs("logs", exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | [%(levelname)s] | %(message)s",
    handlers=[
        logging.StreamHandler(),                 # print to console
        logging.FileHandler("logs/train.log", mode='w')    # save to file
    ]
)

logger = logging.getLogger()
## END LOGGING


# --------------------- Parsing Arguments -------------------------
def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", type=str, default='config/main.yml')
    # Overides config files
    ap.add_argument("--epochs", type=int, default=None, help='Epochs (overrides config file)')
    ap.add_argument('--batch_size', type=int, default=None, help='Batch size (overrides config file)') 
    ap.add_argument('--lr', type=float, default=None, help='Learning rate (overrides config file)')
    ap.add_argument('--save_dir', type=str, default=None, help='Save directory (overrides config file)')
    ap.add_argument('--seed', type=int,default=None, help='Random seed (overrides config file)')
    ap.add_argument('--decay', type=int,default=None, help='decay')

    return ap.parse_args()
# loads config file or uses defualt config file
def load_config(config_path):
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    defualt_config = os.path.join(BASE_DIR, 'config', 'main.yml')
    if not os.path.exists(config_path):
        config_path = os.path.join(BASE_DIR, config_path)
    
    # if still doesnt exist, switch to defualt
    if not os.path.exists(config_path):  # if config file doesnt exist
        logger.info(f"Config file '{config_path}' not found. You might be using relative path, switch to absolute full path")
        logger.info(f"Switching to  default configurations located at {defualt_config}.")
        config_path = defualt_config
    try:
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
    except Exception as e:
        logging.error(f"Failed to load config file: {e}. Please have a yml file containing default configs")
        raise
    
    return config

def overide_configs(args, config):
    if args.epochs is not None:
        config['training']['epochs'] = args.epochs
        logger.info(f"[CLI Override] epochs: {args.epochs}")

    if args.batch_size is not None:
        config['training']['batch_size'] = args.batch_size
        logger.info(f"[CLI Override] batch_size: {args.batch_size}")

    if args.lr is not None:
        config['training']['learning_rate'] = args.lr
        logger.info(f"[CLI Override] learning_rate: {args.lr}")

    if args.decay is not None:
        config['training']['weight_decay'] = args.decay
        logger.info(f"[CLI Override] weight_decay: {args.decay}")

    if args.seed is not None:
        config['reproducibility']['seed'] = args.seed
        logger.info(f"[CLI Override] seed: {args.seed}")

    if args.save_dir is not None:
        config['checkpoint']['save_dir'] = args.save_dir
        logger.info(f"[CLI Override] save_dir: {args.save_dir}")

   

    return config

ap = parse_args()
config = overide_configs(ap, load_config(ap.config))
PAD = config['special_tokens']['pad']

# ---------------- Device ----------------
def pick_device():
    if torch.backends.mps.is_available():
        return "mps"  # Apple Silicon / Metal
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"

DEVICE = pick_device()
logger.info(f"Using device:{DEVICE}")

# ---------------- Reproducibility ----------------
def seed_everything(seed: int = 42):
    import random
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

seed_everything(config["reproducibility"]["seed"])

# ----------------  Preprocessing ------------------
# Word Level Preprocessing Step
def preprocessing_step(train_path, val_path, test_path, max_len, special_tokens, min_freq=0):
    def read_file(path):
        with open(path, 'r', encoding='utf-8') as f:
            return f.readlines()

    try:
        # format of files: each line is "word1 word2 ..."
        train_data = read_file(train_path)
        val_data   = read_file(val_path)
        test_data = read_file(test_path)
    except Exception as e:
        print(f"Error loading dataset: {e}")
    

    def yield_tokens(data, max_len=max_len):
        for sentences in data:
            tokens = sentences.split()[:max_len]
            yield tokens

    def build_vocab_from_iterator(iterator, min_freq, specials=None):
        counter = Counter()
        for tokens in iterator:
            counter.update(tokens)
        
        vocab = {}
        if specials:
            for i, token in enumerate(specials):
                vocab[token] = i
        
        idx = len(vocab)
        for token, freq in counter.items():
            if freq >= min_freq and token not in vocab:
                vocab[token] = idx
                idx += 1
        
        return vocab


    word2idx = build_vocab_from_iterator(yield_tokens(train_data, max_len), specials=special_tokens, min_freq=min_freq)

    raw_train_tokens = list(yield_tokens(train_data, max_len))
    raw_valid_tokens = list(yield_tokens(val_data, max_len))
    raw_test_tokens = list(yield_tokens(test_data, max_len))

    return raw_train_tokens, raw_valid_tokens, raw_test_tokens, word2idx


# bpe preprocessing
def preprocessing_step_bpe(train_path, val_path, test_path, max_len, special_tokens, min_freq=0, vocab_size=10000):
    def read_file(path):
        with open(path, 'r', encoding='utf-8') as f:
            return f.readlines()
    try:
        train_data = read_file(train_path)
        val_data = read_file(val_path)
        test_data = read_file(test_path)
    except Exception as e:
        print(f"Error loading dataset: {e}")
        raise

    # Initialize BPE tokenizer
    tokenizer = Tokenizer(BPE(unk_token=special_tokens[2]))  # unk_token 
    tokenizer.pre_tokenizer = Whitespace()

    # Create trainer with special tokens
    trainer = BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=special_tokens,
        min_frequency=min_freq
    )

    # Train tokenizer on training data
    tokenizer.train_from_iterator(
        train_data,
        trainer=trainer
    )

    def yield_tokens(data, max_len=max_len):
        for sentence in data:
            yield tokenizer.encode(sentence).tokens[:max_len]

    raw_train_tokens = list(yield_tokens(train_data, max_len))
    raw_val_tokens = list(yield_tokens(val_data, max_len))
    raw_test_tokens = list(yield_tokens(test_data, max_len))

    word2idx = tokenizer.get_vocab()


    return raw_train_tokens, raw_val_tokens, raw_test_tokens, word2idx, tokenizer

   
raw_train_tokens, raw_valid_tokens, raw_test_tokens, word2idx, tokenizer = \
        preprocessing_step_bpe(config['data']['train_path'], 
                           config['data']['valid_path'], 
                           config['data']['test_path'], 
                           config['data']['max_len'], 
                           [config['special_tokens']['pad'], config['special_tokens']['eos'], config['special_tokens']['unk']], 
                           config['data']['min_freq'],
                           vocab_size=config['model']['vocab_size'])
    


logger.info(
    "Dataset preprocessing completed | vocab=%d | train=%d | valid=%d | test=%d",
    len(word2idx),
    len(raw_train_tokens),
    len(raw_valid_tokens),
    len(raw_test_tokens),
)

# Updating the the configurations file
config['model']['vocab_size'] = len(word2idx)
logger.info(
    "Updated config | vocab_size=%s | max_len=%s (includes <eos>)",
    config['model']['vocab_size'],
    config['model']['max_len']
)

# ----------------- DATASET CLASSES --------------
def numericalize(toks, vocab, eos, unk):
    unk_idx = vocab[unk]
    ids = [vocab.get(t, unk_idx) for t in toks] + [vocab[eos]]
    return torch.tensor(ids, dtype=torch.long)

class PTB_Dataset(Dataset):
    def __init__(self, raw_tokens, vocab, pad, eos, unk):
        self.raw_tokens = raw_tokens
        self.vocab = vocab
        self.pad = pad
        self.eos = eos
        self.unk = unk

    def __len__(self):
        # Should return number of sentences
        return len(self.raw_tokens)
    
    def __getitem__(self, idx):
        tokens = self.raw_tokens[idx]
        tokens = numericalize(tokens, self.vocab, self.eos, self.unk)

        inputs  = tokens[:-1]   # x[:, 0:T-1]
        targets = tokens[1:]

        return inputs, targets

# Creates (B, T) Padded Sequences
def pad_sequence(sequences, max_len, pad_value): 
    batch_size = len(sequences)
    padded = torch.full((batch_size, max_len), pad_value, dtype=torch.long)

    for i, seq in enumerate(sequences):
        seq_len = len(seq)
        padded[i, :seq_len] = seq
    return padded

def collate_fn(batch):
    src_seqs, tgt_seqs = zip(*batch)

    src_lens = torch.tensor([len(s) for s in src_seqs])
    tgt_lens = torch.tensor([len(t) for t in tgt_seqs])

    max_src_len = max(src_lens)
    max_tgt_len = max(tgt_lens)

    padded_src = pad_sequence(src_seqs, max_src_len, word2idx[PAD])
    padded_tgt = pad_sequence(tgt_seqs, max_tgt_len, word2idx[PAD])
    
    return padded_src, padded_tgt


# ----------------- Training Helper functions ------------------- 
def train_epoch(model, dataloader, optimizer, criterion, device,  clip_grad=1.0, warmup_scheduler=None, batches_seen=0):
    model.train()
    total_loss = 0.0
    i  = 0
    for src, tgt in tqdm.tqdm(dataloader,"train"):
        i+=1
        src, tgt = src.to(device), tgt.to(device)
        optimizer.zero_grad()

        logits, _ = model(src)  # (B, T, V)
        B, T, V = logits.shape

        loss = criterion(logits.view(B * T, V), tgt.view(B * T))
        loss.backward()

        if clip_grad is not None:
            nn.utils.clip_grad_norm_(model.parameters(), clip_grad)

        optimizer.step()

        
        total_loss += loss.item()
        
        if warmup_scheduler is not None:
            warmup_scheduler.step()
            
    avg_loss = total_loss / len(dataloader)

    nll = float(avg_loss)
    ppl = float(math.exp(nll))
    return avg_loss, nll, ppl

def evaluate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0.0
    
    with torch.no_grad():
        for src, tgt in tqdm.tqdm(dataloader, "Evaluating"):
            src, tgt = src.to(device), tgt.to(device)
            logits, _ = model(src)  # (B, T, V)
            B, T, V = logits.shape
            
            loss = criterion(logits.view(B * T, V), tgt.view(B * T))
            
        
            total_loss += loss.item()
    
    nll = total_loss / len(dataloader)
    ppl = math.exp(nll)
    return nll, ppl


# ---------------- MODEL CONFIGURATIONS -------------------
# Creating a config class for better readibility in the code
class Config:
    # Preprocessing configurations
    preprocess_max_len = config['data']['max_len']
    min_freq = config['data']['min_freq']
    pad, eos, unk = config['special_tokens']['pad'], config['special_tokens']['eos'], config['special_tokens']['unk']
    

    # Model Parameters
    VOCAB_SIZE = int(config['model']['vocab_size'])   # Will be set during training
    D_MODEL = int(config['model']['d_model'])
    N_HEADS = int(config['model']['n_heads'])
    N_LAYERS = int(config['model']['n_layers'])
    MAX_LEN = int(config['model']['max_len'])  # To account for <eos> tag
    D_FF = int(config['model']['d_ff'])
    DROPOUT = float(config['model']['dropout'])
    PAD_ID = int(word2idx[pad])


    # Hyperparameters
    EPOCHS = int(config['training']['epochs'])
    BATCH_SIZE = int(config['training']['batch_size'])
    LR = float(config['training']['learning_rate'])
    GRAD_CLIP = float(config['training']['grad_clip'])
    WEIGHT_DECAY =  float(config['training']['weight_decay'])
    PATIENCE = int(config['training']['patience'])
    LABEL_SMOOTHING = float(config['training']['label_smoothing'])

    # scheduler
    WARMUP_STEPS = int(config['training']['warmup_steps'])

    # Save file
    SAVE_DIR = config['checkpoint']['save_dir']    


# -------------------------------- Main Training Loop ------------------------------
def main():
    # - - - -- - - - - Configurations and results folders - --  -- -- 
    conf = Config()
    
    save_path = os.path.join(conf.SAVE_DIR)
    os.makedirs(save_path, exist_ok=True)
    logger.info(f"Saving all results, model configurations, results in {save_path}")
    
    import yaml

    with open(os.path.join(save_path, "best_config.yml"), 'w') as f:
        yaml.dump(config, f, default_flow_style=False)
    # - - - - - -  - - - - - - - - - END - - - -- -- - - - -- - - -- - -- 

    # Model Optimizer and Loss function (criterion)
    model = LM_GPT(conf.VOCAB_SIZE, conf.D_MODEL, conf.N_HEADS, conf.N_LAYERS, conf.MAX_LEN, conf.D_FF, conf.DROPOUT, conf.PAD_ID).to(DEVICE)
    optimizer = AdamW(model.parameters(), lr=conf.LR, weight_decay=conf.WEIGHT_DECAY)
    criterion = nn.CrossEntropyLoss(ignore_index=conf.PAD_ID, label_smoothing=conf.LABEL_SMOOTHING)

    # Parameter Calculation
    total_params     = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    logger.info(f"Total params:     {total_params:,}")
    logger.info(f"Trainable params: {trainable_params:,}")
    logger.info(f"Non-trainable:    {total_params - trainable_params:,}")

    # Datasets and Dataloaders
    train_ds = PTB_Dataset(raw_train_tokens,word2idx, pad=conf.pad, eos=conf.eos, unk=conf.unk)
    train_dl = DataLoader(train_ds, batch_size=conf.BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

    valid_ds = PTB_Dataset(raw_valid_tokens,word2idx, pad=conf.pad, eos=conf.eos, unk=conf.unk)
    valid_dl = DataLoader(valid_ds, batch_size=conf.BATCH_SIZE, shuffle=False, collate_fn=collate_fn)

    ## Scheduler (Warmup + Cosine Anealing): Warmup(~0 -> LR) -> Then -> CosineAnealing (LR  -> ~0)
    #  ----------------  Scheduler  -----------------
    warmup_scheduler = LinearLR(
        optimizer, 
        start_factor=1/conf.WARMUP_STEPS, 
        end_factor=1.0, 
        total_iters=conf.WARMUP_STEPS)
    
    epoch_by_warmup = (conf.WARMUP_STEPS // len(train_dl))
    remaining_epochs = conf.EPOCHS - (epoch_by_warmup)
    cosine_scheduler = CosineAnnealingLR(optimizer, T_max=remaining_epochs)
    # ----------------  Scheduler End  -----------------

    # Loss Curve lists
    tl_list, vl_list = [], []
    tp_list, vp_list = [], []

    # Stop loss 
    best_valid_loss  = float("inf")
    patience_counter = 0
    patience = conf.PATIENCE
    best_model_save_path = os.path.join(save_path, "best_model.pt")

    logger.info("======= " + "Starting Training " + ("=" * 60))
    # ## LR
    # from utils import find_optimal_lr
    # optimal_lr = find_optimal_lr(
    #     model=model,
    #     train_dataloader=train_dl,
    #     optimizer=optimizer,
    #     criterion=criterion,
    #     device=DEVICE,
    #     start_lr=1e-6,
    #     end_lr=1e-2,
    #     num_iter=100,
    #     save_path="lr_range_test.png"
    # )
    
    # ## 
    # return

    for epoch in range(conf.EPOCHS):
        if epoch > epoch_by_warmup:
            warmup_scheduler = None
        
        train_loss, train_nll, train_ppl = train_epoch(model, train_dl, optimizer, criterion, DEVICE, warmup_scheduler=warmup_scheduler, clip_grad=conf.GRAD_CLIP)
        valid_nll,  valid_ppl  = evaluate(model,valid_dl,criterion,DEVICE)
        
        # Early Stopping
        if valid_nll < best_valid_loss:
            best_valid_loss  = valid_nll
            patience_counter = 0
            torch.save({
                'model_state_dict': model.state_dict(),
                'hyperparams': {
                    'vocab_size':conf.VOCAB_SIZE,
                    'd_model': conf.D_MODEL,
                    'n_heads': conf.N_HEADS,
                    'n_layers': conf.N_LAYERS,
                    'max_len': conf.MAX_LEN,
                    'd_ff': conf.D_FF,
                    'dropout': conf.DROPOUT,
                    'pad_id': conf.PAD_ID
                },
                'optimizer_state_dict': optimizer.state_dict(),
                'epoch': epoch,
                'valid_loss': best_valid_loss,
            }, best_model_save_path)
            print(f"    At epoch: {epoch+1}, best model saved at {best_model_save_path}")
        else:
            patience_counter += 1

        if patience_counter >= patience:
            print(f"\nEarly stopping triggered after {epoch+1} epochs")
            break
        
        if epoch > epoch_by_warmup:
            cosine_scheduler.step()
            
        # Metricss
        tl_list.append(train_loss); vl_list.append(valid_nll)
        tp_list.append(train_ppl);  vp_list.append(valid_ppl)

        logger.info(
            "Epoch %d/%d | Train Loss=%.4f | Train PPL=%.2f | Valid Loss=%.4f | Valid PPL=%.2f",
            epoch + 1,
            conf.EPOCHS,
            train_loss,
            train_ppl,
            valid_nll,
            valid_ppl,
        )
        
        
    
    # ── Evaluation Test Dataset ─────────────────────────────────────────────────────────────────────
    logger.info("======= " + "Starting Evaluating " + ("=" * 60))

    
    test_ds = PTB_Dataset(raw_test_tokens,word2idx, pad=conf.pad, eos=conf.eos, unk=conf.unk)
    test_dl = DataLoader(test_ds, batch_size=conf.BATCH_SIZE, shuffle=False, collate_fn=collate_fn)

    ckpt = torch.load(best_model_save_path)
    hp   = ckpt['hyperparams']
    model = LM_GPT(**hp).to(DEVICE)
    model.load_state_dict(ckpt['model_state_dict'])
    logger.info(f"Loaded checkpoint from epoch {ckpt['epoch']+1} with valid loss {ckpt['valid_loss']:.4f}")


    criterion = nn.CrossEntropyLoss(ignore_index=hp['pad_id'])
    test_nll, test_ppl = evaluate(model, test_dl, criterion, DEVICE)
    logger.info(f"Test Loss: {test_nll:.4f} || Test PPL: {test_ppl:.2f}")


    # ── Plots ─────────────────────────────────────────────────────────────────────
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 5))

    ax1.plot(tl_list, label='Train Loss', marker='o')
    ax1.plot(vl_list, label='Valid Loss', marker='s')
    ax1.set_xlabel('Epoch'); ax1.set_ylabel('Loss')
    ax1.set_title('Training and Validation Loss', fontweight='bold')
    ax1.legend(); ax1.grid(True, alpha=0.3)

    ax2.plot(tp_list, label='Train PPL', marker='o')
    ax2.plot(vp_list, label='Valid PPL', marker='s')
    ax2.set_xlabel('Epoch'); ax2.set_ylabel('Perplexity')
    ax2.set_title('Training and Validation Perplexity', fontweight='bold')
    ax2.legend(); ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(save_path + '/training_curves.png', dpi=150, bbox_inches='tight')
    plt.show()

    logger.info("Training & Evaluating completed!")

    # ---------------------------- Training and Evaluating Complete ------------------------------------------ #

    # --------------------   ------ JUST  SAVING STUFF HERE ----------------- -------------
    
    results_dict = save_training_results(
        save_path=save_path,
        conf=conf,
        config_dict=config,
        model=model,
        best_model_path=best_model_save_path,
        best_epoch=ckpt["epoch"],
        best_valid_loss=ckpt["valid_loss"],
        test_nll=test_nll,
        test_ppl=test_ppl,
        train_losses=tl_list,
        train_ppls=tp_list,
        valid_losses=vl_list,
        valid_ppls=vp_list,
        raw_train_tokens=raw_train_tokens,
        raw_valid_tokens=raw_valid_tokens,
        raw_test_tokens=raw_test_tokens,
        vocab_size=len(word2idx),
        warmup_steps=conf.WARMUP_STEPS,
    )

    
    tokenizer_file = os.path.join(save_path, "bpe_tokenizer.json")
    tokenizer.save(tokenizer_file)

    logger.info(f"BPE Tokenizer saved to: {tokenizer_file}")

    ## LAST STEP ## 
    # Flush handlers and save log file
    for handler in logger.handlers:
        handler.flush()
    
    temp_log = "logs/train.log"
    final_log = os.path.join(save_path, "train.log")
    if os.path.exists(temp_log):
        import shutil
        shutil.copy(temp_log, final_log)


if __name__ == "__main__":
   main()