Function test_save_load_equality

catboost/python-package/ut/medium/test.py:1357–1398

(task_type)

def test_save_load_equality(task_type):
    output_model_path = test_output_path(OUTPUT_MODEL_PATH)

    def check_equality(model1, model2):
        assert model1 == model2
        for attr in ('best_score_', 'evals_result_', 'best_iteration_'):
            assert compare_with_limited_precision(getattr(model1, attr), getattr(model2, attr))

    def check_load_from_stream(model):
        cb_stream = CatBoost()
        with open(output_model_path, 'rb') as stream:
            cb_stream.load_model(stream=stream)
        check_equality(model, cb_stream)

    def check_load_from_string(model):
        cb_blob = CatBoost()
        cb_blob.load_model(blob=open(output_model_path, 'rb').read())
        check_equality(model, cb_blob)

    def fill_check_model(params, train_file, test_file, cd_file):
        model, _ = fit_from_file(params, train_file, test_file, cd_file)
        model.save_model(fname=output_model_path)
        check_load_from_string(model)
        check_load_from_stream(model)

    fill_check_model({'iterations': 10, 'task_type': task_type, 'gpu_ram_part': TEST_GPU_RAM_PART, 'devices': '0'}, TRAIN_FILE, TEST_FILE, CD_FILE)
    fill_check_model({'loss_function': 'RMSE', 'iterations': 10}, HIGGS_TRAIN_FILE, HIGGS_TEST_FILE, HIGGS_CD_FILE)

    params = {
        'dictionaries': [
            {'dictionary_id': 'UniGram', 'token_level_type': 'Letter', 'occurrence_lower_bound': '1'},
            {'dictionary_id': 'BiGram', 'token_level_type': 'Letter', 'occurrence_lower_bound': '1', 'gram_order': '2'},
            {'dictionary_id': 'Word', 'occurrence_lower_bound': '1'},
        ],
        'feature_calcers': ['NaiveBayes', 'BoW:top_tokens_count=10'],
        'iterations': 10,
        'loss_function': 'MultiClass',
        'task_type': task_type,
        'gpu_ram_part': TEST_GPU_RAM_PART,
        'devices': '0'
    }
    fill_check_model(params, ROTTEN_TOMATOES_TRAIN_FILE, ROTTEN_TOMATOES_TEST_FILE, ROTTEN_TOMATOES_CD_FILE)


def test_load_model_incorrect_argument(task_type):