src/mlpipeline/automl/automl_factory.py

from typing import Union
from mlpipeline.automl.autogluon_trainer import AutoGluonTrainer
from mlpipeline.automl.flaml_trainer import FLAMLTrainer
from mlpipeline.automl.pycaret_trainer import PyCaretTrainer
from mlpipeline.logging.logger import get_logger

logger = get_logger(__name__)


class AutoMLFactory:
    @staticmethod
    def create_trainer(library: str, config: dict) -> Union[AutoGluonTrainer, FLAMLTrainer, PyCaretTrainer]:
        library = library.lower()
        
        if library == 'autogluon':
            logger.info("Creating AutoGluon trainer")
            return AutoGluonTrainer(config)
        
        elif library == 'flaml':
            logger.info("Creating FLAML trainer")
            return FLAMLTrainer(config)
        
        elif library == 'pycaret':
            logger.info("Creating PyCaret trainer")
            return PyCaretTrainer(config)
        
        else:
            raise ValueError(f"Unknown AutoML library: {library}. Choose from: autogluon, flaml, pycaret")
    
    @staticmethod
    def get_available_libraries():
        return ['autogluon', 'flaml', 'pycaret']