From 329e2a0f650863887e27932ef7d7aac8b2bdf643 Mon Sep 17 00:00:00 2001
From: vmpuri <puri@meta.com>
Date: Wed, 9 Oct 2024 13:13:21 -0700
Subject: [PATCH] Enable hf_transfer for faster download

---
 install/install_requirements.sh | 2 ++
 install/requirements.txt        | 2 +-
 torchchat/cli/download.py       | 3 +++
 3 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/install/install_requirements.sh b/install/install_requirements.sh
index e525434a3..26203367f 100755
--- 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
index 3329563b4..ddb32a38f 100644
--- 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
index 3c6579d6f..14c19f943 100644
--- 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]
 ):