update_metadata_util.py

import numpy as np
import openml


classification_tasks = [
    232,
    236,
    241,
    245,
    253,
    254,
    256,
    258,
    260,
    262,
    267,
    271,
    273,
    275,
    279,
    288,
    336,
    340,
    2119,
    2120,
    2121,
    2122,
    2123,
    2125,
    2356,
    3044,
    3047,
    3048,
    3049,
    3053,
    3054,
    3055,
    75089,
    75092,
    75093,
    75098,
    75100,
    75108,
    75109,
    75112,
    75114,
    75115,
    75116,
    75118,
    75120,
    75121,
    75125,
    75126,
    75129,
    75131,
    75133,
    75134,
    75136,
    75139,
    75141,
    75142,
    75143,
    75146,
    75147,
    75148,
    75149,
    75153,
    75154,
    75156,
    75157,
    75159,
    75161,
    75163,
    75166,
    75169,
    75171,
    75173,
    75174,
    75176,
    75178,
    75179,
    75180,
    75184,
    75185,
    75187,
    75192,
    75195,
    75196,
    75199,
    75210,
    75212,
    75213,
    75215,
    75217,
    75219,
    75221,
    75223,
    75225,
    75232,
    75233,
    75234,
    75235,
    75236,
    75237,
    75239,
    75250,
    126021,
    126024,
    126028,
    126030,
    126031,
    146574,
    146575,
    146576,
    146577,
    146578,
    146583,
    146586,
    146592,
    146593,
    146594,
    146596,
    146597,
    146600,
    146601,
    146602,
    146603,
    146679,
    166859,
    166866,
    166872,
    166875,
    166882,
    166897,
    166905,
    166906,
    166913,
    166915,
    166931,
    166932,
    166944,
    166950,
    166951,
    166953,
    166956,
    166957,
    166958,
    166959,
    166970,
    166996,
    167085,
    167086,
    167087,
    167088,
    167089,
    167090,
    167094,
    167096,
    167097,
    167099,
    167100,
    167101,
    167103,
    167105,
    167106,
    167202,
    167203,
    167204,
    167205,
    168785,
    168791,
    189779,
    189786,
    189828,
    189829,
    189836,
    189840,
    189841,
    189843,
    189844,
    189845,
    189846,
    189857,
    189858,
    189859,
    189863,
    189864,
    189869,
    189870,
    189875,
    189878,
    189880,
    189881,
    189882,
    189883,
    189884,
    189887,
    189890,
    189893,
    189894,
    189899,
    189900,
    189902,
    190154,
    190155,
    190156,
    190157,
    190158,
    190159,
    211720,
    211721,
    211722,
    211723,
    211724,
]
regression_tasks = [
    359997,
    359998,
    359999,
    360000,
    360001,
    360002,
    360003,
    167146,
    360004,
    360005,
    360006,
    360007,
    211696,
    360009,
    360010,
    360011,
    360012,
    360013,
    360014,
    360015,
    360016,
    360017,
    360018,
    360019,
    360020,
    360021,
    360022,
    360023,
    360024,
    360025,
    360026,
    360027,
    360028,
    360029,
    360030,
    360031,
    360032,
    360033,
    360034,
    360035,
    360036,
    360037,
    360038,
    360039,
    360040,
    360041,
    360042,
    360043,
    360044,
    360045,
    360046,
    360047,
    360048,
    360049,
    360050,
    360051,
    360052,
    360053,
    360054,
    360055,
    360056,
    360057,
    360058,
    360059,
    360060,
    360061,
    360062,
    360063,
    360064,
    360066,
    360067,
    360068,
    360069,
    360070,
    360071,
    360072,
    360073,
    360074,
    360075,
    360076,
    360077,
    360078,
    360079,
    360080,
    360081,
    360082,
    360083,
    360084,
    360085,
    360086,
    360087,
    360088,
    360089,
    360090,
    360091,
    360092,
    360093,
    360094,
    360095,
    360096,
    360097,
    360098,
    360100,
    360101,
    360102,
    360103,
    360104,
    360105,
    360106,
    360107,
    360108,
]


def load_task(task_id):
    task = openml.tasks.get_task(task_id)
    X, y = task.get_X_and_y()
    train_indices, test_indices = task.get_train_test_split_indices()
    X_train = X[train_indices]
    y_train = y[train_indices]
    X_test = X[test_indices]
    y_test = y[test_indices]
    dataset = openml.datasets.get_dataset(task.dataset_id)
    _, _, cat, _ = dataset.get_data(target=task.target_name)
    name = dataset.name.lower()
    del _
    del dataset
    cat = {i: "categorical" if c else "numerical" for i, c in enumerate(cat)}

    if isinstance(task, openml.tasks.OpenMLClassificationTask):
        task_type = "classification"
    elif isinstance(task, openml.tasks.OpenMLRegressionTask):
        task_type = "regression"
    else:
        raise ValueError("Unknown task type")

    return X_train, y_train, X_test, y_test, cat, task_type, name