Spaces:
Paused
Paused
File size: 3,962 Bytes
f1b21db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | 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"} |