Added code to generate test data from the same training examples, depending on user's choice