Done

Dockerfile

@@ -32,7 +32,7 @@ RUN wget https://github.com/wkhtmltopdf/packaging/releases/download/0.12.6.1-2/w
 WORKDIR /app
 
 COPY requirements.txt .
-RUN pip install --no-cache-dir gunicorn eventlet huggingface_hub[cli]
+RUN pip install --no-cache-dir gunicorn eventlet "huggingface_hub[cli,hf_transfer]"
 RUN pip install --no-cache-dir -r requirements.txt
 
 COPY . .
@@ -46,6 +46,7 @@ ENV PATH=/home/user/.local/bin:$PATH
 
 ENV PORT=7680
 ENV PYTHONUNBUFFERED=1
+ENV HF_HUB_ENABLE_HF_TRANSFER=1
 
 EXPOSE 7680
 

hf_sync.py

@@ -6,7 +6,7 @@ import json
 import time
 import hashlib
 from datetime import datetime
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi, snapshot_download
 
 # Configuration
 REPO_ID = os.environ.get("DATASET_REPO_ID", "Jaimodiji/Report-Generator-Data")
@@ -192,48 +192,17 @@ def download(revision=None):
         with open(LOCK_FILE, 'w') as f:
             f.write(str(os.getpid()))
         write_status("running", "download", 5, "Preparing initial data sync...")
-
-        temp_dir = os.path.join("/tmp", f"hf_sync_download_{os.getpid()}")
-        if os.path.exists(temp_dir):
-            shutil.rmtree(temp_dir)
-        os.makedirs(temp_dir, exist_ok=True)
-
-        repo_files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset", revision=revision)
-        total_files = max(len(repo_files), 1)
-        write_status("running", "download", 10, "Downloading dataset files...", f"0/{len(repo_files)} files downloaded")
-
-        for index, repo_file in enumerate(repo_files, start=1):
-            hf_hub_download(
-                repo_id=REPO_ID,
-                repo_type="dataset",
-                filename=repo_file,
-                local_dir=temp_dir,
-                token=HF_TOKEN,
-                revision=revision,
-            )
-            progress = 10 + int((index / total_files) * 75)
-            write_status("running", "download", progress, "Downloading dataset files...", f"{index}/{len(repo_files)} files downloaded")
-
-        write_status("running", "download", 90, "Applying downloaded data locally...")
-        protected_names = {os.path.basename(STATUS_FILE)}
-        for entry in os.listdir(DATA_DIR):
-            if entry in protected_names:
-                continue
-            entry_path = os.path.join(DATA_DIR, entry)
-            if os.path.isdir(entry_path):
-                shutil.rmtree(entry_path)
-            else:
-                os.remove(entry_path)
-
-        for entry in os.listdir(temp_dir):
-            src = os.path.join(temp_dir, entry)
-            dst = os.path.join(DATA_DIR, entry)
-            if os.path.isdir(src):
-                shutil.copytree(src, dst, dirs_exist_ok=True)
-            else:
-                shutil.copy2(src, dst)
-
-        shutil.rmtree(temp_dir, ignore_errors=True)
+        write_status("running", "download", 15, "Listing dataset snapshot...", "Fetching repository metadata")
+        write_status("running", "download", 35, "Downloading dataset files...", "Parallel download in progress")
+        snapshot_download(
+            repo_id=REPO_ID,
+            repo_type="dataset",
+            local_dir=DATA_DIR,
+            token=HF_TOKEN,
+            revision=revision,
+            max_workers=8
+        )
+        write_status("running", "download", 90, "Finalizing local dataset...", "Verifying downloaded files")
         write_status("completed", "download", 100, "Initial data sync completed.")
         print("Download successful.")
     except Exception as e: