build(utils): allow to download in safetensors format (#174)

Author: fardeon

utils/download.py

@@ -3,6 +3,7 @@
 This allows memory-constrained CI/CD runners to build container images
 with large bundled models. See ../deployments/bundle/ for examples.
 """
+import os
 import sys
 import tempfile
 
@@ -11,6 +12,11 @@
 if len(sys.argv) < 3:
     sys.exit("usage: python download.py REPO_ID LOCAL_DIR [REVISION]")
 
+if os.getenv("TENSOR_FORMAT") == "safetensors":
+    tensor_format = "*.safetensors"
+else:
+    tensor_format = "*.bin"
+
 with tempfile.TemporaryDirectory() as cache_dir:
     huggingface_hub.snapshot_download(
         repo_id=sys.argv[1],
@@ -19,5 +25,5 @@
         cache_dir=cache_dir,
         local_dir_use_symlinks=False,
         resume_download=True,
-        allow_patterns=["*.bin", "*.json", "*.model", "*.py"],
+        allow_patterns=[tensor_format, "*.json", "*.model", "*.py"],
     )