Added code to train naive bayes