101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from huggingface_hub import HfApi, hf_hub_url
|
|
|
|
from common import load_model_card
|
|
|
|
|
|
def selected_files(repo_id: str, allow_patterns: list[str] | None) -> list[tuple[str, int | None]]:
|
|
api = HfApi()
|
|
info = api.model_info(repo_id=repo_id, files_metadata=True)
|
|
files = []
|
|
for sibling in info.siblings:
|
|
name = sibling.rfilename
|
|
if name == ".gitattributes":
|
|
continue
|
|
if allow_patterns and name not in allow_patterns:
|
|
continue
|
|
files.append((name, sibling.size))
|
|
return files
|
|
|
|
|
|
def run_curl(url: str, output: Path) -> int:
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
command = [
|
|
"curl.exe",
|
|
"-L",
|
|
"-C",
|
|
"-",
|
|
"--retry",
|
|
"10",
|
|
"--retry-all-errors",
|
|
"--retry-delay",
|
|
"2",
|
|
"--connect-timeout",
|
|
"30",
|
|
"--speed-time",
|
|
"60",
|
|
"--speed-limit",
|
|
"1024",
|
|
"-o",
|
|
str(output),
|
|
url,
|
|
]
|
|
return subprocess.run(command, check=False).returncode
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Download files from a Hugging Face model card using curl resume.")
|
|
parser.add_argument("card_id")
|
|
parser.add_argument("--local-dir", required=True)
|
|
parser.add_argument("--allow-file", action="append", dest="allow_files")
|
|
parser.add_argument("--max-files", type=int)
|
|
args = parser.parse_args()
|
|
|
|
card = load_model_card(args.card_id)
|
|
repo_id = card.get("upstream_id")
|
|
if not repo_id:
|
|
print(f"{args.card_id}: upstream_id is required", file=sys.stderr)
|
|
return 1
|
|
|
|
files = selected_files(repo_id, args.allow_files)
|
|
if args.max_files:
|
|
files = files[: args.max_files]
|
|
|
|
local_dir = Path(args.local_dir)
|
|
failed = []
|
|
for name, expected_size in files:
|
|
output = local_dir / name
|
|
current_size = output.stat().st_size if output.exists() else 0
|
|
if expected_size and current_size == expected_size:
|
|
print(f"OK {name}: {current_size}/{expected_size}")
|
|
continue
|
|
if expected_size and current_size > expected_size:
|
|
failed.append((name, current_size, expected_size, "oversized"))
|
|
continue
|
|
|
|
print(f"GET {name}: {current_size}/{expected_size or '?'}")
|
|
url = hf_hub_url(repo_id, name)
|
|
code = run_curl(url, output)
|
|
current_size = output.stat().st_size if output.exists() else 0
|
|
if code != 0 or (expected_size and current_size != expected_size):
|
|
failed.append((name, current_size, expected_size, code))
|
|
|
|
if failed:
|
|
print("Incomplete downloads:", file=sys.stderr)
|
|
for name, current_size, expected_size, code in failed:
|
|
print(f"- {name}: {current_size}/{expected_size or '?'} curl_code={code}", file=sys.stderr)
|
|
return 1
|
|
|
|
print(f"Downloaded {len(files)} file(s) for {args.card_id}.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|