diff --git a/scripts/finetune.py b/scripts/finetune.py
index 77f1b3c78..4f2a9a255 100644
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -306,6 +306,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
 
 def train(
     config: Path = Path("configs/"),
+    prepare_ds_only: bool = False,
     **kwargs,
 ):
     if Path(config).is_dir():
@@ -396,6 +397,10 @@ def train(
             else:
                 dataset.save_to_disk(DEFAULT_DATASET_PREPARED_PATH)
 
+        if prepare_ds_only:
+            logger.info("Finished preparing dataset. Exiting...")
+            return
+
     train_dataset = dataset["train"]
     eval_dataset = dataset["test"]