(parser)
| 25 | |
| 26 | |
| 27 | def add_code_generation_args(parser): |
| 28 | group = parser.add_argument_group(title="code generation") |
| 29 | group.add_argument( |
| 30 | "--num-layers", |
| 31 | type=int, |
| 32 | default=39, |
| 33 | ) |
| 34 | group.add_argument( |
| 35 | "--hidden-size", |
| 36 | type=int, |
| 37 | default=5120, |
| 38 | ) |
| 39 | group.add_argument( |
| 40 | "--num-attention-heads", |
| 41 | type=int, |
| 42 | default=40, |
| 43 | ) |
| 44 | group.add_argument( |
| 45 | "--padded-vocab-size", |
| 46 | type=int, |
| 47 | default=52224, |
| 48 | ) |
| 49 | group.add_argument( |
| 50 | "--max-position-embeddings", |
| 51 | type=int, |
| 52 | default=2048, |
| 53 | ) |
| 54 | group.add_argument( |
| 55 | "--tokenizer-path", |
| 56 | type=str, |
| 57 | default="./tokenizer", |
| 58 | ) |
| 59 | group.add_argument( |
| 60 | "--example-path", |
| 61 | type=str, |
| 62 | default="./", |
| 63 | ) |
| 64 | group.add_argument( |
| 65 | "--load", |
| 66 | type=str, |
| 67 | ) |
| 68 | group.add_argument( |
| 69 | "--state-dict-path", |
| 70 | type=str, |
| 71 | ) |
| 72 | group.add_argument( |
| 73 | "--micro-batch-size", |
| 74 | type=int, |
| 75 | default=1, |
| 76 | ) |
| 77 | group.add_argument( |
| 78 | "--quantize", |
| 79 | action="store_true", |
| 80 | ) |
| 81 | |
| 82 | return parser |
| 83 | |
| 84 |