train_model.py

"""
Script to train and save the predictive maintenance model
Run this script before using the Streamlit app for the first time
"""

from preprocessing import DataPreprocessor
from model import PredictiveMaintenanceModel
import pickle

def main():
    print("="*60)
    print("Training Predictive Maintenance Model")
    print("="*60)
    
    # Initialize preprocessor
    print("\n1. Loading and preprocessing data...")
    preprocessor = DataPreprocessor('ai4i2020.csv')
    
    # Prepare data
    X_train, X_test, y_train, y_test, feature_columns = preprocessor.prepare_data()
    
    # Train model
    print("\n2. Training model...")
    model = PredictiveMaintenanceModel()
    model.train(X_train, y_train)
    
    # Evaluate model
    print("\n3. Evaluating model...")
    results = model.evaluate(X_test, y_test)
    
    # Save model and preprocessor
    print("\n4. Saving model and preprocessor...")
    model.save_model('predictive_maintenance_model.pkl')
    
    with open('preprocessor.pkl', 'wb') as f:
        pickle.dump(preprocessor, f)
    
    print("\n5. Saving feature columns...")
    with open('feature_columns.pkl', 'wb') as f:
        pickle.dump(feature_columns, f)
    
    print("\n" + "="*60)
    print("Model training complete!")
    print("="*60)
    print("\nModel files saved:")
    print("  - predictive_maintenance_model.pkl")
    print("  - preprocessor.pkl")
    print("  - feature_columns.pkl")
    print("\nYou can now run the Streamlit app: streamlit run app.py")

if __name__ == "__main__":
    main()