normalize pretraining_dataset configuration

axolotl-ai-cloud · Jan 20, 2024 · a787d57 · a787d57
1 parent 585c11d
commit a787d57
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/axolotl/utils/config.py b/src/axolotl/utils/config.py
@@ -155,6 +155,9 @@ def normalize_config(cfg):
     if isinstance(cfg.learning_rate, str):
         cfg.learning_rate = float(cfg.learning_rate)
 
+    if isinstance(cfg.pretraining_dataset, dict):
+        cfg.pretraining_dataset = [cfg.pretraining_dataset]
+
     log_gpu_memory_usage(LOG, "baseline", cfg.device)
 
 

diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py
@@ -71,10 +71,7 @@ def prepare_dataset(cfg, tokenizer):
     else:
         path = cfg.pretraining_dataset
         name = None
-        if isinstance(cfg.pretraining_dataset, dict):
-            path = cfg.pretraining_dataset["path"]
-            name = cfg.pretraining_dataset["name"]
-        elif isinstance(cfg.pretraining_dataset, list) and isinstance(
+        if isinstance(cfg.pretraining_dataset, list) and isinstance(
             cfg.pretraining_dataset[0], dict
         ):
             path = cfg.pretraining_dataset[0]["path"]