scripts/export_to_hub.py

# hf_custom_proj/scripts/export_to_hub.py
"""
개발 프로젝트(hf_custom_proj/)에서 학습 산출물은 artifacts/에 보관하고,
배포(Hub/로컬 from_pretrained)용 repo 루트는 dist/my-mnist-hf/ 로 구성합니다.

목표 (당신이 요구한 "HF 스타일"을 만족)
- 개발: src/ 레이아웃 유지 (pip install -e . / python -m examples.*)
- 배포: dist/ 는 Hub / local from_pretrained 모두 동작
- 임시처방 X: dist 구조를 HF dynamic module 로딩 규칙에 "정합"되게 생성

중요 포인트 (현재 당신 상황의 핵심)
- preprocessor_config.json의 auto_map이
    "AutoImageProcessor": "image_processing_my_mnist.MyMNISTImageProcessor"
  처럼 "패키지명 없이" 저장되는 경우가 있음.
- 이 경우 transformers는 dist 루트에서 image_processing_my_mnist.py 를 찾는다.
- 따라서 dist 루트에 configuration/modeling/image_processing *.py 를 flat하게 둔다.
"""

from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path


# -----------------------------
# path utils
# -----------------------------
def project_root_from_this_file() -> Path:
    return Path(__file__).resolve().parents[1]


def rmtree_if_exists(p: Path) -> None:
    if p.exists():
        shutil.rmtree(p)


def copytree(src: Path, dst: Path) -> None:
    if not src.exists():
        raise FileNotFoundError(f"Missing source path: {src}")
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(src, dst)


def copy2(src: Path, dst: Path) -> None:
    if not src.exists():
        raise FileNotFoundError(f"Missing file: {src}")
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)


def copy_if_exists(src: Path, dst: Path) -> None:
    if src.exists():
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)


def write_text(dst: Path, text: str) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_text(text, encoding="utf-8")


# -----------------------------
# artifacts validate
# -----------------------------
def find_model_file(artifact_dir: Path) -> Path:
    st = artifact_dir / "model.safetensors"
    pt = artifact_dir / "pytorch_model.bin"
    if st.exists():
        return st
    if pt.exists():
        return pt
    raise FileNotFoundError(
        f"Missing model file in artifacts: expected '{st.name}' or '{pt.name}'"
    )


def validate_artifacts(artifact_dir: Path) -> tuple[Path, Path, Path]:
    if not artifact_dir.exists():
        raise FileNotFoundError(
            f"Artifacts directory does not exist: {artifact_dir}\n"
            "먼저 학습 후 아래처럼 저장하세요:\n"
            "  save_dir = 'artifacts/my_mnist'\n"
            "  trainer.save_model(save_dir)\n"
            "  processor.save_pretrained(save_dir)\n"
        )

    config = artifact_dir / "config.json"
    preproc = artifact_dir / "preprocessor_config.json"
    model = find_model_file(artifact_dir)

    missing = []
    if not config.exists():
        missing.append("config.json")
    if not preproc.exists():
        missing.append("preprocessor_config.json")
    if missing:
        raise FileNotFoundError(
            "Artifacts 폴더에 학습 산출물이 없습니다(또는 불완전합니다).\n"
            f"  artifacts: {artifact_dir}\n"
            "  missing:\n"
            + "\n".join([f"    - {m}" for m in missing])
            + "\n\n"
            "먼저 학습 후 아래처럼 저장하세요:\n"
            "  save_dir = 'artifacts/my_mnist'\n"
            "  trainer.save_model(save_dir)\n"
            "  processor.save_pretrained(save_dir)\n"
        )

    return config, model, preproc


# -----------------------------
# dist build
# -----------------------------
def prepare_dist_repo_root(
    *,
    root: Path,
    dist_repo_dir: Path,
    artifact_dir: Path,
    package_name: str,
    copy_scripts_into_dist: bool,
    requirements_text: str | None,
) -> None:
    """
    dist repo 루트를 생성합니다.

    - src/<package_name>/ 안의 "배포에 필요한 모듈 파일"을 dist 루트로 flat 복사
      (현재 auto_map이 'image_processing_my_mnist.MyMNISTImageProcessor' 형태로 저장되는 상황을
       가장 확실히 만족시키는 방식)
    - examples/는 dist에 같이 넣어 테스트 편의 제공(원치 않으면 지워도 됨)
    - artifacts의 config/model/preprocessor_config를 dist 루트로 복사
    """
    rmtree_if_exists(dist_repo_dir)
    dist_repo_dir.mkdir(parents=True, exist_ok=True)

    # 1) 코드: src/<package_name> 에서 필요한 py를 dist 루트로 "flat" 복사
    src_pkg = root / "src" / package_name
    if not src_pkg.exists():
        raise FileNotFoundError(f"Missing src package dir: {src_pkg}")

    required_py = [
        "configuration_my_mnist.py",
        "modeling_my_mnist.py",
        "image_processing_my_mnist.py",
    ]
    for fname in required_py:
        copy2(src_pkg / fname, dist_repo_dir / fname)

    # __init__.py는 필수는 아니지만 같이 두면 디버깅에 유리
    if (src_pkg / "__init__.py").exists():
        copy2(src_pkg / "__init__.py", dist_repo_dir / "__init__.py")

    # 2) 예제/스크립트 복사 (편의용)
    copytree(root / "examples", dist_repo_dir / "examples")
    if copy_scripts_into_dist:
        # dist 안의 scripts는 "배포물"이라기보다 "참고용"입니다.
        copytree(root / "scripts", dist_repo_dir / "scripts")

    # 3) 메타 파일
    copy_if_exists(root / "README.md", dist_repo_dir / "README.md")
    copy_if_exists(root / "LICENSE", dist_repo_dir / "LICENSE")
    copy_if_exists(root / ".gitignore", dist_repo_dir / ".gitignore")
    copy_if_exists(root / "pyproject.toml", dist_repo_dir / "pyproject.toml")

    # 4) requirements.txt
    req_dst = dist_repo_dir / "requirements.txt"
    if requirements_text is not None:
        write_text(req_dst, requirements_text)
    else:
        if (root / "requirements.txt").exists():
            copy2(root / "requirements.txt", req_dst)
        else:
            write_text(
                req_dst,
                "\n".join(["torch", "torchvision", "transformers", "evaluate", "numpy", "Pillow", ""]),
            )

    # 5) 학습 산출물 복사
    config, model, preproc = validate_artifacts(artifact_dir)
    copy2(config, dist_repo_dir / "config.json")
    copy2(preproc, dist_repo_dir / "preprocessor_config.json")
    copy2(model, dist_repo_dir / model.name)

    # 안내 메모(선택)
    write_text(
        dist_repo_dir / "DIST_NOTE.txt",
        "\n".join(
            [
                "This folder is a Hub/local from_pretrained() compatible repo root.",
                "Dynamic module loading expects flat *.py modules at repo root (per auto_map).",
                "",
            ]
        ),
    )


# -----------------------------
# Hub upload (optional)
# -----------------------------
def upload_to_hub(
    *,
    dist_repo_dir: Path,
    repo_id: str,
    private: bool,
    repo_type: str,
    commit_message: str,
) -> None:
    try:
        from huggingface_hub import HfApi
        from huggingface_hub.utils import HfHubHTTPError
        from huggingface_hub.hf_api import HfFolder
    except Exception as e:
        raise RuntimeError(
            "huggingface_hub가 필요합니다. 설치 후 다시 실행하세요:\n"
            "  pip install -U huggingface_hub\n"
        ) from e

    token = HfFolder.get_token()
    if not token:
        raise RuntimeError(
            "Hugging Face 인증 토큰이 없습니다.\n"
            "다음 중 하나를 수행하세요:\n"
            "  1) 터미널에서: huggingface-cli login\n"
            "  2) 환경변수로: export HF_TOKEN=... (Colab/CI 포함)\n"
            "  3) 파이썬에서: from huggingface_hub import login; login('HF_TOKEN')\n"
        )

    api = HfApi()

    try:
        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True, private=private)
    except HfHubHTTPError as e:
        raise RuntimeError(
            f"Repo 생성에 실패했습니다: {repo_id} (repo_type={repo_type})\n"
            "권한(organization repo 여부), repo_id 오타, 토큰 권한을 확인하세요."
        ) from e

    try:
        api.upload_folder(
            repo_id=repo_id,
            repo_type=repo_type,
            folder_path=str(dist_repo_dir),
            commit_message=commit_message,
        )
    except HfHubHTTPError as e:
        raise RuntimeError(
            f"업로드에 실패했습니다: {repo_id} (repo_type={repo_type})\n"
            "토큰 권한, 대용량 파일, 네트워크 상태를 확인하세요."
        ) from e


# -----------------------------
# CLI
# -----------------------------
def build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Build dist repo root (Hub-ready) from src/<pkg> + artifacts, and optionally push to Hub."
    )

    p.add_argument("--push", action="store_true", help="지정 시 Hub에 업로드까지 수행합니다.")
    p.add_argument("--repo-id", default=None, help='Hub repo id. 예: "YOUR_ID/my-mnist-hf" ( --push일 때 필수 )')
    p.add_argument("--private", action="store_true", help="Hub repo를 private로 생성( --push일 때만 의미 )")
    p.add_argument("--repo-type", default="model", help='Hub repo type. 기본: "model"')
    p.add_argument("--commit-message", default="Release custom MNIST model", help="Hub 커밋 메시지")

    p.add_argument("--artifact-dir", default="artifacts/my_mnist", help="학습 산출물 폴더. 기본: artifacts/my_mnist")
    p.add_argument("--dist-dir", default="dist/my-mnist-hf", help="배포용 dist 폴더. 기본: dist/my-mnist-hf")

    p.add_argument(
        "--package-name",
        default="my_mnist_hf",
        help="src/ 아래의 패키지 폴더명. 기본: my_mnist_hf",
    )

    p.add_argument("--no-copy-scripts", action="store_true", help="dist에 scripts/ 복사를 생략")
    p.add_argument(
        "--requirements",
        default=None,
        help="requirements.txt 내용을 문자열로 지정(파일 생성). 미지정이면 루트 requirements.txt 복사 또는 기본값 생성.",
    )
    return p


def main() -> None:
    args = build_argparser().parse_args()

    root = project_root_from_this_file()
    artifact_dir = root / args.artifact_dir
    dist_repo_dir = root / args.dist_dir

    prepare_dist_repo_root(
        root=root,
        dist_repo_dir=dist_repo_dir,
        artifact_dir=artifact_dir,
        package_name=str(args.package_name),
        copy_scripts_into_dist=not bool(args.no_copy_scripts),
        requirements_text=args.requirements,
    )

    print("\n[OK] dist 폴더 생성 완료")
    print(f" - artifacts : {artifact_dir}")
    print(f" - dist repo : {dist_repo_dir}")
    print(f" - package  : {args.package_name}")

    if not args.push:
        print("\n[Info] --push를 지정하지 않아 Hub 업로드는 생략했습니다(local-only).")
        print("\n로컬 로드 예시:")
        print("  from transformers import AutoModelForImageClassification, AutoImageProcessor")
        print(f"  p = '{args.dist_dir}'")
        print("  processor = AutoImageProcessor.from_pretrained(p, trust_remote_code=True)")
        print("  model = AutoModelForImageClassification.from_pretrained(p, trust_remote_code=True)")
        return

    if not args.repo_id:
        raise ValueError("--push를 사용하려면 --repo-id를 반드시 지정해야 합니다.")

    upload_to_hub(
        dist_repo_dir=dist_repo_dir,
        repo_id=str(args.repo_id),
        private=bool(args.private),
        repo_type=str(args.repo_type),
        commit_message=str(args.commit_message),
    )

    print("\n[OK] Hub 업로드 완료")
    print(f" - hub repo : {args.repo_id}")
    print("\nHub 로드 예시:")
    print("  from transformers import AutoModelForImageClassification, AutoImageProcessor")
    print(f"  repo_id = '{args.repo_id}'")
    print("  processor = AutoImageProcessor.from_pretrained(repo_id, trust_remote_code=True)")
    print("  model = AutoModelForImageClassification.from_pretrained(repo_id, trust_remote_code=True)")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n오류: {e}\n", file=sys.stderr)
        raise