Enable hf_transfer for faster download

pytorch · Oct 9, 2024 · 329e2a0 · 329e2a0
1 parent cfeafce
commit 329e2a0
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/install/install_requirements.sh b/install/install_requirements.sh
@@ -106,3 +106,5 @@ fi
   set -x
   $PIP_EXECUTABLE install evaluate=="0.4.3" lm-eval=="0.4.2" psutil=="6.0.0"
 )
+
+export HF_HUB_ENABLE_HF_TRANSFER=1
diff --git a/install/requirements.txt b/install/requirements.txt
@@ -1,7 +1,7 @@
 # Requires python >=3.10
 
 # Hugging Face download
-huggingface_hub
+huggingface_hub[hf_transfer]
 
 # GGUF import
 gguf

diff --git a/torchchat/cli/download.py b/torchchat/cli/download.py
@@ -22,6 +22,9 @@
 # Both $HF_HOME and $HUGGINGFACE_HUB_CACHE are valid environment variables for the same directory.
 HUGGINGFACE_HOME_PATH = Path(os.environ.get("HF_HOME", os.environ.get("HUGGINGFACE_HUB_CACHE", os.path.expanduser("~/.cache/huggingface/hub"))))
 
+if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", None) is None:
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
 def _download_hf_snapshot(
     model_config: ModelConfig, hf_token: Optional[str]
 ):
-Original file line number
+Diff line change
@@ Expand Up / @@ -106,3 +106,5 @@ fi @@
       set -x
       $PIP_EXECUTABLE install evaluate=="0.4.3" lm-eval=="0.4.2" psutil=="6.0.0"
     )
+    export HF_HUB_ENABLE_HF_TRANSFER=1