Unnecessary level of abstraction about palma HOT 1 OPEN

from flaml import AutoML
from datetime import datetime
from palma.utils.utils import get_hash
from palma import logger

import logging
from contextlib import contextmanager

@contextmanager
def disable_logging():
    logging.disable()
    yield
    logging.basicConfig(level=logging.DEBUG)


class AbstractOptimizer:
    def __init__(self, **parameters):
        self.date = datetime.now()
        self.run_id = get_hash(date=self.date)
        self.parameters = parameters


class FlamlOptimizer(AbstractOptimizer):
    def __init__(self, **parameters):
        super().__init__(**parameters)
        self.optimizer = AutoML(**parameters)

    def find_best_model(self, project):
        X = project.X.iloc[project.validation_strategy.train_index]
        y = project.y.iloc[project.validation_strategy.train_index]
        splitter = project.validation_strategy
        split_type = None if splitter is None else splitter.splitter
        groups = None if splitter is None else splitter.groups
        groups = groups if groups is None else groups[splitter.train_index]
        with disable_logging():
            self.optimizer.fit(
                    X_train=pd.DataFrame(X.values, index=range(len(X))),
                    y_train=pd.Series(y.values, index=range(len(X))),
                    split_type=split_type, groups=groups,
                    mlflow_logging=False,
                    task=project.problem
                    )
        best_model = self.optimizer.model.model
        logger.logger.log_artifact(best_model, self.run_id)
        return best_model

optim = FlamlOptimizer(time_budget=2)
model = optim.find_best_model(project)