Chaobin Tang (唐超斌)

# We are using numpy for computation, and Python NLTK for some utilities 

# A word on NLTK
# NLTK perhaps remains the go-to recommendations for many people when it comes to NLP.
# There are a couple of things people liked about this library:
# 1. It has an abundant collection of corpora, installed with nltk.download()
# 2. It has an abundant collection of algorithms for many NLP routines,
#   such as tokenization, stemming, lemmatization, etc.
# 3. It supports a lot of languages.
#
# However, NLTK falls short when it comes to performance. If you look at
# the source code implementation of algorithms such as that of some
# tokenizers in nltk, those are pure Python code.
# When we are processing a large amount of text, or that our model is
# training with online data when data has to be feed-in and
# processed in real time, the heavy cost of pure Python
# code easily becomes bottleneck. With text processing,
# a lot of str objects are being created and manipulated.
# These operations have extensive memory footprints, also because str objects
# are immutable in Python, the many intermediate operations over
# a str objects cast away the the object soon it finishes off with it,
# causing the garbage collector to be very busy while
# keeps the system busy alocating memory in the heap.
#
# For this reason, efforts such as spacy (https://spacy.io/) emerged.
# Spacy is implemented in Cython, and claims to be the-fastest
# in several preprocessing algorithms. The point here is
# that even though Python may be the most popular choice in NLP,
# one should be aware of Python's shortcomings and
# also know that people developed solutions for it.

import numpy as np
import nltk
print("numpy: ", np.__version__)
print("nltk: ", nltk.__version__)

# labeled email samples from http://csmining.org/index.php/spam-email-datasets-.html
import os
path_emails = '/Users/cbt/Projects/isaac/data'
path_label = os.path.join(path_emails, 'SPAMTrain.label')
IS_SPAM = 0
NO_SPAM = 1

import glob

import mailparser

def extract_text_from_email(eml_file):
    mail = mailparser.parse_from_file(eml_file)
    return os.linesep.join([mail.subject, mail.body])

def parse_emails(src, saveto):
    emls = glob.glob(os.path.join(path_emls, '*.eml'))
    texts = ((f, extract_text_from_email(f)) for f in emls)
    for (f, text) in texts:
        name = os.path.basename(f)
        name, ext = os.path.splitext(name)
        name = '%s.msg' % name
        dest = os.path.join(saveto, name)       
        with open(dest, 'w') as out:
            out.write(text)

# a set utility functions
def head(sequence_or_map):
    if isinstance(sequence_or_map, list):
        return sequence_or_map[0]
    elif isinstance(sequence_or_map, dict):
        return next(iter(sequence_or_map.items()))

# read each email into words and do the following:
# - remove stopwords
# - remove punctuation
# - possibly remove high frequency words

import re
import glob
import codecs
import string
import math
from collections import Counter

from bs4 import BeautifulSoup
from nltk.tokenize.casual import casual_tokenize
from nltk.corpus import stopwords
from nltk.probability import FreqDist
from nltk.stem import WordNetLemmatizer
from nltk.stem import SnowballStemmer


STOPWORDS = set(stopwords.words('english'))
EXTRA_STOPWORDS = set(["i'm", "even", "it's", "using", "also", "don't", "isn't",
                      "i've", "within", "without"])
STOPWORDS = set(STOPWORDS) | EXTRA_STOPWORDS
PUNCT = set(list(string.punctuation))
REMOVED_WORDS = (STOPWORDS | PUNCT)
REMOVED_WORDS = (REMOVED_WORDS | set(list(string.ascii_lowercase)) | set(list(string.digits)))
REMOVED_WORDS |= set(['...', '..', ])


from nltk.stem.api import StemmerI
from nltk.compat import python_2_unicode_compatible


@python_2_unicode_compatible
class SimpleStemmer(StemmerI):
    
    IRREGULAR_FORMS = {
        'email': ['e-mail']
    }
    
    def __init__(self):

        self.pool = {}
        for word in self.IRREGULAR_FORMS:
            for form in self.IRREGULAR_FORMS[word]:
                self.pool[form] = word

    def stem(self, word):
        return self.pool.get(word, word)

    def __repr__(self):
        return '<%s>'.format(self.__class__.__name__)

LEMMATIZER = WordNetLemmatizer()
STEMMERS = [SnowballStemmer('english'), SimpleStemmer()]

class Vocabulary(Counter):
    
    def size(self):
        return len(self)
    
    def drop(self, threshold=None, keep=(), remove=(), pattern=None):
        if threshold:
            try:
                lt, gt = threshold
            except ValueError as e:
                raise ValueError('threshold must be a 2-item tuple, %s given' % str(threshold))
        else:
            lt, gt = -1, math.inf
            
        deletion = []
        for (k, count) in self.items():
            if (
                (count <= lt or count >= gt) or
                (k in remove) or
                (pattern and pattern.match(k))
            ) and not (k in keep):
                deletion.append(k)
        size = self.size()
        for k in deletion:
            del self[k]
        print('purged %d word' % (size - self.size()))
    
    def search(self, pattern):
        matches = []
        for (word, count) in self.items():
            if pattern.match(word):
                matches.append(word)
        return matches
    
def load(filename):
    with codecs.open(filename, "r",
                     encoding='utf-8', errors='replace') as raw:
        text = raw.read()
    return text

def read_raw():
    '''
    The reason this is used over the corpus reader in NLP
    is to allow dealing with unicode error.
    '''
    texts = {}
    files = glob.glob('%s/*.msg' % path_emails)
    for fname in files:
        text = load(fname)
        texts[os.path.basename(fname)] = text
    return texts

def read_labels():
    labels = {}
    with open(path_label) as label_file:
        for line in label_file:
            int_label, filename = line.split(' ')
            filename = filename.strip()
            filename = filename.replace('eml', 'msg')
            labels[filename] = int(int_label)
    return labels

def remove_stopwords(words):
    _words = []
    
    for word in words: # O(n)
        _word = word.lower()
        if not _word in REMOVED_WORDS: # O(1)
            _words.append(_word)
    return _words

def clean_html(text):
    soup = BeautifulSoup(text, 'html5lib')
    return soup.get_text()

def stem(word):
    _word = word
    for stemmer in STEMMERS:
        _word = stemmer.stem(word)
    return _word

def tokenize(text):
    # some sanitization
    text = clean_html(text)
    # tokenize, preserving URL
    # There are a variety of tokenizers in NLP, and there is no "BEST" tokenizer.
    # A choice is usually made considering the specific application.
    words = [stem(LEMMATIZER.lemmatize(w.lower()))
             for w in casual_tokenize(text, preserve_case=False)
             if not w.lower() in set()]
    # remove stopwords and some punctuations (not all)
    words = remove_stopwords(words)
    return words

def parse_emails(texts):
    named_emails = []
    for (fname, text) in texts.items():
        words = tokenize(text)
        named_emails.append((fname, words))
    print('%d emails parsed' % len(named_emails))
    return named_emails

def merge(parsed_emails, labels):
    data = []
    for (name, words) in parsed_emails:
        data.append((words, labels[name]))
    return data

from collections import Counter


class ModelUntrained(Exception): pass
class ModelNotReady(Exception): pass
class Classifier(object):
    
    def __init__(self):
        self.likelihood = None
        self.num_classes = None
        self.priors = None
        
        self.V_C = None
        self.V = None
        
        self.num_trainings = 0
        self.distribution = {}
        
        self._prior_updated = False
        self._likelihood_updated = False
    
    def stable(self):
        return all((
                self._prior_updated, self._likelihood_updated))
    
    def require_update(self):
        self._prior_updated = self._likelihood_updated = False
    
    def update_vocabulary(self, training_set, drop_threshold=None,
                          keep=(), remove=(), pattern=None):
        self.num_trainings += len(training_set)
        label_count = {}
        # initialize vocabulary
        self.V_C = self.V_C or {}
        self.V = self.V or Vocabulary()

        # build vocabulary of all classes and of each class
        for (words, label) in training_set:
            v_c = self.V_C.setdefault(label, Vocabulary())
            v_c.update(words)
            self.V.update(words)
            self.distribution[label] = (self.distribution.setdefault(label, 0) + 1)
        
        for (_, v_c) in self.V_C.items():
            v_c.drop(threshold=drop_threshold, keep=keep, remove=remove, pattern=pattern)
        self.V.drop(threshold=drop_threshold, keep=keep, remove=remove, pattern=pattern)
        print('updated vocabulary size:', self.V.size())
        self.require_update()
    
    def update_prior(self):
        self.priors = self.priors or {}
        # calculate the priors
        for (label, count) in self.distribution.items():
            self.priors[label] = math.log(count / self.num_trainings)
        self.num_classes = len(self.priors)
        self._prior_updated = True

    def update_likelihood(self):
        # calculate the probability of word in class-vocabulary
        # against all-vocabulary
        self.likelihood = self.likelihood or {}
        for (c, v_c) in self.V_C.items():
            likelihood = self.likelihood.setdefault(c, {})
            for (word, count) in v_c.items():
                # Lapalace smoothing (additive smoothing)
                p = (count + 1) / (v_c.size() + self.V.size())
                likelihood[word] = math.log(p)
        self._likelihood_updated = True
                
    def update(self):
        self.update_prior()
        self.update_likelihood()
    
    def preprocess(self, filename):
        text = load(filename)
        words = tokenize(text)
        return words
    
    def train(self, training_set):
        self.update_vocabulary(training_set)
        self.update()
        
    def classify(self, words):
        if not self.stable(): raise ModelNotReady()
        if isinstance(words, str): # filename
            words = self.preprocess(words)
        probabilities = np.zeros(
            self.num_classes,
            dtype=np.float32)
        classes = []
        for (i, (label, v_c)) in enumerate(self.likelihood.items()):
            classes.append(label)
            prior = self.priors[label]
            probabilities[i] = prior
            for word in words:
                p = v_c.get(word)
                if p is None: continue
                probabilities[i] += (probabilities[i] + p)
        return classes[probabilities.argmax()]
    

def cross_validate(testing_set, labels, classifier):
    error = 0
    for (i, (name, words)) in enumerate(testing_set):
        label = labels[name]
        prediction = classifier.classify(words)
        if label != prediction:
            error += 1
    return error / len(testing_set)

# The I/O operation that loads data into Python
texts = read_raw()
labels = read_labels()

# The preprocessing of all training data
# This is the most computation-expensive part of the process. But this can be carried out
# offline.
parsed_emails = parse_emails(texts)

# split the training set
SPLIT = int(len(parsed_emails) * (0.9))
training_set = parsed_emails[:SPLIT]
testing_set = parsed_emails[SPLIT:]

PTN_REMOVE = re.compile(r'_+|^\w\d$|^\d+px$|^#\w+$')
REMOVE = set(('wa', '�', '=d', 'font-size', 'ha', 'color', 'br',
         'text', 'text-decoration', 'de', 'font-weight',
         'content-transfer-encoding', 'le', 'content-type',
         'sans-serif', '-8859-1', '•', 'bb',
         'quoted-printable', 'font-family', 'plain',
         'charset', 'spam', 'ba', 'dqo', 'bc', 'ab',
         'af', '©', 'line-height', '..', 'email', ';}', 'year', 'list'))
classifier = Classifier()
classifier.update_vocabulary(merge(training_set, labels),
                             drop_threshold=(10, math.inf),
                             keep=('free',),
                             remove=REMOVE,
                             pattern=PTN_REMOVE
                            )
classifier.update()

cross_validate(testing_set, labels, classifier)

# The word usage in the spams:
classifier.V_C[IS_SPAM].most_common(100)

classifier.V_C[NO_SPAM].most_common(100)