train_catt_whisper.py

import torch
import pickle
from eo_pl import TashkeelModel
from tashkeel_tokenizer import TashkeelTokenizer
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger
from tashkeel_dataset import AudioDataset, AudioDatasetParquet, PrePaddingAudioDataLoader
import argparse
from pytorch_lightning.callbacks import EarlyStopping


def freeze(model):
    for param in model.parameters():
        param.requires_grad = False

def unfreeze(model):
    for param in model.parameters():
        param.requires_grad = True

if __name__ == '__main__':
    # Model's Configs
    dl_num_workers = 2
    batch_size = 32
    max_seq_len = 1024
    threshold = 0

    # Catt Whisper best checkpoint
    pretrained_mlm_pt = "models/catt_whisper_base_model_v1_epoch_26_with_spec_augment.pt"
    tokenizer = TashkeelTokenizer()

    parser = argparse.ArgumentParser()
    parser.add_argument("--train_path", type=str, required=True)
    parser.add_argument("--val_path", type=str, required=True)
    args = parser.parse_args()

    speech_model_name = 'base'
    n_mels = 80

#    speech_model_name = 'large-v3'
#    n_mels = 128

    print('Creating KSSA Audio Train Dataset...')
    train_audio_dataset = AudioDataset(args.train_path, tokenizer, max_seq_len, tashkeel_to_text_ratio_threshold=threshold, n_mels=n_mels, augment=True)
    train_audio_dataloader = PrePaddingAudioDataLoader(tokenizer, train_audio_dataset, batch_size=batch_size, num_workers=dl_num_workers, shuffle=True)

    print('Creating KSAA Audio Validation Dataset...')
    val_audio_dataset = AudioDataset(args.val_path, tokenizer, max_seq_len, tashkeel_to_text_ratio_threshold=threshold, n_mels=n_mels, augment=False)
    val_audio_dataloader = PrePaddingAudioDataLoader(tokenizer, val_audio_dataset, batch_size=batch_size, num_workers=dl_num_workers, shuffle=False)


    print('Creating Model...')
    model = TashkeelModel(tokenizer, max_seq_len=max_seq_len, n_layers=6, learnable_pos_emb=False, speech_model_name=speech_model_name)

    # Use the pretrained weights of Catt Whisper best checkpoint to initialize the model
    if not pretrained_mlm_pt is None:
        missing = model.load_state_dict(torch.load(pretrained_mlm_pt), strict=False)
        print(f'Missing layers: {missing}')

    
    freeze(model)   #Freeze the entire encoder first
    unfreeze(model.transformer.decoder)   #Unfreeze only the final classification head
    unfreeze(model.transformer.encoder.layers[5])    #Unfreeze last encoder layer (index 5)

    dirpath = 'catt_whisper_model_v1/' # Whisper Base + spec_augment

    checkpoint_callback = ModelCheckpoint(dirpath=dirpath, save_top_k=10, save_last=True,
                                          monitor='val_der',
                                          filename='catt_whisper_model_v1-{epoch:02d}-{val_loss:.5f}-{val_der:.5f}')
    

    early_stop_callback = EarlyStopping(
        monitor      = "val_der",           # metric to monitor
        min_delta    = 0.001,               # minimum improvement required
        patience     = 3,                   # how many epochs with no improvement before stop
        verbose      = True,                # shows "EarlyStopping counter: X/Y" in console
        mode         = "min",               # "min" for loss, "max" for accuracy/f1
        strict       = True,                # crash if metric not found (good for debugging)
    )

    print('Creating Trainer...')

    logs_path = f'{dirpath}/logs'

    print('#'*100)
    print(model)
    print('#'*100)


    trainer = Trainer(
        accelerator="cuda",
        devices=-1,
        max_epochs=20,
        callbacks=[TQDMProgressBar(refresh_rate=1), checkpoint_callback, early_stop_callback],
        logger=CSVLogger(save_dir=logs_path),
        strategy="auto"
        )

    trainer.fit(model, train_audio_dataloader, val_audio_dataloader)