Function load_simple_dataset_as_lists

catboost/python-package/ut/medium/test.py:304–317

(is_test)

# returns (features_data, labels)
def load_simple_dataset_as_lists(is_test):
    features_data = []
    labels = []
    with open(TEST_FILE if is_test else TRAIN_FILE) as data_file:
        for line in data_file:
            elements = line[:-1].split('\t')
            features_data.append([])
            for column_idx, element in enumerate(elements):
                if column_idx == TARGET_IDX:
                    labels.append(element)
                else:
                    features_data[-1].append(element if column_idx in CAT_COLUMNS else float(element))

    return features_data, labels


# Test cases begin here ########################################################