def parse_args()

in optimum/commands/neuron/cache.py [0:0]

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

    def parse_args(parser: "ArgumentParser"):
        parser.add_argument("-m", "--model", type=str, required=True, help="The name of model or path of the model.")
        parser.add_argument("--task", type=str, required=True, help="The task for which the model should be compiled.")

        # Shapes
        parser.add_argument(
            "--train_batch_size",
            type=int,
            required=True,
            help="The batch size to use during the model compilation for training.",
        )

        parser.add_argument(
            "--eval_batch_size",
            type=int,
            default=None,
            help="The batch size to use during model compilation for evaluation.",
        )

        sequence_length_group = parser.add_mutually_exclusive_group()

        sequence_length_group.add_argument(
            "--sequence_length", type=int, help="The sequence length of the model during compilation."
        )

        seq2seq_sequence_length_group = sequence_length_group.add_argument_group()
        seq2seq_sequence_length_group.add_argument(
            "--encoder_sequence_length",
            type=int,
            help="The sequence length of the encoder part of the model during compilation.",
        )
        seq2seq_sequence_length_group.add_argument(
            "--decoder_sequence_length",
            type=int,
            help="The sequence length of the decoder part of the model during compilation.",
        )

        parser.add_argument(
            "--gradient_accumulation_steps", type=int, default=1, help="The number of gradient accumulation steps.."
        )

        parser.add_argument(
            "--precision",
            choices=["fp", "bf16"],
            type=str,
            required=True,
            help="The precision to use during the model compilation.",
        )
        parser.add_argument(
            "--num_cores",
            choices=list(range(1, 33)),
            type=int,
            required=True,
            help="The number of neuron cores to use during compilation.",
        )
        parser.add_argument(
            "--example_dir", type=str, default=None, help="Path to where the example scripts are stored."
        )
        parser.add_argument(
            "--max_steps", type=int, default=10, help="The maximum number of steps to run compilation for."
        )