#!/bin/bash
# DCU DeepSeek-R1-Distill-Qwen-7B Model Install Bash
# Run with bash `curl -fsSL https://static.ksuser.cn/dcu/model.sh | bash -s --`
# env.sh Needed
# By KsuserKqy
# 2025-10-12

set -e  # 遇到错误立即退出

echo "=== 使用Conda环境下载DeepSeek-R1-Distill-Qwen-7B模型 ==="

# 检查Conda是否安装
if ! command -v conda &> /dev/null; then
    echo "错误: 未找到Conda，请先安装Conda"
    exit 1
fi

# 创建或激活Conda环境
ENV_NAME="dcu"
PYTHON_VERSION="3.10.18"
# 示例：在 Dockerfile 或终端中执行
echo "source /opt/conda/etc/profile.d/conda.sh" >> ~/.bash_profile
if conda env list | grep -q "$ENV_NAME"; then
    echo "激活现有Conda环境: $ENV_NAME"
    source /opt/conda/etc/profile.d/conda.sh
    # source /opt/conda/etc/profile.d/conda.sh
    # conda activate "$ENV_NAME"
    /opt/conda/bin/conda activate dcu
else
    source /opt/conda/etc/profile.d/conda.sh
    # echo "创建新的Conda环境: $ENV_NAME (Python $PYTHON_VERSION)"
    # source /opt/conda/etc/profile.d/conda.sh
    # conda create -n "$ENV_NAME" python="$PYTHON_VERSION" -y
    # conda activate "$ENV_NAME"
    /opt/conda/bin/conda activate dcu
fi
# 创建模型存储目录
MODEL_DIR="./models/DeepSeek-R1-Distill-Qwen-7B"
mkdir -p "$MODEL_DIR"
echo "模型将下载到: $(realpath "$MODEL_DIR")"

# 安装必要的Python包
echo "=== 安装依赖包 ==="
pip install --upgrade pip
pip install modelscope transformers torch accelerate sentencepiece
pip install protobuf scipy numpy tqdm

# 使用ModelScope下载模型
echo "=== 使用ModelScope下载模型 ==="
cat > download_model.py << 'EOF'
import os
import sys
from modelscope import snapshot_download
from pathlib import Path

def download_model():
    model_id = 'deepseek-ai/DeepSeek-R1-Distill-Qwen-7B'
    model_dir = './models/DeepSeek-R1-Distill-Qwen-7B'
    
    # 确保目录存在
    Path(model_dir).mkdir(parents=True, exist_ok=True)
    
    print(f"开始下载模型: {model_id}")
    print(f"下载目录: {os.path.abspath(model_dir)}")
    
    try:
        # 下载模型
        model_path = snapshot_download(
            model_id=model_id,
            cache_dir=model_dir,
            revision='master',
            ignore_file_pattern=[".*bin"],  # 可选：忽略某些文件以加快下载
            resume_download=True  # 支持断点续传
        )
        
        print(f"✅ 模型下载完成！")
        print(f"模型路径: {model_path}")
        return model_path
        
    except Exception as e:
        print(f"❌ 下载失败: {e}")
        # 尝试使用Hugging Face作为备选
        return download_huggingface_backup()

def download_huggingface_backup():
    """备选方案：使用Hugging Face下载"""
    print("尝试使用Hugging Face下载...")
    
    from transformers import AutoTokenizer
    import torch
    
    model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
    local_dir = "./models/DeepSeek-R1-Distill-Qwen-7B"
    
    try:
        # 下载tokenizer
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            cache_dir=local_dir,
            trust_remote_code=True
        )
        
        print("✅ Tokenizer下载成功")
        
        # 尝试下载模型配置（不加载权重）
        from transformers import AutoConfig
        config = AutoConfig.from_pretrained(
            model_name,
            cache_dir=local_dir,
            trust_remote_code=True
        )
        
        print("✅ 模型配置下载成功")
        return local_dir
        
    except Exception as e:
        print(f"❌ Hugging Face下载也失败: {e}")
        raise

if __name__ == "__main__":
    try:
        model_path = download_model()
        sys.exit(0)
    except Exception as e:
        print(f"所有下载方法都失败: {e}")
        sys.exit(1)
EOF

# 运行下载脚本
python download_model.py

# 验证下载结果
echo "=== 验证下载结果 ==="
cat > verify_download.py << 'EOF'
import os
import sys
from pathlib import Path

def verify_model():
    model_dir = "./models/DeepSeek-R1-Distill-Qwen-7B"
    model_path = Path(model_dir)
    
    if not model_path.exists():
        print(f"错误: 模型目录不存在: {model_dir}")
        return False
    
    # 检查必要的文件
    required_patterns = [
        "config.json",
        "*.safetensors",  # 权重文件可能是safetensors格式
        "tokenizer.json",
        "tokenizer_config.json"
    ]
    
    print("模型目录内容:")
    found_files = []
    for file in model_path.iterdir():
        print(f"  - {file.name}")
        found_files.append(file.name)
    
    # 检查是否有权重文件
    has_weights = any(file.suffix in ['.bin', '.safetensors'] for file in model_path.iterdir())
    
    if not has_weights:
        print("警告: 未找到模型权重文件")
        return False
    
    # 检查文件大小
    total_size = 0
    for file in model_path.iterdir():
        if file.is_file():
            size = file.stat().st_size / (1024**3)  # GB
            total_size += size
            print(f"  {file.name}: {size:.2f} GB")
    
    print(f"模型总大小: {total_size:.2f} GB")
    
    # 检查是否有大文件（模型权重）
    if total_size < 1.0:  # 如果小于1GB，可能下载不完整
        print("警告: 模型文件可能不完整")
        return False
    
    return True

if __name__ == "__main__":
    success = verify_model()
    sys.exit(0 if success else 1)
EOF

python verify_download.py

echo "=== 下载完成！ ==="
echo "模型路径: ./models/DeepSeek-R1-Distill-Qwen-7B"
echo "Conda环境名称: $ENV_NAME"
echo "使用以下命令激活环境: conda activate $ENV_NAME"

# 创建环境配置脚本
cat > setup_environment.sh << 'EOF'
#!/bin/bash
# 激活Conda环境并设置路径
conda activate deepseek_quant
export MODEL_PATH="./models/DeepSeek-R1-Distill-Qwen-7B"
echo "环境已设置，模型路径: $MODEL_PATH"
EOF

chmod +x setup_environment.sh

echo "=== 使用以下命令设置环境: source setup_environment.sh ==="