setup.py

"""
Setup script for Hugging Face Spaces deployment
Downloads data and trains model if not present
"""
import os
import sys

def setup():
    """Initialize data and model for the application"""
    print("="*80)
    print("EXOPLANET HUNTER AI - SETUP")
    print("="*80)

    # Check if model exists
    model_exists = os.path.exists('model/exoplanet_classifier.pkl')
    data_exists = os.path.exists('data/kepler_koi.csv')

    if not data_exists:
        print("\n📥 Downloading NASA datasets...")
        print("This may take a few minutes...")
        try:
            import download_data
            download_data.main()
            print("✅ Data downloaded successfully!")
        except Exception as e:
            print(f"❌ Error downloading data: {e}")
            sys.exit(1)
    else:
        print("\n✅ Data already exists, skipping download")

    # Always retrain model on first deployment to ensure version compatibility
    if not model_exists:
        print("\n🤖 Training machine learning model...")
        print("This ensures compatibility with the deployment environment...")
        print("This may take a few minutes...")
        try:
            import train_model
            train_model.main()
            print("✅ Model trained successfully!")
        except Exception as e:
            print(f"❌ Error training model: {e}")
            print(f"Error details: {str(e)}")
            import traceback
            traceback.print_exc()
            sys.exit(1)
    else:
        print("\n✅ Model already exists, skipping training")
        print("   (If you encounter pickle errors, delete model/ and restart)")

    print("\n" + "="*80)
    print("SETUP COMPLETE! Starting application...")
    print("="*80 + "\n")

if __name__ == "__main__":
    setup()