Update app.py

app.py

@@ -1,51 +1,95 @@
 import gradio as gr
 import pandas as pd
-import joblib
+import numpy as np
+import mlflow
+from sklearn.preprocessing import StandardScaler
+import sklearn
 
-# Load the model
-model_path = "model.pkl"
-model = joblib.load(model_path)
+print(f"Prediction environment scikit-learn version: {sklearn.__version__}")
 
-# Define prediction function
-def predict_stroke(age, avg_glucose_level, bmi, gender, hypertension, heart_disease, ever_married, work_type, residence_type, smoking_status):
-    # Prepare input data
-    input_data = pd.DataFrame([{
-        "age": age,
-        "avg_glucose_level": avg_glucose_level,
-        "bmi": bmi,
-        "gender": gender,
-        "hypertension": hypertension,
-        "heart_disease": heart_disease,
-        "ever_married": ever_married,
-        "work_type": work_type,
-        "Residence_type": residence_type,
-        "smoking_status": smoking_status
-    }])
-    # Make predictions
-    prediction = model.predict(input_data)[0]
-    probability = model.predict_proba(input_data)[0][1]
-    return f"Prediction: {'Stroke' if prediction == 1 else 'No Stroke'} (Probability of Stroke: {probability:.2f})"
+# Load model from MLflow artifacts
+model_path = "metadata/mlflow/mlartifacts/0951b451e9554321adaebc8f9f15ac8c/artifacts/train/model/artifacts/sk_model/model.pkl"
+loaded_model = mlflow.sklearn.load_model(model_path)
 
-# Define Gradio interface
-demo = gr.Interface(
-    fn=predict_stroke,
+def preprocess_input(data_dict):
+    """Preprocess input data to match the training format"""
+    df = pd.DataFrame([data_dict])
+    
+    # Numeric features
+    numeric_features = ['age', 'avg_glucose_level', 'bmi']
+    
+    # Scale numeric features
+    scaler = StandardScaler()
+    df[numeric_features] = scaler.fit_transform(df[numeric_features])
+    
+    # Create dummy variables for categorical features
+    df = pd.get_dummies(df, columns=['gender', 'hypertension', 'heart_disease', 
+                                    'ever_married', 'work_type', 'Residence_type', 
+                                    'smoking_status'])
+    
+    # Ensure all expected columns are present
+    expected_columns = [
+        'num__age', 'num__avg_glucose_level', 'num__bmi',
+        'cat__gender_Male', 'cat__gender_Other', 'cat__hypertension_1',
+        'cat__heart_disease_1', 'cat__ever_married_Yes',
+        'cat__work_type_Never_worked', 'cat__work_type_Private',
+        'cat__work_type_Self-employed', 'cat__work_type_children',
+        'cat__Residence_type_Urban', 'cat__smoking_status_formerly smoked',
+        'cat__smoking_status_never smoked', 'cat__smoking_status_smokes'
+    ]
+    
+    for col in expected_columns:
+        if col not in df.columns:
+            df[col] = 0
+            
+    return df[expected_columns]
+
+def predict(gender, age, hypertension, ever_married, work_type, heart_disease, 
+           avg_glucose_level, bmi, smoking_status, Residence_type):
+    """Make prediction using the loaded model"""
+    # Create input dictionary
+    input_data = {
+        'gender': gender,
+        'age': age,
+        'hypertension': 1 if hypertension == 'Yes' else 0,
+        'heart_disease': 1 if heart_disease == 'Yes' else 0,
+        'ever_married': ever_married,
+        'work_type': work_type,
+        'Residence_type': Residence_type,
+        'avg_glucose_level': avg_glucose_level,
+        'bmi': bmi,
+        'smoking_status': smoking_status
+    }
+    
+    # Preprocess the input
+    processed_input = preprocess_input(input_data)
+    
+    # Use the loaded model
+    try:
+        prediction = loaded_model.predict_proba(processed_input)[0][1]
+        return f"The probability of stroke is {prediction:.2%}"
+    except Exception as e:
+        return f"Error making prediction: {str(e)}"
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict,
     inputs=[
-        gr.Number(label="Age"),
+        gr.Radio(choices=['Female', 'Male'], label="Gender"),
+        gr.Slider(minimum=0, maximum=100, label="Age"),
+        gr.Radio(choices=['Yes', 'No'], label="Hypertension"),
+        gr.Radio(choices=['Yes', 'No'], label="Ever Married"),
+        gr.Radio(choices=['Private', 'Self-employed', 'Govt_job', 'children', 'Never_worked'], label="Work Type"),
+        gr.Radio(choices=['Yes', 'No'], label="Heart Disease"),
         gr.Number(label="Average Glucose Level"),
-        gr.Number(label="BMI"),
-        gr.Radio(["Male", "Female", "Other"], label="Gender"),
-        gr.Radio([0, 1], label="Hypertension (0 = No, 1 = Yes)"),
-        gr.Radio([0, 1], label="Heart Disease (0 = No, 1 = Yes)"),
-        gr.Radio(["Yes", "No"], label="Ever Married"),
-        gr.Dropdown(["Private", "Self-employed", "Govt_job", "children", "Never_worked"], label="Work Type"),
-        gr.Radio(["Urban", "Rural"], label="Residence Type"),
-        gr.Dropdown(["never smoked", "formerly smoked", "smokes"], label="Smoking Status"),
+        gr.Slider(minimum=10, maximum=50, label="BMI"),
+        gr.Radio(choices=['formerly smoked', 'never smoked', 'smokes', 'Unknown'], label="Smoking Status"),
+        gr.Radio(choices=['Urban', 'Rural'], label="Residence Type")
     ],
-    outputs="text",
-    title="Stroke Prediction",
-    description="Predict the likelihood of a stroke based on health and lifestyle inputs.",
+    outputs='text',
+    title='Stroke Probability Predictor',
+    description='Predicts the probability of having a stroke based on input features.'
 )
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()