File size: 3,233 Bytes
22cec44
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
#!/usr/bin/env python
# coding=utf-8

"""

Script to install requirements in the correct order for the Phi-4 training project.

This ensures base requirements are installed first, followed by additional requirements.

"""

import os
import sys
import subprocess
import argparse
import logging
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

def install_requirements(include_flash=False):
    """Install requirements in the correct order."""
    current_dir = Path(__file__).parent
    base_req_path = current_dir / "requirements-base.txt"
    main_req_path = current_dir / "requirements.txt"
    flash_req_path = current_dir / "requirements-flash.txt"
    
    if not base_req_path.exists():
        logger.error(f"Base requirements file not found: {base_req_path}")
        return False
    
    if not main_req_path.exists():
        logger.error(f"Main requirements file not found: {main_req_path}")
        return False
    
    logger.info("Installing dependencies in sequential order...")
    
    try:
        # Step 1: Install base requirements
        logger.info(f"Step 1: Installing base requirements from {base_req_path}")
        subprocess.run([sys.executable, "-m", "pip", "install", "-r", str(base_req_path)], 
                      check=True)
        logger.info("Base requirements installed successfully")
        
        # Step 2: Install main requirements
        logger.info(f"Step 2: Installing additional requirements from {main_req_path}")
        subprocess.run([sys.executable, "-m", "pip", "install", "-r", str(main_req_path)], 
                      check=True)
        logger.info("Additional requirements installed successfully")
        
        # Step 3: Optionally install flash-attention
        if include_flash and flash_req_path.exists():
            logger.info(f"Step 3: Installing flash-attention from {flash_req_path}")
            subprocess.run([sys.executable, "-m", "pip", "install", "-r", str(flash_req_path), "--no-build-isolation"], 
                          check=True)
            logger.info("Flash-attention installed successfully")
        elif include_flash:
            logger.warning(f"Flash requirements file not found: {flash_req_path}")
        
        logger.info("All required packages installed successfully!")
        return True
    
    except subprocess.CalledProcessError as e:
        logger.error(f"Error installing dependencies: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Install requirements for Phi-4 training")
    parser.add_argument("--flash", action="store_true", help="Also install flash-attention (optional)")
    args = parser.parse_args()
    
    success = install_requirements(include_flash=args.flash)
    if success:
        logger.info("Installation completed successfully!")
    else:
        logger.error("Installation failed. Please check the logs for details.")
        sys.exit(1)

if __name__ == "__main__":
    main()