"""环境检查与依赖验证脚本"""
import subprocess
import sys


def check_and_install():
    """检查并安装依赖"""
    print("=" * 60)
    print("Qwen3.5-9B 测试环境检查")
    print("=" * 60)

    # 检查 Python 版本
    print(f"\nPython 版本: {sys.version}")

    # 检查 CUDA
    try:
        import torch
        print(f"PyTorch 版本: {torch.__version__}")
        print(f"CUDA 可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"GPU: {torch.cuda.get_device_name(0)}")
            vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
            print(f"VRAM: {vram_gb:.1f} GB")
    except ImportError:
        print("ERROR: PyTorch 未安装")
        sys.exit(1)

    # 安装依赖
    print("\n安装依赖包...")
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", "-r",
        "vsp/qwen3.5-9b/requirements.txt", "-q"
    ])

    # 验证关键包
    packages = ["transformers", "modelscope", "accelerate", "bitsandbytes"]
    for pkg in packages:
        try:
            mod = __import__(pkg)
            ver = getattr(mod, "__version__", "unknown")
            print(f"  {pkg}: {ver}")
        except ImportError:
            print(f"  ERROR: {pkg} 安装失败")

    print("\n环境检查完成！")


if __name__ == "__main__":
    check_and_install()