grosse mise à jour sur le train test eval

src/cuml_trainer.py

@@ -3,10 +3,12 @@
 # ===========================
 
 from abc import ABC, abstractmethod
-from typing import Union
+from typing import Optional, Union, Tuple
 import cupy as cp
 from scipy.sparse import csr_matrix
 import cudf
+from cuml.model_selection import train_test_split
+import logging
 
 from config import Config
 from base_trainer import BaseTrainer
@@ -36,8 +38,21 @@ class CuMLTrainer(BaseTrainer, ABC):
         """
         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).
+        self.classifier: object = None # Déjà dans BaseTrainer, mais redéfini pour clarté
+        self.logger = logging.getLogger(__name__)
+
+        # Attributs pour stocker les données splittées (texte brut)
+        self.X_train_text: Optional[cudf.Series] = None
+        self.X_val_text: Optional[cudf.Series] = None
+        self.X_test_text: Optional[cudf.Series] = None
+        self.y_train: Optional[cp.ndarray] = None
+        self.y_val: Optional[cp.ndarray] = None
+        self.y_test: Optional[cp.ndarray] = None
+
+        # Attributs pour stocker les données vectorisées
+        self.X_train_vec: Optional[Union[cp.ndarray, csr_matrix]] = None
+        self.X_val_vec: Optional[Union[cp.ndarray, csr_matrix]] = None
+        self.X_test_vec: Optional[Union[cp.ndarray, csr_matrix]] = None
 
     @abstractmethod
     def build_components(self) -> None:
@@ -47,66 +62,152 @@ class CuMLTrainer(BaseTrainer, ABC):
         """
         pass
 
-    def train(self) -> None:
+    def _load_and_split_data(self, test_size=0.2, val_size=0.5, random_state=42) -> None:
         """
-        Entraîne le classifieur sur les données vectorisées.
-        Cette implémentation générique fonctionne pour tous les trainers cuML.
+        Charge les données depuis data_path, les sépare en features/labels,
+        et les divise en ensembles train/validation/test (80/10/10 par défaut).
+        Stocke les résultats dans les attributs de l'instance.
         """
-        # Chargement des données
+        if self.X_train_text is not None: # Évite de recharger/resplitter
+            self.logger.info("Données déjà chargées et splittées.")
+            return
 
+        self.logger.info(f"Chargement des données depuis {self.data_path}...")
         data = cudf.read_csv(self.data_path)
 
-        # Identification des colonnes de features (toutes sauf la cible)
+        # Identification et concaténation des features
         feature_columns = [col for col in data.columns if col != self.target_column]
         if not feature_columns:
-            raise ValueError("Aucune colonne de feature trouvée. Le dataset doit contenir au moins une colonne en plus de la colonne cible.")
-
-        # Concaténation des features (même si une seule, pour la cohérence et l'évolutivité)
-        # Convertit en string et ajoute un espace comme séparateur si plusieurs colonnes existent.
+            raise ValueError("Aucune colonne de feature trouvée.")
         texts_concatenated = data[feature_columns].astype(str).agg(' '.join, axis=1)
-        labels = data[self.target_column].astype(self._get_label_dtype()).values # Assurer le bon dtype pour les labels
+        labels = data[self.target_column].astype(self._get_label_dtype()).values
+
+        self.logger.info("Séparation des données en train/validation/test (80/10/10)...")
+        # Premier split: 80% train, 20% temp (pour val+test)
+        X_train, X_temp, y_train, y_temp = train_test_split(
+            texts_concatenated, labels, test_size=test_size, random_state=random_state, stratify=labels
+        )
+
+        # Deuxième split: 50% validation, 50% test sur l'ensemble temp
+        # (val_size=0.5 sur 20% donne 10% du total pour val et 10% pour test)
+        # Utilisation de stratify=y_temp pour maintenir la distribution des classes
+        X_val, X_test, y_val, y_test = train_test_split(
+            X_temp, y_temp, test_size=val_size, random_state=random_state, stratify=y_temp
+        )
 
-        # Vectorisation des textes concaténés
-        X = self.vectorizer.fit_transform(texts_concatenated)
-        X_prepared = self._prepare_input_for_fit(X)
-        # Entraînement du modèle
-        self.classifier.fit(X_prepared, labels)
+        # Stockage des résultats
+        self.X_train_text = X_train
+        self.X_val_text = X_val
+        self.X_test_text = X_test
+        self.y_train = y_train
+        self.y_val = y_val
+        self.y_test = y_test
 
-    def evaluate(self) -> dict:
+        self.logger.info(f"Taille Train: {len(self.X_train_text)}, Val: {len(self.X_val_text)}, Test: {len(self.X_test_text)}")
+
+
+    def train(self) -> None:
         """
-        Évalue le classifieur et calcule les métriques.
-        Cette implémentation générique fonctionne pour tous les trainers cuML.
+        Entraîne le classifieur sur l'ensemble d'entraînement après vectorisation.
+        Vectorise également les ensembles de validation et de test.
         """
-        # Chargement des données (idéalement un jeu de test séparé)
-        data = cudf.read_csv(self.data_path)
+        self._load_and_split_data() # Assure que les données sont chargées et splittées
 
-        # Identification des colonnes de features (toutes sauf la cible)
-        feature_columns = [col for col in data.columns if col != self.target_column]
-        if not feature_columns:
-            raise ValueError("Aucune colonne de feature trouvée pour l'évaluation.")
+        if self.vectorizer is None or self.classifier is None:
+             raise RuntimeError("Les composants (vectorizer, classifier) doivent être construits avant l'entraînement. Appelez build_components().")
 
-        # Concaténation des features
-        texts_concatenated = data[feature_columns].astype(str).agg(' '.join, axis=1)
-        y_true = data[self.target_column].astype(self._get_label_dtype()).values # Assurer le bon dtype
+        self.logger.info("Vectorisation des données textuelles...")
+        # fit_transform sur l'entraînement
+        self.X_train_vec = self.vectorizer.fit_transform(self.X_train_text)
+        # transform sur validation et test
+        self.X_val_vec = self.vectorizer.transform(self.X_val_text)
+        self.X_test_vec = self.vectorizer.transform(self.X_test_text)
+
+        # Préparation pour cuML (conversion en dense si nécessaire)
+        X_train_prepared = self._prepare_input_for_fit(self.X_train_vec)
+
+        self.logger.info("Entraînement du modèle...")
+        self.classifier.fit(X_train_prepared, self.y_train)
+        self.logger.info("Entraînement terminé.")
+
+
+    def evaluate(self, use_validation_set=False) -> dict:
+        """
+        Évalue le classifieur sur l'ensemble de test (par défaut) ou de validation.
+
+        :param use_validation_set: Si True, évalue sur l'ensemble de validation.
+                                   Sinon (défaut), évalue sur l'ensemble de test.
+        :return: Dictionnaire de métriques.
+        """
+        if self.X_test_vec is None or self.y_test is None or self.X_val_vec is None or self.y_val is None:
+            raise RuntimeError("Les données doivent être chargées, splittées et vectorisées avant l'évaluation. Appelez train().")
+        if self.classifier is None:
+             raise RuntimeError("Le classifieur doit être entraîné avant l'évaluation. Appelez train().")
+
+        if use_validation_set:
+            self.logger.info("Évaluation sur l'ensemble de validation...")
+            X_eval_vec = self.X_val_vec
+            y_true = self.y_val
+            dataset_name = "validation"
+        else:
+            self.logger.info("Évaluation sur l'ensemble de test...")
+            X_eval_vec = self.X_test_vec
+            y_true = self.y_test
+            dataset_name = "test"
+
+        # Préparation pour cuML et prédiction
+        X_eval_prepared = self._prepare_input_for_predict(X_eval_vec)
+        y_pred = self.classifier.predict(X_eval_prepared)
+
+        # Essayer de récupérer les probabilités
+        y_proba = None
+        try:
+            # Utilise la méthode _get_positive_probabilities définie dans BaseTrainer
+            # et potentiellement surchargée dans les sous-classes (comme SvmTrainer)
+            y_proba = self._get_positive_probabilities(X_eval_prepared)
+        except NotImplementedError:
+            self.logger.warning("La méthode _get_positive_probabilities n'est pas implémentée pour ce modèle, AUC pourrait être moins précis ou indisponible.")
+        except Exception as e:
+            self.logger.warning(f"Erreur lors de la récupération des probabilités : {e}")
 
-        # Vectorisation et prédiction
-        X = self.vectorizer.transform(texts_concatenated)
-        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
-        )
+        prefix = f"{self.config.model.type.lower()}_{dataset_name}"
+        if self.metrics_calculator is None:
+             # Initialisation par défaut si non fait ailleurs
+             from interfaces.metrics_calculator import DefaultMetricsCalculator # Utiliser l'implémentation par défaut
+             self.metrics_calculator = DefaultMetricsCalculator()
+
+        # Déterminer si binaire ou multiclasse
+        num_classes = len(cp.unique(y_true))
+        self.logger.info(f"Nombre de classes détectées dans y_true ({dataset_name}): {num_classes}")
+
+        if num_classes <= 2:
+             metrics = self.metrics_calculator.calculate_and_log(
+                 y_true=y_true,
+                 y_pred=y_pred,
+                 y_proba=y_proba, # Passer les probas
+                 prefix=prefix
+             )
+        else:
+             metrics = self.metrics_calculator.calculate_and_log_multiclass(
+                 y_true=y_true,
+                 y_pred=y_pred,
+                 y_proba=y_proba, # Passer les probas
+                 prefix=prefix
+             )
 
         # Afficher les résultats
-        print(f"Métriques d'évaluation {prefix}: {metrics}")
-        
+        self.logger.info(f"Métriques d'évaluation ({dataset_name}): {metrics}")
+
         return metrics
 
+    # Note: La méthode optimize_if_needed appelle l'optimiseur qui, à son tour,
+    # peut appeler train() et evaluate(). Il faudra s'assurer que l'optimiseur
+    # utilise evaluate(use_validation_set=True) pour l'évaluation des hyperparamètres.
+    # Cela pourrait nécessiter une modification des classes Optimizer ou de la façon
+    # dont la fonction objectif est définie dans l'optimiseur.
+
     def _prepare_input_for_fit(self, X: Union[cp.ndarray,
                                               csr_matrix]) -> cp.ndarray:
         """

src/interfaces/metrics_calculator.py

@@ -1,8 +1,13 @@
 import cupy as cp
-import numpy as np
-from typing import Dict, Protocol
-import logging
-from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+from typing import Dict, Protocol, Optional
+import warnings
+# Utiliser cuml.metrics
+from cuml.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score
+
+# Ignorer les avertissements (peut nécessiter ajustement si cuML utilise d'autres types)
+# Attention: Masquer tous les warnings peut cacher des problèmes potentiels.
+warnings.filterwarnings("ignore", category=Warning)
+
 
 class MetricsCalculator(Protocol):
     """
@@ -12,98 +17,141 @@ class MetricsCalculator(Protocol):
         self,
         y_true: cp.ndarray,
         y_pred: cp.ndarray,
+        y_proba: Optional[cp.ndarray], # Probabilités pour AUC
         prefix: str
     ) -> Dict[str, float]:
         """
-        Calcule et log les métriques pour un problème binaire.
+        Calcule les métriques pour un problème binaire.
         """
         pass
-    
+
     def calculate_and_log_multiclass(
         self,
         y_true: cp.ndarray,
         y_pred: cp.ndarray,
+        y_proba: Optional[cp.ndarray], # Probabilités (potentiellement pour futures métriques)
         prefix: str
     ) -> Dict[str, float]:
         """
-        Calcule et log les métriques pour un problème multiclasses.
+        Calcule les métriques pour un problème multiclasses.
         """
         pass
 
-logger = logging.getLogger(__name__)
 
 class DefaultMetricsCalculator(MetricsCalculator):
     """
-    Implémentation concrète de MetricsCalculator qui calcule
-    accuracy, f1, precision, recall, et auc-roc.
-    Fonctionne pour binaire ou multiclasses (avec 'ovr' ou 'macro').
+    Implémentation concrète de MetricsCalculator utilisant cuML.
+    Calcule accuracy, et F1/precision/recall pondérés.
+    Calcule AUC-ROC pour les problèmes binaires *uniquement* si les probabilités sont fournies.
+    Ne calcule pas l'AUC-ROC pour les problèmes multiclasses (non supporté par cuml.metrics.roc_auc_score).
+    Retourne NaN pour les métriques non calculables.
     """
 
     def calculate_and_log(
         self,
         y_true: cp.ndarray,
         y_pred: cp.ndarray,
+        y_proba: Optional[cp.ndarray], # Probabilités requises pour AUC
         prefix: str
     ) -> Dict[str, float]:
         """
-        Calcule et log les métriques pour un problème binaire
-        en supposant y_pred est dans {0,1} ou {True,False}.
+        Calcule les métriques pour un problème binaire.
+        Utilise y_proba pour AUC si disponible.
+        Utilise average='weighted' pour precision/recall/f1.
         """
-        y_true_np = cp.asnumpy(y_true)
-        y_pred_np = cp.asnumpy(y_pred)
-
-        acc = accuracy_score(y_true_np, y_pred_np)
-        prec = precision_score(y_true_np, y_pred_np, zero_division=0)
-        rec = recall_score(y_true_np, y_pred_np, zero_division=0)
-        f1 = f1_score(y_true_np, y_pred_np, zero_division=0)
-        
-        # Calcul AUC pour un problème binaire (si y_pred est 0/1)
-        # On treat y_pred_np as our "probabilities" only if truly 0/1. 
-        # In a real pipeline you might store probabilities separately.
+        metrics: Dict[str, float] = {}
+
         try:
-            auc = roc_auc_score(y_true_np, y_pred_np)
-        except ValueError:
-            auc = 0.0
-
-        metrics = {
-            f"{prefix}_accuracy"  : acc,
-            f"{prefix}_precision" : prec,
-            f"{prefix}_recall"    : rec,
-            f"{prefix}_f1"        : f1,
-            f"{prefix}_auc_roc"   : auc
-        }
-        logger.info(f"[{prefix}] Metrics: {metrics}")
-        return metrics
+            # Accuracy
+            acc = accuracy_score(y_true, y_pred)
+            metrics[f"{prefix}_accuracy"] = float(acc)
+
+            # Precision, Recall, F1 (Weighted)
+            prec, rec, f1, _ = precision_recall_fscore_support(
+                y_true, y_pred, average='weighted'
+            )
+            metrics[f"{prefix}_precision_weighted"] = float(prec)
+            metrics[f"{prefix}_recall_weighted"] = float(rec)
+            metrics[f"{prefix}_f1_weighted"] = float(f1)
+
+        except Exception:
+            # En cas d'erreur sur les métriques de base, remplir avec NaN
+            metrics.setdefault(f"{prefix}_accuracy", float('nan'))
+            metrics.setdefault(f"{prefix}_precision_weighted", float('nan'))
+            metrics.setdefault(f"{prefix}_recall_weighted", float('nan'))
+            metrics.setdefault(f"{prefix}_f1_weighted", float('nan'))
+
+        # AUC-ROC (Binary only, requires probabilities)
+        auc: float = float('nan') # Default to NaN
+        if y_proba is not None:
+            try:
+                # Ensure y_true and y_proba have compatible shapes and types
+                if y_true.dtype != cp.int32 and y_true.dtype != cp.int64:
+                     y_true = y_true.astype(cp.int32)
+
+                # roc_auc_score expects probabilities of the positive class
+                if y_proba.ndim == 2 and y_proba.shape[1] == 2:
+                    proba_pos_class = y_proba[:, 1]
+                elif y_proba.ndim == 1:
+                    proba_pos_class = y_proba # Assume already positive class proba
+                else:
+                    # Forme inattendue, ne peut pas calculer l'AUC
+                     raise ValueError("y_proba a une forme inattendue pour le calcul AUC binaire.")
+
+                if proba_pos_class.dtype != cp.float32 and proba_pos_class.dtype != cp.float64:
+                     proba_pos_class = proba_pos_class.astype(cp.float32)
+
+                # Check if y_true contains more than one class before calculating AUC
+                unique_labels = cp.unique(y_true)
+                if len(unique_labels) >= 2:
+                    auc_score = roc_auc_score(y_true, proba_pos_class)
+                    auc = float(auc_score) # Cast to float
+
+            except (ValueError, TypeError, Exception):
+                 # Si une erreur se produit (ex: une seule classe, type incorrect, autre), AUC reste NaN
+                 pass # auc est déjà float('nan')
+
+        metrics[f"{prefix}_auc_roc"] = auc
+        # Ensure all values in the returned dict are standard floats
+        return {k: float(v) for k, v in metrics.items()}
+
 
     def calculate_and_log_multiclass(
         self,
         y_true: cp.ndarray,
         y_pred: cp.ndarray,
+        y_proba: Optional[cp.ndarray], # Gardé pour cohérence d'interface
         prefix: str
     ) -> Dict[str, float]:
         """
-        Calcule et log les métriques pour un problème multiclasses.
-        AUC-ROC en mode 'macro' si possible.
+        Calcule les métriques pour un problème multiclasses.
+        AUC-ROC n'est pas calculé (retourne NaN) car non supporté par cuml.metrics.roc_auc_score.
+        Utilise average='weighted' pour precision/recall/f1.
         """
-        y_true_np = cp.asnumpy(y_true)
-        y_pred_np = cp.asnumpy(y_pred)
-
-        acc = accuracy_score(y_true_np, y_pred_np)
-        prec = precision_score(y_true_np, y_pred_np, average="macro", zero_division=0)
-        rec = recall_score(y_true_np, y_pred_np, average="macro", zero_division=0)
-        f1 = f1_score(y_true_np, y_pred_np, average="macro", zero_division=0)
-
-        # Pour le multiclasses, la roc_auc_score nécessite des scores proba
-        # ou "decision_function" => vous ajusterez selon votre cas.
-        # Ici, on met 0.0 en fallback.
-        auc = 0.0  
-
-        metrics = {
-            f"{prefix}_accuracy"  : acc,
-            f"{prefix}_precision" : prec,
-            f"{prefix}_recall"    : rec,
-            f"{prefix}_f1"        : f1,
-            f"{prefix}_auc_roc"   : auc
-        }
-        logger.info(f"[{prefix}] Multiclass metrics: {metrics}")
-        return metrics
+        metrics: Dict[str, float] = {}
+
+        try:
+            # Accuracy
+            acc = accuracy_score(y_true, y_pred)
+            metrics[f"{prefix}_accuracy"] = float(acc)
+
+            # Precision, Recall, F1 (Weighted)
+            prec, rec, f1, _ = precision_recall_fscore_support(
+                y_true, y_pred, average="weighted"
+            )
+            metrics[f"{prefix}_precision_weighted"] = float(prec)
+            metrics[f"{prefix}_recall_weighted"] = float(rec)
+            metrics[f"{prefix}_f1_weighted"] = float(f1)
+
+        except Exception:
+             # En cas d'erreur sur les métriques de base, remplir avec NaN
+            metrics.setdefault(f"{prefix}_accuracy", float('nan'))
+            metrics.setdefault(f"{prefix}_precision_weighted", float('nan'))
+            metrics.setdefault(f"{prefix}_recall_weighted", float('nan'))
+            metrics.setdefault(f"{prefix}_f1_weighted", float('nan'))
+
+        # AUC Multiclasse non supporté, retourner NaN
+        metrics[f"{prefix}_auc_roc"] = float('nan')
+
+        # Ensure all values in the returned dict are standard floats
+        return {k: float(v) for k, v in metrics.items()}

src/trainers/huggingface/huggingface_transformer_trainer.py

@@ -2,16 +2,94 @@
 # Fichier: trainers/huggingface/huggingface_transformer_trainer.py
 # ============================================================
 
-from typing import Optional
+from typing import Optional, Dict, List, Any
 import cupy as cp
+import numpy as np
 import cudf
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+import torch.nn.functional as F # Pour softmax
+# Utiliser cuml.model_selection
+from cuml.model_selection import train_test_split
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, EvalPrediction
+from datasets import Dataset as HFDataset # Utiliser le type Dataset de Hugging Face pour plus de clarté
+# Utiliser cuml.metrics
+from cuml.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score
+# Importer cupy pour la conversion et argmax
+import cupy as cp
+# numpy n'est plus nécessaire ici
+# import numpy as np
 
 from base_trainer import BaseTrainer
 from config import Config
 
 
+# Fonction pour calculer les métriques en utilisant cuML (sans logs)
+def compute_metrics(p: EvalPrediction) -> Dict[str, float]:
+    logits = p.predictions
+    # Convertir les labels numpy en cupy
+    labels_cp = cp.asarray(p.label_ids)
+    # Obtenir les prédictions en appliquant argmax aux logits avec cupy
+    preds_cp = cp.argmax(cp.asarray(logits), axis=1) # Utilisation de cp.argmax
+
+    metrics: Dict[str, float] = {}
+
+    try:
+        # Accuracy avec cuML
+        acc = accuracy_score(labels_cp, preds_cp)
+        metrics["accuracy"] = float(acc)
+
+        # Precision, Recall, F1 (Weighted) avec cuML
+        prec, rec, f1, _ = precision_recall_fscore_support(
+            labels_cp, preds_cp, average='weighted'
+        )
+        metrics["precision_weighted"] = float(prec)
+        metrics["recall_weighted"] = float(rec)
+        metrics["f1_weighted"] = float(f1)
+
+    except Exception:
+        # Remplir avec NaN si erreur
+        metrics.setdefault("accuracy", float('nan'))
+        metrics.setdefault("precision_weighted", float('nan'))
+        metrics.setdefault("recall_weighted", float('nan'))
+        metrics.setdefault("f1_weighted", float('nan'))
+
+    # Calcul AUC (binaire seulement avec cuML)
+    auc: float = float('nan') # Default NaN
+    num_classes = logits.shape[1]
+
+    if num_classes == 2:
+        try:
+            # Obtenir les probabilités (softmax) et convertir en cupy
+            probas_torch = F.softmax(torch.tensor(logits), dim=-1)
+            probas_cp = cp.asarray(probas_torch)
+
+            # Utiliser les probas de la classe positive
+            proba_pos_class = probas_cp[:, 1]
+
+            # S'assurer que les types sont corrects pour cuML roc_auc_score
+            if labels_cp.dtype != cp.int32 and labels_cp.dtype != cp.int64:
+                labels_cp = labels_cp.astype(cp.int32)
+            if proba_pos_class.dtype != cp.float32 and proba_pos_class.dtype != cp.float64:
+                proba_pos_class = proba_pos_class.astype(cp.float32)
+
+            # Vérifier qu'il y a plus d'une classe dans les labels réels
+            unique_labels = cp.unique(labels_cp)
+            if len(unique_labels) >= 2:
+                 auc_score = roc_auc_score(labels_cp, proba_pos_class)
+                 auc = float(auc_score) # Cast to float
+            # else: # Pas de log
+
+        except (ValueError, TypeError, Exception):
+            # auc reste NaN en cas d'erreur, pas de log
+            pass
+    # else: # Pas de log pour le cas multiclasse
+        # auc reste NaN
+
+    metrics["auc_roc"] = auc
+
+    # Retourner les métriques avec les noms de base
+    return {k: float(v) for k, v in metrics.items()} # Assurer float standard
+
 class HuggingFaceTransformerTrainer(BaseTrainer):
     """
     Entraîneur spécifique Hugging Face, utilisant un tokenizer,
@@ -31,9 +109,13 @@ class HuggingFaceTransformerTrainer(BaseTrainer):
         :param target_column: Nom de la colonne cible dans vos données.
         """
         super().__init__(config, data_path, target_column)
+        super().__init__(config, data_path, target_column)
         self.tokenizer: Optional[AutoTokenizer] = None
         self.model: Optional[AutoModelForSequenceClassification] = None
         self.hf_trainer: Optional[Trainer] = None
+        self.train_dataset: Optional[HFDataset] = None
+        self.eval_dataset: Optional[HFDataset] = None
+        self.test_dataset: Optional[HFDataset] = None
 
     def build_components(self) -> None:
         """
@@ -47,12 +129,16 @@ class HuggingFaceTransformerTrainer(BaseTrainer):
             model_name)
         training_args = self._prepare_training_args()
 
-        # Le HF Trainer a besoin de datasets, qui sont construits
-        # dans le code de train/evaluate ou un data loader.
-        self.hf_trainer = Trainer(model=self.model,
-                                  args=training_args,
-                                  train_dataset=None,
-                                  eval_dataset=None)
+        # Le HF Trainer a besoin de datasets, qui sont construits dans train()
+        # On ajoute compute_metrics ici
+        self.hf_trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=self.train_dataset, # Sera défini dans train()
+            eval_dataset=self.eval_dataset,   # Sera défini dans train()
+            compute_metrics=compute_metrics,
+            tokenizer=self.tokenizer # Ajout du tokenizer pour le padding dynamique si besoin
+        )
 
     def train(self) -> None:
         """
@@ -68,18 +154,46 @@ class HuggingFaceTransformerTrainer(BaseTrainer):
         for col in features_df.columns[1:]:
             texts = texts.str.cat(features_df[col], sep=' ')
         texts_list = texts.to_arrow().to_pylist()
-        # Tokenization
-        encodings = self.tokenizer(texts_list, padding=True, truncation=True, return_tensors="pt")
-        # Création du dataset PyTorch (sous forme de liste de dictionnaires)
-        dataset = [{
-            "input_ids": encodings["input_ids"][i],
-            "attention_mask": encodings["attention_mask"][i],
-            "labels": torch.tensor(cp.asnumpy(labels)[i])
-        } for i in range(len(texts_list))]
-        # Assignation des datasets au Trainer HF
-        self.hf_trainer.train_dataset = dataset
-        self.hf_trainer.eval_dataset = dataset
+        # texts est une cudf.Series, labels est un cp.ndarray
+        # Utiliser cuml.model_selection.train_test_split directement
+        # Premier split: 80% train, 20% temp
+        X_train_text, X_temp_text, y_train, y_temp = train_test_split(
+            texts, labels, test_size=0.2, random_state=42, stratify=labels
+        )
+        # Deuxième split: 50% validation, 50% test sur temp (donne 10% val, 10% test du total)
+        X_val_text, X_test_text, y_val, y_test = train_test_split(
+            X_temp_text, y_temp, test_size=0.5, random_state=42, stratify=y_temp
+        )
+
+        # Fonction pour créer un dataset Hugging Face à partir de cudf.Series et cp.ndarray
+        def create_hf_dataset(text_series: cudf.Series, label_array: cp.ndarray) -> HFDataset:
+            # Convertir en listes Python pour le tokenizer et HF Dataset
+            texts_list = text_series.to_arrow().to_pylist()
+            # Convertir cupy array en numpy puis en liste pour HF Dataset
+            labels_list = cp.asnumpy(label_array).tolist()
+
+            encodings = self.tokenizer(texts_list, padding=True, truncation=True) # Pas de return_tensors="pt" ici
+            # Crée un dictionnaire compatible avec Dataset.from_dict
+            data_dict = {
+                "input_ids": encodings["input_ids"],
+                "attention_mask": encodings["attention_mask"],
+                "labels": labels_list
+            }
+            return HFDataset.from_dict(data_dict)
+
+        # Création des datasets
+        self.train_dataset = create_hf_dataset(X_train_text, y_train)
+        self.eval_dataset = create_hf_dataset(X_val_text, y_val)
+        self.test_dataset = create_hf_dataset(X_test_text, y_test) # Garder pour evaluate()
+
+        # Assignation des datasets au Trainer HF (déjà fait dans build_components mais on réassigne ici)
+        self.hf_trainer.train_dataset = self.train_dataset
+        self.hf_trainer.eval_dataset = self.eval_dataset
+
         # Lancement du fine‑tuning
+        print(f"Starting training with {len(self.train_dataset)} samples.")
+        print(f"Validation during training with {len(self.eval_dataset)} samples.")
+        print(f"Test set prepared with {len(self.test_dataset)} samples.")
         self.hf_trainer.train()
 
     def evaluate(self) -> dict:
@@ -87,10 +201,19 @@ class HuggingFaceTransformerTrainer(BaseTrainer):
         Évalue le modèle Hugging Face; la logique de calcul
         des métriques est en partie assurée par le HF Trainer.
         
-        :return: Dictionnaire contenant les métriques calculées.
+        :return: Dictionnaire contenant les métriques calculées sur l'ensemble de test.
         """
-        # À implémenter
-        return {}
+        if self.hf_trainer is None or self.test_dataset is None:
+            raise ValueError("Trainer or test dataset not initialized. Run train() first.")
+
+        print(f"Evaluating on the test set ({len(self.test_dataset)} samples)...")
+        # Utiliser predict pour obtenir les métriques sur le jeu de test
+        results = self.hf_trainer.predict(self.test_dataset)
+
+        # results.metrics contient déjà les métriques calculées par compute_metrics
+        # sur le test_dataset fourni.
+        print("Evaluation results:", results.metrics)
+        return results.metrics
 
     def _create_torch_dataset(self, texts: cudf.Series,
                               labels: cp.ndarray) -> torch.utils.data.Dataset:
@@ -103,8 +226,9 @@ class HuggingFaceTransformerTrainer(BaseTrainer):
         :return: Un Dataset PyTorch utilisable par Trainer.
         """
         # Implémentation possible : tokenization + construction d'un dataset custom.
+        # Cette méthode n'est plus directement utilisée car on crée les HFDatasets dans train()
         raise NotImplementedError(
-            "La méthode _create_torch_dataset est à implémenter selon vos besoins."
+            "La méthode _create_torch_dataset n'est plus utilisée directement."
         )
 
     def _prepare_training_args(self) -> TrainingArguments:
@@ -124,4 +248,11 @@ class HuggingFaceTransformerTrainer(BaseTrainer):
             warmup_steps=params.get("warmup_steps"),
             weight_decay=params.get("weight_decay"),
             adam_epsilon=params.get("adam_epsilon"),
+            # Ajout de paramètres importants pour l'évaluation
+            evaluation_strategy="epoch", # Évaluer à chaque époque
+            save_strategy="epoch",       # Sauvegarder le modèle à chaque époque
+            load_best_model_at_end=True, # Charger le meilleur modèle à la fin
+            metric_for_best_model="f1",  # Utiliser F1 pour déterminer le meilleur modèle (ou 'accuracy')
+            logging_dir=params.get("logging_dir", "./logs"), # Pour les logs TensorBoard
+            logging_steps=params.get("logging_steps", 10), # Fréquence des logs
         )