3
\$\begingroup\$

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.
See https://nlp.stanford.edu/IR-book/pdf/irbookonlinereading.pdf for the algorithm.
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
def load_data(filepath):
    """Load the sentiment labelled data."""
    # 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:
        print("Please supply a document string.")
        return

    user_doc = argv[0]

    library = load_data(DATAFILE)
    clean_library(library)
    vocabulary, priors, cond_prob = train_categories(library)

    doc_cat = apply_nb(vocabulary, priors, cond_prob, user_doc)
    print(f'"{user_doc}": Category {doc_cat}')

if __name__ == "__main__":
    main(sys.argv[1:])

\$\endgroup\$
  • \$\begingroup\$ Add either clearer docstrings indicating what types are taken in and returned, or new python type annotations. Ex. what does "load_data" return? \$\endgroup\$ – Zachary Vance Nov 4 at 2:32
  • \$\begingroup\$ Define a data type "bayes model" to hold: "vocabulary, priors, cond_prob" \$\endgroup\$ – Zachary Vance Nov 4 at 2:33

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.