File size: 4,931 Bytes
47beeae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
decfb95
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env python
# coding=utf-8

"""
Installation script for Hugging Face Space setup.
This script ensures all dependencies are installed correctly 
during the Space build process.
"""

import unsloth
import os
import sys
import subprocess
import logging
import traceback
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 run_command(cmd, description=""):
    """Run a shell command and log the output."""
    logger.info(f"Running: {description if description else cmd}")
    try:
        process = subprocess.run(
            cmd, 
            shell=True, 
            check=True,
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE,
            text=True
        )
        logger.info(f"Command output: {process.stdout}")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Command failed with exit code {e.returncode}")
        logger.error(f"Error output: {e.stderr}")
        return False

def install_dependencies():
    """Install all required dependencies in the correct order."""
    current_dir = Path(__file__).parent
    req_path = current_dir / "requirements.txt"
    
    if not req_path.exists():
        logger.error(f"Requirements file not found: {req_path}")
        return False
    
    try:
        # Step 1: Upgrade pip
        run_command(f"{sys.executable} -m pip install --upgrade pip", "Upgrading pip")
        
        # Step 2: Install direct torch version for CUDA compatibility
        run_command(
            f"{sys.executable} -m pip install torch>=2.0.0,<2.2.0 --extra-index-url https://download.pytorch.org/whl/cu118",
            "Installing PyTorch with CUDA support"
        )
        
        # Step 3: Install base dependencies
        run_command(
            f"{sys.executable} -m pip install transformers accelerate bitsandbytes peft einops",
            "Installing ML dependencies"
        )
        
        # Step 4: Install unsloth separately
        run_command(
            f"{sys.executable} -m pip install unsloth>=2024.3",
            "Installing Unsloth"
        )
        
        # Step 5: Install all remaining requirements
        run_command(
            f"{sys.executable} -m pip install -r {req_path}",
            "Installing all requirements"
        )
        
        # Verify critical packages
        import_check = verify_imports()
        if not import_check:
            logger.error("Failed to verify critical packages")
            return False
        
        logger.info("All dependencies installed successfully!")
        return True
        
    except Exception as e:
        logger.error(f"Error installing dependencies: {str(e)}")
        traceback.print_exc()
        return False

def verify_imports():
    """Verify that critical packages can be imported."""
    critical_packages = [
        "torch", "transformers", "unsloth", "peft", 
        "gradio", "accelerate", "bitsandbytes"
    ]
    
    success = True
    for package in critical_packages:
        try:
            module = __import__(package)
            version = getattr(module, "__version__", "unknown")
            logger.info(f"Successfully imported {package} (version: {version})")
        except ImportError:
            logger.error(f"CRITICAL: Failed to import {package}")
            success = False
        except Exception as e:
            logger.error(f"Error verifying {package}: {str(e)}")
            success = False
    
    # Check CUDA
    try:
        import torch
        cuda_available = torch.cuda.is_available()
        device_count = torch.cuda.device_count() if cuda_available else 0
        if cuda_available:
            device_name = torch.cuda.get_device_name(0)
            logger.info(f"CUDA available - Devices: {device_count}, Name: {device_name}")
        else:
            logger.warning(f"CUDA not available - This might affect performance")
    except Exception as e:
        logger.error(f"Error checking CUDA: {str(e)}")
    
    return success

def main():
    logger.info("Starting installation for Phi-4 Unsloth Training Space")
    
    try:
        # Install dependencies
        if not install_dependencies():
            logger.error("Failed to install dependencies")
            sys.exit(1)
        
        # Create marker file to show successful installation
        with open("INSTALL_SUCCESS.txt", "w") as f:
            f.write("Installation completed successfully")
        
        logger.info("Installation completed successfully")
        return 0
        
    except Exception as e:
        logger.error(f"Installation failed with error: {str(e)}")
        traceback.print_exc()
        return 1

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