def get_problem_type()

in bayesmark/data.py [0:0]

• 9 lines of code
• 5 McCabe index (conditional complexity)

def get_problem_type(dataset_name):
    """Determine if this dataset is a regression of classification problem.

    Parameters
    ----------
    dataset : str
        Which data set to use, must be key in `DATA_LOADERS` dict, or name of custom csv file.

    Returns
    -------
    problem_type : ProblemType
        `Enum` to indicate if regression of classification data set.
    """
    if dataset_name in DATA_LOADERS:
        _, problem_type = DATA_LOADERS[dataset_name]
        return problem_type

    # Maybe we can come up with a better system, but for now let's use a convention based on the naming of the csv file.
    if dataset_name.startswith("reg-"):
        return ProblemType.reg
    if dataset_name.startswith("clf-"):
        return ProblemType.clf
    assert False, "Can't determine problem type from dataset name."