()
| 23 | |
| 24 | |
| 25 | def parse_args(): |
| 26 | parser = argparse.ArgumentParser(description="Train a Recurrent-Depth Transformer model for multi-hop composition.") |
| 27 | |
| 28 | parser.add_argument('--data_dir', type=str, default='data/multi_hop/', |
| 29 | help='Path to dataset directory') |
| 30 | parser.add_argument('--test_file', type=str, default='test.json', |
| 31 | help='Path to dataset directory') |
| 32 | parser.add_argument('--seed', type=int, default=42) |
| 33 | parser.add_argument('--checkpoint_dir', type=str, |
| 34 | default='checkpoints/multi_hop/r_dyn/', |
| 35 | help='Path to save model checkpoints') |
| 36 | parser.add_argument('--log_file', type=str, |
| 37 | default='results/multi_hop/r_dyn.txt', |
| 38 | help='Path to save training logs') |
| 39 | |
| 40 | parser.add_argument('--num_epochs', type=int, default=100001, help='Number of training epochs') |
| 41 | parser.add_argument('--batch_size', type=int, default=128, help='Batch size for training and evaluation') |
| 42 | parser.add_argument('--max_len', type=int, default=50, help='Maximum number of tokens') |
| 43 | parser.add_argument('--max_hop', type=int, default=40, help='Maximum number of hops') |
| 44 | parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') |
| 45 | parser.add_argument('--weight_decay', type=float, default=0.01, help='Weight decay for optimizer') |
| 46 | parser.add_argument('--warmup_steps', type=int, default=2000, |
| 47 | help='Number of warmup steps for learning rate scheduler') |
| 48 | parser.add_argument('--acc_threshold_curriculum', type=float, default=0.95, |
| 49 | help='Threshold beyond which we move to next stage in curriculum training.') |
| 50 | parser.add_argument('--pred_pos', type=str, choices=['inp_len', 'last_token'], default="last_token", |
| 51 | help='Position for output prediction.') |
| 52 | |
| 53 | parser.add_argument('--resume_training', action='store_true', |
| 54 | help='Resume curriculum from a checkpoint and skip completed hop stages.') |
| 55 | |
| 56 | parser.add_argument('--resume_checkpoint', type=str, default=None, |
| 57 | help='Checkpoint filename inside --checkpoint_dir OR an absolute path.') |
| 58 | |
| 59 | parser.add_argument('--hops_generalized', type=int, default=9, |
| 60 | help='Last hop level already completed (e.g., 10). Training resumes at hops_generalized+1.') |
| 61 | |
| 62 | parser.add_argument('--use_lr_decay', action='store_true', |
| 63 | help='lr decay after each stage.') |
| 64 | parser.add_argument('--stage_lr_base', type=float, default=1e-4, |
| 65 | help='LR for the first curriculum stage (2-hop).') |
| 66 | parser.add_argument('--stage_lr_gamma', type=float, default=0.9, |
| 67 | help='Multiply LR by this factor each new stage.') |
| 68 | parser.add_argument('--stage_lr_min', type=float, default=2e-5, |
| 69 | help='Min LR for any curriculum stage.') |
| 70 | parser.add_argument('--max_grad_norm', type=float, default=0.0, |
| 71 | help='Maximum gradient norm for gradient clipping. Set to 0.0 to disable.') |
| 72 | parser.add_argument('--dropout', type=float, default=0.0) |
| 73 | parser.add_argument('--c_scale', type=float, default=0.0) |
| 74 | parser.add_argument('--input_injection', action='store_true', |
| 75 | help='Enable input injection (adding input embeddings at the start of each recurrence).') |
| 76 | parser.add_argument('--force_grok', action='store_true', |
| 77 | help='Force 1000 epochs on 2-hop curriculum stage.') |
| 78 | |
| 79 | parser.add_argument('--train_mode', type=str, choices=['max-autotune', 'regular'], default='max-autotune', |
| 80 | help='Whether training run was optimized.') |
| 81 | |
| 82 | parser.add_argument('--d_model', type=int, default=768, help='Dimension of model embeddings') |
no outgoing calls
no test coverage detected