train.py

from model import build_transformer
from dataset import BilingualDataset, causal_mask
from config import get_config, get_weights_file_path


import datasets
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import IterableDataset, DataLoader, random_split
from torch.optim.lr_scheduler import LambdaLR

import warnings
from tqdm import tqdm
import os
from pathlib import Path

# Huggingface datasets and tokenizers
from datasets import load_dataset
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.trainers import WordLevelTrainer
from tokenizers.pre_tokenizers import Whitespace

import torchmetrics
import wandb
import accelerate
from torch.utils.tensorboard import SummaryWriter
from safetensors.torch import load_model, save_model
from accelerate import Accelerator
from transformers import GPT2TokenizerFast
import threading


def greedy_decode(model, source, source_mask, tokenizer_tgt, max_len, device):
    sos_idx = tokenizer_tgt.convert_tokens_to_ids('[SOS]')
    eos_idx = tokenizer_tgt.convert_tokens_to_ids('[EOS]')

    # Precompute the encoder output and reuse it for every step
    encoder_output = model.module.encode(source, None)
   
    # Initialize the decoder input with the sos token
    decoder_input = torch.empty(1, 1).fill_(sos_idx).long().to(device)
    while True:
        if decoder_input.size(1) == max_len:
            break

        # build mask for target
        decoder_mask = causal_mask(decoder_input.size(1)).long().to(device)
      

        # calculate output
        out = model.module.decode(encoder_output, source_mask, decoder_input, decoder_mask)
        # print(f'out: {out.shape}')

        # Get next token probabilities with temperature applied
        logits = model.module.project(out[:, -1]) 
        probabilities = F.softmax(logits, dim=-1)

        # Greedily select the next word
        next_word = torch.argmax(probabilities, dim=1)
        
        # Append next word
        decoder_input = torch.cat([decoder_input, next_word.unsqueeze(0)], dim=1)
        # # get next token
        # prob = model.project(out[:, -1])
        # _, next_word = torch.max(prob, dim=1)
        # # print(f'prob: {prob.shape}')
        # decoder_input = torch.cat(
        #     [decoder_input, torch.empty(1, 1).long().fill_(next_word.item()).to(device)], dim=1
        # )

        if next_word.item() == eos_idx:
            break

    return decoder_input.squeeze(0)


def run_validation(model, validation_ds,tokenizer_tgt, max_len, device, print_msg, global_step, num_examples=3):
    model.eval()
    count = 0

    source_texts = []
    expected = []
    predicted = []

    try:
        # get the console window width
        with os.popen('stty size', 'r') as console:
            _, console_width = console.read().split()
            console_width = int(console_width)+_
    except:
        # If we can't get the console width, use 80 as default
        console_width = 80

    with torch.no_grad():
        for batch in validation_ds:
            count += 1
            encoder_input = batch["encoder_input"].to(device) # (b, seq_len)
            encoder_mask = batch["encoder_mask"].to(device) # (b, 1, 1, seq_len)

            # check that the batch size is 1
            assert encoder_input.size(
                0) == 1, "Batch size must be 1 for validation"

            model_out = greedy_decode(model, encoder_input, None, tokenizer_tgt, max_len, device)

            # source_text = batch["src_text"][0]
            target_text = batch["tgt_text"][0]
           
            model_out_text = tokenizer_tgt.decode(model_out.detach().cpu().numpy())

            # source_texts.append(source_text)
            expected.append(target_text)
            predicted.append(model_out_text)
            
            # Print the source, target and model output
            print_msg('-'*console_width)
            # print_msg(f"{f'SOURCE: ':>12}{source_text}")
            print_msg(f"{f'TARGET: ':>12}{target_text}")
            print_msg(f"{f'PREDICTED: ':>12}{model_out_text}")

            if count == num_examples:
                print_msg('-'*console_width)
                break
            
    
    # if writer:
    #     # Evaluate the character error rate
    #     # Compute the char error rate 
    #     metric = torchmetrics.CharErrorRate()
    #     cer = metric(predicted, expected)
    #     writer.add_scalar('validation cer', cer, global_step)
    #     writer.flush()

    #     # Compute the word error rate
    #     metric = torchmetrics.WordErrorRate()
    #     wer = metric(predicted, expected)
    #     writer.add_scalar('validation wer', wer, global_step)
    #     writer.flush()

    #     # Compute the BLEU metric
    #     metric = torchmetrics.BLEUScore()
    #     bleu = metric(predicted, expected)
    #     writer.add_scalar('validation BLEU', bleu, global_step)
    #     writer.flush()

def get_all_sentences(ds):
    for item in ds:
        yield item['text']
def batch_iterator(data):
    for i in range(0, len(data)):
        yield data[i]['text'] 

# Assuming batch_iterator is a function that yields batches
def tqdm_batch_iterator(data, *args, **kwargs):
    for batch in tqdm(batch_iterator(data, *args, **kwargs), total=len(data)):
        yield batch               

def get_or_build_tokenizer(config, ds):
    tokenizer = GPT2TokenizerFast.from_pretrained("openai-community/gpt2", unk_token ='[UNK]', bos_token = '[SOS]', eos_token = '[EOS]' , pad_token = '[PAD]')
    return tokenizer
    # tokenizer_path = Path(config['tokenizer_file'])
    # if not Path.exists(tokenizer_path):
    #     # Most code taken from: https://huggingface.co/docs/tokenizers/quicktour
    #     tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))
    #     tokenizer.pre_tokenizer = Whitespace()
    #     trainer = WordLevelTrainer(special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"], min_frequency=2)
    #     tokenizer.train_from_iterator(get_all_sentences(ds), trainer=trainer)
    #     tokenizer.save(str(tokenizer_path))
    # else:
    #     tokenizer = Tokenizer.from_file(str(tokenizer_path))
    # return tokenizer

def get_ds(config):
    # It only has the train split, so we divide it overselves
    # ds_raw = load_dataset("HausaNLP/HausaVG", split='train+validation+test+challenge_test')
    train_ds_raw = load_dataset("MMInstruction/M3IT", 'coco', split ='train')
    
    val_ds_raw =   load_dataset("MMInstruction/M3IT", 'coco', split ='validation[:2%]')
    
    # ds_raw = load_dataset('opus_books', f"{config['lang_src']}-{config['lang_tgt']}", split='train')

    # Build tokenizers
    
    tokenizer_tgt = get_or_build_tokenizer(config, train_ds_raw,)
    seed = 20  # You can choose any integer as your seed
    torch.manual_seed(seed)
    # # Keep 90% for training, 10% for validation
    # train_ds_size = int(0.9 * len(ds_raw))
    # val_ds_size = len(ds_raw) - train_ds_size
    # train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size])

    train_ds = BilingualDataset(train_ds_raw, tokenizer_tgt,  config['seq_len'])
    val_ds = BilingualDataset(val_ds_raw, tokenizer_tgt, config['seq_len'])
    

    train_dataloader = DataLoader(train_ds,batch_size=config['batch_size'], shuffle=True )
   
    val_dataloader = DataLoader(val_ds, batch_size=1,shuffle=True )

    return train_dataloader, val_dataloader, tokenizer_tgt

def get_model(config, vocab_tgt_len):
    model = build_transformer(vocab_tgt_len, config['seq_len'], d_model=config['d_model'])
    return model

def train_model(config):

    accelerator = Accelerator()
  

    print()
    wandb.login(key = 'c20a1022142595d7d1324fdc53b3ccb34c0ded22')
    wandb.init(project="Vision", name=config['project_name'])

    # Initialize WandB configuration
    wandb.config.epochs = config['num_epochs']
    wandb.config.batch_size = config['batch_size']
    wandb.config.learning_rate = config['lr'] 
    # Define the devic
    # Define the device
    device = accelerator.device
    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device:", device)

    # Make sure the weights folder exists
    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)

    train_dataloader, val_dataloader, tokenizer_tgt = get_ds(config)
    model = get_model(config, len(tokenizer_tgt)).to(device)
   

    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], betas=(0.9, 0.98),eps=1e-9)

    model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, val_dataloader
)

    # If the user specified a model to preload before training, load it
    initial_epoch = 0
    global_step = 0

    def save_models():
        accelerator.save_state(output_dir=f'/kaggle/working/weights/tmodel_00')
        print(f'saving global step {global_step}')
    
    if config['preload']:
        model_filename = get_weights_file_path(config, config['preload'])
        print(f'Preloading model {model_filename}')
        accelerator.load_state(model_filename)
        initial_epoch = 4
   
        # state = torch.load(model_filename)
        # model.load_state_dict(state['model_state_dict'])
        # initial_epoch = state['epoch'] + 1
        # optimizer.load_state_dict(state['optimizer_state_dict'])
        # global_step = state['global_step']

    loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer_tgt.convert_tokens_to_ids('[PAD]'), label_smoothing=0.1).to(device)
   
    for epoch in range(initial_epoch, config['num_epochs']):

        # timer = threading.Timer(5*60, save_models)
        # timer.start()

        model.train()
        batch_iterator = tqdm(train_dataloader, desc=f"Processing Epoch {epoch:02d}")
        
        for batch in batch_iterator:
           
            encoder_input = batch["encoder_input"].to(device) # (b, seq_len)
            decoder_input = batch["decoder_input"].to(device) # (B, seq_len)
            encoder_mask = batch["encoder_mask"].to(device) # (B, 1, 1, seq_len)
            decoder_mask = batch["decoder_mask"].to(device) # (B, 1, seq_len, seq_len)

            # Run the tensors through the encoder, decoder and the projection layer
            encoder_output = model.module.encode(encoder_input, None) # (B, seq_len, d_model)
            decoder_output = model.module.decode(encoder_output, None, decoder_input, decoder_mask) # (B, seq_len, d_model)
            proj_output = model.module.project(decoder_output)
            
             # (B, seq_len, vocab_size)

            # Compare the output with the label
            label = batch["label"].to(device) # (B, seq_len)

            # Compute the loss using a simple cross entropy
            loss = loss_fn(proj_output.view(-1, len(tokenizer_tgt)), label.view(-1))
            batch_iterator.set_postfix({"loss": f"{loss.item():6.3f}"})

            # Log the loss
            wandb.log({"Training Loss": loss.item(), "Global Step": global_step})

            # # Backpropagate the loss
            # loss.backward()
            accelerator.backward(loss)

            # Update the weights
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)

            global_step += 1
            # if global_step == 20000 or global_step == 25000:
            #     print(f'saved state at {global_step}')
            #     accelerator.save_state(output_dir=f'/kaggle/working/weights/tmodel_{epoch:02d}')  
            if global_step == 1000 or global_step == 5000  or global_step == 10000 or global_step == 15000 or global_step == 20000 or global_step == 30000:
                run_validation(model, val_dataloader, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step)
                model.train()
                

        # # Run validation at the end of every epoch
            # Save the model at the end of every epoch
        model_filename = get_weights_file_path(config, f"{epoch:02d}")
        # torch.save({
        #     'epoch': epoch,
        #     'model_state_dict': model.state_dict(),
        #     'optimizer_state_dict': optimizer.state_dict(),
        #     'global_step': global_step
        # }, model_filename)
        # accelerator.save_model(model, model_filename)
    
        accelerator.save_state(output_dir=f'/kaggle/working/weights/tmodel_{epoch:02d}')
        # run_validation(model, val_dataloader, tokenizer_src, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step, writer)
        model.eval()
        eval_loss = 0.0

        #accelerate 
        accurate = 0
        num_elems = 0
        # batch_iterator = tqdm(v_dataloader, desc=f"Processing Epoch {epoch:02d}")
        with torch.no_grad():
            batch_itere = tqdm(val_dataloader, desc=f"Processing loss")
            for batch in batch_itere:
            

                encoder_input = batch['encoder_input'].to(device) # (b, seq_len)
                decoder_input = batch['decoder_input'].to(device) # (B, seq_len)
                encoder_mask = batch['encoder_mask'].to(device) # (B, 1, 1, seq_len)
                decoder_mask = batch['decoder_mask'].to(device) # (B, 1, seq_len, seq_len)

                # Run the tensors through the encoder, decoder and the projection layer
            
                encoder_output = model.module.encode(encoder_input, None) # (B, seq_len, d_model)
                decoder_output = model.module.decode(encoder_output, None, decoder_input, decoder_mask)# (B, seq_len, d_model)
                proj_output = model.module.project(decoder_output)
            
                # (B, seq_len, vocab_size)

                # Compare the output with the label
                # label = batch['label'].to(device) # (B, seq_len)
                proj_output, label = accelerator.gather_for_metrics((
                proj_output, batch["label"]
            ))

                # Compute the loss using a simple cross entropy
                ls = loss_fn(proj_output.view(-1, len(tokenizer_tgt)), label.view(-1))
                batch_itere.set_postfix({"loss": f"{ls.item():6.3f}"})
                eval_loss += ls
                # loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))
           
                
        avg_val_loss = eval_loss / len(val_dataloader)
        accelerator.print(f"Epoch {epoch},Validation Loss: {avg_val_loss})Validation Loss: {avg_val_loss}")
        # print(f'Epoch {epoch},Validation Loss: {avg_val_loss.item()}')
        wandb.log({"Validation Loss": avg_val_loss.item(), "Global Step": global_step})

    
        run_validation(model, val_dataloader, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step)


if __name__ == '__main__':
    warnings.filterwarnings("ignore")
    config = get_config()
    train_model(config)