Upload hf_sync.py with huggingface_hub

hf_sync.py

@@ -1,39 +1,30 @@
 import os
-import subprocess
 import sys
 import shutil
+from huggingface_hub import snapshot_download, HfApi
 
 # Configuration
 REPO_ID = os.environ.get("DATASET_REPO_ID")
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-def run_command(command):
-    print(f"Running: {' '.join(command)}")
-    # Ensure HF_TOKEN is in environment for the command
-    env = os.environ.copy()
-    if HF_TOKEN:
-        env["HF_TOKEN"] = HF_TOKEN
-    
-    result = subprocess.run(command, capture_output=True, text=True, env=env)
-    if result.returncode != 0:
-        print(f"Error: {result.stderr}")
-    else:
-        print(f"Output: {result.stdout}")
-    return result.returncode == 0
-
 def download():
     if not REPO_ID:
         print("DATASET_REPO_ID not set, skipping download.")
         return
     
     print(f"Downloading data from {REPO_ID}...")
-    # hf download REPO_ID --repo-type dataset --local-dir data_repo
-    # Using --local-dir-use-symlinks False to avoid issues in some environments
-    success = run_command(["hf", "download", REPO_ID, "--repo-type", "dataset", "--local-dir", "data_repo", ])
-    if success:
+    try:
+        # snapshot_download is more efficient for many files than the CLI
+        snapshot_download(
+            repo_id=REPO_ID,
+            repo_type="dataset",
+            local_dir="data_repo",
+            token=HF_TOKEN,
+            max_workers=8
+        )
         print("Download successful.")
-    else:
-        print("Download failed or repository is empty.")
+    except Exception as e:
+        print(f"Download failed: {e}")
 
 def upload():
     if not REPO_ID:
@@ -44,22 +35,23 @@ def upload():
         return
 
     print(f"Uploading data to {REPO_ID}...")
-    # hf upload REPO_ID data_repo / --repo-type dataset
-    # We upload the contents of data_repo to the root of the dataset
-    success = run_command(["hf", "upload", REPO_ID, "data_repo", ".", "--repo-type", "dataset"])
-    if success:
+    try:
+        api = HfApi(token=HF_TOKEN)
+        api.upload_folder(
+            folder_path="data_repo",
+            repo_id=REPO_ID,
+            repo_type="dataset",
+            # This handles large folders by committing in chunks if necessary
+        )
         print("Upload successful.")
-    else:
-        print("Upload failed.")
+    except Exception as e:
+        print(f"Upload failed: {e}")
 
 def init_local():
     """Ensure data_repo has the necessary structure if download failed or it's new."""
     os.makedirs("data_repo/output", exist_ok=True)
     os.makedirs("data_repo/processed", exist_ok=True)
     os.makedirs("data_repo/uploads", exist_ok=True)
-    # database.db will be created by the app if it doesn't exist, 
-    # but we should ensure it's in data_repo.
-    # We'll handle this in entrypoint.sh by symlinking.
 
 if __name__ == "__main__":
     if len(sys.argv) < 2:
@@ -74,4 +66,4 @@ if __name__ == "__main__":
     elif action == "init":
         init_local()
     else:
-        print(f"Unknown action: {action}")
+        print(f"Unknown action: {action}")