Pasted

class SpamDetector:
 
    def __init__(self):
        self.p_class = [0, 0]
        self.p_words = [{}, {}]
 
    # typical machine learning model interface
    def fit(self, x, y):
        self.p_class[1] = np.mean(y)
        self.p_class[0] = 1 - self.p_class[1]
        for c in [0, 1]:
            data = x[y == c]
            words = self._count_words(data)
            words_sum = sum(words.values())
            self.p_words[c] = {word: log(count / words_sum) for (word, count) in words.items()}
            
    def predict(self, d):
        proba = self.predict_proba(d)
        return proba[:, 1] < proba[:, 0]
 
    def predict_proba(self, d):
        ret = np.zeros((d.shape[0], 2))
        for i, msg in enumerate(d):
            for c in [0, 1]:
                ret[i, c] = self.p_class[c] * self.__p_message(msg, c)
        return ret
 
    # helper functions
    def _count_words(self, d):
        ret = defaultdict(lambda: 0)
        for line in d:
            for word in line:
                ret[word] += 1
        return ret
 
    def __p_message(self, msg, c):
        d = self.p_words[c]
        return sum(d[s] if s in d else 0 for s in msg)