FineTurnX/quickstart.py

"""
FineTuneX 快速启动脚本

使用方法:
    python quickstart.py

这个脚本将:
1. 检查环境
2. 初始化项目
3. 运行示例微调
"""

import os
import sys
import subprocess


def print_header(text):
    """打印标题"""
    print("\n" + "=" * 60)
    print(text.center(60))
    print("=" * 60 + "\n")


def run_command(command, description):
    """运行命令"""
    print(f"正在 {description}...")
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    
    if result.returncode == 0:
        print(f"✓ {description} 完成")
        if result.stdout:
            print(result.stdout)
        return True
    else:
        print(f"✗ {description} 失败")
        if result.stderr:
            print(result.stderr)
        return False


def check_environment():
    """检查环境"""
    print_header("步骤 1: 检查环境")
    
    # 检查 Python 版本
    python_version = sys.version_info
    if python_version.major < 3 or python_version.minor < 9:
        print("✗ Python 版本过低，需要 3.9+")
        return False
    
    print(f"✓ Python {python_version.major}.{python_version.minor}.{python_version.micro}")
    
    # 检查依赖
    try:
        import torch
        print(f"✓ PyTorch {torch.__version__}")
        
        if torch.cuda.is_available():
            print(f"✓ CUDA 可用：{torch.cuda.get_device_name(0)}")
        else:
            print("⚠ CUDA 不可用，将使用 CPU")
    except ImportError:
        print("✗ PyTorch 未安装")
        return False
    
    return True


def install_dependencies():
    """安装依赖"""
    print_header("步骤 2: 安装依赖")
    
    if os.path.exists("requirements.txt"):
        response = input("是否安装依赖？(y/n): ")
        if response.lower() == 'y':
            run_command("pip install -r requirements.txt", "安装依赖")
    else:
        print("⚠ requirements.txt 不存在")


def initialize_project():
    """初始化项目"""
    print_header("步骤 3: 初始化项目")
    
    if os.path.exists("scripts/init_project.py"):
        run_command("python scripts/init_project.py", "初始化项目")
    else:
        # 手动创建目录
        directories = ["outputs", "data", "logs"]
        for directory in directories:
            os.makedirs(directory, exist_ok=True)
            print(f"✓ 创建目录：{directory}")


def run_example():
    """运行示例"""
    print_header("步骤 4: 运行示例")
    
    print("可运行的示例:")
    print("1. Qwen3.5 微调示例")
    print("2. 环境检查")
    print("3. 运行测试")
    
    choice = input("\n请选择 (1-3): ")
    
    if choice == "1":
        if os.path.exists("examples/qwen3.5_0.8b_finetune.py"):
            print("\n开始运行 Qwen3.5 微调示例...")
            print("注意：这将下载模型并开始训练，可能需要较长时间")
            response = input("是否继续？(y/n): ")
            if response.lower() == 'y':
                run_command("python examples/qwen3.5_0.8b_finetune.py", "运行示例")
        else:
            print("✗ 示例文件不存在")
    
    elif choice == "2":
        if os.path.exists("scripts/check_env.py"):
            run_command("python scripts/check_env.py", "环境检查")
    
    elif choice == "3":
        if os.path.exists("tests/test_all.py"):
            run_command("python tests/test_all.py", "运行测试")
    
    else:
        print("无效选择")


def show_next_steps():
    """显示下一步"""
    print_header("完成!")
    
    print("""
项目已成功设置！接下来你可以:

1. 查看文档:
   - README.md - 项目概述
   - QUICKSTART.md - 快速开始
   - docs/usage.md - 详细使用文档

2. 运行示例:
   python examples/qwen3.5_0.8b_finetune.py

3. 使用自己的数据:
   - 准备数据文件 (JSON 格式)
   - 修改示例脚本中的数据集路径
   - 运行微调

4. 启动 API 服务:
   python scripts/start_api.py

5. 模型推理:
   python scripts/inference.py --model_path ./outputs/xxx --interactive

祝你使用愉快！
""")


def main():
    """主函数"""
    print_header("FineTuneX 快速启动")
    
    # 1. 检查环境
    if not check_environment():
        print("\n请先安装必要的依赖:")
        print("  pip install -r requirements.txt")
        return
    
    # 2. 安装依赖（可选）
    # install_dependencies()
    
    # 3. 初始化项目
    initialize_project()
    
    # 4. 运行示例
    run_example()
    
    # 5. 显示下一步
    show_next_steps()


if __name__ == "__main__":
    main()