Update app.py

app.py

@@ -3,45 +3,64 @@ import pickle
 import numpy as np
 import matplotlib.pyplot as plt
 
-# Load model
+# Load trained model
 with open("DecisionTreeClassifier.pkl", "rb") as file:
     model = pickle.load(file)
 
 # Prediction function
 def predict_churn(age, gender, tenure, usage, support, delay,
                   subscription, contract, spend, interaction):
+    try:
+        # 🔹 Gender encoding
+        gender_val = 1 if gender == "Female" else 0
 
-    # Convert categorical to numeric (must match training!)
-    gender_map = {"Male": 0, "Female": 1}
-    sub_map = {"Basic": 0, "Standard": 1, "Premium": 2}
-    contract_map = {"Monthly": 0, "Quarterly": 1, "Annual": 2}
+        # 🔹 One-hot encoding (MUST match training columns)
+        sub_premium = 1 if subscription == "Premium" else 0
+        sub_standard = 1 if subscription == "Standard" else 0
 
-    input_data = np.array([[ 
-        age,
-        gender_map[gender],
-        tenure,
-        usage,
-        support,
-        delay,
-        sub_map[subscription],
-        contract_map[contract],
-        spend,
-        interaction
-    ]])
+        contract_monthly = 1 if contract == "Monthly" else 0
+        contract_quarterly = 1 if contract == "Quarterly" else 0
 
-    # Prediction
-    pred = model.predict(input_data)[0]
-    prob = model.predict_proba(input_data)[0][1]
+        # 🔹 Input array (order matters!)
+        input_data = np.array([[ 
+            age,
+            gender_val,
+            tenure,
+            usage,
+            support,
+            delay,
+            spend,
+            interaction,
+            sub_premium,
+            sub_standard,
+            contract_monthly,
+            contract_quarterly
+        ]])
 
-    result = "Churn ⚠️" if pred == 1 else "No Churn 🙂"
+        # 🔹 Prediction
+        pred = model.predict(input_data)[0]
+        prob = model.predict_proba(input_data)[0][1]
 
-    # 📊 Graph: Probability Bar Chart
-    fig, ax = plt.subplots()
-    ax.bar(["No Churn", "Churn"], [1-prob, prob])
-    ax.set_title("Churn Probability")
-    ax.set_ylabel("Probability")
+        result = "Churn ⚠️" if pred == 1 else "No Churn 🙂"
 
-    return result, f"{prob*100:.2f}%", fig
+        # 🔹 Risk Level
+        if prob > 0.7:
+            risk = "High Risk 🔴"
+        elif prob > 0.4:
+            risk = "Medium Risk 🟠"
+        else:
+            risk = "Low Risk 🟢"
+
+        # 🔹 Graph
+        fig, ax = plt.subplots()
+        ax.bar(["No Churn", "Churn"], [1 - prob, prob])
+        ax.set_title("Churn Probability")
+        ax.set_ylabel("Probability")
+
+        return result, f"{prob*100:.2f}%", risk, fig
+
+    except Exception as e:
+        return f"Error: {str(e)}", "", "", None
 
 
 # 🎨 Gradio UI
@@ -62,10 +81,11 @@ interface = gr.Interface(
     outputs=[
         gr.Text(label="Prediction"),
         gr.Text(label="Churn Probability"),
+        gr.Text(label="Risk Level"),
         gr.Plot(label="Graph")
     ],
     title="📊 Customer Churn Prediction System",
-    description="Enter customer details to predict churn probability"
+    description="Enter customer details to predict churn probability and risk level"
 )
 
 interface.launch()