new hf_use_auth_token setting so login to hf isn't required

README.md

@@ -207,6 +207,9 @@ datasets:
 dataset_prepared_path: data/last_run_prepared
 # push prepared dataset to hub
 push_dataset_to_hub: # repo path
+# whether to use hf `use_auth_token` for loading datasets. Useful for fetching private datasets
+# required to be true when used in combination with `push_dataset_to_hub`
+hf_use_auth_token: # boolean
 # How much of the dataset to set aside as evaluation. 1 = 100%, 0.50 = 50%, etc
 val_set_size: 0.04
 # Num shards for whole dataset

src/axolotl/utils/data.py

@@ -61,10 +61,11 @@ def load_tokenized_prepared_datasets(
         else Path(default_dataset_prepared_path) / ds_hash
     )
     dataset = None
+    use_auth_token = cfg.hf_use_auth_token
     try:
         if cfg.push_dataset_to_hub:
             dataset = load_dataset(
-                f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True
+                f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=use_auth_token
             )
             dataset = dataset["train"]
     except:
@@ -84,7 +85,7 @@ def load_tokenized_prepared_datasets(
             ds: Union[Dataset, DatasetDict] = None
             ds_from_hub = False
             try:
-                load_dataset(d.path, streaming=True, use_auth_token=True)
+                load_dataset(d.path, streaming=True, use_auth_token=use_auth_token)
                 ds_from_hub = True
             except FileNotFoundError:
                 pass
@@ -100,10 +101,10 @@ def load_tokenized_prepared_datasets(
                         d.path,
                         streaming=False,
                         data_files=d.data_files,
-                        use_auth_token=True,
+                        use_auth_token=use_auth_token,
                     )
                 else:
-                    ds: Dataset = load_dataset(d.path, streaming=False, use_auth_token=True)
+                    ds: Dataset = load_dataset(d.path, streaming=False, use_auth_token=use_auth_token)
             else:
                 fp = hf_hub_download(
                     repo_id=d.path, repo_type="dataset", filename=d.data_files
@@ -274,13 +275,14 @@ def load_prepare_datasets(
         )
 
         dataset = None
+        use_auth_token = cfg.hf_use_auth_token
         try:
             if cfg.push_dataset_to_hub:
                 logging.info(
                     f"Checking for packed prepared dataset from hub... {cfg.push_dataset_to_hub}/{ds_hash}"
                 )
                 dataset = load_dataset(
-                    f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=True
+                    f"{cfg.push_dataset_to_hub}/{ds_hash}", use_auth_token=use_auth_token
                 )
                 dataset = dataset["train"]
         except:

src/axolotl/utils/validation.py

@@ -37,6 +37,9 @@ def validate_config(cfg):
             "`trust_remote_code` is set to true. Please make sure that you reviewed the remote code/model."
         )
 
+    if cfg.push_dataset_to_hub and cfg.hf_use_auth_token is not True:
+        raise ValueError("Require cfg.hf_use_auth_token to be True for push_dataset_to_hub")
+
     # TODO
     # MPT 7b
     # https://github.com/facebookresearch/bitsandbytes/issues/25

tests/test_validation.py

@@ -93,3 +93,29 @@ class ValidationTest(unittest.TestCase):
 
         with pytest.raises(ValueError, match=r".*4bit.*"):
             validate_config(cfg)
+
+    def test_hf_use_auth_token(self):
+        base_cfg = DictDefault(
+            {
+                "push_dataset_to_hub": None,
+                "hf_use_auth_token": None,
+            }
+        )
+
+        cfg = base_cfg | DictDefault(
+            {
+                "push_dataset_to_hub": "namespace/repo",
+            }
+        )
+
+        with pytest.raises(ValueError, match=r".*hf_use_auth_token.*"):
+            validate_config(cfg)
+
+        cfg = base_cfg | DictDefault(
+            {
+                "push_dataset_to_hub": "namespace/repo",
+                "hf_use_auth_token": True,
+            }
+        )
+        validate_config(cfg)
+