Add prediction type specifier

Previous commit forgot to add a prediction type parameter. Unit tests and model used for unit test are changed correspondingly.