REALM arguments

hashed_index.py

@@ -203,7 +203,7 @@ def load_ict_checkpoint():
         iteration = int(f.read().strip())
 
     assert iteration > 0
-    checkpoint_name = get_checkpoint_name(args.load, iteration, False)
+    checkpoint_name = get_checkpoint_name(args.ict_load, iteration, False)
     if mpu.get_data_parallel_rank() == 0:
         print('global rank {} is loading checkpoint {}'.format(
             torch.distributed.get_rank(), checkpoint_name))
@@ -221,7 +221,7 @@ def load_ict_checkpoint():
 def get_ict_dataset():
     args = get_args()
     block_dataset = get_indexed_dataset_(args.data_path, 'mmap', True)
-    titles_dataset = get_indexed_dataset_(args.data_path + '-titles', 'mmap', True)
+    titles_dataset = get_indexed_dataset_(args.titles_data_path, 'mmap', True)
 
     kwargs = dict(
         name='full',

megatron/arguments.py

@@ -219,6 +219,8 @@ def _add_checkpointing_args(parser):
                        help='Do not save current rng state.')
     group.add_argument('--load', type=str, default=None,
                        help='Directory containing a model checkpoint.')
+    group.add_argument('--ict-load', type=str, default=None,
+                       help='Directory containing an ICTBertModel checkpoint')
     group.add_argument('--no-load-optim', action='store_true',
                        help='Do not load optimizer when loading checkpoint.')
     group.add_argument('--no-load-rng', action='store_true',
@@ -294,6 +296,8 @@ def _add_data_args(parser):
 
     group.add_argument('--data-path', type=str, default=None,
                        help='Path to combined dataset to split.')
+    group.add_argument('--titles-data-path', type=str, default=None,
+                       help='Path to titles dataset used for ICT')
     group.add_argument('--split', type=str, default='969, 30, 1',
                        help='Comma-separated list of proportions for training,'
                        ' validation, and test split. For example the split '