Function test_custom_class_labels

catboost/python-package/ut/medium/test.py:2001–2045

(loss_function, label_type, class_count, task_type)

    ]
)
def test_custom_class_labels(loss_function, label_type, class_count, task_type):
    if label_type == 'boolean':
        assert (class_count == 2)
        labels = [False, True]
    elif label_type == 'int':
        labels = {2 : [1, 2], 5 : [1, 3, 2, 10, 20]}[class_count]
    elif label_type == 'string':
        labels = {
            2 : ['Class1', 'Class2'],
            5 : ['Class1', 'Class3', 'Class2', 'Class10', 'Class20']
        }[class_count]

    prng = np.random.RandomState(seed=0)
    train_pool = Pool(prng.random_sample(size=(100, 10)), label=prng.choice(labels, size=100))
    test_features = prng.random_sample(size=(50, 10))
    test_label = prng.choice(labels, size=50)

    classifier = CatBoostClassifier(iterations=2, loss_function=loss_function, thread_count=8, task_type=task_type, gpu_ram_part=TEST_GPU_RAM_PART, devices='0')

    if (loss_function == 'Logloss') and (class_count != 2):
        with pytest.raises(CatBoostError):
            classifier.fit(train_pool)
    else:
        classifier.fit(train_pool)
        output_model_path = test_output_path(OUTPUT_MODEL_PATH)
        classifier.save_model(output_model_path)
        new_classifier = CatBoostClassifier()
        new_classifier.load_model(output_model_path)
        pred = new_classifier.predict_proba(test_features)
        classes = new_classifier.predict(test_features)
        assert pred.shape == (50, class_count)
        assert all(((class1 in labels) for class1 in classes))
        preds_path = test_output_path(PREDS_TXT_PATH)
        np.savetxt(preds_path, np.array(pred), fmt='%.8f')

        score = new_classifier.score(test_features, test_label)

        score_path = test_output_path(SCORE_PATH)
        with open(score_path, 'w') as score_file:
            score_file.write(f'{score:.8f}')

        return [
            local_canonical_file(preds_path, diff_tool=get_limited_precision_dsv_diff_tool(1e-6, False)),
            local_canonical_file(score_path, diff_tool=get_limited_precision_dsv_diff_tool(1e-6, False)),
        ]


def test_multiclass_custom_class_labels_from_files(task_type):