main.py

import json, re
from collections import Counter
import bayes

TRAIN = 'data/train.json'
VALIDATE = 'data/validate.json'

train = json.loads(open(TRAIN).read())
validate = json.loads(open(VALIDATE).read())

def test(dataset, outcomes):
    answers = dict([x.split(" ") for x in open(dataset + "_validate.txt").read().split("\n")[:-1]])

    bayes.model_dict = bayes.train(train[dataset])

    correct_by_outcome = Counter()
    incorrect_by_outcome = Counter()

    for point in validate[dataset]:
        words = set(bayes.tokenize(point['contents']))
        prediction = bayes.predict(outcomes, words)
        answer = answers[point['name']]
        if prediction == answer:
            correct_by_outcome[answer] += 1
        else:
            incorrect_by_outcome[answer] += 1

    print(correct_by_outcome)
    print(incorrect_by_outcome)

test('tweets', set(['positive', 'negative']))
test('emails', set(['spam', 'ham']))