| 38 | |
| 39 | |
| 40 | def parse_args(): |
| 41 | parser = argparse.ArgumentParser( |
| 42 | description= |
| 43 | "Finetune a transformers model on a causal language modeling task") |
| 44 | parser.add_argument('--data_path', |
| 45 | type=str, |
| 46 | required=True, |
| 47 | help='Path to the training dataset.') |
| 48 | parser.add_argument( |
| 49 | "--model_name_or_path", |
| 50 | type=str, |
| 51 | help= |
| 52 | "Path to pretrained model or model identifier from huggingface.co/models.", |
| 53 | required=True, |
| 54 | ) |
| 55 | parser.add_argument( |
| 56 | "--per_device_train_batch_size", |
| 57 | type=int, |
| 58 | default=1, |
| 59 | help="Batch size (per device) for the training dataloader.", |
| 60 | ) |
| 61 | parser.add_argument( |
| 62 | "--max_seq_len", |
| 63 | type=int, |
| 64 | default=512, |
| 65 | help="The maximum sequence length.", |
| 66 | ) |
| 67 | parser.add_argument( |
| 68 | "--learning_rate", |
| 69 | type=float, |
| 70 | default=1e-5, |
| 71 | help= |
| 72 | "Initial learning rate (after the potential warmup period) to use.", |
| 73 | ) |
| 74 | parser.add_argument("--weight_decay", |
| 75 | type=float, |
| 76 | default=0., |
| 77 | help="Weight decay to use.") |
| 78 | parser.add_argument("--num_train_epochs", |
| 79 | type=int, |
| 80 | default=1, |
| 81 | help="Total number of training epochs to perform.") |
| 82 | parser.add_argument( |
| 83 | "--gradient_accumulation_steps", |
| 84 | type=int, |
| 85 | default=1, |
| 86 | help= |
| 87 | "Number of updates steps to accumulate before performing a backward/update pass.", |
| 88 | ) |
| 89 | parser.add_argument( |
| 90 | "--lr_scheduler_type", |
| 91 | type=SchedulerType, |
| 92 | default="cosine", |
| 93 | help="The scheduler type to use.", |
| 94 | choices=[ |
| 95 | "linear", "cosine", "cosine_with_restarts", "polynomial", |
| 96 | "constant", "constant_with_warmup" |
| 97 | ], |