arguments.py

import argparse

def get_train_parser():
    """
    Generate a parameters parser for training.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--node_rank', type=int, default=0, metavar='N', help='node rank.')
    parser.add_argument('--local_rank', type=int, default=0, metavar='N', help='Local process rank.')
    parser.add_argument('--nproc_per_node', type=int, default=1, metavar='N', help='gpu process per node.')
    parser.add_argument("--job_name", default=None,type=str,required=True,help="wandb jobs name")
    parser.add_argument("--base_model_name",default=None,type=str,required=True,help="pretrained model name")
    parser.add_argument("--use_pooler", default=False, action='store_true', help="use pooler output")
    parser.add_argument("--aggregation", default='cls',type=str, required=False, help="dense aggregation strategy")
    parser.add_argument("--data_dir", default=None,type=str,required=False,help="dataset directory")
    parser.add_argument("--qidpidtriples",default=None,type=str,help="msmarco triples ids")
    parser.add_argument("--langs",default=None,type=str,nargs='+',required=True,help="languages")
    parser.add_argument("--output_dir",default=None,type=str,required=True,help="output directory")
    parser.add_argument("--checkpoint",default=None,type=str,help="custom model init checkpoint")
    parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
    parser.add_argument("--batch_size", type=int, default=2, help="batch size")
    parser.add_argument("--num_train_epochs", default=3, type=int, help="total number of training epochs.")
    parser.add_argument("--num_train", default=10000000, type=int, help="num of training triples")
    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
    parser.add_argument("--logging_steps", type=int, default=1000, help="logging every n steps")
    parser.add_argument("--query_maxlen", type=int, default=32, help="max query token length")
    parser.add_argument("--doc_maxlen", type=int, default=180, help="max document token length")
    parser.add_argument("--train_n_passages", type=int, default=2, help="number of passages per query")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=8, help="gradient accumulation")
    parser.add_argument("--domain_accumulation_size", type=int, default=128, help="domain classifier accumulation size")
    parser.add_argument('--normalize', default=False, action='store_true', help="normalize embeddings")
    parser.add_argument("--temperature", default=1.0, type=float, help="temperature for softmax")
    parser.add_argument('--fp16', default=False, action='store_true', help="mixed precision training")
    parser.add_argument('--add_prefix', default=False, action='store_true', help="add input type prefix for mE5 model")
    parser.add_argument("--parallel_collection",default=None,type=str,help="parallel collection")
    return parser

def get_index_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--node_rank', type=int, default=0, metavar='N', help='node rank.')
    parser.add_argument('--local_rank', type=int, default=0, metavar='N', help='Local process rank.')
    parser.add_argument('--nproc_per_node', type=int, default=1, metavar='N', help='gpu process per node.')
    parser.add_argument("--base_model_name",default=None,type=str,required=True,help="pretrained model name")
    parser.add_argument("--num_layers", type=int, default=3, help="num of decoder layer")
    parser.add_argument("--language_list",default=None,type=str,help="language in parallel collection")
    parser.add_argument("--index_name",default=None,type=str,help="index name")
    parser.add_argument("--checkpoint",default=None,type=str,help="custom model init checkpoint")
    parser.add_argument("--fitter_checkpoint",default=None,type=str,help="LEACE checkpoint")
    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
    parser.add_argument("--batch_size", type=int, default=20, help="batch size")
    parser.add_argument("--collection",default=None,type=str,help="collection in tsv file")
    parser.add_argument("--query_maxlen", type=int, default=180, help="max query length")
    parser.add_argument("--doc_maxlen", type=int, default=180, help="max document (passage) token length")
    parser.add_argument("--stride", type=int, default=90, help="overlap between passages")
    parser.add_argument("--chunk_size", type=int, default=100000, help="cache vectors to file")
    parser.add_argument("--output_dir",default=None,type=str,required=True,help="output directory")
    parser.add_argument("--faiss_index", default=None, type=str, help="faiss index file")
    parser.add_argument("--test_queries", default=None, type=str, help="test query file")
    parser.add_argument("--test_qrel", default=None, type=str, help="test qrel file")
    parser.add_argument("--metrics", default=None, type=str, help="trec eval metrics")
    parser.add_argument("--topK", type=int, default=100, help="retrieve top-k document")
    parser.add_argument("--trec_eval",default='/gypsum/work1/allan/zhiqihuang/genmatch/bin/trec_eval', type=str, help="trec eval file")
    parser.add_argument('--save_vectors', default=False, action='store_true', help="save encoding results")
    parser.add_argument("--langs",default=None,type=str,nargs='+',required=True,help="languages")
    parser.add_argument("--use_pooler", default=False, action='store_true', help="use pooler output")
    parser.add_argument('--normalize', default=False, action='store_true', help="normalize embeddings")
    parser.add_argument("--temperature", default=1.0, type=float, help="temperature for softmax")
    parser.add_argument("--aggregation", default='cls',type=str, required=False, help="dense aggregation strategy")
    parser.add_argument('--add_prefix', default=False, action='store_true', help="add input type prefix for mE5 model")
    parser.add_argument('--fp16', default=False, action='store_true', help="mixed precision training")
    return parser