training/run_retriever_no_trainer.py

import argparse
import functools
import logging
import math
from random import choice, randint

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from datasets import load_dataset
from torch.utils import checkpoint
from torch.utils.data import Dataset, RandomSampler, DataLoader, SequentialSampler
from tqdm.auto import tqdm
from transformers import get_scheduler, AutoTokenizer, AdamW, SchedulerType, AutoModelForSequenceClassification

logger = logging.getLogger(__name__)


def get_parser():
    parser = argparse.ArgumentParser(description="Train ELI5 retriever")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="vblagoje/lfqa",
        help="The name of the dataset to use (via the datasets library).",
    )

    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=1024,
    )

    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=1024,
        help="Batch size (per device) for the evaluation dataloader.",
    )

    parser.add_argument(
        "--max_length",
        type=int,
        default=128,
    )

    parser.add_argument(
        "--checkpoint_batch_size",
        type=int,
        default=32,
    )

    parser.add_argument(
        "--pretrained_model_name",
        type=str,
        default="google/bert_uncased_L-8_H-768_A-12",
    )

    parser.add_argument(
        "--model_save_name",
        type=str,
        default="eli5_retriever_model_l-12_h-768_b-512-512",
    )

    parser.add_argument(
        "--learning_rate",
        type=float,
        default=2e-4,
    )

    parser.add_argument(
        "--weight_decay",
        type=float,
        default=0.2,
    )

    parser.add_argument(
        "--log_freq",
        type=int,
        default=500,
        help="Log train/validation loss every log_freq update steps"
    )

    parser.add_argument(
        "--num_train_epochs",
        type=int,
        default=4,
    )

    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )

    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )

    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="linear",  # this is linear with warmup
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )

    parser.add_argument(
        "--num_warmup_steps",
        type=int,
        default=100,
        help="Number of steps for the warmup in the lr scheduler."
    )

    parser.add_argument(
        "--warmup_percentage",
        type=float,
        default=0.08,
        help="Number of steps for the warmup in the lr scheduler."
    )
    return parser


class RetrievalQAEmbedder(torch.nn.Module):
    def __init__(self, sent_encoder):
        super(RetrievalQAEmbedder, self).__init__()
        dim = sent_encoder.config.hidden_size
        self.bert_query = sent_encoder
        self.output_dim = 128
        self.project_query = torch.nn.Linear(dim, self.output_dim, bias=False)
        self.project_doc = torch.nn.Linear(dim, self.output_dim, bias=False)
        self.ce_loss = torch.nn.CrossEntropyLoss(reduction="mean")

    def embed_sentences_checkpointed(self, input_ids, attention_mask, checkpoint_batch_size=-1):
        # reproduces BERT forward pass with checkpointing
        if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
            return self.bert_query(input_ids, attention_mask=attention_mask)[1]
        else:
            # prepare implicit variables
            device = input_ids.device
            input_shape = input_ids.size()
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
            head_mask = [None] * self.bert_query.config.num_hidden_layers
            extended_attention_mask: torch.Tensor = self.bert_query.get_extended_attention_mask(
                attention_mask, input_shape, device
            )

            # define function for checkpointing
            def partial_encode(*inputs):
                encoder_outputs = self.bert_query.encoder(inputs[0], attention_mask=inputs[1], head_mask=head_mask, )
                sequence_output = encoder_outputs[0]
                pooled_output = self.bert_query.pooler(sequence_output)
                return pooled_output

            # run embedding layer on everything at once
            embedding_output = self.bert_query.embeddings(
                input_ids=input_ids, position_ids=None, token_type_ids=token_type_ids, inputs_embeds=None
            )
            # run encoding and pooling on one mini-batch at a time
            pooled_output_list = []
            for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size)):
                b_embedding_output = embedding_output[b * checkpoint_batch_size: (b + 1) * checkpoint_batch_size]
                b_attention_mask = extended_attention_mask[b * checkpoint_batch_size: (b + 1) * checkpoint_batch_size]
                pooled_output = checkpoint.checkpoint(partial_encode, b_embedding_output, b_attention_mask)
                pooled_output_list.append(pooled_output)
            return torch.cat(pooled_output_list, dim=0)

    def embed_questions(self, q_ids, q_mask, checkpoint_batch_size=-1):
        q_reps = self.embed_sentences_checkpointed(q_ids, q_mask, checkpoint_batch_size)
        return self.project_query(q_reps)

    def embed_answers(self, a_ids, a_mask, checkpoint_batch_size=-1):
        a_reps = self.embed_sentences_checkpointed(a_ids, a_mask, checkpoint_batch_size)
        return self.project_doc(a_reps)

    def forward(self, q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=-1):
        device = q_ids.device
        q_reps = self.embed_questions(q_ids, q_mask, checkpoint_batch_size)
        a_reps = self.embed_answers(a_ids, a_mask, checkpoint_batch_size)
        compare_scores = torch.mm(q_reps, a_reps.t())
        loss_qa = self.ce_loss(compare_scores, torch.arange(compare_scores.shape[1]).to(device))
        loss_aq = self.ce_loss(compare_scores.t(), torch.arange(compare_scores.shape[0]).to(device))
        loss = (loss_qa + loss_aq) / 2
        return loss


class ELI5DatasetQARetriever(Dataset):
    def __init__(self, examples_array, extra_answer_threshold=3, min_answer_length=64, training=True, n_samples=None):
        self.data = examples_array
        self.answer_thres = extra_answer_threshold
        self.min_length = min_answer_length
        self.training = training
        self.n_samples = self.data.num_rows if n_samples is None else n_samples

    def __len__(self):
        return self.n_samples

    def make_example(self, idx):
        example = self.data[idx]
        question = example["title"]
        if self.training:
            answers = [a for i, (a, sc) in enumerate(zip(example["answers"]["text"], example["answers"]["score"]))]
            answer_tab = choice(answers).split(" ")
            start_idx = randint(0, max(0, len(answer_tab) - self.min_length))
            answer_span = " ".join(answer_tab[start_idx:])
        else:
            answer_span = example["answers"]["text"][0]
        return question, answer_span

    def __getitem__(self, idx):
        return self.make_example(idx % self.data.num_rows)


def make_qa_retriever_batch(qa_list, tokenizer, max_len=64):
    q_ls = [q for q, a in qa_list]
    a_ls = [a for q, a in qa_list]
    q_toks = tokenizer(q_ls, padding="max_length", max_length=max_len, truncation=True)
    q_ids, q_mask = (
        torch.LongTensor(q_toks["input_ids"]),
        torch.LongTensor(q_toks["attention_mask"])
    )
    a_toks = tokenizer(a_ls, padding="max_length", max_length=max_len, truncation=True)
    a_ids, a_mask = (
        torch.LongTensor(a_toks["input_ids"]),
        torch.LongTensor(a_toks["attention_mask"]),
    )
    return q_ids, q_mask, a_ids, a_mask


def evaluate_qa_retriever(model, data_loader):
    # make iterator
    epoch_iterator = tqdm(data_loader, desc="Iteration", disable=True)
    tot_loss = 0.0
    with torch.no_grad():
        for step, batch in enumerate(epoch_iterator):
            q_ids, q_mask, a_ids, a_mask = batch
            loss = model(q_ids, q_mask, a_ids, a_mask)
            tot_loss += loss.item()
        return tot_loss / (step + 1)


def train(config):
    set_seed(42)
    args = config["args"]
    data_files = {"train": "train.json", "validation": "validation.json", "test": "test.json"}
    eli5 = load_dataset(args.dataset_name, data_files=data_files)

    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
    accelerator = Accelerator()
    # 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,
    )
    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
    logger.info(accelerator.state)

    # prepare torch Dataset objects
    train_dataset = ELI5DatasetQARetriever(eli5['train'], training=True)
    valid_dataset = ELI5DatasetQARetriever(eli5['validation'], training=False)

    tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name)
    base_model = AutoModel.from_pretrained(args.pretrained_model_name)

    model = RetrievalQAEmbedder(base_model)
    no_decay = ['bias', 'LayerNorm.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 = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, weight_decay=args.weight_decay)

    model_collate_fn = functools.partial(make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length)
    train_dataloader = DataLoader(train_dataset, batch_size=args.per_device_train_batch_size,
                                  sampler=RandomSampler(train_dataset), collate_fn=model_collate_fn)

    model_collate_fn = functools.partial(make_qa_retriever_batch, tokenizer=tokenizer, max_len=args.max_length)
    eval_dataloader = DataLoader(valid_dataset, batch_size=args.per_device_eval_batch_size,
                                 sampler=SequentialSampler(valid_dataset), collate_fn=model_collate_fn)

    # train the model
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer,
                                                                              train_dataloader, eval_dataloader)
    # Scheduler and math around the number of training steps.
    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
    else:
        args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    num_warmup_steps = args.num_warmup_steps if args.num_warmup_steps else math.ceil(args.max_train_steps *
                                                                                     args.warmup_percentage)
    scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps,
        num_training_steps=args.max_train_steps,
    )

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

    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}")
    logger.info(f"  Warmup steps = {num_warmup_steps}")
    logger.info(f"  Logging training progress every {args.log_freq} optimization steps")

    loc_loss = 0.0
    current_loss = 0.0
    checkpoint_step = 0

    completed_steps = checkpoint_step
    progress_bar = tqdm(range(args.max_train_steps), initial=checkpoint_step,
                        disable=not accelerator.is_local_main_process)
    for epoch in range(args.num_train_epochs):
        model.train()
        batch = next(iter(train_dataloader))
        for step in range(1000):
        #for step, batch in enumerate(train_dataloader, start=checkpoint_step):
            # model inputs
            q_ids, q_mask, a_ids, a_mask = batch
            pre_loss = model(q_ids, q_mask, a_ids, a_mask, checkpoint_batch_size=args.checkpoint_batch_size)
            loss = pre_loss.sum() / args.gradient_accumulation_steps
            accelerator.backward(loss)
            loc_loss += loss.item()
            if ((step + 1) % args.gradient_accumulation_steps == 0) or (step + 1 == len(train_dataloader)):
                current_loss = loc_loss
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                progress_bar.update(1)
                progress_bar.set_postfix(loss=loc_loss)
                loc_loss = 0
                completed_steps += 1

            if step % (args.log_freq * args.gradient_accumulation_steps) == 0:
                accelerator.wait_for_everyone()
                unwrapped_model = accelerator.unwrap_model(model)
                eval_loss = evaluate_qa_retriever(unwrapped_model, eval_dataloader)
                logger.info(f"Train loss {current_loss} , eval loss {eval_loss}")
                if args.wandb and accelerator.is_local_main_process:
                    import wandb
                    wandb.log({"loss": current_loss, "eval_loss": eval_loss, "step": completed_steps})

            if completed_steps >= args.max_train_steps:
                break

        logger.info("Saving model {}".format(args.model_save_name))
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        accelerator.save(unwrapped_model.state_dict(), "{}_{}.bin".format(args.model_save_name, epoch))
        eval_loss = evaluate_qa_retriever(unwrapped_model, eval_dataloader)
        logger.info("Evaluation loss epoch {:4d}: {:.3f}".format(epoch, eval_loss))


if __name__ == "__main__":
    parser = get_parser()
    parser.add_argument(
        "--wandb",
        action="store_true",
        help="Whether to use W&B logging",
    )
    main_args, _ = parser.parse_known_args()
    config = {"args": main_args}
    if main_args.wandb:
        import wandb
        wandb.init(project="Retriever")

    train(config=config)