import argparse

def parse_args():
    parser = argparse.ArgumentParser(description="AIO pipeline")
    parser.add_argument(
        "--model_name",
        type=str,
        default="vicuna",
        help="Name of the white-box LLM",
    )
    parser.add_argument(
        "--task",
        type=str,
        default=None,
        help="The name of the dataset to use (via the datasets library).",
    )
    parser.add_argument(
        "--n_prompt_tokens",
        type=int,
        default=5,
        help="The number of prompt tokens."
    )
    parser.add_argument(
        "--HF_cache_dir",
        type=str,
        default="PATH",
        help="Your vicuna directory"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Set the seed."    
    )
    parser.add_argument(
        "--nu",
        type=float,
        default=0.1,
        help="Set the parameter nu."    
    )
    parser.add_argument(
        "--lamdba",
        type=float,
        default=0.1,
        help="Set the lamdba parameter."    
    )
    parser.add_argument(
        "--n_init",
        type=int,
        default=40,
        help="Set the number of initialization points."    
    )
    parser.add_argument(
        "--n_domain",
        type=int,
        default=10000,
        help="Set the number of domain."    
    )
    parser.add_argument(
        "--total_iter",
        type=int,
        default=165,
        help="Set the number of total queries."    
    )
    parser.add_argument(
        "--local_training_iter",
        type=int,
        default=30,
        help="Set the number of total queries."    
    )
    parser.add_argument(
        "--random_proj",
        type=str,
        default='uniform',
        help="Set the projection method."    
    )
    parser.add_argument(
        "--intrinsic_dim",
        type=int,
        default=100,
        help="Set the number of intrinsic dim."    
    )
    parser.add_argument(
        "--n_eval",
        type=int,
        default=1000,
        help="Set the number of domains to be evaluated at each ucb iteration."    
    )
    parser.add_argument(
        "--name",
        type=str,
        default="iter165_gpt-0301",
        help="Set the name of the experiments."    
    )
    parser.add_argument(
        "--gpt",
        type=str,
        default="gpt-3.5-turbo",
        help="Which version of gpt to use."    
    )
    parser.add_argument(
        "--init_scale",
        type=float,
        default=1,
        help="Which scale to use."    
    )
    parser.add_argument(
        "--pooling",
        type=str,
        default="last",
        help="Which pooling method to use."    
    )
    
    # ------------------------------------------------------
    
    parser.add_argument(
        "--white_box_LLM_eval_HF_cache_dir",
        type=str,
        default=None,
        help="Whether to evaluate performance with white-box LLM."    
    )

    parser.add_argument(
        "--few_shot_sample_num",
        type=int,
        default=20,
        help="Number of few-shot samples."    
    )
    
    parser.add_argument(
        "--evaluation_sample_num",
        type=int,
        default=100,
        help="Number of validation and testing samples."    
    )

    ###
    parser.add_argument(
        "--gen_demo_num_comparison",
        type=int,
        default=0,
        help="Number of samples included in the instruction."    
    )
    
    ###
    parser.add_argument(
        "--training_module_name",
        nargs="*",
        type=str,
        default=None,
        help="Which modules to fine-tune."  
    )

    ###
    parser.add_argument(
        "--max_seq_length",
        type=int,
        default=8000,
        help="Number of validation and testing samples."    
    )

    ###
    parser.add_argument(
        "--num_mezo_steps",
        type=int,
        default=-1,
        help="Number of MeZO steps."    
    )

    ###
    parser.add_argument(
        "--training_profile_index",
        type=int,
        default=1,
        help="Index of AIO training profile."    
    )

    ###
    parser.add_argument(
        "--use_baseline_method_name",
        type=str,
        default=None,
        help="Whether to use baseline, and its name"    
    )

    ###
    parser.add_argument('--zero_shot_evaluation_flag', action='store_true')

    # ------------------------------------------------------

    args = parser.parse_args()
    return args