project0802shenji/src/model/downloader.py

# -*- coding: utf-8 -*-
"""
模型下载器模块

负责从ModelScope或HuggingFace下载模型。
"""

import os
from pathlib import Path
from typing import Optional

from ..config import Config
from ..utils.logger import get_logger


class ModelDownloader:
    """模型下载器"""
    
    def __init__(self, config: Config = None, model_key: str = None):
        self.config = config or Config()
        
        # 设置当前模型
        if model_key:
            self.config.set_current_model(model_key)
        
        # 获取模型配置
        self.model_config = self.config.get_current_model_config()
        self.logger = get_logger(self.__class__.__name__)
        
        # 设置环境变量
        self.config.setup_environment()
    
    def download_model(self, model_name: Optional[str] = None, 
                      force_download: bool = False) -> Path:
        """下载模型"""
        if model_name is None:
            model_name = self.config.get_current_model_key()
        
        model_path = self.get_model_path(model_name)
        
        # 检查模型是否已存在
        if self.check_model_exists(model_name) and not force_download:
            self.logger.info(f"模型已存在: {model_path}")
            return model_path
        
        self.logger.info(f"开始下载模型: {model_name}")
        
        try:
            if self.config.USE_MODELSCOPE:
                return self._download_from_modelscope(model_name, model_path)
            else:
                return self._download_from_huggingface(model_name, model_path)
        except Exception as e:
            self.logger.error(f"模型下载失败: {e}")
            # 如果ModelScope失败，尝试HuggingFace
            if self.config.USE_MODELSCOPE:
                self.logger.info("尝试从HuggingFace下载...")
                return self._download_from_huggingface(model_name, model_path)
            raise
    
    def _download_from_modelscope(self, model_name: str, model_path: Path) -> Path:
        """从ModelScope下载模型"""
        # 首先尝试git clone方式下载
        try:
            return self._download_from_modelscope_git(model_name, model_path)
        except Exception as git_error:
            self.logger.warning(f"Git下载失败: {git_error}，尝试SDK方式")
            
        # 如果git方式失败，回退到SDK方式
        try:
            return self._download_from_modelscope_sdk(model_name, model_path)
        except Exception as sdk_error:
            self.logger.error(f"SDK下载也失败: {sdk_error}")
            raise
    
    def _download_from_modelscope_git(self, model_name: str, model_path: Path) -> Path:
        """使用git clone从ModelScope下载模型"""
        import subprocess
        import shutil
        
        # 获取ModelScope下载ID
        download_id = self.config.get_model_id_for_download('modelscope', model_name)
        
        if not download_id:
            # 回退到原有映射逻辑
            modelscope_names = {
                "Qwen/Qwen2-1.5B": "qwen/Qwen2-1.5B",
                "Qwen/Qwen2-1.5B-Instruct": "qwen/Qwen2-1.5B-Instruct",
                "Qwen/Qwen2.5-1.5B-Instruct": "qwen/Qwen2.5-1.5B-Instruct"
            }
            download_id = modelscope_names.get(model_name, model_name)
        
        ms_model_name = download_id
        git_url = f"https://www.modelscope.cn/{ms_model_name}.git"
        
        self.logger.info(f"使用git clone从ModelScope下载: {git_url}")
        
        # 检查git是否可用
        try:
            subprocess.run(["git", "--version"], check=True, capture_output=True)
        except (subprocess.CalledProcessError, FileNotFoundError):
            raise Exception("Git未安装或不可用")
        
        # 确保目标目录存在
        model_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 如果目标目录已存在，先删除
        if model_path.exists():
            shutil.rmtree(model_path)
        
        # 执行git clone
        try:
            cmd = ["git", "clone", git_url, str(model_path)]
            result = subprocess.run(
                cmd, 
                check=True, 
                capture_output=True, 
                text=True,
                timeout=1800  # 30分钟超时
            )
            
            self.logger.info(f"Git clone成功: {model_path}")
            return model_path
            
        except subprocess.TimeoutExpired:
            raise Exception("Git clone超时")
        except subprocess.CalledProcessError as e:
            raise Exception(f"Git clone失败: {e.stderr}")
    
    def _download_from_modelscope_sdk(self, model_name: str, model_path: Path) -> Path:
        """使用SDK从ModelScope下载模型"""
        try:
            # 设置SSL配置以解决证书验证问题
            import ssl
            import urllib3
            import requests
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
            
            # 彻底禁用SSL验证
            original_context = ssl._create_default_https_context
            ssl._create_default_https_context = ssl._create_unverified_context
            
            # Monkey patch urllib3 to disable SSL verification
            import urllib3.util.ssl_
            original_ssl_wrap_socket = urllib3.util.ssl_.ssl_wrap_socket
            
            def patched_ssl_wrap_socket(*args, **kwargs):
                kwargs['cert_reqs'] = ssl.CERT_NONE
                kwargs['check_hostname'] = False
                return original_ssl_wrap_socket(*args, **kwargs)
            
            urllib3.util.ssl_.ssl_wrap_socket = patched_ssl_wrap_socket
            
            try:
                from modelscope import snapshot_download
                
                # 获取ModelScope下载ID
                download_id = self.config.get_model_id_for_download('modelscope', model_name)
                
                if not download_id:
                    # 回退到原有映射逻辑
                    modelscope_names = {
                        "Qwen/Qwen2-1.5B": "qwen/Qwen2-1.5B",
                        "Qwen/Qwen2-1.5B-Instruct": "qwen/Qwen2-1.5B-Instruct",
                        "Qwen/Qwen2.5-1.5B-Instruct": "qwen/Qwen2.5-1.5B-Instruct"
                    }
                    download_id = modelscope_names.get(model_name, model_name)
                
                ms_model_name = download_id
                
                self.logger.info(f"使用SDK从ModelScope下载: {ms_model_name}")
                
                downloaded_path = snapshot_download(
                    model_id=ms_model_name,
                    cache_dir=str(self.config.MODELS_DIR),
                    local_dir=str(model_path)
                )
                
                self.logger.info(f"SDK下载完成: {downloaded_path}")
                return Path(downloaded_path)
                
            finally:
                # 恢复原始函数
                ssl._create_default_https_context = original_context
                urllib3.util.ssl_.ssl_wrap_socket = original_ssl_wrap_socket
            
        except ImportError:
            self.logger.error("ModelScope未安装，请安装: pip install modelscope")
            raise
    
    def _download_from_huggingface(self, model_name: str, model_path: Path) -> Path:
        """从HuggingFace下载模型"""
        try:
            from transformers import AutoTokenizer, AutoModelForCausalLM
            
            self.logger.info(f"从HuggingFace下载: {model_name}")
            
            # 下载分词器
            tokenizer = AutoTokenizer.from_pretrained(
                model_name,
                cache_dir=str(self.config.CACHE_DIR),
                trust_remote_code=True
            )
            tokenizer.save_pretrained(str(model_path))
            
            # 下载模型
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                cache_dir=str(self.config.CACHE_DIR),
                trust_remote_code=True
            )
            model.save_pretrained(str(model_path))

            self.logger.info(f"模型下载完成: {model_path}")
            return model_path
            
        except Exception as e:
            self.logger.error(f"HuggingFace下载失败: {e}")
            raise
    
    def check_model_exists(self, model_name: Optional[str] = None) -> bool:
        """检查模型是否存在"""
        if model_name is None:
            model_name = self.config.get_current_model_key()
        
        model_path = self.get_model_path(model_name)
        
        # 检查关键文件是否存在
        required_files = ["config.json", "tokenizer.json"]
        
        if not model_path.exists():
            return False
        
        for file_name in required_files:
            if not (model_path / file_name).exists():
                return False
        
        return True
    
    def get_model_path(self, model_name: Optional[str] = None) -> Path:
        """获取模型路径"""
        if model_name is None:
            model_name = self.model_config.name
        
        # 尝试使用配置中的路径方法
        if hasattr(self.config, 'get_model_path'):
            return Path(self.config.get_model_path(model_name))
        
        # 优先检查原始名称（去掉组织前缀）
        simple_name = model_name.split("/")[-1] if "/" in model_name else model_name
        simple_path = self.config.MODELS_DIR / simple_name
        
        # 如果简单名称的目录存在，使用它
        if simple_path.exists():
            return simple_path
        
        # 否则使用下划线替换的名称
        return self.config.MODELS_DIR / model_name.replace("/", "_")
    
    def get_model_info(self, model_name: Optional[str] = None) -> dict:
        """获取模型信息"""
        model_path = self.get_model_path(model_name)
        
        if not self.check_model_exists(model_name):
            return {"exists": False, "path": str(model_path)}
        
        # 计算模型大小
        total_size = 0
        file_count = 0
        
        for file_path in model_path.rglob("*"):
            if file_path.is_file():
                total_size += file_path.stat().st_size
                file_count += 1
        
        return {
            "exists": True,
            "path": str(model_path),
            "size_mb": total_size / (1024 * 1024),
            "file_count": file_count
        }