in notebooks/text-classification/scripts/train.py [0:0]
def parse_args():
"""Parse the arguments."""
parser = argparse.ArgumentParser()
# add model id and dataset path argument
parser.add_argument(
"--model_id",
type=str,
default="bert-base-uncased",
help="Model id to use for training.",
)
parser.add_argument(
"--output_dir",
type=str,
default=None,
help="Hugging Face Repository id for uploading models",
)
parser.add_argument(
"--repository_id",
type=str,
default=None,
help="Hugging Face Repository id for uploading models",
)
parser.add_argument("--epochs", type=int, default=3, help="Number of epochs to train for.")
parser.add_argument("--max_steps", type=int, default=-1, help="Number of steps to train for.")
parser.add_argument(
"--per_device_train_batch_size",
type=int,
default=8,
help="Batch size to use for training.",
)
parser.add_argument(
"--per_device_eval_batch_size",
type=int,
default=8,
help="Batch size to use for validation.",
)
parser.add_argument(
"--train_max_length",
type=int,
default=128,
help="Maximum length of tokens to be used for training.",
)
parser.add_argument(
"--learning_rate",
type=float,
default=5e-5,
help="Learning rate to use for training.",
)
parser.add_argument("--seed", type=int, default=42, help="Seed to use for training.")
parser.add_argument(
"--hf_token",
type=str,
default=HfFolder.get_token(),
help="Token to use for uploading models to Hugging Face Hub.",
)
args = parser.parse_known_args()
return args