Merge pull request #573 from microsoft/daden/bertsumext

Daden/bertsumext improvement of bertsum extractive summarization

Author: saidbleik

examples/text_summarization/extractive_summarization_cnndm_distributed_train.py

@@ -14,7 +14,7 @@
     sys.path.insert(0, nlp_path)
 
 
-from utils_nlp.dataset.cnndm import CNNDMSummarizationDataset
+from utils_nlp.dataset.cnndm import CNNDMBertSumProcessedData, CNNDMSummarizationDataset
 from utils_nlp.models.transformers.extractive_summarization import (
     ExtractiveSummarizer,
     ExtSumProcessedData,
@@ -47,6 +47,12 @@
 parser.add_argument("--encoder", type=str.lower, default='transformer', 
                     choices=['baseline', 'classifier', 'transformer', 'rnn'],
                     help="Encoder types in the extractive summarizer.")
+parser.add_argument(
+    "--max_pos_length",
+    type=int,
+    default=512,
+    help="maximum input length in terms of input token numbers in training",
+)
 parser.add_argument("--learning_rate", type=float, default=1e-3,
                     help="Learning rate.")
 parser.add_argument("--batch_size", type=int, default=3000,
@@ -86,8 +92,8 @@ def main():
     os.makedirs(args.cache_dir, exist_ok=True)
 
     ngpus_per_node = torch.cuda.device_count()
-
-    summarizer = ExtractiveSummarizer(args.model_name, args.encoder, args.cache_dir)
+    processor = ExtSumProcessor(model_name=args.model_name)
+    summarizer = ExtractiveSummarizer(processor, args.model_name, args.encoder, args.max_pos_length, args.cache_dir)
 
     mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, summarizer, args))
 
@@ -109,9 +115,17 @@ def main_worker(local_rank, ngpus_per_node, summarizer, args):
         rank=rank,
     )
 
-    train_dataset, test_dataset = ExtSumProcessedData().splits(root=args.data_dir)
+    if local_rank not in [-1, 0]:
+         torch.distributed.barrier()
+    
+    download_path = CNNDMBertSumProcessedData.download(local_path=args.data_dir)
+    train_dataset, test_dataset = ExtSumProcessedData().splits(root=args.data_dir, train_iterable=True)
+    
+    if local_rank in [-1, 0]:
+        torch.distributed.barrier()
+    
     # total number of steps for training
-    MAX_STEPS = 1e3
+    MAX_STEPS = 1e2
     # number of steps for warm up
     WARMUP_STEPS = 5e2
     if args.quick_run.lower() == "false":
@@ -137,17 +151,18 @@ def main_worker(local_rank, ngpus_per_node, summarizer, args):
         verbose=True,
         report_every=REPORT_EVERY,
         clip_grad_norm=False,
-        local_rank=rank,
+        local_rank=local_rank,
         save_every=save_every,
-        world_size=world_size
+        world_size=world_size,
+        rank=rank,
+        use_preprocessed_data=True
     )
 
     end = time.time()
     print("rank {0}, duration {1:.6f}s".format(rank, end - start))
-    if rank in [-1, 0]:
-        summarizer.save_model(os.path.join(args.output_dir, args.model_filename))
+    torch.distributed.barrier()
+    if local_rank in [-1, 0]:
         prediction = summarizer.predict(test_dataset, num_gpus=ngpus_per_node, batch_size=128)
-
         def _write_list_to_file(list_items, filename):
             with open(filename, "w") as filehandle:
                 # for cnt, line in enumerate(filehandle):