| 39 | |
| 40 | |
| 41 | def add_code_generation_args(parser): |
| 42 | group = parser.add_argument_group(title="code generation") |
| 43 | group.add_argument( |
| 44 | "--num-layers", |
| 45 | type=int, |
| 46 | default=39, |
| 47 | ) |
| 48 | group.add_argument( |
| 49 | "--hidden-size", |
| 50 | type=int, |
| 51 | default=5120, |
| 52 | ) |
| 53 | group.add_argument( |
| 54 | "--num-attention-heads", |
| 55 | type=int, |
| 56 | default=40, |
| 57 | ) |
| 58 | group.add_argument( |
| 59 | "--padded-vocab-size", |
| 60 | type=int, |
| 61 | default=52224, |
| 62 | ) |
| 63 | group.add_argument( |
| 64 | "--max-position-embeddings", |
| 65 | type=int, |
| 66 | default=2048, |
| 67 | ) |
| 68 | group.add_argument( |
| 69 | "--temperature", |
| 70 | type=float, |
| 71 | default=1.0, |
| 72 | help="Sampling temperature.", |
| 73 | ) |
| 74 | group.add_argument( |
| 75 | "--greedy", |
| 76 | action="store_true", |
| 77 | default=False, |
| 78 | help="Use greedy sampling.", |
| 79 | ) |
| 80 | group.add_argument( |
| 81 | "--top-p", |
| 82 | type=float, |
| 83 | default=0.0, |
| 84 | help="Top p sampling.", |
| 85 | ) |
| 86 | group.add_argument( |
| 87 | "--top-k", |
| 88 | type=int, |
| 89 | default=0, |
| 90 | help="Top k sampling.", |
| 91 | ) |
| 92 | group.add_argument( |
| 93 | "--out-seq-length", |
| 94 | type=int, |
| 95 | default=2048, |
| 96 | help="Size of the output generated text.", |
| 97 | ) |
| 98 | group.add_argument( |