def test_multiclass_prediction(predict_mock, dask_client): predicted_probabilities = [[0, 0, 0.99], [0, 0.99, 0], [0.99, 0, 0], [0, 0.99, 0], [0, 0, 0.99]] predicted_indexes = [2, 1, 0, 1, 2] expected_result = ['c', 'b', 'a', 'b', 'c'] predict_mock.return_value = np.array(predicted_probabilities) classifier = AutoMLClassifier( time_left_for_this_task=1, per_run_time_limit=1, dask_client=dask_client, ) classifier.InputValidator = InputValidator(is_classification=True) classifier.InputValidator.target_validator.fit( pd.DataFrame(expected_result, dtype='category'), ) classifier.InputValidator._is_fitted = True actual_result = classifier.predict([None] * len(predicted_indexes)) np.testing.assert_array_equal(expected_result, actual_result)
def test_multilabel_prediction(predict_mock, dask_client): predicted_probabilities = [[0.99, 0], [0.99, 0], [0, 0.99], [0.99, 0.99], [0.99, 0.99]] predicted_indexes = np.array([[1, 0], [1, 0], [0, 1], [1, 1], [1, 1]]) predict_mock.return_value = np.array(predicted_probabilities) classifier = AutoMLClassifier( time_left_for_this_task=1, per_run_time_limit=1, dask_client=dask_client, ) classifier.InputValidator = InputValidator(is_classification=True) classifier.InputValidator.target_validator.fit( pd.DataFrame(predicted_indexes, dtype='int64'), ) classifier.InputValidator._is_fitted = True assert classifier.InputValidator.target_validator.type_of_target == 'multilabel-indicator' actual_result = classifier.predict([None] * len(predicted_indexes)) np.testing.assert_array_equal(predicted_indexes, actual_result)