# Naive Bayes Classifier for sentiment labelled documents

In order to continue improving my Python knowledge, I have implemented a naïve Bayes classifier as described in "An introduction to Information Retrieval". I would be very interested which parts could be improved, be it e.g. coding style or use of data structures.

"""Implementation of a naive Bayes classifier based on sentiment labelled sentences.
The dataset was obtained from
https://archive.ics.uci.edu/ml/datasets/Sentiment+Labelled+Sentences"""
import sys
import re
import math
from collections import Counter
from stop_words import get_stop_words

# PARAMETERS
DATAFILE = "data\\imdb_labelled.txt"

# FUNCTIONS
# A library is a list of categories, which label a list of documents
library = [[],[]]

# The textfile is formatted as document (string), TAB, category (int), NL
with open(filepath, 'r') as file:
for line in file:
document, category = line.split('\t')
library[int(category)].append(document)

return library

def clean_library(library):
"""Clean documents in the library array."""
for i, category in enumerate(library):
for j, document in enumerate(category):
cleaned_doc = clean_document(document)
library[i][j] = cleaned_doc

def clean_document(document):
"""Clean a document from stop words, numbers and various other
characters and return a list of all words."""
stop_words = get_stop_words('en')

new_doc = document.strip().lower()
new_doc = re.sub("[-0-9.,!;:\\/()\"&]", "", new_doc)
new_doc = new_doc.split()
new_doc = [word for word in new_doc if word not in stop_words]

return new_doc

def train_categories(library):
"""Calculate probabilities for the naive Bayes classifier and
return the vocabulary with conditional probabilities and the priors."""
total_docs = sum((len(category) for category in library))
vocabulary = [word for category in library
for document in category
for word in document]

cond_prob = []
prior = []

for category in library:
# Prior probability
total_cat_docs = len(category)
prior.append(total_cat_docs / total_docs)

# Conditional probabilities
text = [word for document in category for word in document]

word_count = Counter(text)
total_word_count = sum(word_count.values())

cat_cond_prob = {}

for word in vocabulary:
cat_cond_prob[word] = (word_count[word] + 1) / (total_word_count + 1)

cond_prob.append(cat_cond_prob)

return (vocabulary, prior, cond_prob)

def apply_nb(vocabulary, priors, cond_prob, document):
"""Apply the naive Bayes classification to a document in order
to retrieve its category."""
prepared_doc = clean_document(document)
prepared_doc = [word for word in prepared_doc if word in vocabulary]

score = [math.log(prior) for prior in priors]

for cat, cat_cond_prob in enumerate(cond_prob):
score[cat] = sum((math.log(cat_cond_prob[word]) for word in prepared_doc))

return score.index(max(score))

def main(argv):
"""Train a naive Bayes classifier and apply it to a user-supplied string."""
if len(argv) == 0:
return

user_doc = argv