()
| 12 | |
| 13 | |
| 14 | def parse_args(): |
| 15 | parser = argparse.ArgumentParser( |
| 16 | description="Evaluate a RecurrentDepthTransformer model with adaptive recurrence based on Cosine Similarity.") |
| 17 | |
| 18 | parser.add_argument('--data_dir', type=str, default='data/multi_hop') |
| 19 | parser.add_argument('--test_file', type=str, default='test.json') |
| 20 | parser.add_argument('--checkpoint_dir', type=str, |
| 21 | default='checkpoints/multi_hop/r_dyn/') |
| 22 | parser.add_argument('--model_name', type=str, default='checkpoint_epoch_5388.pt') |
| 23 | parser.add_argument('--output_file', type=str, |
| 24 | default='outputs/multi_hop/r_dyn_adaptive.json', |
| 25 | help='Path to save the output JSON file with results.') |
| 26 | parser.add_argument('--pred_pos', type=str, choices=['inp_len', 'last_token'], default="last_token", |
| 27 | help='Prediction position.') |
| 28 | parser.add_argument('--input_injection', action='store_true', |
| 29 | help='Enable input injection (adding input embeddings at the start of each recurrence).') |
| 30 | |
| 31 | parser.add_argument('--batch_size', type=int, default=1024) |
| 32 | parser.add_argument('--max_hop', type=int, default=40) |
| 33 | parser.add_argument('--max_len', type=int, default=50) |
| 34 | parser.add_argument('--seed', type=int, default=42) |
| 35 | |
| 36 | parser.add_argument('--max_recurrence', type=int, default=16, |
| 37 | help='Maximum number of recurrent iterations to allow.') |
| 38 | |
| 39 | parser.add_argument('--d_model', type=int, default=768) |
| 40 | parser.add_argument('--num_recurrent_layers', type=int, default=4) |
| 41 | parser.add_argument('--num_heads', type=int, default=12) |
| 42 | parser.add_argument('--positional_embedding_type', type=str, default='none') |
| 43 | |
| 44 | parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu') |
| 45 | |
| 46 | return parser.parse_args() |
| 47 | |
| 48 | |
| 49 | if __name__ == '__main__': |
no outgoing calls
no test coverage detected