add data pipeline, baseline, and XGBoost model

notebooks/01_eda.ipynb

@@ -0,0 +1,48 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5589eb3",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31mFailed to start the Kernel. \n",
+      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
+     ]
+    }
+   ],
+   "source": [
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "import seaborn as sns\n",
+    "\n",
+    "df = pd.read_csv('../data/Wellbeing_and_lifestyle_data_Kaggle.csv')\n",
+    "print(df.shape)\n",
+    "print(df.dtypes)\n",
+    "df.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "01e7fd72",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/baseline.py

@@ -0,0 +1,27 @@
+import numpy as np
+from sklearn.dummy import DummyClassifier
+from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
+import sys
+sys.path.insert(0, '.')
+from src.data_loader import load_data, preprocess, split_and_scale
+
+def evaluate(model, X, y, split_name):
+    preds = model.predict(X)
+    proba = model.predict_proba(X)[:, 1]
+    acc = accuracy_score(y, preds)
+    f1 = f1_score(y, preds, zero_division=0)
+    auc = roc_auc_score(y, proba)
+    print(f"{split_name} — Accuracy: {acc:.3f} | F1: {f1:.3f} | AUC: {auc:.3f}")
+
+if __name__ == '__main__':
+    df = load_data()
+    X, y, feature_cols = preprocess(df)
+    X_train, X_val, X_test, y_train, y_val, y_test, scaler = split_and_scale(X, y)
+
+    # Majority class baseline
+    dummy = DummyClassifier(strategy='most_frequent')
+    dummy.fit(X_train, y_train)
+    print("=== Baseline (majority class) ===")
+    evaluate(dummy, X_train, y_train, "Train")
+    evaluate(dummy, X_val, y_val, "Val")
+    evaluate(dummy, X_test, y_test, "Test")

src/data_loader.py

@@ -0,0 +1,53 @@
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+
+def load_data(path='data/Wellbeing_and_lifestyle_data_Kaggle.csv'):
+    df = pd.read_csv(path)
+    return df
+
+def preprocess(df):
+    # Drop timestamp
+    df = df.drop(columns=['Timestamp'])
+    
+    # Encode categorical columns
+    df['GENDER'] = df['GENDER'].map({'Female': 0, 'Male': 1})
+    age_map = {'Less than 20': 0, '21 to 35': 1, '36 to 50': 2, '51 or more': 3}
+    df['AGE'] = df['AGE'].map(age_map)
+    
+    # Keep only numeric columns
+    df = df.apply(pd.to_numeric, errors='coerce')
+    
+    # Drop rows with nulls
+    df = df.dropna()
+    
+    # Create burnout label: bottom 25% of work-life balance = high burnout risk
+    threshold = df['WORK_LIFE_BALANCE_SCORE'].quantile(0.25)
+    df['BURNOUT_RISK'] = (df['WORK_LIFE_BALANCE_SCORE'] <= threshold).astype(int)
+    
+    # Features and target
+    feature_cols = [c for c in df.columns if c not in ['WORK_LIFE_BALANCE_SCORE', 'BURNOUT_RISK']]
+    X = df[feature_cols]
+    y = df['BURNOUT_RISK']
+    
+    return X, y, feature_cols
+
+def split_and_scale(X, y):
+    X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=42)
+    X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42)
+    
+    scaler = StandardScaler()
+    X_train = scaler.fit_transform(X_train)
+    X_val = scaler.transform(X_val)
+    X_test = scaler.transform(X_test)
+    
+    print(f"Train: {X_train.shape}, Val: {X_val.shape}, Test: {X_test.shape}")
+    return X_train, X_val, X_test, y_train, y_val, y_test, scaler
+
+if __name__ == '__main__':
+    df = load_data()
+    X, y, feature_cols = preprocess(df)
+    X_train, X_val, X_test, y_train, y_val, y_test, scaler = split_and_scale(X, y)
+    print(f"Burnout rate: {y.mean():.2%}")
+    print(f"Features: {feature_cols}")

src/train_xgboost.py

@@ -0,0 +1,50 @@
+import numpy as np
+import joblib
+from xgboost import XGBClassifier
+from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
+import sys
+sys.path.insert(0, '.')
+from src.data_loader import load_data, preprocess, split_and_scale
+
+def evaluate(model, X, y, split_name):
+    preds = model.predict(X)
+    proba = model.predict_proba(X)[:, 1]
+    acc = accuracy_score(y, preds)
+    f1 = f1_score(y, preds)
+    auc = roc_auc_score(y, proba)
+    print(f"{split_name} — Accuracy: {acc:.3f} | F1: {f1:.3f} | AUC: {auc:.3f}")
+    return acc, f1, auc
+
+if __name__ == '__main__':
+    df = load_data()
+    X, y, feature_cols = preprocess(df)
+    X_train, X_val, X_test, y_train, y_val, y_test, scaler = split_and_scale(X, y)
+
+    model = XGBClassifier(
+        n_estimators=200,
+        max_depth=6,
+        learning_rate=0.1,
+        subsample=0.8,
+        colsample_bytree=0.8,
+        use_label_encoder=False,
+        eval_metric='logloss',
+        early_stopping_rounds=20,
+        random_state=42
+    )
+
+    model.fit(
+        X_train, y_train,
+        eval_set=[(X_val, y_val)],
+        verbose=10
+    )
+
+    print("\n=== XGBoost Results ===")
+    evaluate(model, X_train, y_train, "Train")
+    evaluate(model, X_val, y_val, "Val")
+    evaluate(model, X_test, y_test, "Test")
+
+    # Save model and scaler
+    joblib.dump(model, 'models/xgboost_model.pkl')
+    joblib.dump(scaler, 'models/scaler.pkl')
+    joblib.dump(feature_cols, 'models/feature_cols.pkl')
+    print("\nModel saved to models/")