()
| 11 | |
| 12 | |
| 13 | def parse_args(): |
| 14 | parser = argparse.ArgumentParser( |
| 15 | description="Evaluate a RecurrentDepthTransformer model across a range of recurrence depths.") |
| 16 | |
| 17 | parser.add_argument('--data_dir', type=str, default='data/multi_hop') |
| 18 | parser.add_argument('--test_file', type=str, default='test.json') |
| 19 | parser.add_argument('--checkpoint_dir', type=str, |
| 20 | default='checkpoints/multi_hop/r_dyn/') |
| 21 | parser.add_argument('--model_name', type=str, default='checkpoint_epoch_5388.pt') |
| 22 | parser.add_argument('--output_file', type=str, |
| 23 | default='outputs/multi_hop/r_dyn.json', |
| 24 | help='Path to save the output JSON file with results.') |
| 25 | parser.add_argument('--input_injection', action='store_true', |
| 26 | help='Enable input injection (adding input embeddings at the start of each recurrence).') |
| 27 | parser.add_argument('--pred_pos', type=str, choices=['inp_len', 'last_token'], default="last_token", |
| 28 | help='Prediction position.') |
| 29 | |
| 30 | parser.add_argument( |
| 31 | "--recurrence_range", |
| 32 | type=int, |
| 33 | nargs="+", |
| 34 | # default=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], |
| 35 | default=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, |
| 36 | 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40], |
| 37 | help="List of recurrence depths" |
| 38 | ) |
| 39 | |
| 40 | parser.add_argument('--batch_size', type=int, default=1024) |
| 41 | parser.add_argument('--max_hop', type=int, default=40) |
| 42 | parser.add_argument('--max_len', type=int, default=50) |
| 43 | parser.add_argument('--seed', type=int, default=42) |
| 44 | |
| 45 | parser.add_argument('--d_model', type=int, default=768) |
| 46 | parser.add_argument('--num_recurrent_layers', type=int, default=4) |
| 47 | parser.add_argument('--num_heads', type=int, default=12) |
| 48 | parser.add_argument('--positional_embedding_type', type=str, default='none') |
| 49 | |
| 50 | parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') |
| 51 | |
| 52 | return parser.parse_args() |
| 53 | |
| 54 | |
| 55 | if __name__ == '__main__': |
no outgoing calls
no test coverage detected