Direction9/model/download_from_modelscope.py

#!/usr/bin/env python
"""
从魔搭（ModelScope）下载 LLM 模型

支持模型列表（示例）：
- Qwen2.5-7B-Instruct（已有，可重新下载）
- Qwen2.5-14B-Instruct
- Qwen2.5-32B-Instruct
- Qwen3.5-7B-Instruct（如果可用）
- Llama-3-8B-Instruct
- Llama-3.1-8B-Instruct

用法:
    # 下载 Qwen2.5-7B-Instruct
    python model/download_from_modelscope.py --model qwen2.5-7b

    # 下载 Llama-3-8B-Instruct
    python model/download_from_modelscope.py --model llama3-8b

    # 下载指定模型（自定义）
    python model/download_from_modelscope.py --model-id Qwen/Qwen2.5-7B-Instruct
"""

import os
import sys
import argparse
from pathlib import Path
from typing import Optional


# 模型映射表
MODEL_MAPPING = {
    # Qwen 系列
    "qwen2.5-7b": "Qwen/Qwen2.5-7B-Instruct",
    "qwen2.5-14b": "Qwen/Qwen2.5-14B-Instruct",
    "qwen2.5-32b": "Qwen/Qwen2.5-32B-Instruct",
    "qwen2.5-72b": "Qwen/Qwen2.5-72B-Instruct",
    "qwen3.5-7b": "Qwen/Qwen3.5-7B-Instruct",
    "qwen3.5-14b": "Qwen/Qwen3.5-14B-Instruct",
    "qwen3.5-32b": "Qwen/Qwen3.5-32B-Instruct",

    # Llama 系列
    "llama3-8b": "LLM-Research/Meta-Llama-3-8B-Instruct",
    "llama3-70b": "LLM-Research/Meta-Llama-3-70B-Instruct",
    "llama3.1-8b": "LLM-Research/Meta-Llama-3.1-8B-Instruct",
    "llama3.1-70b": "LLM-Research/Meta-Llama-3.1-70B-Instruct",
    "llama3.2-3b": "LLM-Research/Llama-3.2-3B-Instruct",  # 魔搭特供版本

    # Mistral 系列（已有）
    "mistral-7b": "LLM-Research/Mistral-7B-v0.3",

    # Gemma 系列
    "gemma-2b": "LLM-Research/gemma-2b-it",
    "gemma-7b": "LLM-Research/gemma-7b-it",
    "gemma-9b": "LLM-Research/gemma-2-9b-it",
}


def check_modelscope_installed():
    """检查 modelscope 是否安装"""
    try:
        from modelscope import snapshot_download
        return True
    except ImportError:
        print("错误：未安装 modelscope 库")
        print("请先运行以下命令安装：")
        print("  pip install modelscope")
        print("  或")
        print("  pip install modelscope -c https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.txt")
        return False


def download_model(model_id: str, output_dir: str, revision: str = "master") -> bool:
    """
    从魔搭下载模型

    Args:
        model_id: 模型 ID，如 "Qwen/Qwen2.5-7B-Instruct"
        output_dir: 输出目录
        revision: 模型版本

    Returns:
        下载是否成功
    """
    from modelscope import snapshot_download
    import inspect

    print(f"开始下载模型：{model_id}")
    print(f"输出目录：{output_dir}")
    print(f"版本：{revision}")
    print("-" * 50)

    try:
        # 检查 snapshot_download 支持的参数
        sig = inspect.signature(snapshot_download)
        params = {
            'model_id': model_id,
            'revision': revision,
            'local_dir': output_dir,
        }

        # 根据版本决定是否使用 symlinks 参数
        if 'local_dir_use_symlinks' in sig.parameters:
            params['local_dir_use_symlinks'] = False
        elif 'use_symlinks' in sig.parameters:
            params['use_symlinks'] = False

        # 使用 snapshot_download 下载模型
        model_dir = snapshot_download(**params)

        print(f"\n下载完成!")
        print(f"模型路径：{model_dir}")

        # 列出下载的文件
        print("\n下载的文件:")
        for root, dirs, files in os.walk(output_dir):
            level = root.replace(output_dir, '').count(os.sep)
            indent = ' ' * 2 * level
            print(f'{indent}{os.path.basename(root)}/')
            subindent = ' ' * 2 * (level + 1)
            for file in files[:10]:  # 只显示前 10 个文件
                print(f'{subindent}{file}')

        return True

    except Exception as e:
        print(f"\n下载失败：{e}")
        return False


def get_model_info(model_name: str) -> Optional[dict]:
    """获取模型信息"""
    if model_name in MODEL_MAPPING:
        return {
            "name": model_name,
            "model_id": MODEL_MAPPING[model_name],
            "status": "已支持"
        }
    return None


def list_available_models():
    """列出所有可用模型"""
    print("\n可用模型列表:")
    print("-" * 60)
    print(f"{'简写':<20} {'模型 ID':<35} {'状态'}")
    print("-" * 60)

    for name, model_id in MODEL_MAPPING.items():
        print(f"{name:<20} {model_id:<35} 已支持")

    print("-" * 60)
    print("\n使用方法:")
    print("  python model/download_from_modelscope.py --model <简写>")
    print("  python model/download_from_modelscope.py --model-id <完整模型 ID>")
    print()


def main():
    parser = argparse.ArgumentParser(
        description="从魔搭（ModelScope）下载 LLM 模型",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 列出所有可用模型
  python model/download_from_modelscope.py --list

  # 下载 Qwen2.5-7B-Instruct
  python model/download_from_modelscope.py --model qwen2.5-7b

  # 下载 Llama-3-8B-Instruct
  python model/download_from_modelscope.py --model llama3-8b

  # 使用完整模型 ID 下载
  python model/download_from_modelscope.py --model-id Qwen/Qwen2.5-7B-Instruct
        """
    )

    parser.add_argument(
        "--model",
        type=str,
        choices=list(MODEL_MAPPING.keys()),
        help="使用简写下载模型（如 qwen2.5-7b, llama3-8b）"
    )

    parser.add_argument(
        "--model-id",
        type=str,
        help="使用完整模型 ID 下载（如 Qwen/Qwen2.5-7B-Instruct）"
    )

    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="模型输出目录（默认：model/weights/{model_name}）"
    )

    parser.add_argument(
        "--revision",
        type=str,
        default="master",
        help="模型版本（默认：master）"
    )

    parser.add_argument(
        "--list",
        action="store_true",
        help="列出所有可用模型"
    )

    args = parser.parse_args()

    # 列出模型
    if args.list:
        list_available_models()
        return 0

    # 检查参数
    if not args.model and not args.model_id:
        print("错误：请指定 --model 或 --model-id")
        print("使用 --list 查看可用模型")
        return 1

    # 检查安装
    if not check_modelscope_installed():
        return 1

    # 确定模型 ID
    if args.model:
        model_id = MODEL_MAPPING[args.model]
        model_name = args.model
    else:
        model_id = args.model_id
        model_name = model_id.split("/")[-1]

    # 确定输出目录
    if args.output_dir:
        output_dir = args.output_dir
    else:
        output_dir = os.path.join("model/weights", model_name)

    # 创建目录
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    # 检查是否已存在
    config_file = os.path.join(output_dir, "config.json")
    if os.path.exists(config_file):
        print(f"警告：模型目录已存在：{output_dir}")
        response = input("是否覆盖下载？(y/N): ")
        if response.lower() != 'y':
            print("已取消")
            return 0

    # 下载模型
    success = download_model(model_id, output_dir, args.revision)

    if success:
        print("\n" + "=" * 50)
        print("下载成功!")
        print(f"\n下一步：更新 extract_representations.py 中的模型配置")
        print(f"添加新模型：{model_name}")
        print(f"模型路径：{output_dir}")
        print("=" * 50)
        return 0
    else:
        print("\n下载失败，请检查网络连接或模型 ID 是否正确")
        return 1


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