-added check for test set input parameter
-remove unneeded comment
-TODO: evaluate and test extended prediction function