Update app.py

app.py

@@ -1,4 +1,4 @@
-
+# app.py
 
 import os
 import joblib
@@ -7,63 +7,25 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import shap
-from sklearn.compose import ColumnTransformer
-from sklearn.ensemble import RandomForestClassifier
-from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import OneHotEncoder, StandardScaler
+from sklearn.pipeline import Pipeline # Keep for type hinting/structure
 
 # =====================================================================================
-# PART 1: MODEL CREATION AND LOADING (Self-contained for Hugging Face)
-# This part creates, trains, and saves a mock model if one doesn't exist.
-# This ensures the app is fully reproducible in any environment.
+# PART 1: MODEL LOADING
+# This section now directly loads the model file you have uploaded to the repository.
 # =====================================================================================
 
-MODEL_FILE = "machine_failure_model.joblib"
-
-def create_and_train_model():
-    """Creates, trains, and saves a mock model pipeline."""
-    # Mock data that resembles the predictive maintenance dataset
-    mock_features = pd.DataFrame({
-        'Type': ['L', 'M', 'H', 'L', 'M', 'H', 'L', 'M', 'H', 'L'],
-        'Air temperature [K]': [298.1, 298.2, 298.3, 298.4, 299.0, 299.5, 300.1, 301.0, 302.5, 303.0],
-        'Process temperature [K]': [308.6, 308.7, 308.8, 308.9, 309.1, 309.8, 310.5, 311.0, 312.0, 313.5],
-        'Rotational speed [rpm]': [1551, 1428, 1455, 1600, 1750, 2000, 2200, 2500, 2850, 1300],
-        'Torque [Nm]': [42.8, 46.3, 40.0, 50.1, 55.2, 60.0, 65.5, 70.0, 75.0, 35.0],
-        'Tool wear [min]': [0, 5, 10, 15, 25, 50, 80, 120, 180, 210]
-    })
-    # Mock target: 0 = No Failure, 1 = Failure
-    mock_target = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 0])
-
-    # Define preprocessing steps for different column types
-    numeric_features = ['Air temperature [K]', 'Process temperature [K]', 'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]']
-    categorical_features = ['Type']
-
-    preprocessor = ColumnTransformer(
-        transformers=[
-            ('num', StandardScaler(), numeric_features),
-            ('cat', OneHotEncoder(handle_unknown='ignore'), categorical_features)
-        ])
-
-    # Create the full pipeline with preprocessing and a classifier
-    model_pipeline = Pipeline(steps=[
-        ('preprocessor', preprocessor),
-        ('classifier', RandomForestClassifier(n_estimators=50, random_state=42))
-    ])
-
-    # Train the model
-    model_pipeline.fit(mock_features, mock_target)
-
-    # Save the trained model to a file
-    joblib.dump(model_pipeline, MODEL_FILE)
-    print(f"Model trained and saved to {MODEL_FILE}")
-    return model_pipeline
-
-# Check if the model file exists; if not, create it.
-if not os.path.exists(MODEL_FILE):
-    loaded_model = create_and_train_model()
-else:
+# <<< CHANGE 1: Update the filename to match your uploaded model.
+MODEL_FILE = "machine_failure_prediction_model.joblib"
+
+# <<< CHANGE 2: Remove the model creation logic. We will now directly load your file.
+# The app assumes this file exists in your Hugging Face repository.
+try:
     loaded_model = joblib.load(MODEL_FILE)
-    print(f"Model loaded from {MODEL_FILE}")
+    print(f"Successfully loaded model from {MODEL_FILE}")
+except FileNotFoundError:
+    print(f"Error: Model file not found at {MODEL_FILE}. Make sure the file is uploaded to the repository.")
+    # You might want to display an error in the Gradio app itself if the model fails to load.
+    loaded_model = None # Set to None to handle the error case
 
 
 # =====================================================================================
@@ -72,6 +34,9 @@ else:
 
 def predict_failure(Type, air_temperature, process_temperature, rotational_speed, torque, tool_wear):
     """Predicts machine failure and calculates SHAP values using the loaded model."""
+    if loaded_model is None:
+        return "Error: Model not loaded.", None
+
     input_data = pd.DataFrame({
         'Type': [Type], 'Air temperature [K]': [air_temperature],
         'Process temperature [K]': [process_temperature], 'Rotational speed [rpm]': [rotational_speed],
@@ -100,7 +65,7 @@ def predict_failure(Type, air_temperature, process_temperature, rotational_speed
 
 def generate_shap_plot(shap_values, feature_names, base_value):
     """Generates a SHAP waterfall plot for the Gradio interface."""
-    plt.close('all')  # Ensure plots don't stack in memory
+    plt.close('all')
     explanation = shap.Explanation(
         values=shap_values, base_values=base_value, feature_names=feature_names
     )
@@ -111,11 +76,21 @@ def generate_shap_plot(shap_values, feature_names, base_value):
 
 def predict_and_generate_plot(Type, air_temperature, process_temperature, rotational_speed, torque, tool_wear):
     """Wrapper function that connects the backend prediction to the frontend plot."""
-    probability, shap_values, feature_names, base_value = predict_failure(
+    result = predict_failure(
         Type, air_temperature, process_temperature, rotational_speed, torque, tool_wear
     )
-    shap_plot = generate_shap_plot(shap_values, feature_names, base_value)
-    return f"{probability:.2%}", shap_plot # Format probability as percentage
+
+    if isinstance(result, tuple):
+        probability, shap_values, feature_names, base_value = result
+        shap_plot = generate_shap_plot(shap_values, feature_names, base_value)
+        return f"{probability:.2%}", shap_plot
+    else:
+        # Handle the case where the model failed to load
+        error_message = result
+        empty_plot = plt.figure()
+        plt.text(0.5, 0.5, 'Error: Model could not be loaded.', horizontalalignment='center', verticalalignment='center')
+        return error_message, empty_plot
+
 
 # Define the Gradio interface layout and components
 with gr.Blocks(theme=gr.themes.Soft()) as iface_with_shap:
@@ -137,7 +112,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface_with_shap:
             probability_output = gr.Textbox(label="Probability of Machine Failure")
             plot_output = gr.Plot(label="Feature Contribution to Failure (SHAP Waterfall Plot)")
 
-    # Connect the inputs to the function and outputs
+    # This makes the app load the first prediction on startup
+    iface_with_shap.load(
+        fn=predict_and_generate_plot,
+        inputs=inputs,
+        outputs=[probability_output, plot_output]
+    )
+    
+    # This connects the UI changes to the prediction function
     for input_comp in inputs:
         input_comp.change(
             fn=predict_and_generate_plot,
@@ -147,4 +129,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as iface_with_shap:
 
 # Launch the application
 if __name__ == "__main__":
-    iface_with_shap.launch(debug=True)
+    iface_with_shap.launch()