src/cuml_trainer.py

# ===========================
# Fichier: cuml_trainer.py
# ===========================

from abc import ABC, abstractmethod
from typing import Union
import cupy as cp
from scipy.sparse import csr_matrix

from config import Config
from base_trainer import BaseTrainer
from interfaces.vectorizer import Vectorizer


class CuMLTrainer(BaseTrainer, ABC):
    """
    Classe abstraite, héritée de BaseTrainer, représentant un entraîneur
    basé sur la librairie cuML. Elle ajoute notamment le concept de vectoriseur
    et force le passage de la matrice d'entrée en cupy.ndarray pour la plupart
    des opérations.

    Attributs:
        vectorizer (Vectorizer): Objet responsable de la vectorisation du texte.
    """

    def __init__(self, config: Config, data_path: str,
                 target_column: str) -> None:
        """
        Initialise un CuMLTrainer avec la configuration.
        Appelle également le constructeur de BaseTrainer.

        :param config: Configuration globale du système.
        :param data_path: Chemin vers le fichier de données.
        :param target_column: Nom de la colonne cible dans les données.
        """
        super().__init__(config, data_path, target_column)
        self.vectorizer: Vectorizer = None
        # self.classifier est déjà défini dans BaseTrainer.
        # On suppose que 'classifier' sera un modèle cuML (cuml.Base).

    @abstractmethod
    def build_components(self) -> None:
        """
        Méthode abstraite. Instancie concrètement le vectorizer et le classifieur,
        selon la configuration (ex. 'svm', 'random_forest', etc.).
        """
        pass

    def train(self) -> None:
        """
        Entraîne le classifieur sur les données vectorisées.
        Cette implémentation générique fonctionne pour tous les trainers cuML.
        """
        # Chargement des données
        import cudf
        data = cudf.read_csv(self.data_path)
        
        # Séparation des textes et des étiquettes
        texts = data.drop(columns=[self.target_column])
        labels = data[self.target_column].values
        
        # Vectorisation des textes
        X = self.vectorizer.fit_transform(texts)
        X_prepared = self._prepare_input_for_fit(X)
        
        # Entraînement du modèle
        self.classifier.fit(X_prepared, labels)

    def evaluate(self) -> None:
        """
        Évalue le classifieur et calcule les métriques.
        Cette implémentation générique fonctionne pour tous les trainers cuML.
        """
        # Chargement des données (idéalement un jeu de test séparé)
        import cudf
        data = cudf.read_csv(self.data_path)
        
        # Séparation des textes et des étiquettes
        texts = data.drop(columns=[self.target_column])
        y_true = data[self.target_column].values
        
        # Vectorisation et prédiction
        X = self.vectorizer.transform(texts)
        X_prepared = self._prepare_input_for_predict(X)
        y_pred = self.classifier.predict(X_prepared)
        
        # Calcul et logging des métriques
        prefix = self.config.model.type.lower()
        metrics = self.metrics_calculator.calculate_and_log(
            y_true=y_true,
            y_pred=y_pred,
            prefix=prefix
        )
        
        # Afficher les résultats
        print(f"Métriques d'évaluation {prefix}: {metrics}")

    def _prepare_input_for_fit(self, X: Union[cp.ndarray,
                                              csr_matrix]) -> cp.ndarray:
        """
        Convertit, si nécessaire, la matrice en cupy.ndarray pour l'entraînement.

        :param X: Données d'entraînement (cupy.ndarray ou scipy.sparse.csr_matrix).
        :return: Données converties en cupy.ndarray, pour compatibilité cuML.
        """
        if isinstance(X, csr_matrix):
            return cp.asarray(X.toarray())
        return X  # c'est déjà cupy.ndarray

    def _prepare_input_for_predict(
            self, X: Union[cp.ndarray, csr_matrix]) -> cp.ndarray:
        """
        Convertit, si nécessaire, la matrice en cupy.ndarray pour la prédiction.

        :param X: Données de prédiction (cupy.ndarray ou scipy.sparse.csr_matrix).
        :return: Données converties en cupy.ndarray, pour compatibilité cuML.
        """
        if isinstance(X, csr_matrix):
            return cp.asarray(X.toarray())
        return X