Update main.py

main.py

@@ -107,6 +107,7 @@ def _train(config, logger, tokenizer,
       and utils.fsspec_exists(
         config.checkpointing.resume_ckpt_path)):
     ckpt_path = config.checkpointing.resume_ckpt_path
+    print(f"CKPT PATH: {ckpt_path}")
   else:
     ckpt_path = None
 
@@ -120,7 +121,6 @@ def _train(config, logger, tokenizer,
   #   config, tokenizer)
   train_dataset = load_from_disk('/home/tc415/discrete-diffusion-guidance/dataset/3000_400k/train')
   val_dataset = load_from_disk('/home/tc415/discrete-diffusion-guidance/dataset/3000_400k/val')
-  test_dataset = load_from_disk('/home/tc415/discrete-diffusion-guidance/dataset/3000_400k/test')
 
   data_module = dataloader.CustomDataModule(train_dataset, val_dataset, test_dataset, tokenizer, config, batch_size=config.loader.batch_size)
   train_ds = data_module.train_dataloader()
@@ -236,6 +236,32 @@ def _ppl_eval(config, tokenizer):
   ppl = eval_utils.compute_ppl(pretrained, valid_ds)
   print(f"PPL: {ppl:0.3f}")
 
+def _test(config, logger, tokenizer):
+
+    test_dataset = load_from_disk('/home/tc415/discrete-diffusion-guidance/dataset/3000_400k/test')
+    data_module = dataloader.CustomDataModule(None, None, test_dataset=test_dataset, tokenizer=tokenizer, config=config, batch_size=config.loader.batch_size)
+    test_ds = data_module.test_dataloader()
+
+    model = diffusion.Diffusion.load_from_checkpoint(config.eval.checkpoint_path, tokenizer=tokenizer, config=config, logger=False)
+    model.eval()
+
+    # Create a test trainer (without training)
+    trainer = hydra.utils.instantiate(
+        config.trainer,
+        default_root_dir=os.getcwd(),
+        # logger=wandb_logger,
+        strategy=hydra.utils.instantiate(config.strategy),
+        callbacks=[]  # No need for callbacks during testing
+    )
+
+    # Test the model
+    results = trainer.test(model, test_ds)
+
+    # Log or print test results
+    print(f"Test results: {results}")
+
+    return results
+
 
 @hydra.main(version_base=None, config_path='configs',
             config_name='config')
@@ -254,6 +280,8 @@ def main(config):
   elif 'train' in config.mode:
     _train(config, logger, tokenizer,
            train_classifier='classifier' in config.mode)
+  elif 'test' in config.mode:
+    _test(config, logger, tokenizer)
   else:
     raise NotImplementedError(f"Mode {config.mode} not implemented.")