Fix providing glue in seq2seq eval (#4843)

* Fix providing glue in seq2seq eval

Signed-off-by: MaximumEntropy <[email protected]>

* Fix

Signed-off-by: MaximumEntropy <[email protected]>

* Style

Signed-off-by: MaximumEntropy <[email protected]>

Signed-off-by: MaximumEntropy <[email protected]>

Author: MaximumEntropy

examples/nlp/language_modeling/megatron_t5_seq2seq_eval.py

@@ -20,7 +20,13 @@
 
 from nemo.collections.nlp.models.language_modeling.megatron_finetune_model import MegatronT5FinetuneModel
 from nemo.collections.nlp.models.language_modeling.megatron_glue_model import MegatronT5GLUEModel
-from nemo.collections.nlp.parts.nlp_overrides import GradScaler, MegatronHalfPrecisionPlugin, NLPDDPStrategy
+from nemo.collections.nlp.models.language_modeling.megatron_t0_model import MegatronT0Model
+from nemo.collections.nlp.parts.nlp_overrides import (
+    GradScaler,
+    MegatronHalfPrecisionPlugin,
+    NLPDDPStrategy,
+    NLPSaveRestoreConnector,
+)
 from nemo.core.config import hydra_runner
 from nemo.utils import logging
 from nemo.utils.exp_manager import StatelessTimer, exp_manager
@@ -71,15 +77,14 @@ def main(cfg) -> None:
     # NOTE: Only data can be overriden here since this the file being restored here should already correspond to a GLUE/XNLI finetuned model.
     OmegaConf.set_struct(t5_cfg, True)
     with open_dict(t5_cfg):
-        t5_cfg.masked_softmax_fusion = False
         t5_cfg.precision = cfg.trainer.precision
         # Overwrite data configs
-        if cfg.model.data.validation_ds.src_file_name is not None:
+        if cfg.model.data.validation_ds.get('src_file_name', None) is not None:
             logging.info(
                 'Found validation_ds.src_file_name in the config file. Overriding the finetuned model config file with the values from the new config file.'
             )
             t5_cfg.data.validation_ds.src_file_name = cfg.model.data.validation_ds.src_file_name
-        if cfg.model.data.validation_ds.tgt_file_name is not None:
+        if cfg.model.data.validation_ds.get('tgt_file_name', None) is not None:
             logging.info(
                 'Found validation_ds.tgt_file_name in the config file. Overriding the finetuned model config file with the values from the new config file.'
             )
@@ -88,9 +93,28 @@ def main(cfg) -> None:
         t5_cfg.data.validation_ds.micro_batch_size = cfg.model.data.validation_ds.micro_batch_size
         t5_cfg.data.validation_ds.global_batch_size = cfg.model.data.validation_ds.global_batch_size
 
-        model = MegatronT5FinetuneModel.restore_from(
-            restore_path=cfg.model.restore_from_path, trainer=trainer, override_config_path=t5_cfg
-        )
+        if hasattr(cfg.model.data.validation_ds, 'task_name'):
+            model = MegatronT5GLUEModel.restore_from(
+                restore_path=cfg.model.restore_from_path,
+                trainer=trainer,
+                override_config_path=t5_cfg,
+                save_restore_connector=NLPSaveRestoreConnector(),
+            )
+        elif hasattr(cfg.model.data.validation_ds, 'file_names'):
+            model = MegatronT0Model.restore_from(
+                restore_path=cfg.model.restore_from_path,
+                trainer=trainer,
+                override_config_path=t5_cfg,
+                save_restore_connector=NLPSaveRestoreConnector(),
+            )
+        else:
+            model = MegatronT5FinetuneModel.restore_from(
+                restore_path=cfg.model.restore_from_path,
+                trainer=trainer,
+                override_config_path=t5_cfg,
+                save_restore_connector=NLPSaveRestoreConnector(),
+            )
+
     model.freeze()
     trainer.validate(model)
     if hasattr(cfg.model.data, 'test_ds'):