From 23938015c8fe503aa9b5e021d73ceb0f6819516a Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Sat, 15 Apr 2023 16:30:55 -0400
Subject: [PATCH] prepare datasets only flag

---
 scripts/finetune.py | 5 +++++
 1 file changed, 5 insertions(+)

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"]