Add model list and download progress logging

download_model.py

@@ -4,7 +4,7 @@ import os
 import shutil
 from pathlib import Path
 
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi, snapshot_download
 
 
 def require_env(name: str) -> str:
@@ -14,36 +14,35 @@ def require_env(name: str) -> str:
     return value
 
 
-def quant_rank(path: Path) -> tuple[int, int, str]:
-    name = path.name.lower()
-    size = path.stat().st_size
+def quant_rank(file_name: str) -> tuple[int, str]:
+    name = file_name.lower()
 
     q8_tokens = ("q8_k_l", "q8_k_m", "q8_k_s", "q8_k", "q8_0", "q8_1")
     if any(token in name for token in q8_tokens):
-        return (0, -size, path.name)
+        return (0, file_name)
 
     if "f16" in name or "bf16" in name:
-        return (1, -size, path.name)
+        return (1, file_name)
 
     if any(token in name for token in ("q6_k", "q6_0", "q6_1", "q6")):
-        return (2, -size, path.name)
+        return (2, file_name)
 
     if any(token in name for token in ("q5_k_m", "q5_k_s", "q5_k", "q5_0", "q5_1", "q5")):
-        return (3, -size, path.name)
+        return (3, file_name)
 
     if any(token in name for token in ("q4_k_m", "q4_k_s", "q4_k", "q4_0", "q4_1", "q4")):
-        return (4, -size, path.name)
+        return (4, file_name)
 
     if any(token in name for token in ("q3_k_m", "q3_k_s", "q3_k", "q3_0", "q3_1", "q3")):
-        return (5, -size, path.name)
+        return (5, file_name)
 
     if any(token in name for token in ("q2_k", "q2_0", "q2_1", "q2")):
-        return (6, -size, path.name)
+        return (6, file_name)
 
     if "f32" in name:
-        return (7, -size, path.name)
+        return (7, file_name)
 
-    return (8, -size, path.name)
+    return (8, file_name)
 
 
 def main() -> None:
@@ -67,26 +66,28 @@ def main() -> None:
     download_dir.mkdir(parents=True, exist_ok=True)
 
     api = HfApi(token=hf_token)
-    print(f"listing model files: {model_name}")
+    print(f"listing model files: {model_name}", flush=True)
     files = api.list_repo_files(repo_id=model_name, repo_type="model")
-    candidates = [Path(file_name) for file_name in files if file_name.lower().endswith(".gguf")]
+    candidates = [file_name for file_name in files if file_name.lower().endswith(".gguf")]
     if not candidates:
         raise SystemExit(f"no gguf files found in repository: {model_name}")
 
+    print("available gguf files:", flush=True)
+    for candidate in candidates:
+        print(f"- {candidate}", flush=True)
+
     selected = min(candidates, key=quant_rank)
-    print(f"selected model: {selected.name}")
-
-    print(f"downloading model: {model_name}/{selected.as_posix()}")
-    downloaded_path = Path(
-        hf_hub_download(
-            repo_id=model_name,
-            filename=selected.as_posix(),
-            token=hf_token,
-            repo_type="model",
-            cache_dir=os.environ.get("HF_HOME", "/data/hf-cache"),
-            local_dir=str(download_dir),
-        )
+    print(f"selected model: {selected}", flush=True)
+
+    print(f"downloading model: {model_name}/{selected}", flush=True)
+    snapshot_download(
+        repo_id=model_name,
+        token=hf_token,
+        allow_patterns=[selected],
+        local_dir=str(download_dir),
+        cache_dir=os.environ.get("HF_HOME", "/data/hf-cache"),
     )
+    downloaded_path = download_dir / selected
 
     if target_model.exists():
         target_model.unlink()