Upload hf_account_collateral.py with huggingface_hub

hf_account_collateral.py

@@ -15,7 +15,6 @@ from pathlib import Path
 from typing import Dict, List
 
 import requests
-from huggingface_hub import HfApi
 
 
 TEXT_EXTENSIONS = {
@@ -27,6 +26,7 @@ TEXT_EXTENSIONS = {
 MAX_TEXT_FILE_BYTES = int(os.environ.get("HF_COLLATERAL_MAX_TEXT_FILE_BYTES", "1000000"))
 FILE_FETCH_TIMEOUT = int(os.environ.get("HF_COLLATERAL_FILE_TIMEOUT_SECONDS", "10"))
 REPO_TREE_TIMEOUT = int(os.environ.get("HF_COLLATERAL_TREE_TIMEOUT_SECONDS", "10"))
+ACCOUNT_LIST_TIMEOUT = int(os.environ.get("HF_COLLATERAL_ACCOUNT_TIMEOUT_SECONDS", "10"))
 
 
 def _sha256_text(value: str) -> str:
@@ -84,31 +84,40 @@ def _repo_tree_files(repo_id: str, repo_type: str) -> List[str]:
     return sorted(files)
 
 
-def _repo_records(api: HfApi, owner: str) -> List[Dict]:
+def _account_items(path: str, owner: str) -> List[Dict]:
+    response = requests.get(
+        f"https://huggingface.co/api/{path}",
+        params={"author": owner, "limit": 1000},
+        timeout=ACCOUNT_LIST_TIMEOUT,
+    )
+    response.raise_for_status()
+    return response.json()
+
+
+def _repo_records(owner: str) -> List[Dict]:
     records = []
-    for repo_type, iterator in (
-        ("model", api.list_models(author=owner, limit=1000)),
-        ("dataset", api.list_datasets(author=owner, limit=1000)),
-        ("space", api.list_spaces(author=owner, limit=1000)),
+    for repo_type, path in (
+        ("model", "models"),
+        ("dataset", "datasets"),
+        ("space", "spaces"),
     ):
-        for item in iterator:
-            repo_id = getattr(item, "id", None)
+        for item in _account_items(path, owner):
+            repo_id = item.get("id") or item.get("name")
             if not repo_id:
                 continue
             records.append({
                 "repo_type": repo_type,
                 "repo_id": repo_id,
-                "likes": getattr(item, "likes", 0) or 0,
-                "downloads": getattr(item, "downloads", 0) or 0,
-                "last_modified": str(getattr(item, "last_modified", "") or getattr(item, "lastModified", "") or ""),
-                "sdk": getattr(item, "sdk", None),
+                "likes": item.get("likes", 0) or 0,
+                "downloads": item.get("downloads", 0) or 0,
+                "last_modified": str(item.get("lastModified") or item.get("updatedAt") or ""),
+                "sdk": item.get("sdk"),
             })
     return sorted(records, key=lambda item: (item["repo_type"], item["repo_id"]))
 
 
 def scan_hf_account_collateral(owner: str) -> Dict:
-    api = HfApi()
-    repos = _repo_records(api, owner)
+    repos = _repo_records(owner)
     scanned_repos = []
     total_files = 0
     total_text_files = 0