def add_training_args()

in elq/common/params.py [0:0]

• 143 lines of code
• 1 McCabe index (conditional complexity)

    def add_training_args(self, args=None):
        """
        Add model training args.
        """
        parser = self.add_argument_group("Model Training Arguments")
        parser.add_argument(
            "--evaluate", action="store_true", help="Whether to run evaluation."
        )
        parser.add_argument(
            "--output_eval_file",
            default=None,
            type=str,
            help="The txt file where the the evaluation results will be written.",
        )
        parser.add_argument(
            "--train_batch_size", default=8, type=int, 
            help="Total batch size for training."
        )
        parser.add_argument(
            "--eval_batch_size", default=8, type=int,
            help="Total batch size for evaluation.",
        )
        parser.add_argument("--max_grad_norm", default=1.0, type=float)
        parser.add_argument(
            "--learning_rate",
            default=3e-5,
            type=float,
            help="The initial learning rate for Adam.",
        )
        parser.add_argument(
            "--num_train_epochs",
            default=1,
            type=int,
            help="Number of training epochs.",
        )
        parser.add_argument(
            "--print_interval", type=int, default=5, 
            help="Interval of loss printing",
        )
        parser.add_argument(
           "--eval_interval",
            type=int,
            default=40,
            help="Interval for evaluation during training",
        )
        parser.add_argument(
            "--save_interval", type=int, default=1, 
            help="Interval for model saving"
        )
        parser.add_argument(
            "--warmup_proportion",
            default=0.1,
            type=float,
            help="Proportion of training to perform linear learning rate warmup for. "
            "E.g., 0.1 = 10% of training.",
        )
        parser.add_argument(
            "--gradient_accumulation_steps",
            type=int,
            default=1,
            help="Number of updates steps to accumualte before performing a backward/update pass.",
        )
        parser.add_argument(
            "--type_optimization",
            type=str,
            default="all_encoder_layers",
            help="Which type of layers to optimize in BERT",
        )
        parser.add_argument(
            "--shuffle", type=bool, default=False, 
            help="Whether to shuffle train data",
        )
        # TODO DELETE LATER!!!
        parser.add_argument(
            "--start_idx",
            default=None,
            type=int,
        )
        parser.add_argument(
            "--end_idx",
            default=None,
            type=int,
        )
        parser.add_argument(
            "--last_epoch",
            default=0,
            type=int,
            help="Epoch to restore from when pretraining",
        )
        parser.add_argument(
            "--path_to_trainer_state",
            default=None,
            type=str,
            required=False,
            help="The full path to the last checkpoint's training state to load.",
        )
        parser.add_argument(
            '--dont_distribute_train_samples',
            default=False,
            action="store_true",
            help="Don't distribute all training samples across the epochs (go through all samples every epoch)",
        )
        parser.add_argument(
            "--freeze_cand_enc",
            default=False,
            action="store_true",
            help="Freeze the candidate encoder",
        )
        parser.add_argument(
            "--load_cand_enc_only",
            default=False,
            action="store_true",
            help="Only load the candidate encoder from saved model path",
        )
        parser.add_argument(
            "--cand_enc_path",
            default="models/all_entities_large.t7",
            type=str,
            required=False,
            help="Filepath to the saved entity encodings.",
        )
        parser.add_argument(
            "--cand_token_ids_path",
            default="models/entity_token_ids_128.t7",
            type=str,
            required=False,
            help="Filepath to the saved tokenized entity descriptions.",
        )
        parser.add_argument(
            "--index_path",
            default="models/faiss_hnsw_index.pkl",
            type=str,
            required=False,
            help="Filepath to the HNSW index for adversarial training.",
        )
        parser.add_argument(
            "--adversarial_training",
            default=False,
            action="store_true",
            help="Do adversarial training (only takes effect if `freeze_cand_enc` is set)",
        )
        parser.add_argument(
            "--get_losses",
            default=False,
            action="store_true",
            help="Get losses during evaluation",
        )