Hugging Face에서 파일을 다운로드하기 위한 간단한 도구

Source: Dev.to

Source: …

hf_get_from_url.py – Hugging Face 저장소에서 파일 다운로드

머신러닝 모델을 다루는 개발자들은 종종 Hugging Face 저장소에서 파일을 받아와야 합니다. Hugging Face 웹사이트에서 링크를 제공하지만, URL과 경로를 직접 처리하는 것은 번거로울 수 있습니다. 스크립트 **hf_get_from_url.py**는 다양한 입력 형식을 해석하고 Hugging Face CLI(hf)를 호출해 요청된 파일을 손쉽게 가져올 수 있도록 해줍니다.

스크립트가 처리할 수 있는 입력

입력을 파싱한 뒤, 스크립트는 Hugging Face 명령줄 인터페이스를 사용해 지정된 파일을 로컬 디렉터리로 다운로드합니다.

주요 기능

• 유연한 입력 파싱 – 전체 URL, 축약형 URL, 그리고 단순 owner/repo/path 문자열을 모두 받아들입니다.
• 드라이‑런 모드 – --dry-run 옵션을 사용하면 실제 다운로드 없이 실행될 정확한 명령을 출력합니다.
• 사용자 지정 로컬 디렉터리 – --flatten-localdir 옵션을 사용하면 파일이 슬래시 대신 하이픈(owner-repo)을 사용한 폴더에 저장됩니다.
• 오류 처리 – hf 명령이 설치되어 있는지 확인하고, 문제가 발생했을 때 명확한 오류 메시지를 제공합니다.

사용 예시

# 전체 URL을 사용해 다운로드
python hf_get_from_url.py "https://huggingface.co/owner/repo/blob/main/file.gguf"

# repo‑style 경로를 사용해 다운로드
python hf_get_from_url.py owner/repo/file.gguf

# 다운로드 없이 명령만 미리 보기
python hf_get_from_url.py --dry-run "huggingface.co/owner/repo/blob/main/file.gguf"

위 예시들은 Hugging Face에서 모델, 데이터셋, 설정 파일 등을 자주 다운로드하는 사람들의 작업 흐름을 어떻게 간소화할 수 있는지 보여줍니다.

소스 코드

#!/usr/bin/env python3
"""
hf_get_from_url.py
Download files or directories from Hugging Face Hub using huggingface_hub API.

Usage:
  python hf_get_from_url.py [--dry-run] [--flatten-localdir]  [ ...]

Examples:
  python hf_get_from_url.py "https://huggingface.co/owner/repo/blob/main/path/to/file.gguf"
  python hf_get_from_url.py owner/repo/path/to/file.gguf
  python hf_get_from_url.py --dry-run "huggingface.co/owner/repo/blob/main/models"

Notes:
  - Requires: pip install huggingface_hub
  - Authentication (if needed) is taken from env HF_TOKEN or huggingface_hub.login()
"""

from __future__ import annotations

import argparse
import sys
import re
from urllib.parse import urlparse, unquote
from typing import Optional, Tuple, List

# huggingface_hub API
try:
    from huggingface_hub import hf_hub_download, snapshot_download
except Exception:
    hf_hub_download = None
    snapshot_download = None

# ----------------------------------------------------------------------
# Regex patterns
# ----------------------------------------------------------------------
RE_BLOB_RESOLVE = re.compile(
    r'^(?:https?://)?(?:www\.)?huggingface\.co/'
    r'(?P[^/]+/[^/]+)/(?:blob|resolve)/'
    r'(?P[^/]+)/(?P.+)$'
)

RE_NO_PREFIX = re.compile(
    r'^(?P[^/]+/[^/]+)/(?:blob|resolve)/'
    r'(?P[^/]+)/(?P.+)$'
)

RE_SIMPLE = re.compile(
    r'^(?P[^/]+/[^/]+)(?:/(?P.+))?$'
)

# ----------------------------------------------------------------------
# Input parser
# ----------------------------------------------------------------------
def parse_input(s: str) -> Optional[Tuple[str, Optional[str], str]]:
    """
    Parse input and return (repo, revision, path)

    revision may be None (meaning default branch)
    """
    s = s.strip()
    s = unquote(s.split('?', 1)[0].split('#', 1)[0]).rstrip('/')

    # 1) Explicit Hugging Face URL (blob/resolve)
    m = RE_BLOB_RESOLVE.match(s)
    if m:
        return m.group('repo'), m.group('rev'), m.group('path')

    # 2) huggingface.co/... without scheme
    if s.startswith('huggingface.co/'):
        candidate = s[len('huggingface.co/'):].lstrip('/')
        m2 = RE_NO_PREFIX.match(candidate)
        if m2:
            return m2.group('repo'), m2.group('rev'), m2.group('path')

        m2 = RE_SIMPLE.match(candidate)
        if m2 and m2.group('path'):
            return m2.group('repo'), None, m2.group('path')

    # 3) Generic URL parse
    try:
        p = urlparse(s)
    except Exception:
        p = None

    if p and p.netloc and 'huggingface' in p.netloc:
        parts = p.path.lstrip('/').split('/')
        if len(parts) >= 5 and parts[2] in ('blob', 'resolve'):
            repo = f"{parts[0]}/{parts[1]}"
            rev = parts[3]
            path = '/'.join(parts[4:])
            return repo, rev, path
        elif len(parts) >= 3:
            repo = f"{parts[0]}/{parts[1]}"
            path = '/'.join(parts[2:])
            return repo, None, path

    # 4) Direct repo/path
    m3 = RE_NO_PREFIX.match(s)
    if m3:
        return m3.group('repo'), m3.group('rev'), m3.group('path')

    m4 = RE_SIMPLE.match(s)
    if m4 and m4.group('path'):
        return m4.group('repo'), None, m4.group('path')

    return None

# ----------------------------------------------------------------------
# Download logic
# ----------------------------------------------------------------------
def run_hf_download_api(
    repo: str,
    path: str,
    rev: Optional[str],
    local_dir: Optional[str],
    dry_run: bool,
) -> int:
    """
    Try single‑file download first; if it fails, fall back to snapshot_download
    (directory or pattern).
    """
    if hf_hub_download is None or snapshot_download is None:
        print(
            "Error: huggingface_hub is not installed. "
            "Please run `pip install huggingface_hub`.",
            file=sys.stderr,
        )
        return 2

    if local_dir is None:
        local_dir = repo

    rev_disp = rev if rev else "default

    print(f"> (api) download {repo}@{rev_disp} {path} -> local_dir={local_dir}")

    if dry_run:
        return 0

    # ---- Try as single file ----
    try:
        local_path = hf_hub_download(
            repo_id=repo,
            filename=path,
            revision=rev,
            local_dir=local_dir,
        )
        print(f"✔ Downloaded file to {local_path}")
        return 0
    except Exception as e_file:
        print(f"⚠ Single‑file download failed: {e_file}", file=sys.stderr)

    # ---- Fallback: download whole repo or sub‑directory ----
    try:
        snapshot_download(
            repo_id=repo,
            revision=rev,
            local_dir=local_dir,
            allow_patterns=[f"{path}*"],
        )
        print(f"✔ Snapshot downloaded to {local_dir}")
        return 0
    except Exception as e_snap:
        print(f"❌ Snapshot download failed: {e_snap}", file=sys.stderr)
        return 1

# ----------------------------------------------------------------------
# CLI entry point
# ----------------------------------------------------------------------
def main(argv: List[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Download files from Hugging Face Hub using flexible input formats."
    )
    parser.add_argument(
        "inputs",
        nargs="+",
        help="URL, repo/path, or shortened Hugging Face reference.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show the commands that would be run without downloading.",
    )
    parser.add_argument(
        "--flatten-localdir",
        action="store_true",
        help="Replace '/' with '-' in the local directory name.",
    )
    args = parser.parse_args(argv)

    exit_code = 0
    for inp in args.inputs:
        parsed = parse_input(inp)
        if not parsed:
            print(f"❌ Could not parse input: {inp}", file=sys.stderr)
            exit_code = 1
            continue

        repo, rev, path = parsed
        local_dir = repo.replace("/", "-") if args.flatten_localdir else None

        rc = run_hf_download_api(
            repo=repo,
            path=path,
            rev=rev,
            local_dir=local_dir,
            dry_run=args.dry_run,
        )
        if rc != 0:
            exit_code = rc

    return exit_code

if __name__ == "__main__":
    sys.exit(main())

이 도구는 모델, 데이터셋 또는 설정 파일을 Hugging Face에서 자주 다운로드하는 모든 사용자의 작업 흐름을 간소화합니다. 다양한 입력 형식을 허용하고 명확한 피드백을 제공함으로써, 프로세스를 더 빠르고 오류가 적게 만들 수 있습니다.

filename = path,
revision = rev,
local_dir = local_dir,
)
print(f"Downloaded file: {local_path}")
return 0
except Exception as e:
    print(
        f"hf_hub_download failed: {e}. "
        "Trying snapshot_download for directory/pattern...",
        file=sys.stderr,
    )

# ---- Fallback: directory or glob ----
allow_pattern = path.rstrip("/") + "/*"

try:
    repo_local_dir = snapshot_download(
        repo_id=repo,
        revision=rev,
        local_dir=local_dir,
        allow_patterns=[allow_pattern],
    )
    print(f"Snapshot downloaded into: {repo_local_dir}")
    return 0
except Exception as e:
    print(f"snapshot_download failed: {e}", file=sys.stderr)
    return 3

# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main(argv: List[str]) -> int:
    parser = argparse.ArgumentParser(
        description="Download files or directories from Hugging Face Hub"
    )
    parser.add_argument(
        "inputs",
        nargs="+",
        help="Hugging Face URL or /path",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print actions without executing",
    )
    parser.add_argument(
        "--hf-cmd",
        default="hf",
        help="(ignored, kept for compatibility)",

    parser.add_argument(
        "--flatten-localdir",
        action="store_true",
        help="Replace '/' with '-' in local directory name",
    )
    args = parser.parse_args(argv)

    any_failed = False

    for s in args.inputs:
        parsed = parse_input(s)
        if not parsed:
            print(f"Failed to parse input: {s}", file=sys.stderr)
            any_failed = True
            continue

        repo, rev, path = parsed
        if not path:
            print(f"No file path extracted for input: {s}", file=sys.stderr)
            any_failed = True
            continue

        local_dir = repo.replace("/", "-") if args.flatten_localdir else repo

        rc = run_hf_download_api(
            repo=repo,
            path=path,
            rev=rev,
            local_dir=local_dir,
            dry_run=args.dry_run,
        )
        if rc != 0:
            any_failed = True

    return 1 if any_failed else 0

if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))