"""Utility methods for EvalML's model families."""fromevalml.model_family.model_familyimportModelFamily
[docs]defhandle_model_family(model_family):"""Handles model_family by either returning the ModelFamily or converting from a string. Args: model_family (str or ModelFamily): Model type that needs to be handled. Returns: ModelFamily Raises: KeyError: If input is not a valid model family. ValueError: If input is not a string or ModelFamily object. """ifisinstance(model_family,str):try:tpe=ModelFamily[model_family.upper()]returntpeexceptKeyError:raiseKeyError("Model family '{}' does not exist".format(model_family))ifisinstance(model_family,ModelFamily):returnmodel_familyraiseValueError("`handle_model_family` was not passed a str or ModelFamily object")