app.py

from fastapi import FastAPI
from pydantic import BaseModel
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import numpy as np
import joblib
import os

# ------------------------------------------------------------------
# Initialisation de l'API
# ------------------------------------------------------------------
app = FastAPI(
    title="IBM Attrition API",
    description="API qui prédit si un employé va quitter l'entreprise.",
    version="1.0",
)


# ------------------------------------------------------------------
# Définition du schéma des données d'entrée (avec Pydantic)
# ------------------------------------------------------------------
class Employee(BaseModel):
    age: int
    monthly_income: float
    years_at_company: int
    job_satisfaction: int  # de 1 à 4
    work_life_balance: int  # de 1 à 4
    overtime: int  # 0 ou 1


# ------------------------------------------------------------------
# Entraînement (ou chargement) du modèle au démarrage de l'API
# ------------------------------------------------------------------
MODEL_PATH = "attrition_model.joblib"
SCALER_PATH = "attrition_scaler.joblib"


def train_dummy_model():
    """Entraîne un mini modèle de régression logistique sur des données synthétiques."""
    rng = np.random.default_rng(seed=42)
    n = 500

    # Génération de données synthétiques cohérentes.
    age = rng.integers(22, 60, n)
    income = rng.normal(5000, 2000, n).clip(1000, 20000)
    years = rng.integers(0, 30, n)
    satisfaction = rng.integers(1, 5, n)
    wlb = rng.integers(1, 5, n)
    overtime = rng.integers(0, 2, n)

    # Règle simple : risque d'attrition élevé si jeune, faible satisfaction, overtime.
    risk = (
        (60 - age) * 0.02
        + (5 - satisfaction) * 0.3
        + (5 - wlb) * 0.2
        + overtime * 0.5
        - years * 0.05
    )
    proba = 1 / (1 + np.exp(-risk))
    attrition = (rng.random(n) < proba).astype(int)

    X = np.column_stack([age, income, years, satisfaction, wlb, overtime])
    y = attrition

    scaler = StandardScaler().fit(X)
    X_scaled = scaler.transform(X)

    model = LogisticRegression().fit(X_scaled, y)

    joblib.dump(model, MODEL_PATH)
    joblib.dump(scaler, SCALER_PATH)
    return model, scaler


# On entraîne le modèle au démarrage si le fichier n'existe pas.
if os.path.exists(MODEL_PATH) and os.path.exists(SCALER_PATH):
    model = joblib.load(MODEL_PATH)
    scaler = joblib.load(SCALER_PATH)
else:
    model, scaler = train_dummy_model()


# ------------------------------------------------------------------
# Endpoints de l'API
# ------------------------------------------------------------------
@app.get("/")
def root():
    """Page d'accueil avec les instructions."""
    return {
        "message": "IBM Attrition API",
        "endpoints": {
            "/predict": "POST avec les caractéristiques d'un employé pour prédire l'attrition",
            "/docs": "Interface Swagger pour tester l'API",
        },
    }


@app.post("/predict")
def predict(employee: Employee):
    """Prédit la probabilité que l'employé quitte l'entreprise."""
    features = np.array([[
        employee.age,
        employee.monthly_income,
        employee.years_at_company,
        employee.job_satisfaction,
        employee.work_life_balance,
        employee.overtime,
    ]])
    features_scaled = scaler.transform(features)

    proba = float(model.predict_proba(features_scaled)[0, 1])
    prediction = int(proba >= 0.5)

    return {
        "attrition_probability": round(proba, 3),
        "will_leave": bool(prediction),
        "interpretation": (
            "L'employé est à risque de quitter l'entreprise."
            if prediction
            else "L'employé devrait rester."
        ),
    }


@app.get("/health")
def health():
    """Endpoint de santé pour vérifier que l'API tourne."""
    return {"status": "ok"}