J'aime 0
Fork 0

Dernière activité 2 days ago

KarelWintersky a révisé ce gist 2 days ago. Aller à la révision

1 file changed, 428 insertions

download_hf_model.py(fichier créé)

@@ -0,0 +1,428 @@
1 + #!/usr/bin/env python3
2 + """
3 + Download all files from a HuggingFace model repository with retries,
4 + resume support, and integrity verification.
5 +
6 + Usage:
7 + python3 download_hf_model.py <repo_id> [repo_id ...]
8 +
9 + Examples:
10 + python3 download_hf_model.py ai-sage/GigaAM-Multilingual
11 + python3 download_hf_model.py meta-llama/Llama-3-8B-Instruct
12 +
13 + Each repo is downloaded into a directory named by replacing '/' with '.'
14 + in the repo_id (e.g. ai-sage.GigaAM-Multilingual).
15 +
16 + Adapted from download_lilith_q8.py.
17 + """
18 +
19 + import argparse
20 + import hashlib
21 + import importlib
22 + import os
23 + import subprocess
24 + import sys
25 + import tempfile
26 + import time
27 + from pathlib import Path
28 +
29 + MAX_RETRIES = 10
30 + RETRY_DELAY_BASE = 15 # seconds, doubles each retry
31 +
32 + BASE_DIR = Path(__file__).resolve().parent
33 + RUNTIME_DIR = BASE_DIR / ".hf_download_runtime"
34 + HF_HOME_DIR = RUNTIME_DIR / "hf-home"
35 + HF_HUB_CACHE_DIR = HF_HOME_DIR / "hub"
36 + HF_ASSETS_CACHE_DIR = HF_HOME_DIR / "assets"
37 + XDG_CACHE_HOME_DIR = RUNTIME_DIR / "xdg-cache"
38 + PIP_CACHE_DIR = RUNTIME_DIR / "pip-cache"
39 + PYTHON_VENDOR_DIR = RUNTIME_DIR / "python-vendor"
40 + TMP_DIR = RUNTIME_DIR / "tmp"
41 +
42 +
43 + def prepare_runtime_dirs() -> None:
44 + for path in (
45 + RUNTIME_DIR,
46 + HF_HOME_DIR,
47 + HF_HUB_CACHE_DIR,
48 + HF_ASSETS_CACHE_DIR,
49 + XDG_CACHE_HOME_DIR,
50 + PIP_CACHE_DIR,
51 + PYTHON_VENDOR_DIR,
52 + TMP_DIR,
53 + ):
54 + path.mkdir(parents=True, exist_ok=True)
55 +
56 + os.environ["HF_HOME"] = str(HF_HOME_DIR)
57 + os.environ["HF_HUB_CACHE"] = str(HF_HUB_CACHE_DIR)
58 + os.environ["HF_ASSETS_CACHE"] = str(HF_ASSETS_CACHE_DIR)
59 + os.environ["XDG_CACHE_HOME"] = str(XDG_CACHE_HOME_DIR)
60 + os.environ["PIP_CACHE_DIR"] = str(PIP_CACHE_DIR)
61 + os.environ["HF_HUB_DISABLE_XET"] = "1"
62 + os.environ["HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY"] = "1"
63 + os.environ["TMPDIR"] = str(TMP_DIR)
64 + os.environ["TMP"] = str(TMP_DIR)
65 + os.environ["TEMP"] = str(TMP_DIR)
66 + tempfile.tempdir = str(TMP_DIR)
67 +
68 + if str(PYTHON_VENDOR_DIR) not in sys.path:
69 + sys.path.insert(0, str(PYTHON_VENDOR_DIR))
70 +
71 +
72 + prepare_runtime_dirs()
73 +
74 +
75 + def parse_major_minor(version: str) -> tuple[int, int]:
76 + parts = []
77 + for chunk in version.split("."):
78 + digits = "".join(ch for ch in chunk if ch.isdigit())
79 + if digits:
80 + parts.append(int(digits))
81 + if len(parts) == 2:
82 + break
83 + while len(parts) < 2:
84 + parts.append(0)
85 + return parts[0], parts[1]
86 +
87 +
88 + def install_huggingface_hub() -> None:
89 + print("Installing newer huggingface_hub into the local runtime dir...")
90 + subprocess.check_call(
91 + [
92 + sys.executable,
93 + "-m",
94 + "pip",
95 + "install",
96 + "--upgrade",
97 + "--target",
98 + str(PYTHON_VENDOR_DIR),
99 + "--cache-dir",
100 + str(PIP_CACHE_DIR),
101 + "huggingface_hub>=1.0,<2",
102 + ]
103 + )
104 +
105 +
106 + try:
107 + import huggingface_hub as _huggingface_hub
108 +
109 + if parse_major_minor(_huggingface_hub.__version__) < (1, 0):
110 + raise ImportError(f"Found old huggingface_hub {_huggingface_hub.__version__}")
111 +
112 + from huggingface_hub import HfApi, hf_hub_url
113 + from huggingface_hub.utils import RepositoryNotFoundError
114 + except Exception:
115 + install_huggingface_hub()
116 + importlib.invalidate_caches()
117 + for name in list(sys.modules):
118 + if name == "huggingface_hub" or name.startswith("huggingface_hub."):
119 + del sys.modules[name]
120 +
121 + from huggingface_hub import HfApi, hf_hub_url
122 + from huggingface_hub.utils import RepositoryNotFoundError
123 +
124 +
125 + def load_token() -> str | None:
126 + token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
127 + if token:
128 + return token
129 +
130 + token_files = [
131 + HF_HOME_DIR / "token",
132 + Path.home() / ".cache" / "huggingface" / "token",
133 + ]
134 + for token_file in token_files:
135 + if token_file.exists():
136 + return token_file.read_text().strip()
137 + return None
138 +
139 +
140 + TOKEN = load_token()
141 + if not TOKEN:
142 + print("HF token not found. Proceeding without authentication (repo may be public).")
143 + else:
144 + print(f"Using HF token: {TOKEN[:8]}...")
145 +
146 +
147 + def repo_to_dirname(repo_id: str) -> str:
148 + return repo_id.replace("/", ".")
149 +
150 +
151 + def get_dir_size(path: Path) -> int:
152 + return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
153 +
154 +
155 + def expected_size_for(sibling) -> int | None:
156 + size = getattr(sibling, "size", None)
157 + if size is not None:
158 + return size
159 + if sibling.lfs and sibling.lfs.get("size") is not None:
160 + return sibling.lfs["size"]
161 + return None
162 +
163 +
164 + def format_size(n: int | None) -> str:
165 + if n is None:
166 + return "unknown size"
167 + if n >= 1e9:
168 + return f"{n / 1e9:.2f} GB"
169 + if n >= 1e6:
170 + return f"{n / 1e6:.2f} MB"
171 + return f"{n / 1e3:.2f} KB"
172 +
173 +
174 + def quarantine_corrupt_files(paths: list[Path], download_dir: Path) -> None:
175 + if not paths:
176 + return
177 + quarantine_dir = download_dir / ".corrupt"
178 + quarantine_dir.mkdir(exist_ok=True)
179 + stamp = int(time.time())
180 + for path in paths:
181 + target = quarantine_dir / f"{path.name}.{stamp}.corrupt"
182 + suffix = 1
183 + while target.exists():
184 + target = quarantine_dir / f"{path.name}.{stamp}.{suffix}.corrupt"
185 + suffix += 1
186 + path.replace(target)
187 + print(f"Moved corrupt file aside: {target}")
188 +
189 +
190 + def download_file(rel_path: str, dest: Path, repo_id: str) -> None:
191 + url = hf_hub_url(repo_id=repo_id, filename=rel_path)
192 + cmd = [
193 + "curl",
194 + "-L",
195 + "--fail",
196 + "--show-error",
197 + "--connect-timeout",
198 + "30",
199 + "--progress-bar",
200 + "-C",
201 + "-",
202 + "-o",
203 + str(dest),
204 + url,
205 + ]
206 + if TOKEN:
207 + cmd[1:1] = ["-H", f"Authorization: Bearer {TOKEN}"]
208 + subprocess.check_call(cmd)
209 +
210 +
211 + def download_repo_with_retries(repo_id: str, download_dir: Path) -> Path:
212 + """Download all files from a repo with curl, resume, and exponential backoff."""
213 + for attempt in range(1, MAX_RETRIES + 1):
214 + try:
215 + print(f"\n{'=' * 60}")
216 + print(f"Download attempt {attempt}/{MAX_RETRIES}")
217 + print(f"{'=' * 60}")
218 +
219 + api = HfApi(token=TOKEN)
220 + repo_info = api.repo_info(repo_id=repo_id, token=TOKEN, files_metadata=True)
221 + siblings = [s for s in (repo_info.siblings or []) if s.rfilename]
222 +
223 + if attempt == 1:
224 + # Disk space check
225 + needed = 0
226 + for s in siblings:
227 + exp = expected_size_for(s)
228 + dest = download_dir / s.rfilename
229 + have = dest.stat().st_size if dest.exists() else 0
230 + if exp is not None:
231 + needed += max(exp - have, 0)
232 + try:
233 + stat = os.statvfs(download_dir)
234 + free = stat.f_bavail * stat.f_frsize
235 + print(f"Disk check: need ~{format_size(needed)} more, free ~{format_size(free)} on destination.")
236 + if needed > free:
237 + print("WARNING: not enough free space for the remaining download.")
238 + except OSError:
239 + pass
240 +
241 + siblings = sorted(
242 + siblings,
243 + key=lambda s: expected_size_for(s) or 0,
244 + reverse=True,
245 + )
246 +
247 + size_before = get_dir_size(download_dir) if download_dir.exists() else 0
248 + t0 = time.time()
249 +
250 + for index, sibling in enumerate(siblings, start=1):
251 + rel_path = sibling.rfilename
252 + dest = download_dir / rel_path
253 + dest.parent.mkdir(parents=True, exist_ok=True)
254 +
255 + expected_size = expected_size_for(sibling)
256 + current_size = dest.stat().st_size if dest.exists() else 0
257 +
258 + if expected_size is not None and dest.exists() and current_size == expected_size:
259 + print(
260 + f"[{index}/{len(siblings)}] Already present: {rel_path} "
261 + f"({format_size(expected_size)})"
262 + )
263 + continue
264 +
265 + if expected_size is not None and dest.exists() and current_size > expected_size:
266 + quarantine_corrupt_files([dest], download_dir)
267 + current_size = 0
268 +
269 + if current_size > 0:
270 + print(
271 + f"[{index}/{len(siblings)}] Resuming: {rel_path} "
272 + f"from {format_size(current_size)} / {format_size(expected_size)}"
273 + )
274 + else:
275 + print(f"[{index}/{len(siblings)}] Downloading: {rel_path} ({format_size(expected_size)})")
276 +
277 + download_file(rel_path, dest, repo_id)
278 +
279 + if expected_size is not None:
280 + actual_size = dest.stat().st_size
281 + if actual_size != expected_size:
282 + raise ValueError(
283 + f"Size mismatch after download for {rel_path}: "
284 + f"expected {expected_size}, got {actual_size}"
285 + )
286 +
287 + elapsed = time.time() - t0
288 + size_after = get_dir_size(download_dir)
289 + downloaded = max(size_after - size_before, 0)
290 + speed = downloaded / elapsed if elapsed > 0 else 0
291 +
292 + print(f"\nDownload complete: {download_dir}")
293 + print(f" Downloaded: {format_size(downloaded)} in {elapsed:.0f}s")
294 + print(f" Avg speed: {speed / 1e6:.1f} MB/s")
295 + return download_dir
296 +
297 + except KeyboardInterrupt:
298 + print("\nInterrupted by user.")
299 + sys.exit(1)
300 +
301 + except (RepositoryNotFoundError, FileNotFoundError) as exc:
302 + print(f"\nFatal error: {exc}")
303 + sys.exit(1)
304 +
305 + except Exception as exc:
306 + delay = RETRY_DELAY_BASE * (2 ** (attempt - 1))
307 + print(f"\nError on attempt {attempt}: {exc}")
308 + if attempt < MAX_RETRIES:
309 + print(f"Retrying in {delay}s...")
310 + time.sleep(delay)
311 + else:
312 + print("All retries exhausted.")
313 + sys.exit(1)
314 +
315 +
316 + def verify_integrity(repo_id: str, download_dir: Path) -> tuple[bool, list[Path]]:
317 + """Verify all downloaded files against repo metadata."""
318 + print(f"\n{'=' * 60}")
319 + print("Verifying file integrity...")
320 + print(f"{'=' * 60}")
321 +
322 + api = HfApi(token=TOKEN)
323 + try:
324 + repo_info = api.repo_info(repo_id=repo_id, token=TOKEN, files_metadata=True)
325 + except Exception as exc:
326 + print(f"Could not fetch repo metadata for verification: {exc}")
327 + return False, []
328 +
329 + siblings = [s for s in (repo_info.siblings or []) if s.rfilename]
330 + errors = []
331 + checked = 0
332 + corrupt_files: list[Path] = []
333 +
334 + for sibling in siblings:
335 + if not sibling.rfilename:
336 + continue
337 +
338 + file_path = download_dir / sibling.rfilename
339 + if not file_path.exists():
340 + errors.append(f"MISSING: {sibling.rfilename}")
341 + continue
342 +
343 + expected_sha = None
344 + hash_type = None
345 + if sibling.lfs and sibling.lfs.get("sha256"):
346 + expected_sha = sibling.lfs["sha256"]
347 + hash_type = "sha256"
348 + elif hasattr(sibling, "blob_id") and sibling.blob_id:
349 + expected_sha = sibling.blob_id
350 + hash_type = "git-sha1"
351 + else:
352 + continue
353 +
354 + checked += 1
355 + size = file_path.stat().st_size
356 + print(f" Checking {sibling.rfilename} ({format_size(size)})...", end=" ", flush=True)
357 +
358 + if hash_type == "sha256":
359 + digest = hashlib.sha256()
360 + with open(file_path, "rb") as handle:
361 + while chunk := handle.read(8 * 1024 * 1024):
362 + digest.update(chunk)
363 + actual = digest.hexdigest()
364 + else:
365 + digest = hashlib.sha1()
366 + digest.update(f"blob {size}\0".encode())
367 + with open(file_path, "rb") as handle:
368 + while chunk := handle.read(8 * 1024 * 1024):
369 + digest.update(chunk)
370 + actual = digest.hexdigest()
371 +
372 + if actual == expected_sha:
373 + print("OK")
374 + else:
375 + print("MISMATCH")
376 + corrupt_files.append(file_path)
377 + errors.append(
378 + f"HASH MISMATCH: {sibling.rfilename}\n"
379 + f" expected: {expected_sha}\n"
380 + f" actual: {actual}"
381 + )
382 +
383 + print(f"\nChecked {checked} file(s).")
384 +
385 + if errors:
386 + print("\nIntegrity errors:")
387 + for err in errors:
388 + print(f" - {err}")
389 + return False, corrupt_files
390 +
391 + print("All files verified successfully.")
392 + return True, []
393 +
394 +
395 + def main() -> None:
396 + parser = argparse.ArgumentParser(
397 + description="Download files from HuggingFace model repositories with retry and verification."
398 + )
399 + parser.add_argument(
400 + "repos",
401 + nargs="+",
402 + help="HuggingFace repo IDs to download (e.g. ai-sage/GigaAM-Multilingual)",
403 + )
404 + args = parser.parse_args()
405 +
406 + for repo_id in args.repos:
407 + dirname = repo_to_dirname(repo_id)
408 + download_dir = BASE_DIR / dirname
409 + download_dir.mkdir(parents=True, exist_ok=True)
410 +
411 + print(f"\n{'#' * 60}")
412 + print(f"# Repo: {repo_id}")
413 + print(f"# Target: {download_dir}")
414 + print(f"{'#' * 60}")
415 +
416 + download_repo_with_retries(repo_id, download_dir)
417 +
418 + ok, corrupt_files = verify_integrity(repo_id, download_dir)
419 + if not ok:
420 + quarantine_corrupt_files(corrupt_files, download_dir)
421 + print(f"\nIntegrity check FAILED for {repo_id}. Re-run the script to resume.")
422 + sys.exit(1)
423 +
424 + print(f"\nDone. Model saved to: {download_dir}")
425 +
426 +
427 + if __name__ == "__main__":
428 + main()
Plus récent Plus ancien

Propulsé par Opengist ⋅ Load: 31ms⋅

Français