Spaces:
Sleeping
Sleeping
File size: 3,929 Bytes
19d70f4 | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | from pathlib import Path
from typing import Dict, Any, Optional
import pandas as pd
from pycaret.classification import setup as classification_setup, compare_models as classification_compare, finalize_model as classification_finalize, save_model as classification_save, load_model as classification_load
from pycaret.regression import setup as regression_setup, compare_models as regression_compare, finalize_model as regression_finalize, save_model as regression_save, load_model as regression_load
from mlpipeline.logging.logger import get_logger
logger = get_logger(__name__)
class PyCaretTrainer:
def __init__(self, config: Dict[str, Any]):
self.config = config
self.model: Optional[Any] = None
self.is_classification = None
def train(self, train_data: pd.DataFrame, target_column: str, model_path: Path) -> Dict[str, float]:
logger.info("Starting PyCaret training")
if train_data[target_column].dtype == 'object' or train_data[target_column].nunique() < 20:
self.is_classification = True
setup_fn = classification_setup
compare_fn = classification_compare
finalize_fn = classification_finalize
save_fn = classification_save
else:
self.is_classification = False
setup_fn = regression_setup
compare_fn = regression_compare
finalize_fn = regression_finalize
save_fn = regression_save
exp = setup_fn(
data=train_data,
target=target_column,
session_id=self.config.get('session_id', 42),
fold=self.config.get('fold', 5),
verbose=self.config.get('verbose', False),
use_gpu=self.config.get('use_gpu', False),
)
best_model = compare_fn(
n_select=self.config.get('n_select', 5),
verbose=self.config.get('verbose', False),
)
if self.config.get('tuning', {}).get('enabled', True):
from pycaret.classification import tune_model as classification_tune
from pycaret.regression import tune_model as regression_tune
tune_fn = classification_tune if self.is_classification else regression_tune
best_model = tune_fn(
best_model,
n_iter=self.config.get('tuning', {}).get('n_iter', 10),
optimize=self.config.get('tuning', {}).get('optimize', 'Accuracy'),
)
self.model = finalize_fn(best_model)
save_fn(self.model, str(model_path / 'model'))
from pycaret.classification import pull as classification_pull
from pycaret.regression import pull as regression_pull
pull_fn = classification_pull if self.is_classification else regression_pull
results = pull_fn()
metrics = {
'score': float(results.iloc[0]['Mean']) if not results.empty else 0.0,
}
logger.info(f"PyCaret training completed. Score: {metrics['score']}")
return metrics
def predict(self, data: pd.DataFrame) -> pd.Series:
if self.model is None:
raise ValueError("Model not trained. Call train() first.")
from pycaret.classification import predict_model as classification_predict
from pycaret.regression import predict_model as regression_predict
predict_fn = classification_predict if self.is_classification else regression_predict
predictions = predict_fn(self.model, data=data)
return predictions.iloc[:, -1]
def load(self, model_path: Path):
logger.info(f"Loading PyCaret model from {model_path}")
load_fn = classification_load if self.is_classification else regression_load
self.model = load_fn(str(model_path / 'model'))
return self |