download_model.py

#!/usr/bin/env python3
"""
Docker 빌드 시 Hugging Face 모델 다운로드 스크립트
"""

import os
import logging
from transformers import AutoTokenizer, AutoModelForCausalLM
from pathlib import Path

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def download_model():
    """Hugging Face Hub에서 모델 다운로드"""
    
    model_name = "gbrabbit/lily-math-model"
    cache_dir = "/app/cache/transformers"
    
    logger.info(f"📥 모델 다운로드 시작: {model_name}")
    
    try:
        # 캐시 디렉토리 생성
        Path(cache_dir).mkdir(parents=True, exist_ok=True)
        
        # 토크나이저 다운로드
        logger.info("🔤 토크나이저 다운로드 중...")
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=True,
            cache_dir=cache_dir
        )
        logger.info("✅ 토크나이저 다운로드 완료")
        
        # 모델 다운로드 (가중치만)
        logger.info("🧠 모델 다운로드 중...")
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            trust_remote_code=True,
            cache_dir=cache_dir,
            torch_dtype="auto",  # 메모리 절약
            low_cpu_mem_usage=True
        )
        logger.info("✅ 모델 다운로드 완료")
        
        # 메모리 정리
        del model
        del tokenizer
        
        logger.info("🎉 모델 다운로드 및 캐시 완료")
        
    except Exception as e:
        logger.error(f"❌ 모델 다운로드 실패: {e}")
        raise

if __name__ == "__main__":
    download_model()