Subbing Variable into a ML model

Viewed 33

I want a function to take in a model parameter as a variable and sub that parameter into the ML model.

Example is:

def create_model(model_class):
    return sk.linear_model.{model_class}

I think it should be done using classes but i'm pretty lost.

I've tried creating a class but cannot get the output to be the ML model object.

1 Answers

Not sure why, but I think you're looking for something like this?

from sklearn import linear_model

def create_model(model_style: str, **kwargs):
    """Model builder

    Args:
        model_style: which type of linear model you want to use
        kwargs: kwargs for initializing the model_style

    Returns:
        sklearn linear model class
    """
    valid_models = vars(linear_model)['__all__']
    if model_style not in valid_models:
        raise ValueError("model_style is not valid")
    else:
        return getattr(linear_model, model_style)(**kwargs)

ridge = create_model('Ridge')
foo = create_model('foo')  # error
Related