train.py

#!/usr/bin/env python
# coding=utf-8
# Copyright The HuggingFace Team and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import math
import os
import sys

import datasets
import numpy as np
import torch
import transformers
from aac_metrics import evaluate
from accelerate import Accelerator, DistributedDataParallelKwargs
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from omegaconf import OmegaConf
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import (
    AutoTokenizer,
    BartConfig,
    get_inverse_sqrt_schedule,
    get_scheduler,
)

from data.collator import DataCollatorForEnClapBart
from data.preprocess import Preprocessor
from modeling.enclap_bart import EnClapBartForConditionalGeneration

logger = get_logger(__name__)
metric_list = ["meteor", "spider"]


def main():
    # Load Configuration
    cfg_path = sys.argv[1]
    args = OmegaConf.load(cfg_path)

    # Initialize Logging
    accelerator_log_kwargs = {}
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    if args.with_tracking:
        accelerator_log_kwargs["log_with"] = args.report_to
        accelerator_log_kwargs["project_dir"] = args.output_dir

    # Initialize Accelerator
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        split_batches=args.split_batches,
        kwargs_handlers=[ddp_kwargs],
        **accelerator_log_kwargs,
    )
    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
        with open(os.path.join(args.output_dir, "args.yaml"), "w") as f:
            OmegaConf.save(args, f)
    accelerator.wait_for_everyone()

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    file_handler = logging.FileHandler(os.path.join(args.output_dir, "train_log.txt"))
    logger.logger.addHandler(file_handler)
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_warning()
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    # Get the datasets
    data_files = {}
    data_files_eval = {}
    if args.train_file is not None:
        data_files["train"] = args.train_file
    if args.validation_file is not None:
        data_files_eval["validation"] = args.validation_file

    extension = args.train_file.split(".")[-1]
    raw_datasets = load_dataset(extension, data_files=data_files)
    raw_datasets_eval = load_dataset(extension, data_files=data_files_eval)

    # Load pretrained model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)
    if args.config_name_or_path is not None:
        config = BartConfig.from_pretrained(args.config_name_or_path)
    else:
        config = None

    if args.model_name_or_path is not None:
        if config is None:
            model = EnClapBartForConditionalGeneration.from_pretrained(
                args.model_name_or_path
            )
        else:
            model = EnClapBartForConditionalGeneration.from_pretrained(
                args.model_name_or_path, config=config
            )
    else:
        model = EnClapBartForConditionalGeneration(config=config)

    # Set the generation config
    if args.val_max_target_length is None:
        args.val_max_target_length = args.max_target_length

    # Set max encodec length based on the shape of the positional encoding
    max_encodec_length = model.config.max_position_embeddings - 2
    label_pad_token_id = (
        -100 if args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    )
    preprocessor = Preprocessor(
        args.encodec_base_path,
        args.clap_base_path,
        tokenizer,
        model.config.max_position_embeddings,
        args.encodec_masking_prob,
        args.encodec_masking_span,
        label_pad_token_id,
        model.config.encodec_vocab_size,
        args.eval_num_captions,
    )

    with accelerator.main_process_first():
        train_dataset = raw_datasets["train"].map(
            preprocessor.preprocess_train,
            num_proc=args.preprocessing_num_workers,
            load_from_cache_file=not args.overwrite_cache,
            desc="Running tokenizer on dataset",
        )
        train_dataset.set_format(
            "pt",
            columns=[
                "input_ids",
                "attention_mask",
                "clap",
                "labels",
                "decoder_attention_mask",
            ],
        )

        # Temporarily set max_target_length for validation.
        eval_dataset = raw_datasets_eval["validation"].map(
            preprocessor.preprocess_eval,
            num_proc=args.preprocessing_num_workers,
            load_from_cache_file=not args.overwrite_cache,
            desc="Running tokenizer on dataset",
        )
        eval_dataset.set_format(
            "pt",
            columns=["input_ids", "attention_mask", "clap"],
            output_all_columns=True,
        )

    train_data_collator = DataCollatorForEnClapBart(
        tokenizer=tokenizer,
        model=model,
        return_tensors="pt",
        label_pad_token_id=label_pad_token_id,
        max_length=max_encodec_length,
        encodec_masking_prob=args.encodec_masking_prob,
        encodec_masking_span=args.encodec_masking_span,
    )
    valid_data_collator = DataCollatorForEnClapBart(
        tokenizer=tokenizer,
        model=model,
        return_tensors="pt",
        label_pad_token_id=label_pad_token_id,
        max_length=max_encodec_length,
    )

    train_dataloader = DataLoader(
        train_dataset,
        shuffle=True,
        collate_fn=train_data_collator,
        batch_size=args.per_device_train_batch_size,
    )
    eval_dataloader = DataLoader(
        eval_dataset,
        collate_fn=valid_data_collator,
        batch_size=args.per_device_eval_batch_size,
    )

    # Optimizer
    # Split weights in two groups, one with weight decay and the other not.
    no_decay = ["bias", "LayerNorm.weight", "layer_norm.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": args.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 = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)

    # Scheduler and math around the number of training steps.
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps
    )
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    if args.lr_scheduler_type == "inverse_sqrt" and hasattr(args, "time_scale"):
        lr_scheduler = get_inverse_sqrt_schedule(
            optimizer=optimizer,
            num_warmup_steps=args.num_warmup_steps,
            timescale=args.time_scale,
        )
    else:
        lr_scheduler = get_scheduler(
            name=args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=args.num_warmup_steps,
            num_training_steps=args.max_train_steps,
        )

    # Prepare everything with our `accelerator`.
    (
        model,
        optimizer,
        train_dataloader,
        eval_dataloader,
        lr_scheduler,
    ) = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps
    )
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    # Figure out how many steps we should save the Accelerator states
    checkpointing_steps = args.checkpointing_steps
    if checkpointing_steps is not None and checkpointing_steps.isdigit():
        checkpointing_steps = int(checkpointing_steps)

    # The trackers initializes automatically on the main process.
    if args.with_tracking:
        accelerator.init_trackers(args.logging_dir)

    # Train!
    total_batch_size = (
        args.per_device_train_batch_size
        * accelerator.num_processes
        * args.gradient_accumulation_steps
    )

    if args.split_batches:
        total_batch_size = int(total_batch_size / accelerator.num_processes)

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(
        f"  Instantaneous batch size per device = {args.per_device_train_batch_size}"
    )
    logger.info(
        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
    )
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")

    completed_steps = 0
    starting_epoch = 0
    # Potentially load in the weights and states from a previous save
    if not args.overwrite_output_dir and os.path.exists(
        os.path.join(args.output_dir, "checkpoints")
    ):
        if args.resume_from_checkpoint is not None:
            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
            accelerator.load_state(args.resume_from_checkpoint)
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = [
                f
                for f in os.scandir(os.path.join(args.output_dir, "checkpoints"))
                if f.is_dir()
            ]
            dirs.sort(key=os.path.getctime)
            path = dirs[
                -1
            ].name  # Sorts folders by date modified, most recent checkpoint is the last
            accelerator.print(f"Resumed from checkpoint: {dirs[-1]}")
            accelerator.load_state(dirs[-1])
        # Extract `epoch_{i}` or `step_{i}`
        training_difference = os.path.splitext(path)[0]

        if "epoch" in training_difference:
            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
            resume_step = None
            completed_steps = starting_epoch * num_update_steps_per_epoch
        else:
            # need to multiply `gradient_accumulation_steps` to reflect real steps
            resume_step = (
                int(training_difference.replace("step_", ""))
                * args.gradient_accumulation_steps
            )
            starting_epoch = resume_step // len(train_dataloader)
            resume_step -= starting_epoch * len(train_dataloader)
            completed_steps = resume_step // args.gradient_accumulation_stepp

    # update the progress_bar if load from checkpoint
    if args.with_tracking:
        total_loss = 0
        logging_loss = 0
        before_epoch_loss = 0

        if args.encodec_masking_prob > 0:
            total_encodec_loss = 0
            logging_encodec_loss = 0
            before_epoch_encodec_loss = 0

    for epoch in range(starting_epoch, args.num_train_epochs):
        model.train()
        if (
            args.resume_from_checkpoint
            and epoch == starting_epoch
            and resume_step is not None
        ):
            # We skip the first `n` batches in the dataloader when resuming from a checkpoint
            active_dataloader = accelerator.skip_first_batches(
                train_dataloader, resume_step
            )
        else:
            active_dataloader = train_dataloader
        logger.info(f"***** Running epoch {epoch} *****")
        epoch_iterator = tqdm(
            active_dataloader,
            desc="Training",
            disable=not accelerator.is_local_main_process,
            dynamic_ncols=True,
            colour="CYAN",
        )
        for step, batch in enumerate(epoch_iterator):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                # We keep track of the loss at each epoch
                if args.with_tracking:
                    total_loss += outputs.lm_loss.item()
                    if args.encodec_masking_prob > 0:
                        if outputs.encodec_loss is not None:
                            total_encodec_loss += outputs.encodec_loss.item()
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(
                        model.parameters(), max_norm=args.max_grad_norm
                    )
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

                # Checks if the accelerator has performed an optimization step behind the scenes
                if accelerator.sync_gradients:
                    completed_steps += 1
                    # Add loss information to tqdm
                    epoch_iterator.set_postfix(loss=total_loss / completed_steps)

                    if completed_steps % args.logging_steps == 0:
                        train_log = {
                            "train/learning_rate": lr_scheduler.get_last_lr()[0]
                        }
                        train_log["train/loss"] = (
                            total_loss - logging_loss
                        ) / args.logging_steps
                        logging_loss = total_loss
                        if args.encodec_masking_prob > 0:
                            train_log["train/encodec_loss"] = (
                                total_encodec_loss - logging_encodec_loss
                            ) / args.logging_steps
                            logging_encodec_loss = total_encodec_loss
                        accelerator.log(train_log, step=completed_steps)

            if isinstance(checkpointing_steps, int):
                if completed_steps % checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps }"
                    if args.output_dir is not None:
                        output_dir = os.path.join(
                            args.output_dir, "checkpoints", output_dir
                        )
                    accelerator.save_state(output_dir)

            if completed_steps >= args.max_train_steps:
                break

        model.eval()
        gen_kwargs = {
            "max_length": args.val_max_target_length,
        }
        predictions = []
        references = []
        eval_iterator = tqdm(
            eval_dataloader,
            desc="Validation",
            disable=not accelerator.is_local_main_process,
            dynamic_ncols=True,
            colour="MAGENTA",
        )
        for step, batch in enumerate(eval_iterator):
            # Drop the padded samples of the last batch of dataloader
            # try:
            #    if accelerator.gradient_state.end_of_dataloader and accelerator.gradient_state.remainder > 0:
            #        batch = batch[:accelerator.gradient_state.remainder]
            # except:
            #    pass

            with torch.no_grad():
                batch["input_ids"] = batch["input_ids"].cuda()
                batch["clap"] = batch["clap"].cuda()
                batch["attention_mask"] = batch["attention_mask"].cuda()
                batch["eos_mask"] = batch["eos_mask"].cuda()

                generated_tokens = accelerator.unwrap_model(model).generate(
                    batch["input_ids"],
                    clap=batch["clap"],
                    attention_mask=batch["attention_mask"],
                    eos_mask=batch["eos_mask"],
                    **gen_kwargs,
                )

                generated_tokens = accelerator.pad_across_processes(
                    generated_tokens, dim=1, pad_index=tokenizer.pad_token_id
                )
                generated_tokens = generated_tokens.cpu().numpy()
                captions = batch["captions"]

                if isinstance(generated_tokens, tuple):
                    generated_tokens = generated_tokens[0]
                decoded_preds = tokenizer.batch_decode(
                    generated_tokens, skip_special_tokens=True
                )

                predictions.extend(decoded_preds)
                references.extend(captions)

        logger.info("Evaluating predictions...")
        result = evaluate(predictions, references, metrics=metric_list)

        # Gather Result
        result = {k: v.cuda() for k, v in result[0].items()}
        result = accelerator.gather_for_metrics(result)
        # Log the average of metrics among the processes
        if accelerator.num_processes > 1:
            result = {f"eval/{k}": round(v.mean().item(), 4) for k, v in result.items()}
        else:
            result = {f"eval/{k}": round(v.item(), 4) for k, v in result.items()}
        logger.info(result)

        if args.with_tracking:
            result["train/epoch_train_loss"] = (total_loss - before_epoch_loss) / len(
                train_dataloader
            )
            result["train/steps"] = completed_steps
            before_epoch_loss = total_loss
            if args.encodec_masking_prob > 0:
                result["train/epoch_encodec_loss"] = (
                    total_encodec_loss - before_epoch_encodec_loss
                ) / len(train_dataloader)
                before_epoch_encodec_loss = total_encodec_loss
            accelerator.log(result, step=epoch)

        if args.checkpointing_steps == "epoch":
            output_dir = f"epoch_{epoch}"
            if args.output_dir is not None:
                output_dir = os.path.join(args.output_dir, "checkpoints", output_dir)
            accelerator.save_state(output_dir)
            if accelerator.is_main_process:
                unwrapped_model = accelerator.unwrap_model(model)
                unwrapped_model.config.save_pretrained(output_dir)

    if args.output_dir is not None:
        save_dir = os.path.join(args.output_dir, "final")
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        unwrapped_model.save_pretrained(
            save_dir,
            is_main_process=accelerator.is_main_process,
            save_function=accelerator.save,
        )
        if accelerator.is_main_process:
            tokenizer.save_pretrained(save_dir)


if __name__ == "__main__":
    main()