src/modeling/train.py

"""Module d'entraînement des modèles de machine learning.
Ce module contient une fonction pour entraîner un modèle de machine learning à partir d'une pipeline contenant le préprocesseur et le modèle. La fonction gère les exceptions potentielles lors de l'entraînement et retourne la pipeline avec le modèle entraîné.
"""

def train_model(pipeline, X_train, y_train):
    """
    Entraîne le modèle contenu dans la pipeline.
    args:
        pipeline : sklearn.pipeline.Pipeline
        X_train : pd.DataFrame
        y_train : pd.Series ou np.array
    returns:
        pipeline : sklearn.pipeline.Pipeline, avec le modèle entraîné
    """
    try:
        pipeline.fit(X_train, y_train)
        return pipeline

    except Exception as e:
        print(f"Erreur lors de l'entraînement du modèle : {e}")
        return None