setup_inference.py

#!/usr/bin/env python3
"""
Setup script for Math-Eval inference environment.
Handles both API-based and open-source model requirements.
"""

import subprocess
import sys
import argparse
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def run_command(cmd, description):
    """Run a command and handle errors."""
    logger.info(f"Running: {description}")
    try:
        result = subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
        if result.stdout:
            logger.info(result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed: {description}")
        if e.stderr:
            logger.error(e.stderr)
        return False

def install_base_requirements():
    """Install base requirements for inference."""
    logger.info("Installing base inference requirements...")
    cmd = f"{sys.executable} -m pip install pandas pillow tqdm numpy"
    return run_command(cmd, "Installing base packages")

def install_api_requirements():
    """Install requirements for API-based models."""
    logger.info("Installing API model requirements...")
    cmd = f"{sys.executable} -m pip install openai google-generativeai"
    return run_command(cmd, "Installing API model packages")

def install_opensource_requirements():
    """Install requirements for open-source models."""
    logger.info("Installing open-source model requirements...")
    
    # Check if CUDA is available
    try:
        import torch
        cuda_available = torch.cuda.is_available()
        logger.info(f"CUDA available: {cuda_available}")
    except ImportError:
        cuda_available = False
        logger.info("PyTorch not installed yet, assuming CUDA setup")
    
    # Install PyTorch with appropriate CUDA support
    if cuda_available or True:  # Assume CUDA for now
        torch_cmd = f"{sys.executable} -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118"
    else:
        torch_cmd = f"{sys.executable} -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu"
    
    if not run_command(torch_cmd, "Installing PyTorch"):
        return False
    
    # Install transformers and related packages
    transformers_cmd = f"{sys.executable} -m pip install transformers>=4.30.0 accelerate>=0.20.0"
    if not run_command(transformers_cmd, "Installing Transformers"):
        return False
    
    # Install model-specific requirements
    model_specific_cmd = f"{sys.executable} -m pip install qwen-vl-utils einops bitsandbytes"
    if not run_command(model_specific_cmd, "Installing model-specific packages"):
        return False
    
    # Optional: Install flash attention for better performance
    flash_attn_cmd = f"{sys.executable} -m pip install flash-attn --no-build-isolation"
    run_command(flash_attn_cmd, "Installing flash attention (optional)")
    
    return True

def setup_inference_environment():
    """Setup the complete inference environment."""
    logger.info("Setting up Math-Eval inference environment...")
    
    # Always install base requirements
    if not install_base_requirements():
        logger.error("Failed to install base requirements")
        return False
    
    logger.info("✅ Base requirements installed successfully")
    return True

def validate_installation(model_type):
    """Validate that the installation was successful."""
    logger.info(f"Validating {model_type} installation...")
    
    try:
        import pandas
        import PIL
        import tqdm
        import numpy
        logger.info("✅ Base packages validated")
    except ImportError as e:
        logger.error(f"❌ Base package validation failed: {e}")
        return False
    
    if model_type in ['api', 'all']:
        try:
            import openai
            import google.generativeai
            logger.info("✅ API packages validated")
        except ImportError as e:
            logger.error(f"❌ API package validation failed: {e}")
            return False
    
    if model_type in ['opensource', 'all']:
        try:
            import torch
            import transformers
            logger.info(f"✅ Open-source packages validated (CUDA: {torch.cuda.is_available()})")
        except ImportError as e:
            logger.error(f"❌ Open-source package validation failed: {e}")
            return False
    
    return True

def main():
    parser = argparse.ArgumentParser(description="Setup Math-Eval inference environment")
    parser.add_argument('--model_type', type=str, choices=['api', 'opensource', 'all'], 
                       default='all', help='Type of models to setup for')
    parser.add_argument('--validate_only', action='store_true',
                       help='Only validate existing installation')
    args = parser.parse_args()
    
    if args.validate_only:
        success = validate_installation(args.model_type)
        sys.exit(0 if success else 1)
    
    # Setup base environment
    if not setup_inference_environment():
        logger.error("Failed to setup inference environment")
        sys.exit(1)
    
    # Install specific model requirements
    if args.model_type in ['api', 'all']:
        if install_api_requirements():
            logger.info("✅ API model requirements installed")
        else:
            logger.error("❌ Failed to install API model requirements")
    
    if args.model_type in ['opensource', 'all']:
        if install_opensource_requirements():
            logger.info("✅ Open-source model requirements installed")
        else:
            logger.error("❌ Failed to install open-source model requirements")
    
    # Validate installation
    if validate_installation(args.model_type):
        logger.info("🎉 Inference environment setup completed successfully!")
        logger.info("Next steps:")
        logger.info("1. Configure your API keys in inference_config.json")
        logger.info("2. Run inference with: python run_inference.py --help")
    else:
        logger.error("❌ Installation validation failed")
        sys.exit(1)

if __name__ == "__main__":
    main()