further refactor of naive bayes code