Spaces:
Sleeping
Sleeping
File size: 2,693 Bytes
bf5fb5f 8c7aba9 bed4774 bf5fb5f bed4774 bf5fb5f 8c7aba9 bf5fb5f 8c7aba9 bf5fb5f 8c7aba9 bf5fb5f 8c7aba9 |
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 |
# ===========================================
# Fichier: utilities/cuml_pyfunc_wrapper.py
# ===========================================
from typing import Any, Dict, Optional
import pandas as pd
import numpy as np
import cudf
import cupy as cp
import pickle
import os
import mlflow
from interfaces.vectorizer import Vectorizer
class CuMLPyFuncWrapper(mlflow.pyfunc.PythonModel):
"""
Classe wrapper pour intégration de modèles cuML dans MLflow PyFunc,
permettant le chargement et l'inférence.
"""
def __init__(self, vectorizer: Vectorizer, classifier: object) -> None:
"""
Initialise le wrapper avec un vectorizer cuDF/cupy et le classifieur cuML.
:param vectorizer: Instance implémentant l'interface Vectorizer.
:param classifier: Modèle cuML déjà entraîné (ex. SVC, RandomForest, etc.).
"""
self.vectorizer: Vectorizer = vectorizer
self.classifier: object = classifier
def load_context(self, context: Dict[str, Any]) -> None:
"""
Chargé par MLflow PyFunc lors du rechargement du modèle,
par exemple pour initialiser l'environnement.
:param context: Contexte de chargement contenant d'éventuelles informations
sur l'environnement ou l'emplacement d'artefacts.
"""
# Charger le vectorizer depuis les artefacts
vectorizer_path = os.path.join(context.artifacts["vectorizer"])
with open(vectorizer_path, "rb") as f:
self.vectorizer = pickle.load(f)
# Charger le classifier depuis les artefacts
classifier_path = os.path.join(context.artifacts["classifier"])
with open(classifier_path, "rb") as f:
self.classifier = pickle.load(f)
def predict(self, context: Dict[str, Any],
model_input: pd.DataFrame) -> np.ndarray:
"""
Fonction de prédiction, appelée par MLflow PyFunc.
Convertit model_input en cudf, vectorise, puis appelle le modèle cuML.
:param context: Contexte d'exécution éventuel.
:param model_input: Données d'entrée sous forme de DataFrame pandas.
:return: Un vecteur numpy des prédictions.
"""
# Convertir le DataFrame pandas en DataFrame cuDF
cudf_input = cudf.DataFrame(model_input)
# Vectoriser les données d'entrée
X_vectorized = self.vectorizer.transform(cudf_input)
# Effectuer la prédiction avec le classifieur cuML
predictions_gpu = self.classifier.predict(X_vectorized)
# Convertir les prédictions GPU en array numpy pour MLflow
return cp.asnumpy(predictions_gpu)
|