train_dense_encoder_with_llm.py

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


"""
 Pipeline to train DPR Biencoder with a LLM
"""

import argparse
import glob
import logging
import math
import os
import random
import time

import torch
torch.multiprocessing.set_sharing_strategy('file_system')
import transformers
import wandb

from typing import Tuple
from torch import nn
from torch import Tensor as T
from torch.cuda.amp import autocast
from torch.utils.checkpoint import get_device_states, set_device_states

from dpr.models import init_biencoder_components
from dpr.models.biencoder import (
    BiEncoder, BiEncoderKLLoss, BiEncoderInBatchKLLoss,
    BiEncoderNllLoss, BiEncoderBatch, BiEncoderBatchWithScore
)
from dpr.options import (
    add_encoder_params, add_training_params, setup_args_gpu, set_seed,
    print_args, get_encoder_params_state, add_tokenizer_params,
    set_encoder_params_from_state
)
from dpr.utils.data_utils import (
    ShardedDataIterator, Tensorizer,
    read_data_from_json_files
)
from dpr.utils.dist_utils import all_gather_list
from dpr.utils.model_utils import (
    setup_for_distributed_mode, move_to_device, get_schedule_linear,
    CheckpointState, get_model_file, get_model_obj,
    load_states_from_checkpoint
)

logger = logging.getLogger()
logger.setLevel(logging.INFO)
if (logger.hasHandlers()):
    logger.handlers.clear()
console = logging.StreamHandler()
logger.addHandler(console)
transformers.logging.set_verbosity_error()


class RandContext:
    def __init__(self, *tensors):
        self.fwd_cpu_state = torch.get_rng_state()
        self.fwd_gpu_devices, self.fwd_gpu_states = get_device_states(*tensors)

    def __enter__(self):
        self._fork = torch.random.fork_rng(
            devices=self.fwd_gpu_devices,
            enabled=True
        )
        self._fork.__enter__()
        torch.set_rng_state(self.fwd_cpu_state)
        set_device_states(self.fwd_gpu_devices, self.fwd_gpu_states)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._fork.__exit__(exc_type, exc_val, exc_tb)
        self._fork = None


class BiEncoderTrainer(object):
    """
    BiEncoder training pipeline component. Can be used to initiate or resume training and
    validate the trained model using either binary classification's NLL loss or average rank
    of the question's gold passages across dataset provided pools of negative passages.
    For full IR accuracy evaluation, please see generate_dense_embeddings.py
    and dense_retriever.py CLI tools.
    """

    def __init__(self, args):
        self.args = args
        self.shard_id = args.local_rank if args.local_rank != -1 else 0
        self.distributed_factor = args.distributed_world_size or 1

        logger.info("***** Initializing components for training *****")

        # if model file is specified, encoder parameters from saved state should be used for
        # initialization
        model_file = get_model_file(self.args, self.args.checkpoint_file_name)
        saved_state = None
        if model_file:
            saved_state = load_states_from_checkpoint(model_file)
            set_encoder_params_from_state(saved_state.encoder_params, args)

        tensorizer, model, optimizer = init_biencoder_components(args.encoder_model_type, args)

        model, optimizer = setup_for_distributed_mode(model, optimizer, args.device, args.n_gpu,
                                                      args.local_rank,
                                                      args.fp16,
                                                      args.fp16_opt_level)
        self.biencoder = model
        self.optimizer = optimizer
        self.tensorizer = tensorizer
        self.start_epoch = 0
        self.start_batch = 0
        self.total_data_seen = 0
        self.scheduler_state = None
        self.best_validation_result = None
        self.best_cp_name = None
        self.restart = self.args.restart
        if saved_state:
            self._load_saved_state(saved_state)

        self.scaler = torch.cuda.amp.GradScaler() if self.args.fp16 else None

    def get_data_iterator(self, path: str, batch_size: int, shuffle=True,
                          shuffle_seed: int = 0,
                          offset: int = 0, upsample_rates: list = None) -> ShardedDataIterator:
        data_files = glob.glob(path)
        data = read_data_from_json_files(data_files, upsample_rates)

        logger.info('Total cleaned data size: {}'.format(len(data)))

        return ShardedDataIterator(
            data, shard_id=self.shard_id,
            num_shards=self.distributed_factor,
            batch_size=batch_size, shuffle=shuffle, shuffle_seed=shuffle_seed, offset=offset,
            strict_batch_size=True,  # this is not really necessary, one can probably disable it
        )

    def run_train(self):
        args = self.args
        upsample_rates = None
        if args.train_files_upsample_rates is not None:
            upsample_rates = eval(args.train_files_upsample_rates)

        train_iterator = self.get_data_iterator(args.train_file, args.batch_size,
                                                shuffle=True,
                                                shuffle_seed=args.seed, offset=self.start_batch,
                                                upsample_rates=upsample_rates)

        logger.info("  Total iterations per epoch=%d", train_iterator.max_iterations)
        updates_per_epoch = train_iterator.max_iterations // args.gradient_accumulation_steps
        total_updates = max(
            updates_per_epoch * (args.num_train_epochs - self.start_epoch - 1),
            0
        ) + (train_iterator.max_iterations - self.start_batch) // args.gradient_accumulation_steps
        logger.info(" Total updates=%d", total_updates)
        warmup_steps = args.warmup_steps
        scheduler = get_schedule_linear(self.optimizer, warmup_steps, total_updates)

        if self.scheduler_state and self.restart is False:
            logger.info("Loading scheduler state %s", self.scheduler_state)
            scheduler.load_state_dict(self.scheduler_state)

        if args.gradient_checkpointing:
            logger.info("  Gradient checkpointing is enabled.")
            model_to_set = get_model_obj(self.biencoder)
            model_to_set.question_model.gradient_checkpointing_enable()
            model_to_set.ctx_model.gradient_checkpointing_enable()

        if args.wandb_project is not None and args.local_rank in [-1, 0]:
            model_to_get = get_model_obj(self.biencoder)
            wandb.watch(model_to_get.question_model, log_freq=100, idx=0)
            wandb.watch(model_to_get.ctx_model, log_freq=100, idx=1)

        eval_step = math.ceil(train_iterator.max_iterations / args.eval_per_epoch)
        logger.info("  Eval step = %d", eval_step)
        logger.info("***** Training *****")

        if args.restart:
            self.validate_and_save(0, 0, scheduler, do_save=False)

        for epoch in range(self.start_epoch, int(args.num_train_epochs)):
            logger.info("***** Epoch %d *****", epoch)
            self._train_epoch(scheduler, epoch, eval_step, train_iterator)

        if args.num_train_epochs == 0:
            _ = self._save_checkpoint(scheduler, 0, 0)

        if args.local_rank in [-1, 0]:
            logger.info('Training finished. Best validation checkpoint %s', self.best_cp_name)

    def validate_and_save(self, epoch: int, iteration: int, scheduler, do_save: bool = True):
        args = self.args
        # for distributed mode, save checkpoint for only one process
        save_cp = args.local_rank in [-1, 0]

        if epoch == args.val_av_rank_start_epoch:
            self.best_validation_result = None

        validation_nll = self.validate_nll()
        validation_rank = None
        if epoch >= args.val_av_rank_start_epoch:
            validation_rank = self.validate_average_rank()
        validation_loss = validation_rank or validation_nll

        if args.wandb_project is not None and args.local_rank in [-1, 0]:
            wandb.log({'val/nll': validation_nll, "val/rank": validation_rank})

        if save_cp and do_save:
            cp_name = self._save_checkpoint(scheduler, epoch, iteration)
            logger.info('Saved checkpoint to %s', cp_name)

            if validation_loss < (self.best_validation_result or validation_loss + 1):
                self.best_validation_result = validation_loss
                self.best_cp_name = cp_name
                logger.info('New Best validation checkpoint %s', cp_name)

    def validate_nll(self) -> float:
        logger.info('NLL validation ...')
        args = self.args
        self.biencoder.eval()
        data_iterator = self.get_data_iterator(args.dev_file, args.dev_batch_size, shuffle=False)

        total_loss = 0.0
        start_time = time.time()
        total_correct_predictions = 0
        num_hard_negatives = args.hard_negatives
        num_other_negatives = args.other_negatives
        log_result_step = args.log_batch_step
        batches = 0
        for i, samples_batch in enumerate(data_iterator.iterate_data()):
            biencoder_input = BiEncoder.create_biencoder_input(
                samples_batch,
                self.tensorizer,
                True,
                num_hard_negatives,
                num_other_negatives,
                shuffle=False
            )

            loss, correct_cnt = _do_biencoder_fwd_pass_gold(
                self.biencoder, biencoder_input, self.tensorizer, args
            )
            total_loss += loss.item()
            total_correct_predictions += correct_cnt
            batches += 1
            if (i + 1) % log_result_step == 0:
                logger.info('Eval step: %d , used_time=%f sec., loss=%f ',
                            i, time.time() - start_time, loss.item())

        total_loss = total_loss / batches
        total_samples = batches * args.dev_batch_size * self.distributed_factor
        correct_ratio = float(total_correct_predictions / total_samples)
        logger.info('NLL Validation: loss = %f. correct prediction ratio  %d/%d ~  %f', total_loss,
                    total_correct_predictions,
                    total_samples,
                    correct_ratio
                    )
        return total_loss

    def validate_average_rank(self) -> float:
        """
        Validates biencoder model using each question's gold passage's rank across the set
        of passages from the dataset. It generates vectors for specified amount of negative
        passages from each question (see --val_av_rank_xxx params) and stores them in RAM
        as well as question vectors. Then the similarity scores are calculted for the entire
        num_questions x (num_questions x num_passages_per_question) matrix and sorted per
        quesrtion. Each question's gold passage rank in that sorted list of scores is
        averaged across all the questions.
        :return: averaged rank number
        """
        logger.info('Average rank validation ...')

        args = self.args
        self.biencoder.eval()
        distributed_factor = self.distributed_factor

        data_iterator = self.get_data_iterator(args.dev_file, args.dev_batch_size, shuffle=False)

        sub_batch_size = args.val_av_rank_bsz
        sim_score_f = BiEncoderNllLoss.get_similarity_function()
        q_represenations = []
        ctx_represenations = []
        positive_idx_per_question = []

        num_hard_negatives = args.val_av_rank_hard_neg
        num_other_negatives = args.val_av_rank_other_neg

        log_result_step = args.log_batch_step

        for i, samples_batch in enumerate(data_iterator.iterate_data()):
            # samples += 1
            if len(q_represenations) > args.val_av_rank_max_qs / distributed_factor:
                break

            biencoder_input = BiEncoder.create_biencoder_input(
                samples_batch,
                self.tensorizer,
                True,
                num_hard_negatives,
                num_other_negatives,
                shuffle=False
            )

            biencoder_input = BiEncoderBatch(**move_to_device(biencoder_input._asdict(), args.device))
            total_ctxs = len(ctx_represenations)
            ctxs_ids = biencoder_input.context_ids
            ctxs_segments = biencoder_input.ctx_segments
            bsz = ctxs_ids.size(0)

            # split contexts batch into sub batches since it is supposed to be too large to be
            # processed in one batch
            for j, batch_start in enumerate(range(0, bsz, sub_batch_size)):

                q_ids, q_segments = None, None
                if j == 0:
                    q_ids = biencoder_input.question_ids
                    q_segments = biencoder_input.question_segments

                if j == 0 and args.n_gpu > 1 and q_ids.size(0) == 1:
                    # if we are in DP (but not in DDP) mode, all model input tensors should have
                    # batch size >1 or 0, otherwise the other input tensors will be split but
                    # only the first split will be called
                    continue

                ctx_ids_batch = ctxs_ids[batch_start:batch_start + sub_batch_size]
                ctx_seg_batch = ctxs_segments[batch_start:batch_start + sub_batch_size]

                q_attn_mask = self.tensorizer.get_attn_mask(q_ids)
                ctx_attn_mask = self.tensorizer.get_attn_mask(ctx_ids_batch)
                
                with torch.no_grad():
                    if args.fp16:
                        with autocast():
                            q_dense, ctx_dense = self.biencoder(
                                q_ids, q_segments, q_attn_mask,
                                ctx_ids_batch, ctx_seg_batch, ctx_attn_mask
                            )
                    else:
                        q_dense, ctx_dense = self.biencoder(
                            q_ids, q_segments, q_attn_mask,
                            ctx_ids_batch, ctx_seg_batch, ctx_attn_mask
                        )

                if q_dense is not None:
                    q_represenations.extend(q_dense.cpu().split(1, dim=0))

                ctx_represenations.extend(ctx_dense.cpu().split(1, dim=0))

            batch_positive_idxs = biencoder_input.is_positive
            positive_idx_per_question.extend([total_ctxs + v for v in batch_positive_idxs])

            if (i + 1) % log_result_step == 0:
                logger.info('Av.rank validation: step %d, computed ctx_vectors %d, q_vectors %d', i,
                            len(ctx_represenations), len(q_represenations))

        ctx_represenations = torch.cat(ctx_represenations, dim=0)
        q_represenations = torch.cat(q_represenations, dim=0)

        logger.info('Av.rank validation: total q_vectors size=%s', q_represenations.size())
        logger.info('Av.rank validation: total ctx_vectors size=%s', ctx_represenations.size())

        q_num = q_represenations.size(0)
        assert q_num == len(positive_idx_per_question)

        scores = sim_score_f(q_represenations, ctx_represenations)
        values, indices = torch.sort(scores, dim=1, descending=True)

        rank = 0
        for i, idx in enumerate(positive_idx_per_question):
            # aggregate the rank of the known gold passage in the sorted results for each question
            gold_idx = (indices[i] == idx).nonzero()
            rank += gold_idx.item()

        if distributed_factor > 1:
            # each node calcuated its own rank, exchange the information between node and calculate
            # the "global" average rank
            # NOTE: the set of passages is still unique for every node
            eval_stats = all_gather_list([rank, q_num], max_size=100)
            for i, item in enumerate(eval_stats):
                remote_rank, remote_q_num = item
                if i != args.local_rank:
                    rank += remote_rank
                    q_num += remote_q_num

        av_rank = float(rank / q_num)
        logger.info('Av.rank validation: average rank %s, total questions=%d', av_rank, q_num)
        return av_rank

    def _train_epoch(self, scheduler, epoch: int, eval_step: int,
                     train_data_iterator: ShardedDataIterator):
        args = self.args
        rolling_train_loss = 0.0
        epoch_loss = 0

        log_result_step = args.log_batch_step
        rolling_loss_step = args.train_rolling_loss_step
        num_contexts = args.num_contexts
        seed = args.seed
        self.biencoder.train()
        epoch_batches = train_data_iterator.max_iterations
        data_iteration = 0

        for i, samples_batch in enumerate(train_data_iterator.iterate_data(epoch=epoch)):

            # to be able to resume shuffled ctx- pools
            data_iteration = train_data_iterator.get_iteration()
            random.seed(seed + epoch + data_iteration)
            biencoder_batch = BiEncoder.create_biencoder_input_with_score(
                samples_batch,
                self.tensorizer,
                True,
                num_contexts=num_contexts,
                shuffle=not args.no_sample_contexts
            )

            if args.grad_cache:
                loss = _do_biencoder_fwd_bwd_pass_cached(
                    self.biencoder, biencoder_batch, self.tensorizer, args, self)
            else:
                loss = _do_biencoder_fwd_pass(self.biencoder, biencoder_batch, self.tensorizer, args)
                if args.fp16:
                    self.scaler.scale(loss).backward()
                else:
                    loss.backward()

            epoch_loss += loss.item()
            rolling_train_loss += loss.item()
            self.total_data_seen += len(biencoder_batch.question_ids)

            if (i + 1) % args.gradient_accumulation_steps == 0:
                if self.args.fp16:
                    self.scaler.unscale_(self.optimizer)
                if args.max_grad_norm > 0:
                    torch.nn.utils.clip_grad_norm_(self.biencoder.parameters(), args.max_grad_norm)
                if self.args.fp16:
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    self.optimizer.step()

                scheduler.step()
                self.biencoder.zero_grad()

                if args.wandb_project is not None and args.local_rank in [-1, 0]:
                    wandb.log({
                        "train/loss": loss.item(),
                        "train/lr": self.optimizer.param_groups[0]['lr'],
                        "train/data_seen": self.total_data_seen
                    })

            if i % log_result_step == 0:
                lr = self.optimizer.param_groups[0]['lr']
                logger.info('Epoch: %d: Step: %d/%d, loss=%f, lr=%f', epoch, data_iteration,
                            epoch_batches, loss.item(), lr)

            if (i + 1) % rolling_loss_step == 0:
                logger.info('Train batch %d', data_iteration)
                latest_rolling_train_av_loss = rolling_train_loss / rolling_loss_step
                logger.info('Avg. loss per last %d batches: %f', rolling_loss_step,
                            latest_rolling_train_av_loss)
                rolling_train_loss = 0.0

            if (i + 1) % eval_step == 0:
                logger.info('Validation: Epoch: %d Step: %d/%d', epoch, data_iteration,
                            epoch_batches)
                self.validate_and_save(
                    epoch,
                    train_data_iterator.get_iteration(),
                    scheduler,
                    do_save=False
                )
                self.biencoder.train()

        self.validate_and_save(epoch, data_iteration, scheduler)

        epoch_loss = (epoch_loss / epoch_batches) if epoch_batches > 0 else 0
        logger.info('Av Loss per epoch=%f', epoch_loss)

    def _save_checkpoint(self, scheduler, epoch: int, offset: int) -> str:
        args = self.args
        model_to_save = get_model_obj(self.biencoder)
        cp = os.path.join(
            args.output_dir,
            args.checkpoint_file_name + '.' + str(epoch) + ('.' + str(offset) if offset > 0 else '')
        )

        meta_params = get_encoder_params_state(args)

        state = CheckpointState(model_to_save.state_dict(),
                                self.optimizer.state_dict(),
                                scheduler.state_dict(),
                                offset,
                                epoch, meta_params
                                )
        torch.save(state._asdict(), cp)
        logger.info('Saved checkpoint at %s', cp)
        return cp

    def _load_saved_state(self, saved_state: CheckpointState):
        epoch = saved_state.epoch
        offset = saved_state.offset
        if offset == 0:  # epoch has been completed
            epoch += 1
        logger.info('Loading checkpoint @ batch=%s and epoch=%s', offset, epoch)

        if self.restart is False:
            self.start_epoch = epoch
            self.start_batch = offset

        model_to_load = get_model_obj(self.biencoder)
        logger.info('Loading saved model state ...')
        # set strict=False if you use extra projection
        model_to_load.load_state_dict(saved_state.model_dict)

        if self.restart is False and saved_state.optimizer_dict:
            logger.info('Loading saved optimizer state ...')
            self.optimizer.load_state_dict(saved_state.optimizer_dict)

        if self.restart is False and saved_state.scheduler_dict:
            self.scheduler_state = saved_state.scheduler_dict


def _gather_vectors(
    args, local_q_vector, local_ctx_vectors, local_ctx_scores
) -> Tuple[T, bool]:
    """
    Gathers the representations across all the nodes.
    params:
        args: arguments
        local_q_vector: local question vector (batch_size, vector_dim)
        local_ctx_vectors: local context vectors (batch_size, n_contexts, vector_dim)
        local_ctx_scores: local context scores (batch_size, n_contexts)
    """

    distributed_world_size = args.distributed_world_size or 1
    if distributed_world_size > 1:
        batch_size, n_contexts, vector_dim = local_ctx_vectors.size()
        # q_vector_to_send = torch.empty_like(local_q_vector).cpu().copy_(local_q_vector).detach_().view(-1)
        # ctx_vector_to_send = torch.empty_like(local_ctx_vectors).cpu().copy_(
        #     local_ctx_vectors).detach_().view(-1)
        # ctx_scores_to_send = torch.empty_like(local_ctx_scores).cpu().copy_(local_ctx_scores).detach_().view(-1)

        q_vector_to_send = local_q_vector.flatten()
        ctx_vector_to_send = local_ctx_vectors.flatten()
        ctx_scores_to_send = local_ctx_scores.flatten()

        q_vector_size = batch_size * vector_dim
        ctx_vector_size = batch_size * n_contexts * vector_dim

        concat_tensor = torch.cat([q_vector_to_send, ctx_vector_to_send, ctx_scores_to_send])
        gathered_tensor = [torch.empty_like(concat_tensor) for _ in range(distributed_world_size)]

        torch.distributed.all_gather(gathered_tensor, concat_tensor)

        global_q_vector = []
        global_ctxs_vector = []
        global_ctxs_score = []

        for i, tensor in enumerate(gathered_tensor):
            if i != args.local_rank:
                q_vector = tensor[:q_vector_size].view(batch_size, vector_dim)
                ctx_vectors = tensor[q_vector_size:q_vector_size + ctx_vector_size].view(
                    batch_size, n_contexts, vector_dim)
                ctx_scores = tensor[q_vector_size + ctx_vector_size:].view(batch_size, n_contexts)
                global_q_vector.append(q_vector.to(local_q_vector.device))
                global_ctxs_vector.append(ctx_vectors.to(local_q_vector.device))
                global_ctxs_score.append(ctx_scores.to(local_q_vector.device))
            else:
                global_q_vector.append(local_q_vector)
                global_ctxs_vector.append(local_ctx_vectors)
                global_ctxs_score.append(local_ctx_scores)

        global_q_vector = torch.cat(global_q_vector, dim=0)
        global_ctxs_vector = torch.cat(global_ctxs_vector, dim=0)
        global_ctxs_score = torch.cat(global_ctxs_score, dim=0)
    else:
        global_q_vector = local_q_vector
        global_ctxs_vector = local_ctx_vectors
        global_ctxs_score = local_ctx_scores

    return global_q_vector, global_ctxs_vector, global_ctxs_score


def _do_biencoder_fwd_bwd_pass_cached(
    model: nn.Module,
    input: BiEncoderBatchWithScore,
    tensorizer: Tensorizer,
    args,
    trainer: BiEncoderTrainer,
) -> Tuple[torch.Tensor, int]:
    input = BiEncoderBatchWithScore(**move_to_device(input._asdict(), args.device))

    reshaped = False
    if len(input.context_ids.size()) == 3:
        reshaped = True
        input.context_ids = input.context_ids.view(-1, input.context_ids.size(-1))
        input.ctx_segments = input.ctx_segments.view(-1, input.ctx_segments.size(-1))

    q_attn_mask = tensorizer.get_attn_mask(input.question_ids)
    ctx_attn_mask = tensorizer.get_attn_mask(input.context_ids)

    if model.training:
        q_id_chunks = input.question_ids.split(args.q_chunk_size)
        q_seg_chunks = input.question_segments.split(args.q_chunk_size)
        q_attn_mask_chunks = q_attn_mask.split(args.q_chunk_size)

        context_id_chunks = input.context_ids.split(args.ctx_chunk_size)
        ctx_segments_chunks = input.ctx_segments.split(args.ctx_chunk_size)
        ctx_attn_mask_chunks = ctx_attn_mask.split(args.ctx_chunk_size)

        all_q_reps = []
        all_ctx_reps = []

        q_rnds = []
        c_rnds = []

        for id_chunk, seg_chunk, attn_chunk in zip(
                q_id_chunks, q_seg_chunks, q_attn_mask_chunks):
            q_rnds.append(RandContext(id_chunk, seg_chunk, attn_chunk))
            with torch.no_grad():
                if args.fp16:
                    with autocast():
                        q_chunk_reps: T = model(
                            id_chunk, seg_chunk, attn_chunk,
                            None, None, None,
                        )[0]
                else:
                    q_chunk_reps: T = model(
                        id_chunk, seg_chunk, attn_chunk,
                        None, None, None,
                    )[0]
            all_q_reps.append(q_chunk_reps)
        all_q_reps = torch.cat(all_q_reps)

        for id_chunk, seg_chunk, attn_chunk in zip(
                context_id_chunks, ctx_segments_chunks, ctx_attn_mask_chunks):
            c_rnds.append(RandContext(id_chunk, seg_chunk, attn_chunk))
            with torch.no_grad():
                if args.fp16:
                    with autocast():
                        ctx_chunk_reps: T = model(
                            None, None, None,
                            id_chunk, seg_chunk, attn_chunk
                        )[1]
                else:
                    ctx_chunk_reps: T = model(
                        None, None, None,
                        id_chunk, seg_chunk, attn_chunk
                    )[1]
            all_ctx_reps.append(ctx_chunk_reps)
        all_ctx_reps = torch.cat(all_ctx_reps)

        if reshaped:
            all_ctx_reps = all_ctx_reps.view(all_q_reps.size(0), -1, all_ctx_reps.size(-1))

        all_q_reps = all_q_reps.float().detach().requires_grad_()
        all_ctx_reps = all_ctx_reps.float().detach().requires_grad_()

        ctx_scores = input.context_scores
        if args.inbatch_negative:
            loss_function = BiEncoderInBatchKLLoss()
            all_q_reps, all_ctx_reps, ctx_scores = _gather_vectors(args, all_q_reps, all_ctx_reps, ctx_scores)
        else:
            loss_function = BiEncoderKLLoss()

        loss = loss_function.calc(all_q_reps, all_ctx_reps, ctx_scores, args.temperature)

        if args.n_gpu > 1:
            loss = loss.mean()
        if args.gradient_accumulation_steps > 1:
            loss = loss / args.gradient_accumulation_steps
        loss.backward()

        q_grads = all_q_reps.grad.split(args.q_chunk_size)
        ctx_grads = all_ctx_reps.grad.view(-1, all_ctx_reps.size(-1)).split(args.ctx_chunk_size)

        for id_chunk, seg_chunk, attn_chunk, grad, rnd in zip(
                q_id_chunks, q_seg_chunks, q_attn_mask_chunks, q_grads, q_rnds):
            with rnd:
                if args.fp16:
                    with autocast():
                        ctx_chunk_reps: T = model(
                            id_chunk, seg_chunk, attn_chunk,
                            None, None, None,
                        )[0]
                        surrogate = torch.dot(ctx_chunk_reps.flatten().float(), grad.flatten())
                else:
                    ctx_chunk_reps: T = model(
                        id_chunk, seg_chunk, attn_chunk,
                        None, None, None,
                    )[0]
                    surrogate = torch.dot(ctx_chunk_reps.flatten().float(), grad.flatten())
            surrogate = surrogate * (trainer.distributed_factor / 8.)
            if args.fp16:
                trainer.scaler.scale(surrogate).backward()
            else:
                surrogate.backward()

        for id_chunk, seg_chunk, attn_chunk, grad, rnd in zip(
                context_id_chunks, ctx_segments_chunks, ctx_attn_mask_chunks, ctx_grads, c_rnds):
            with rnd:
                if args.fp16:
                    with autocast():
                        ctx_chunk_reps: T = model(
                            None, None, None,
                            id_chunk, seg_chunk, attn_chunk
                        )[1]
                        surrogate = torch.dot(ctx_chunk_reps.flatten().float(), grad.flatten())
                else:
                    ctx_chunk_reps: T = model(
                        None, None, None,
                        id_chunk, seg_chunk, attn_chunk
                    )[1]
                    surrogate = torch.dot(ctx_chunk_reps.flatten().float(), grad.flatten())
            surrogate = surrogate * (trainer.distributed_factor / 8.)
            if args.fp16:
                trainer.scaler.scale(surrogate).backward()
            else:
                surrogate.backward()

        return loss

    else:
        with torch.no_grad():
            if args.fp16:
                with autocast():
                    model_out = model(
                        input.question_ids, input.question_segments, q_attn_mask,
                        input.context_ids, input.ctx_segments, ctx_attn_mask
                    )
            else:
                model_out = model(
                    input.question_ids, input.question_segments, q_attn_mask,
                    input.context_ids, input.ctx_segments, ctx_attn_mask
                )

        q_vectors, ctx_vectors = model_out

        if reshaped:
            ctx_vectors = ctx_vectors.view(q_vectors.size(0), -1, ctx_vectors.size(-1))

        ctx_scores = input.context_scores
        if args.inbatch_negative:
            loss_function = BiEncoderInBatchKLLoss()
            q_vectors, ctx_vectors, ctx_scores = _gather_vectors(args, q_vectors, ctx_vectors, ctx_scores)
        else:
            loss_function = BiEncoderKLLoss()

        loss = loss_function.calc(q_vectors, ctx_vectors, ctx_scores, args.temperature)

        if args.n_gpu > 1:
            loss = loss.mean()
        if args.gradient_accumulation_steps > 1:
            loss = loss / args.gradient_accumulation_steps

        return loss


def _do_biencoder_fwd_pass(
    model: nn.Module,
    input: BiEncoderBatchWithScore,
    tensorizer: Tensorizer,
    args
) -> torch.Tensor:
    input = BiEncoderBatchWithScore(**move_to_device(input._asdict(), args.device))

    reshaped = False
    if len(input.context_ids.size()) == 3:
        reshaped = True
        input.context_ids = input.context_ids.view(-1, input.context_ids.size(-1))
        input.ctx_segments = input.ctx_segments.view(-1, input.ctx_segments.size(-1))

    q_attn_mask = tensorizer.get_attn_mask(input.question_ids)
    ctx_attn_mask = tensorizer.get_attn_mask(input.context_ids)

    if model.training:
        if args.fp16:
            with autocast():
                model_out = model(
                    input.question_ids, input.question_segments, q_attn_mask,
                    input.context_ids, input.ctx_segments, ctx_attn_mask
                )
        else:
            model_out = model(
                input.question_ids, input.question_segments, q_attn_mask,
                input.context_ids, input.ctx_segments, ctx_attn_mask
            )
    else:
        with torch.no_grad():
            if args.fp16:
                with autocast():
                    model_out = model(
                        input.question_ids, input.question_segments, q_attn_mask,
                        input.context_ids, input.ctx_segments, ctx_attn_mask
                    )
            else:
                model_out = model(
                    input.question_ids, input.question_segments, q_attn_mask,
                    input.context_ids, input.ctx_segments, ctx_attn_mask
                )

    q_vectors, ctx_vectors = model_out

    if reshaped:
        ctx_vectors = ctx_vectors.view(q_vectors.size(0), -1, ctx_vectors.size(-1))

    ctx_scores = input.context_scores
    if args.inbatch_negative:
        loss_function = BiEncoderInBatchKLLoss()
        q_vectors, ctx_vectors, ctx_scores = _gather_vectors(args, q_vectors, ctx_vectors, ctx_scores)
    else:
        loss_function = BiEncoderKLLoss()

    loss = loss_function.calc(q_vectors, ctx_vectors, ctx_scores, args.temperature)

    if args.n_gpu > 1:
        loss = loss.mean()
    if args.gradient_accumulation_steps > 1:
        loss = loss / args.gradient_accumulation_steps

    return loss


def _calc_loss(args, loss_function, local_q_vector, local_ctx_vectors, local_positive_idxs,
               local_hard_negatives_idxs: list = None,
               ) -> Tuple[T, bool]:
    """
    Calculates In-batch negatives schema loss and supports to run it in DDP mode by exchanging the
    representations across all the nodes.
    """
    distributed_world_size = args.distributed_world_size or 1
    if distributed_world_size > 1:
        q_vector_to_send = torch.empty_like(local_q_vector).cpu().copy_(local_q_vector).detach_()
        ctx_vector_to_send = torch.empty_like(local_ctx_vectors).cpu().copy_(
            local_ctx_vectors).detach_()

        global_question_ctx_vectors = all_gather_list(
            [q_vector_to_send, ctx_vector_to_send, local_positive_idxs, local_hard_negatives_idxs],
            max_size=args.global_loss_buf_sz)

        global_q_vector = []
        global_ctxs_vector = []

        # ctxs_per_question = local_ctx_vectors.size(0)
        positive_idx_per_question = []
        hard_negatives_per_question = []

        total_ctxs = 0

        for i, item in enumerate(global_question_ctx_vectors):
            q_vector, ctx_vectors, positive_idx, hard_negatives_idxs = item

            if i != args.local_rank:
                global_q_vector.append(q_vector.to(local_q_vector.device))
                global_ctxs_vector.append(ctx_vectors.to(local_q_vector.device))
                positive_idx_per_question.extend([v + total_ctxs for v in positive_idx])
                hard_negatives_per_question.extend([[v + total_ctxs for v in l]
                                                    for l in hard_negatives_idxs])
            else:
                global_q_vector.append(local_q_vector)
                global_ctxs_vector.append(local_ctx_vectors)
                positive_idx_per_question.extend([v + total_ctxs for v in local_positive_idxs])
                hard_negatives_per_question.extend([[v + total_ctxs for v in l]
                                                    for l in local_hard_negatives_idxs])
            total_ctxs += ctx_vectors.size(0)

        global_q_vector = torch.cat(global_q_vector, dim=0)
        global_ctxs_vector = torch.cat(global_ctxs_vector, dim=0)

    else:
        global_q_vector = local_q_vector
        global_ctxs_vector = local_ctx_vectors
        positive_idx_per_question = local_positive_idxs
        hard_negatives_per_question = local_hard_negatives_idxs

    loss, is_correct = loss_function.calc(
        global_q_vector, global_ctxs_vector, positive_idx_per_question, hard_negatives_per_question
    )

    return loss, is_correct


def _do_biencoder_fwd_pass_gold(
    model: nn.Module,
    input: BiEncoderBatch,
    tensorizer: Tensorizer,
    args
) -> Tuple[torch.Tensor, int]:
    input = BiEncoderBatch(**move_to_device(input._asdict(), args.device))

    q_attn_mask = tensorizer.get_attn_mask(input.question_ids)
    ctx_attn_mask = tensorizer.get_attn_mask(input.context_ids)

    if model.training:
        if args.fp16:
            with autocast():
                model_out = model(
                    input.question_ids, input.question_segments, q_attn_mask, input.context_ids,
                    input.ctx_segments, ctx_attn_mask
                )
        else:
            model_out = model(
                input.question_ids, input.question_segments, q_attn_mask, input.context_ids,
                input.ctx_segments, ctx_attn_mask
            )
    else:
        with torch.no_grad():
            if args.fp16:
                with autocast():
                    model_out = model(
                        input.question_ids, input.question_segments, q_attn_mask, input.context_ids,
                        input.ctx_segments, ctx_attn_mask
                    )
            else:
                model_out = model(
                    input.question_ids, input.question_segments, q_attn_mask, input.context_ids,
                    input.ctx_segments, ctx_attn_mask
                )

    local_q_vector, local_ctx_vectors = model_out

    loss_function = BiEncoderNllLoss()

    loss, is_correct = _calc_loss(
        args, loss_function, local_q_vector, local_ctx_vectors,
        input.is_positive, input.hard_negatives
    )

    is_correct = is_correct.sum().item()

    if args.n_gpu > 1:
        loss = loss.mean()
    if args.gradient_accumulation_steps > 1:
        loss = loss / args.gradient_accumulation_steps

    return loss, is_correct


def main():
    parser = argparse.ArgumentParser()

    add_encoder_params(parser)
    add_training_params(parser)
    add_tokenizer_params(parser)

    # biencoder specific training features
    parser.add_argument(
        "--eval_per_epoch",
        default=1,
        type=int,
        help="How many times it evaluates on dev set per epoch and saves a checkpoint"
    )

    parser.add_argument(
        "--global_loss_buf_sz",
        type=int,
        default=150000,
        help='Buffer size for distributed mode representations al gather operation. \
             Increase this if you see errors like "encoded data exceeds max_size ..."'
    )
    parser.add_argument("--shuffle_positive_ctx", action='store_true')

    # input/output src params
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        help="The output directory where the model checkpoints will be written or resumed from"
    )

    # data handling parameters
    parser.add_argument("--hard_negatives", default=1, type=int,
                        help="amount of hard negative ctx per question")
    parser.add_argument("--other_negatives", default=0, type=int,
                        help="amount of 'other' negative ctx per question")
    parser.add_argument("--num_contexts", default=0, type=int,
                        help="amount of 'other' negative ctx per question")
    parser.add_argument("--inbatch_negative", action="store_true",
                        help="set to true to enable using in-batch negatives for the KL loss")
    parser.add_argument("--no_sample_contexts", action="store_true",
                        help="set to true to disable sampling contexts from the candidate list")
    parser.add_argument("--train_files_upsample_rates", type=str,
                        help="list of up-sample rates per each train file. Example: [1,2,1]")
    parser.add_argument("--temperature", type=float, default=1.0, help="temperature for softmax")

    # parameters for Av.rank validation method
    parser.add_argument("--val_av_rank_start_epoch", type=int, default=10000,
                        help="Av.rank validation: the epoch from which to enable this validation")
    parser.add_argument(
        "--val_av_rank_hard_neg",
        type=int,
        default=30,
        help="Av.rank validation: how many hard negatives to take from each question pool"
    )
    parser.add_argument(
        "--val_av_rank_other_neg",
        type=int,
        default=30,
        help="Av.rank validation: how many 'other' negatives to take from each question pool"
    )
    parser.add_argument("--val_av_rank_bsz", type=int, default=128,
                        help="Av.rank validation: batch size to process passages")
    parser.add_argument("--val_av_rank_max_qs", type=int, default=10000,
                        help="Av.rank validation: max num of questions")
    parser.add_argument(
        '--checkpoint_file_name',
        type=str,
        default='dpr_biencoder',
        help="Checkpoints file prefix"
    )

    parser.add_argument(
        "--restart",
        action='store_true',
        help="set true if you want to reset the optimizer and schedular states when you \
             start fine-tuning from pre-trained models"
    )
    args = parser.parse_args()

    if args.gradient_accumulation_steps < 1:
        raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
            args.gradient_accumulation_steps))

    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)

    if args.wandb_project is not None and args.local_rank in [-1, 0]:
        wandb.init(project=args.wandb_project, name=args.wandb_name, config=args)

    setup_args_gpu(args)
    set_seed(args)
    print_args(args)

    trainer = BiEncoderTrainer(args)

    if args.train_file is not None:
        trainer.run_train()
    elif args.model_file and args.dev_file:
        logger.info(
            "No train files are specified. Run 2 types of validation for specified model file")
        trainer.validate_nll()
        trainer.validate_average_rank()
    else:
        logger.warning(
            "Neither train_file or (model_file & dev_file) parameters are specified. Nothing to do."
        )


if __name__ == "__main__":
    main()