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)
