fix test

Author: daden-ms

tests/unit/test_abstractive_summarization_seq2seq.py

@@ -84,7 +84,9 @@ def s2s_test_data():
 
 @pytest.mark.gpu
 def test_S2SAbstractiveSummarizer(s2s_test_data, tmp):
-    processor = S2SAbsSumProcessor(cache_dir=tmp)
+    cache_dir = tmp
+    model_dir = tmp
+    processor = S2SAbsSumProcessor(cache_dir=cache_dir)
     train_dataset = processor.s2s_dataset_from_json_or_file(
         s2s_test_data["train_ds"], train_mode=True
     )
@@ -95,14 +97,14 @@ def test_S2SAbstractiveSummarizer(s2s_test_data, tmp):
         max_seq_length=MAX_SEQ_LENGTH,
         max_source_seq_length=MAX_SOURCE_SEQ_LENGTH,
         max_target_seq_length=MAX_TARGET_SEQ_LENGTH,
-        cache_dir=tmp,
+        cache_dir=cache_dir,
     )
 
     # test fit and predict
-    abs_summarizer.fit(
+    global_step = abs_summarizer.fit(
         train_dataset,
         per_gpu_batch_size=TRAIN_PER_GPU_BATCH_SIZE,
-        save_model_to_dir=tmp,
+        save_model_to_dir=model_dir,
     )
     abs_summarizer.predict(
         test_dataset,
@@ -112,12 +114,12 @@ def test_S2SAbstractiveSummarizer(s2s_test_data, tmp):
 
     # test load model from local disk
     abs_summarizer_loaded = S2SAbstractiveSummarizer(
-        load_model_from_dir=tmp,
-        model_file_name="model.1.bin",
+        load_model_from_dir=model_dir,
+        model_file_name="model.{}.bin".format(global_step),
         max_seq_length=MAX_SEQ_LENGTH,
         max_source_seq_length=MAX_SOURCE_SEQ_LENGTH,
         max_target_seq_length=MAX_TARGET_SEQ_LENGTH,
-        cache_dir=tmp,
+        cache_dir=cache_dir,
     )
 
     abs_summarizer_loaded.predict(
@@ -130,10 +132,10 @@ def test_S2SAbstractiveSummarizer(s2s_test_data, tmp):
     abs_summarizer.fit(
         train_dataset,
         per_gpu_batch_size=TRAIN_PER_GPU_BATCH_SIZE,
-        save_model_to_dir=tmp,
-        recover_step=1,
-        recover_dir=tmp,
-        max_steps=4,
+        save_model_to_dir=model_dir,
+        recover_step=global_step,
+        recover_dir=model_dir,
+        max_steps=global_step + 3,
     )
 
     abs_summarizer.predict(

utils_nlp/models/transformers/abstractive_summarization_seq2seq.py

@@ -524,7 +524,7 @@ def __init__(
             + self.max_target_seq_length,
         )
         logger.info("Model config for seq2seq: %s", str(config))
-    
+
         self.model = model_class.from_pretrained(
             model_to_load,
             config=config,
@@ -732,11 +732,12 @@ def fit(
         )
 
         if save_model_to_dir is not None and local_rank in [-1, 0]:
-            self.save_model(save_model_to_dir, global_step, fp16)
+            self.save_model(save_model_to_dir, global_step - 1, fp16)
 
         # release GPU memories
         self.model.cpu()
         torch.cuda.empty_cache()
+        return global_step - 1
 
     def predict(
         self,
@@ -896,8 +897,7 @@ def collate_fn(input_batch):
                 is_roberta=is_roberta,
                 no_segment_embedding=no_segment_embedding
             )
-        # print(self._bert_model_name)
-        # print(type(bert_config))
+
         model = BertForSeq2SeqDecoder.from_pretrained(
             self._bert_model_name,
             bert_config,
@@ -955,7 +955,6 @@ def collate_fn(input_batch):
             batch_size=batch_size,
             collate_fn=collate_fn,
         )
-        print(device)
         for batch, buf_id in tqdm(
             test_dataloader, desc="Evaluating", disable=not verbose
         ):