DeepSeek-R1 로드

 

import logging
from typing import Optional, Union, Dict, Any, Tuple

# transformers 라이브러리에서 필요한 클래스와 함수를 불러옵니다.
from transformers import (
    AutoModel,          # 사전 학습된 모델을 쉽게 불러올 수 있는 클래스
    AutoTokenizer,      # 사전 학습된 토크나이저를 쉽게 불러올 수 있는 클래스
    AutoConfig,         # 모델의 설정(config)을 불러오는 클래스
    PreTrainedModel,    # 사전 학습된 모델의 기본 클래스
    PreTrainedTokenizer # 사전 학습된 토크나이저의 기본 클래스
)

# 로깅 설정: 현재 모듈 이름으로 로거를 생성합니다.
logger = logging.getLogger(__name__)

def load_model_with_quantization_fallback(
    model_name: str = "deepseek-ai/DeepSeek-R1",  # 기본 모델 이름을 설정합니다.
    trust_remote_code: bool = True,                # 원격 저장소의 커스텀 코드를 신뢰할지 여부를 지정합니다.
    device_map: Optional[Union[str, Dict[str, Any]]] = "auto",  # 모델을 어느 장치에서 로드할지 설정합니다. ('auto'이면 자동으로 할당)
    **kwargs                                      # 추가적으로 전달할 인자들
) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
    """
    Hugging Face Hub에서 모델과 토크나이저를 로드하는 함수입니다.
    
    기본적으로 모델을 로드하는데, 만약 '알 수 없는 양자화(quantization) 타입'으로 인해 ValueError가 발생하면,
    설정(config)에서 'quantization_config'를 제거한 후 모델을 다시 로드합니다.
    
    :param model_name: 사전 학습된 모델의 경로나 식별자.
    :param trust_remote_code: 모델 리포지토리의 커스텀 코드를 신뢰할지 여부.
    :param device_map: 모델 로드시 사용할 장치 맵 (예: 'auto').
    :param kwargs: 추가적으로 AutoModel.from_pretrained에 전달할 인자들.
    :return: (모델, 토크나이저) 튜플.
    """
    
    try:
        # 첫 번째 시도로 모델을 로드합니다.
        logger.info(f"원격 코드 신뢰({trust_remote_code})로 '{model_name}' 모델을 로드합니다.")
        model = AutoModel.from_pretrained(
            model_name,
            trust_remote_code=trust_remote_code,
            device_map=device_map,
            **kwargs
        )
        # 모델과 함께 사용할 토크나이저도 로드합니다.
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=trust_remote_code
        )
        logger.info("원래 설정으로 모델을 성공적으로 로드했습니다.")
        # 성공적으로 로드한 모델과 토크나이저를 반환합니다.
        return model, tokenizer

    except ValueError as e:
        # ValueError가 발생했을 때, 에러 메시지에 'Unknown quantization type'이라는 문자열이 포함되어 있는지 확인합니다.
        if "Unknown quantization type" in str(e):
            logger.warning(
                "알 수 없는 양자화 타입이 감지되었습니다. "
                "'quantization_config'를 제거한 후 재시도합니다..."
            )
            
            try:
                # 모델의 설정(config)을 다시 로드합니다.
                config = AutoConfig.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code
                )
                # 로드된 설정에 'quantization_config' 속성이 있는지 확인합니다.
                if hasattr(config, "quantization_config"):
                    # 해당 속성이 존재하면 삭제합니다.
                    delattr(config, "quantization_config")
                    logger.info("'quantization_config'가 설정에서 제거되었습니다.")
                
                # 수정된 설정(config)을 사용하여 모델을 재로드합니다.
                model = AutoModel.from_pretrained(
                    model_name,
                    config=config,  # 수정된 config를 전달합니다.
                    trust_remote_code=trust_remote_code,
                    device_map=device_map,
                    **kwargs
                )
                # 토크나이저도 다시 로드합니다.
                tokenizer = AutoTokenizer.from_pretrained(
                    model_name,
                    trust_remote_code=trust_remote_code
                )
                logger.info("양자화 설정 없이 모델을 성공적으로 로드했습니다.")
                # 재시도에 성공하면 모델과 토크나이저를 반환합니다.
                return model, tokenizer

            except Exception as inner_e:
                # 내부 재시도 과정에서 다른 예외가 발생한 경우, 에러를 로깅하고 예외를 다시 발생시킵니다.
                logger.error(f"양자화 없이 모델 로드 실패: {inner_e}")
                raise inner_e
        
        # 만약 발생한 ValueError가 양자화 문제와 관련이 없다면, 에러를 로깅하고 다시 예외를 발생시킵니다.
        logger.error(f"모델 로딩 중 예상치 못한 ValueError 발생: {e}")
        raise

    except Exception as e:
        # ValueError 외에 다른 종류의 예외가 발생한 경우, 이를 로깅한 후 예외를 발생시킵니다.
        logger.error(f"예상치 못한 오류로 모델 로드 실패: {e}")
        raise

주요 동작 설명

1. 첫 번째 시도 (try 블록):
   • 주어진 model_name과 옵션을 사용하여 모델과 토크나이저를 로드합니다.
   • 성공하면 로드된 모델과 토크나이저를 반환합니다.
2. ValueError 예외 처리:
   • 만약 ValueError가 발생하고 에러 메시지에 "Unknown quantization type"이 포함되어 있으면,
     • 모델의 설정(config)을 재로드한 후 quantization_config 속성을 제거합니다.
     • 수정된 설정을 사용하여 모델을 다시 로드합니다.
     • 재시도에 성공하면 수정된 모델과 토크나이저를 반환합니다.
   • 다른 ValueError가 발생하면 에러를 기록하고 예외를 발생시킵니다.
3. 그 외의 예외 처리:
   • HTTP, I/O, 또는 기타 예상치 못한 오류가 발생하면 에러를 기록한 후 예외를 발생시킵니다.