update pr

Author: rasbt

litgpt/scripts/download.py

@@ -3,6 +3,7 @@
 import os
 from concurrent.futures import ProcessPoolExecutor
 from contextlib import contextmanager
+import importlib.util
 from pathlib import Path
 from typing import List, Optional, Tuple
 
@@ -56,6 +57,8 @@ def download_from_hub(
         return
 
     from huggingface_hub import snapshot_download
+    if importlib.util.find_spec("hf_transfer") is None:
+        print("It is recommended to install hf_transfer for faster checkpoint download speeds: `pip install hf_transfer`")
 
     download_files = ["tokenizer*", "generation_config.json", "config.json"]
     if not tokenizer_only:
@@ -70,22 +73,14 @@ def download_from_hub(
         else:
             raise ValueError(f"Couldn't find weight files for {repo_id}")
 
-    # Get and set env variable to improve download speed
-    user_env_value = os.environ.get("HF_HUB_ENABLE_HF_TRANSFER")
-
-    if user_env_value is None:
-        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-        print("Setting HF_HUB_ENABLE_HF_TRANSFER=1 by default")
-
     import huggingface_hub._snapshot_download as download
     import huggingface_hub.constants as constants
 
-    previous_flag = constants.HF_HUB_ENABLE_HF_TRANSFER # this may be redundant
-    
-    if _HF_TRANSFER_AVAILABLE and not previous_flag:
-         print("Setting HF_HUB_ENABLE_HF_TRANSFER=1")
-         constants.HF_HUB_ENABLE_HF_TRANSFER = True
-         download.HF_HUB_ENABLE_HF_TRANSFER = True
+    previous = constants.HF_HUB_ENABLE_HF_TRANSFER
+    if _HF_TRANSFER_AVAILABLE and not previous:
+        print("Setting HF_HUB_ENABLE_HF_TRANSFER=1")
+        constants.HF_HUB_ENABLE_HF_TRANSFER = True
+        download.HF_HUB_ENABLE_HF_TRANSFER = True
 
     directory = checkpoint_dir / repo_id
     with gated_repo_catcher(repo_id, access_token):
@@ -96,8 +91,8 @@ def download_from_hub(
             token=access_token,
         )
 
-    constants.HF_HUB_ENABLE_HF_TRANSFER = previous_flag
-    download.HF_HUB_ENABLE_HF_TRANSFER = previous_flag
+    constants.HF_HUB_ENABLE_HF_TRANSFER = previous
+    download.HF_HUB_ENABLE_HF_TRANSFER = previous
 
     if convert_checkpoint and not tokenizer_only:
         print("Converting checkpoint files to LitGPT format.")