retrain / retrain.py
gdleds's picture
modif pour version sickit-learn
3e142f4
Raw
History Blame Contribute Delete
2.58 kB
import pandas as pd
import mlflow
import joblib
import io
import boto3
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.metrics import f1_score, recall_score, precision_score
from xgboost import XGBClassifier
import psycopg2
from sqlalchemy import create_engine
import os
def run_retraining():
db_user = os.getenv("DB_USER")
db_password = os.getenv("DB_PASS")
db_host = os.getenv("DB_HOST")
db_name = os.getenv("DB_NAME")
engine = create_engine(f"postgresql+psycopg2://{db_user}:{db_password}@{db_host}/{db_name}")
df = pd.read_sql("SELECT * FROM fraud_transaction", engine)
df=pd.DataFrame(df)
X=df.drop(columns=["is_fraud", "is_fraud_pred", "timestamp","id", "inserted_at"])
y=df["is_fraud"]
X_train, X_test, y_train,y_test=train_test_split(X,y, test_size=0.2, random_state=42, stratify=y)
numeric_features = X.select_dtypes(include=['int64', 'float64', "int32"]).columns
categoric_features = X.select_dtypes(include=['object']).columns
numeric_transformer = StandardScaler()
categorical_transformer = OneHotEncoder()
preprocessor = ColumnTransformer(
transformers=[
('num', numeric_transformer, numeric_features),
('cat', categorical_transformer, categoric_features)])
xgb = XGBClassifier(n_estimators=1200,max_depth=10, learning_rate=0.2, random_state=42)
xgb_final = Pipeline(steps=[
('preprocessor', preprocessor),
('xgboost_best', xgb)
])
mlflow.set_tracking_uri("https://gdleds-mlflow-fraud.hf.space")
mlflow.set_experiment("fraud_detection")
with mlflow.start_run() as run:
xgb_final.fit(X_train, y_train)
y_pred = xgb_final.predict(X_test)
f1 = f1_score(y_test, y_pred)
mlflow.log_metric("F1", f1)
mlflow.log_metric("precision", precision_score(y_test, y_pred))
mlflow.log_metric("recall", recall_score(y_test, y_pred))
# Sauvegarde modèle dans S3
s3 = boto3.client("s3")
bucket = os.getenv("S3_BUCKET")
key = f"mlflow/models/xgboost_model_{run.info.run_id}.joblib"
buffer = io.BytesIO()
joblib.dump(xgb_final, buffer)
s3.put_object(
Bucket=bucket,
Key=key,
Body=buffer.getvalue()
)
return {
"status": "success",
"run_id": run.info.run_id,
"f1_score": f1,
"model_path": key
}