pass the chat_template value correctly

truefoundry · Jun 12, 2024 · 872152d · 872152d
1 parent 063d828
commit 872152d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/data_utils.py b/data_utils.py
@@ -88,9 +88,11 @@ def dataset_uri_to_axolotl_datasources(
         datasources = []
         if os.path.isdir(uri):
             for filepath in find_all_jsonl_files(uri):
-                datasources.append(_make_dataset_file_source(path=filepath, dataset_type=dataset_type))
+                datasources.append(
+                    _make_dataset_file_source(path=filepath, dataset_type=dataset_type, chat_template=chat_template)
+                )
         else:
-            datasources = [_make_dataset_file_source(path=uri, dataset_type=dataset_type)]
+            datasources = [_make_dataset_file_source(path=uri, dataset_type=dataset_type, chat_template=chat_template)]
         return datasources
     else:
         raise ValueError("Unsupported data uri or path does not exist: {uri}")

diff --git a/train.py b/train.py
@@ -176,9 +176,10 @@ def make_axolotl_config(config_base, kwargs, timestamp=None):
         set_cfg_option_if_auto(cfg, "unsloth_lora_qkv", use_unsloth)
         set_cfg_option_if_auto(cfg, "unsloth_lora_o", use_unsloth)
 
-        model_type = getattr(model_hf_config, "model_type", None)
-        chat_template = MODEL_TYPE_TO_CHAT_TEMPLATE.get(model_type, "chatml")
-        set_cfg_option_if_auto(cfg, "chat_template", chat_template)
+        if cfg.chat_template == "auto":
+            model_type = getattr(model_hf_config, "model_type", None)
+            chat_template = MODEL_TYPE_TO_CHAT_TEMPLATE.get(model_type, "chatml")
+            set_cfg_option_if_auto(cfg, "chat_template", chat_template)
 
         if cfg.datasets == "auto":
             if not cfg.train_data_uri:
@@ -187,7 +188,7 @@ def make_axolotl_config(config_base, kwargs, timestamp=None):
                 uri=cfg.train_data_uri,
                 download_dir=cfg.data_dir,
                 dataset_type=cfg.dataset_type,
-                chat_template=chat_template,
+                chat_template=cfg.chat_template,
             )
         if cfg.test_datasets == "auto":
             if cfg.val_data_uri and str(cfg.val_data_uri).lower() != "na":