upload_to_hf.py

#!/usr/bin/env python3
"""
Hugging Face Hub 업로드 스크립트
- 병합된 EEVE 모델을 Hugging Face Hub에 업로드
- 모델 카드(README.md) 포함
"""

import os
import argparse
from pathlib import Path
from huggingface_hub import HfApi, create_repo, upload_folder


def upload_model_to_hub(
    model_dir: str,
    repo_id: str,
    token: str = None,
    private: bool = False,
    commit_message: str = "Upload EEVE Korean Custom model"
):
    """
    모델을 Hugging Face Hub에 업로드
    
    Args:
        model_dir: 업로드할 모델 디렉토리 경로
        repo_id: Hugging Face 리포지토리 ID (username/model-name)
        token: Hugging Face API 토큰 (None이면 환경변수 또는 CLI 로그인 사용)
        private: Private 리포지토리로 생성할지 여부
        commit_message: 커밋 메시지
    """
    
    print("\n" + "="*80)
    print(" Hugging Face Hub 업로드")
    print("="*80)
    print(f"📁 모델 디렉토리: {model_dir}")
    print(f"🎯 리포지토리: {repo_id}")
    print(f"🔒 공개 여부: {'Private' if private else 'Public'}")
    print("="*80 + "\n")
    
    # HfApi 초기화
    api = HfApi(token=token)
    
    # 1. 리포지토리 생성 (이미 있으면 스킵)
    print("1️⃣ 리포지토리 생성 중...")
    try:
        repo_url = create_repo(
            repo_id=repo_id,
            token=token,
            private=private,
            exist_ok=True,
            repo_type="model"
        )
        print(f"✓ 리포지토리: {repo_url}\n")
    except Exception as e:
        print(f"⚠️  리포지토리가 이미 존재하거나 오류 발생: {e}\n")
    
    # 2. 파일 업로드
    print("2️⃣ 모델 파일 업로드 중...")
    print("   ⏱️  이 작업은 시간이 걸립니다 (모델 크기: ~20GB)...\n")
    
    try:
        upload_folder(
            repo_id=repo_id,
            folder_path=model_dir,
            token=token,
            commit_message=commit_message,
            repo_type="model"
        )
        print("✓ 업로드 완료!\n")
    except Exception as e:
        print(f"❌ 업로드 실패: {e}")
        raise
    
    # 3. 완료
    print("="*80)
    print("✅ 업로드 성공!")
    print("="*80)
    print(f"\n🔗 모델 페이지: https://huggingface.co/{repo_id}")
    print(f"📖 사용 방법:")
    print(f"""
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("{repo_id}")
tokenizer = AutoTokenizer.from_pretrained("{repo_id}")
""")
    print("="*80 + "\n")


def main():
    parser = argparse.ArgumentParser(
        description="EEVE 모델을 Hugging Face Hub에 업로드",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
사용 예시:

# 기본 사용 (public)
python upload_to_hf.py --repo-id username/model-name

# Private 리포지토리로 업로드
python upload_to_hf.py --repo-id username/model-name --private

# 토큰 직접 지정
python upload_to_hf.py --repo-id username/model-name --token hf_xxxxx

# 다른 디렉토리에서 업로드
python upload_to_hf.py --repo-id username/model-name --model-dir /path/to/model

전체 옵션:
python upload_to_hf.py \\
    --repo-id MyeongHo0621/EEVE-Korean-Custom-10.8B \\
    --model-dir /home/work/tesseract/eeve_hf_upload \\
    --private \\
    --commit-message "Initial upload: checkpoint-500 merged"

주의사항:
1. 먼저 Hugging Face에 로그인해야 합니다:
   huggingface-cli login
   또는
   hf auth login

2. 토큰은 Write 권한이 있어야 합니다:
   https://huggingface.co/settings/tokens

3. 업로드 시간: ~20-30분 (네트워크 속도에 따라)
        """
    )
    
    parser.add_argument(
        "--repo-id",
        type=str,
        required=True,
        help="Hugging Face 리포지토리 ID (username/model-name)"
    )
    
    parser.add_argument(
        "--model-dir",
        type=str,
        default="/home/work/tesseract/eeve_hf_upload",
        help="업로드할 모델 디렉토리 (기본: eeve_hf_upload)"
    )
    
    parser.add_argument(
        "--token",
        type=str,
        default=None,
        help="Hugging Face API 토큰 (선택, 없으면 CLI 로그인 사용)"
    )
    
    parser.add_argument(
        "--private",
        action="store_true",
        help="Private 리포지토리로 생성"
    )
    
    parser.add_argument(
        "--commit-message",
        type=str,
        default="Upload EEVE Korean Custom model (checkpoint-500 merged)",
        help="커밋 메시지"
    )
    
    args = parser.parse_args()
    
    # 모델 디렉토리 확인
    model_dir = Path(args.model_dir)
    if not model_dir.exists():
        print(f"❌ 오류: 모델 디렉토리를 찾을 수 없습니다: {model_dir}")
        return 1
    
    # 필수 파일 확인
    required_files = ["config.json", "tokenizer_config.json"]
    missing_files = [f for f in required_files if not (model_dir / f).exists()]
    if missing_files:
        print(f"❌ 오류: 필수 파일이 없습니다: {missing_files}")
        return 1
    
    # 업로드 실행
    try:
        upload_model_to_hub(
            model_dir=str(model_dir),
            repo_id=args.repo_id,
            token=args.token,
            private=args.private,
            commit_message=args.commit_message
        )
        print("✅ 모든 작업이 완료되었습니다!")
        return 0
        
    except Exception as e:
        print(f"\n❌ 오류 발생: {e}")
        import traceback
        traceback.print_exc()
        return 1


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