options.py

"""
Command line arguments utils
"""

import argparse
import logging
import os
import random
import socket

import numpy as np
import torch

logger = logging.getLogger()


def add_tokenizer_params(parser: argparse.ArgumentParser):
    parser.add_argument(
        "--do_lower_case",
        action="store_true",
        help="Whether to lower case the input text. True for uncased models, False for cased models.",
    )


def add_encoder_params(parser: argparse.ArgumentParser):
    """
    Common parameters to initialize an encoder-based model
    """
    parser.add_argument(
        "--pretrained_model_cfg",
        default=None,
        type=str,
        help="config name for model initialization",
    )
    parser.add_argument(
        "--encoder_model_type",
        default=None,
        type=str,
        help="model type. One of [hf_bert, pytext_bert, fairseq_roberta]",
    )
    parser.add_argument(
        "--decoder_name",
        default="google/bert_uncased_L-2_H-768_A-12",
        type=str,
        help="model type. One of [hf_bert, pytext_bert, fairseq_roberta]",
    )
    parser.add_argument(
        "--pretrained_file",
        type=str,
        help="Some encoders need to be initialized from a file",
    )
    parser.add_argument(
        "--model_file",
        default=None,
        type=str,
        help="Saved bi-encoder checkpoint file to initialize the model",
    )
    parser.add_argument(
        "--projection_dim",
        default=0,
        type=int,
        help="Extra linear layer on top of standard bert/roberta encoder",
    )
    parser.add_argument(
        "--sequence_length",
        type=int,
        default=512,
        help="Max length of the encoder input sequence",
    )
    parser.add_argument(
        "--weight_sharing",
        action="store_true",
        help="Whether to share weights between question encoder and ctx encoder"
    )


def add_training_params(parser: argparse.ArgumentParser):
    """
    Common parameters for training
    """
    add_cuda_params(parser)
    parser.add_argument(
        "--train_file", default=None, type=str, help="File pattern for the train set"
    )
    parser.add_argument("--dev_file", default=None, type=str, help="")

    parser.add_argument(
        "--batch_size", default=2, type=int, help="Amount of questions per batch"
    )
    parser.add_argument(
        "--dev_batch_size",
        type=int,
        default=4,
        help="amount of questions per batch for dev set validation",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="random seed for initialization and dataset shuffling",
    )

    parser.add_argument(
        "--adam_eps", default=1e-8, type=float, help="Epsilon for Adam optimizer."
    )
    parser.add_argument(
        "--adam_betas",
        default="(0.9, 0.999)",
        type=str,
        help="Betas for Adam optimizer.",
    )

    parser.add_argument(
        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
    )
    parser.add_argument("--log_batch_step", default=100, type=int, help="")
    parser.add_argument("--train_rolling_loss_step", default=100, type=int, help="")
    parser.add_argument("--weight_decay", default=0.0, type=float, help="")
    parser.add_argument(
        "--learning_rate",
        default=1e-5,
        type=float,
        help="The initial learning rate for Adam.",
    )

    parser.add_argument(
        "--warmup_steps", default=100, type=int, help="Linear warmup over warmup_steps."
    )
    parser.add_argument("--dropout", default=0.1, type=float, help="")

    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(
        "--num_train_epochs",
        default=3.0,
        type=float,
        help="Total number of training epochs to perform.",
    )


def add_cuda_params(parser: argparse.ArgumentParser):
    parser.add_argument(
        "--no_cuda", action="store_true", help="Whether not to use CUDA when available"
    )
    parser.add_argument(
        "--local_rank",
        type=int,
        default=-1,
        help="local_rank for distributed training on gpus",
    )
    parser.add_argument(
        "--fp16",
        action="store_true",
        help="Whether to use 16-bit float precision instead of 32-bit",
    )

    parser.add_argument(
        "--fp16_opt_level",
        type=str,
        default="O1",
        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
        "See details at https://nvidia.github.io/apex/amp.html",
    )


def add_reader_preprocessing_params(parser: argparse.ArgumentParser):
    parser.add_argument(
        "--gold_passages_src",
        type=str,
        help="File with the original dataset passages (json format). Required for train set",
    )
    parser.add_argument(
        "--gold_passages_src_dev",
        type=str,
        help="File with the original dataset passages (json format). Required for dev set",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=16,
        help="number of parallel processes to binarize reader data",
    )


def get_encoder_checkpoint_params_names():
    return [
        "do_lower_case",
        "pretrained_model_cfg",
        "encoder_model_type",
        "pretrained_file",
        "projection_dim",
        "sequence_length",
        "weight_sharing",
    ]


def get_encoder_params_state(args):
    """
     Selects the param values to be saved in a checkpoint, so that a trained model faile can be used for downstream
     tasks without the need to specify these parameter again
    :return: Dict of params to memorize in a checkpoint
    """
    params_to_save = get_encoder_checkpoint_params_names()

    r = {}
    for param in params_to_save:
        r[param] = getattr(args, param)
    return r


def set_encoder_params_from_state(state, args):
    if not state:
        return
    params_to_save = get_encoder_checkpoint_params_names()

    override_params = [
        (param, state[param])
        for param in params_to_save
        if param in state and state[param]
    ]
    for param, value in override_params:
        if hasattr(args, param):
            logger.warning(
                "Overriding args parameter value from checkpoint state. Param = %s, value = %s",
                param,
                value,
            )
        setattr(args, param, value)
    return args


def set_seed(args):
    seed = args.seed
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(seed)


def setup_args_gpu(args):
    """
    Setup arguments CUDA, GPU & distributed training
    """

    if args.local_rank == -1 or args.no_cuda:  # single-node multi-gpu (or cpu) mode
        device = torch.device(
            "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"
        )
        args.n_gpu = torch.cuda.device_count()
    else:  # distributed mode
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl")
        args.n_gpu = 1
    args.device = device
    ws = os.environ.get("WORLD_SIZE")

    args.distributed_world_size = int(ws) if ws else 1

    logger.info(
        "Initialized host %s as d.rank %d on device=%s, n_gpu=%d, world size=%d",
        socket.gethostname(),
        args.local_rank,
        device,
        args.n_gpu,
        args.distributed_world_size,
    )
    logger.info("16-bits training: %s ", args.fp16)


def print_args(args, output_dir=None):
    logger.info(" **************** CONFIGURATION **************** ")
    for key, val in sorted(vars(args).items()):
        keystr = "{}".format(key) + (" " * (30 - len(key)))
        logger.info("%s -->   %s", keystr, val)
    logger.info(" **************** CONFIGURATION **************** ")

    if output_dir is not None:
        with open(os.path.join(output_dir, "args.txt"), "w") as f:
            for key, val in sorted(vars(args).items()):
                keystr = "{}".format(key) + (" " * (30 - len(key)))
                f.write(f"{keystr}   {val}\n")