def test_save_nn_model():
"""Test save nn model"""
classifier = ReviewClassifier('nn')
classifier.train('test-reviews.csv')
classifier.save_model()
assert os.path.exists('trained_nn_model.tar')
os.remove('trained_nn_model.tar')
def test_save_svm_model():
"""Test save svm model"""
classifier = ReviewClassifier('svm')
classifier.train('test-reviews.csv')
classifier.save_model()
assert os.path.exists('trained_svm_model.pickle')
os.remove('trained_svm_model.pickle')
def main():
""" Main function.
"""
# Review sentiment classifier
review_classifier = ReviewClassifier(
'nb') # Try 'nb', 'dt', 'rf', and 'nn'
review_classifier.train('citalopram-reviews.csv')
review_classifier.save_model()