#!/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()