args.py

import os
import sys
import argparse

xnli_langs="en fr es de el bg ru tr ar vi th zh hi sw ur"
pawsx_langs="en de es fr ja ko zh"
panx_langs="ar he vi id jv ms tl eu ml ta te af nl en de el bn hi mr ur fa fr it pt es bg ru ja ka ko th sw yo my zh kk tr et fi hu"
conll_langs="en es nl de ar fi"
multiatis_langs="en de es fr hi ja pt tr zh"


def cross_lingual_params(parser):
    group = parser.add_argument_group('Cross-lingual params.')
    group.add_argument("--dev_lang",
                       default="en", 
                       type=str, 
                       help="Name of the development language, seperated by `;`. (model tuned by this language dev. set). Value-type: (str)")
    group.add_argument("--src_lang",
                       default="en", 
                       type=str, 
                       help="Name of the source languages, seperated by `;`. Value-type: (str)")
    group.add_argument("--tgt_lang",
                       default="en", 
                       type=str, 
                       help="Name of the tgt language, seperated by `;`. Value-type: (str)")


def xnli_data_params(parser):
    
    assert len(xnli_langs.split()) == 15
    default_train_data, default_dev_data, default_test_data = [], [], []
    for lg in xnli_langs.split():
        default_train_data.append(
            "data/xnli/XNLI-MT-1.0/multinli/multinli.train.{}.tsv;utf-8-sig;{}".format(lg, lg)
        )
        default_dev_data.append(
            "data/xnli/XNLI-1.0/xnli.dev.tsv;utf-8-sig;{}".format(lg)
        )
        default_test_data.append(
            "data/xnli/XNLI-1.0/xnli.test.tsv;utf-8-sig;{}".format(lg)
        )
        

    group = parser.add_argument_group('Dataset params.')
    group.add_argument("--train",
                        nargs='*',
                        default=default_train_data, 
                        help="Train set location. Value-type: list(string)")
    group.add_argument("--dev",
                        nargs='*',
                        default=default_dev_data, 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--test",
                        nargs='*',
                        default=default_test_data, 
                        help="Test set location. Value-type: list(string)")
    group.add_argument("--external_data",
                        nargs='*',
                        default=[], 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--label",
                       default=None, 
                       type=str, 
                       help="Path where label file is saved.")


def pawsx_data_params(parser):

    assert len(pawsx_langs.split()) == 7
    default_train_data, default_dev_data, default_test_data = [], [], []
    for lg in pawsx_langs.split():
        file_name = "train"
        if lg != "en":
            file_name = "translated_train"
        default_train_data.append(
            "data/pawsx/{}/{}.tsv;utf-8-sig;{}".format(lg, file_name, lg)
        )
        default_dev_data.append(
            "data/pawsx/{}/dev_2k.tsv;utf-8-sig;{}".format(lg, lg)
        )
        default_test_data.append(
           "data/pawsx/{}/test_2k.tsv;utf-8-sig;{}".format(lg, lg)
        )
        

    group = parser.add_argument_group('Dataset params.')
    group.add_argument("--train",
                        nargs='*',
                        default=default_train_data, 
                        help="Train set location. Value-type: list(string)")
    group.add_argument("--dev",
                        nargs='*',
                        default=default_dev_data, 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--test",
                        nargs='*',
                        default=default_test_data, 
                        help="Test set location. Value-type: list(string)")
    group.add_argument("--external_data",
                        nargs='*',
                        default=[], 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--label",
                       default=None, 
                       type=str, 
                       help="Path where label file is saved.")


def panx_data_params(parser, model_type, tok_name):
    
    assert len(panx_langs.split()) == 40
    default_train_data, default_dev_data, default_test_data = [], [], []
    for lg in panx_langs.split():
        default_train_data.append(
            "data/panx/{}.train.{}.{}.tok;utf-8-sig;{}".format(lg, model_type, tok_name, lg)
        )
        default_dev_data.append(
            "data/panx/{}.dev.{}.{}.tok;utf-8-sig;{}".format(lg, model_type, tok_name, lg)
        )
        default_test_data.append(
            "data/panx/{}.test.{}.{}.tok;utf-8-sig;{}".format(lg, model_type, tok_name, lg)
        )

    group = parser.add_argument_group('Dataset params.')
    group.add_argument("--train",
                        nargs='*',
                        default=default_train_data, 
                        help="Train set location. Value-type: list(string)")
    group.add_argument("--dev",
                        nargs='*',
                        default=default_dev_data, 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--test",
                        nargs='*',
                        default=default_test_data, 
                        help="Test set location. Value-type: list(string)")
    group.add_argument("--external_data",
                        nargs='*',
                        default=[], 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--label",
                       default=None, 
                       type=str, 
                       help="Path where label file is saved.")


def conll_data_params(parser):
    
    assert len(conll_langs.split()) == 6
    default_train_data, default_dev_data, default_test_data = [], [], []
    for lg in conll_langs.split():
        encoding = 'latin-1' if lg == 'de' else 'utf-8' 
        default_train_data.append(
            "data/conll_ner/{}/{}.train.iob2;{};{}".format(lg, lg, encoding, lg)
        )
        default_dev_data.append(
            "data/conll_ner/{}/{}.testa.iob2;{};{}".format(lg, lg, encoding, lg)
        )
        default_test_data.append(
            "data/conll_ner/{}/{}.testb.iob2;{};{}".format(lg, lg, encoding, lg)
        )
        

    group = parser.add_argument_group('Dataset params.')
    group.add_argument("--train",
                        nargs='*',
                        default=default_train_data, 
                        help="Train set location. Value-type: list(string)")
    group.add_argument("--dev",
                        nargs='*',
                        default=default_dev_data, 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--test",
                        nargs='*',
                        default=default_test_data, 
                        help="Test set location. Value-type: list(string)")
    group.add_argument("--external_data",
                        nargs='*',
                        default=[], 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--label",
                       default=None, 
                       type=str, 
                       help="Path where label file is saved.")

def multiatis_ic_data_params(parser):
    
    assert len(xnli_langs.split()) == 15
    default_train_data, default_dev_data, default_test_data = [], [], []
    for lg in multiatis_langs.split():
        default_train_data.append(
            "data/multiatis/data/MultiATIS++/data/train_dev_test/train_{}.ic.tsv;utf-8-sig;{}".format(lg.upper(), lg)
        )
        default_dev_data.append(
            "data/multiatis/data/MultiATIS++/data/train_dev_test/dev_{}.ic.tsv;utf-8-sig;{}".format(lg.upper(), lg)
        )
        default_test_data.append(
            "data/multiatis/data/MultiATIS++/data/train_dev_test/test_{}.ic.tsv;utf-8-sig;{}".format(lg.upper(), lg)
        )

    group = parser.add_argument_group('Dataset params.')
    group.add_argument("--train",
                        nargs='*',
                        default=default_train_data, 
                        help="Train set location. Value-type: list(string)")
    group.add_argument("--dev",
                        nargs='*',
                        default=default_dev_data, 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--test",
                        nargs='*',
                        default=default_test_data, 
                        help="Test set location. Value-type: list(string)")
    group.add_argument("--external_data",
                        nargs='*',
                        default=[], 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--label",
                       default=None, 
                       type=str, 
                       help="Path where label file is saved.")


def model_params(parser):
    group = parser.add_argument_group('Model params.')
    group.add_argument("--model_type", 
                        default="bert", 
                        type=str, 
                        help="Model type selected in the list. Value-type: string")
    group.add_argument("--model_name_or_path", 
                        default="bert-base-multilingual-cased", 
                        type=str, 
                        help="Path to pre-trained model or shortcut name selected in the list. Value-type: string")
    group.add_argument("--config_name", 
                        default="", 
                        type=str,
                        help="Pretrained config name or path if not the same as model_name. Value-type: string")
    group.add_argument("--tokenizer_name", 
                        default="bert-base-multilingual-cased", 
                        type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name. Value-type: string")
    parser.add_argument("--do_lower_case", 
                        default=0, 
                        type=int, 
                        help="Set this flag if you are using an uncased model.")    
    group.add_argument("--dropout", 
                        default=.1, 
                        type=float,
                        help="Dropout value of the hidden representation of the LM. Value-type: float")
    group.add_argument("--max_seq_length", 
                        default=512, 
                        type=int,
                        help="The maximum total input sequence length after tokenization. Sequences longer "
                             "than this will be truncated, sequences shorter will be padded. Value-type: int")


def logistics_params(parser):
    group = parser.add_argument_group('Logistics params.')
    group.add_argument("--task_name", 
                        default="temp-task", 
                        type=str, 
                        help="task name. Value-type: path(string)",
                        choices=["xnli", "pawsx", "panx", "conll", "multiatis_ic", "temp-task"])
    group.add_argument("--output_dir", 
                        default="./dumped", 
                        type=str, 
                        help="The output directory where the model predictions and checkpoints will be written. Value-type: path(string)")
    group.add_argument("--cache_dir", 
                        default=None, 
                        help="Where do you want to store the pre-trained models downloaded from s3. Value-type: path(string)")
    group.add_argument("--per_gpu_train_batch_size", 
                        type=int, 
                        default=4,
                        help="Batch size per GPU/CPU for evaluation. Value-type: int")
    group.add_argument("--eval_single_model", 
                        type=int,
                        default=0,
                        help="Evaluate a single model.")
    group.add_argument("--per_gpu_eval_batch_size", 
                        type=int, 
                        default=32,
                        help="Batch size per GPU/CPU for evaluation. Value-type: int")
    group.add_argument("--gradient_accumulation_steps", 
                        type=int, 
                        default=4,
                        help="Number of updates steps to accumulate before performing a backward/update pass. Value-type: int")
    group.add_argument("--logging_steps", 
                        type=int, 
                        default=50,
                        help="Log every X updates steps. Value-type: int")
    group.add_argument("--save_steps", 
                        type=int, 
                        default=50,
                        help="Save checkpoint every X updates steps. Value-type: int")
    group.add_argument("--eval_steps", 
                        type=int, 
                        default=250,
                        help="Do a evaluation on dev dataset to select best model. Value-type: int")
    group.add_argument("--sampling_penalty", 
                        type=float, 
                        default=.5,
                        help="Save checkpoint every X updates steps. Value-type: float")
    group.add_argument("--eval_all_checkpoints", 
                        action="store_true",
                        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
    group.add_argument("--no_cuda", 
                        action="store_true",
                        help="Avoid using CUDA when available")
    group.add_argument("--overwrite_output_dir", 
                        action="store_true",
                        help="Overwrite the content of the output directory")
    group.add_argument("--overwrite_cache", 
                        action="store_true",
                        help="Overwrite the cached training and evaluation sets")
    group.add_argument("--seed", 
                        type=int, 
                        default=1234,
                        help="random seed for initialization, Value-type: int")
    group.add_argument("--resume_training", 
                        action="store_true",
                        help="Resume training from a step, Value-type: int")
    group.add_argument("--logger_id", 
                        default=None,
                        help="logger id, if None, by default gpu_id is selected. Value-type: str")
    group.add_argument("--load_args", 
                        default=None,
                        help="load_args from file. Value-type: str")
    group.add_argument("--disable_tqdm", 
                        action="store_true",
                        help="Disable tqdm in simulation. Value-type: str")
    group.add_argument("--overwrite_num_of_label", 
                        default=None,
                        type=int, 
                        help="Overwrite num of label while loading. Value-type: str")


def training_params(parser):
    group = parser.add_argument_group('Training params.')
    group.add_argument("--do_train", 
                        action="store_true",
                        help="Whether to run training.")
    group.add_argument("--do_eval", 
                        action="store_true",
                        help="Whether to run eval on the dev set.")
    group.add_argument("--do_transductive_eval", 
                        action="store_true",
                        help="Whether to run transductive eval or not.")
    group.add_argument("--do_maml_train", 
                        action="store_true",
                        help="Whether to run transductive eval or not.")
    group.add_argument("--do_few_shot_benchamrk", 
                        action="store_true",
                        help="Whether to run transductive eval or not.")
    group.add_argument("--evaluate_during_training", 
                        action="store_true",
                        help="Whether to run evaluation during training at each logging step.")
    group.add_argument("--evaluate_test_on_best_dev", 
                        action="store_true",
                        help="Evaluate test score when best dev found.")
    group.add_argument("--model_selection_metric", 
                        default="acc",
                        type=str,
                        help="Whether to run evaluation during training at each logging step.")
    group.add_argument("--dev_metric_comp", 
                        default="larger",
                        type=str,
                        help="Whether to run evaluation during training at each logging step.",
                        choices=["larger", "smaller"])
    group.add_argument("--n_best_dev", 
                        default=3,
                        type=int,
                        help="N Number of best mode will be selected.")
    group.add_argument("--learning_rate", 
                        default=2e-5, 
                        type=float,
                        help="The initial learning rate for Adam. Value-type: float")
    group.add_argument("--weight_decay", 
                        default=0.01, 
                        type=float,
                        help="Weight decay if we apply some. Value-type: float")
    group.add_argument("--adam_epsilon", 
                        default=1e-8, 
                        type=float,
                        help="Epsilon for Adam optimizer. Value-type: float")
    group.add_argument("--max_grad_norm", 
                        default=1.0, 
                        type=float,
                        help="Max gradient norm. Value-type: float")
    group.add_argument("--num_train_epochs", 
                        default=3.0, 
                        type=float,
                        help="Total number of training epochs to perform. Value-type: float")
    group.add_argument("--max_steps", 
                        default=0, 
                        type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs. Value-type: int")
    group.add_argument("--warmup_steps", 
                        default=-1, 
                        type=int,
                        help="Linear warmup over percentage of batch sample, if negative value (<0), it converts warmup_percentage into warmup_steps. Value-type: int")
    group.add_argument("--warmup_percentage", 
                        default=-1, 
                        type=float,
                        help="Percentage of training steps that will be used for warmup. Value-type: float")
    group.add_argument("--train_data_percentage", 
                        default=100, 
                        type=float,
                        help="Percentage of training data that will be selected. Value-type: float")
    group.add_argument("--conf_penalty", 
                        default=0, 
                        type=int,
                        help="Add a NegEntropy term with loss if not zero. Value-type: int")
    group.add_argument("--marginal_entropy", 
                        default=0, 
                        type=int,
                        help="Add a Marginal Entropy term with loss if not zero. Value-type: int")


def dist_params(parser):
    group = parser.add_argument_group('Distributed params.')
    group.add_argument("--fp16", 
                        action="store_true",
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    group.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")
    group.add_argument("--local_rank", 
                        type=int, 
                        default=-1,
                        help="For distributed training: local_rank")
    group.add_argument("--server_ip", 
                        type=str, 
                        default="", 
                        help="For distant debugging.")
    group.add_argument("--server_port", 
                        type=str, 
                        default="", 
                        help="For distant debugging.")


def inference_params(parser):
    group = parser.add_argument_group('Inference params.')
    group.add_argument('--lmd', 
                        default=0, 
                        type=float,
                        help='weight for Laplacian')
    group.add_argument('--knn', 
                        default=3, 
                        type=int,
                        help='knn for affinity')
    group.add_argument('--lshot', 
                        action='store_true',
                        help='enable LaplacianShot.')
    group.add_argument('--tune-lmd', 
                        default = 0,  
                        type=int,
                        help='Tune Lambda on Validation set.')
    group.add_argument('--proto-rect', 
                        default = 1,  
                        type=int,
                        help='Prototype rectification')
    group.add_argument('--plot-converge', 
                        action='store_true',
                        help='plot the energy in each bound updates.')
    group.add_argument('--cache_logit', 
                        action='store_true',
                        help='Cache logit, lm_output of the train, dev, trest dataset.')
    group.add_argument('--benchmark_transductive', 
                        action='store_true',
                        help='Benchmark regular inference with laplacian-shot.')


def meta_params(parser):
    group = parser.add_argument_group('Meta params.')
    group.add_argument("--support_set",
                        nargs='*',
                        default=[], 
                        help="Validation set location. Value-type: list(string)")
    group.add_argument("--cross_task_name",
                        type=str,
                        default=None, 
                        help="Name of the cross task. Value-type: list(string)")
    group.add_argument("--shot",
                        type=int,
                        default=5, 
                        help="Number of sample for a class. Value-type: list(string)")
    group.add_argument("--val_shot",
                        type=int,
                        default=0, 
                        help="Number of sample for a class. Value-type: list(string)")
    group.add_argument("--support_split",
                        type=str,
                        default="dev", 
                        help="Support set comes from a split. Value-type: list(string)")
    parser.add_argument('--meta_test_iter', type=int, default=10000,
                        help='number of iterations for meta test')
    parser.add_argument('--benchmarks', 
                        nargs='*',
                        default=["zero_shot", "knn", "finetuning"], 
                        help='Types of benchmark to be evaluated.')
    parser.add_argument('--fs_finetune_lr', 
                        nargs='*',
                        default=[7.5e-06], 
                        help='Types of benchmark to be evaluated.')
    # 5e-6, 7.5e-06, 5e-05, 1e-5
    #, 16, 32
    parser.add_argument('--fs_finetune_batch', 
                        nargs='*',
                        default=[4], 
                        help='Types of benchmark to be evaluated.')
    parser.add_argument('--fs_grad_acc_step', 
                        nargs='*',
                        default=[4], 
                        help='Types of benchmark to be evaluated.')
    # "constant"
    parser.add_argument('--fs_finetune_lr_scheduler', 
                        nargs='*',
                        default=["linear_decay"], 
                        help='Types of benchmark to be evaluated.')
    # , "head"
    parser.add_argument('--fs_finetune_type', 
                        nargs='*',
                        default=["full", "head"], 
                        help='Types of benchmark to be evaluated.')
    parser.add_argument('--fewshot_benchmark_epoch', 
                        type=int, 
                        default=5,
                        help='number of epoch in episodic supevised training to compare with lshot, used for benchmarking l-shot.')
    parser.add_argument('--fs_finetune_cf', 
                        nargs='*',
                        default=[0], 
                        help='number of epoch in episodic supevised training to compare with lshot, used for benchmarking l-shot.')
    parser.add_argument('--fs_finetune_transductive', 
                        nargs='*',
                        default=[0, 2], 
                        type=int,
                        help='Should we perform transductive finetuning.')
    parser.add_argument('--meta_val_way', 
                        type=int, 
                        default=3,
                        help='number of ways for meta val/test')
    parser.add_argument('--meta_val_query', 
                        type=int, 
                        default=15,
                        help='number of queries for meta val/test')
    parser.add_argument('--lshot_remap', 
                        action='store_true',
                        help='Perform a hungerian based remapping after prediction.')
    parser.add_argument('--train_mean', 
                        action='store_true',
                        help='No train mean for knn.')

    
def find_task():
    for i, name in enumerate(sys.argv):
        if name == "--task_name":
            return sys.argv[i+1]
    return "temp-task"

def find_model_type():
    for i, name in enumerate(sys.argv):
        if name == "--model_type":
            return sys.argv[i+1]
    raise NotImplementedError

def find_tokenizer_name():
    for i, name in enumerate(sys.argv):
        if name == "--tokenizer_name":
            return sys.argv[i+1]
    raise NotImplementedError

def load_args():
    
    parser = argparse.ArgumentParser("Few-Shot Contextual Cross-lingual Adaptation.")

    task_name = find_task()
    model_type = find_model_type()
    tok_name = find_tokenizer_name()

    if task_name == "xnli":
        xnli_data_params(parser)
    elif task_name == "pawsx":
        pawsx_data_params(parser)
    elif task_name == "panx":
        panx_data_params(parser, model_type, tok_name)
    elif task_name == "conll":
        conll_data_params(parser)
    elif task_name == "multiatis_ic":
        multiatis_ic_data_params(parser)

    cross_lingual_params(parser)
    model_params(parser)
    logistics_params(parser)
    training_params(parser)
    dist_params(parser)
    inference_params(parser)
    meta_params(parser)

    args = parser.parse_args()
    args = args

    if args.task_name == "temp-task":
        raise NotImplementedError("Please provide a valid task name.")
    return args