python自己实现朴素贝叶斯分类模型

import numpy as np
x = [['happy','new','year','every','day','sunny',],
     ['sunny','happy','slow','great','cool'],
     ['sad','bad','no'],
     ['sad','hard','worry'],
     ['happy','love','warm','dream','sweat'],
     ['hard','cold','slow']]
y = [0,0,1,1,0,1]
def vocabVec(x,y):
    vocabSet = set()
    for i in x:
        vocabSet = vocabSet|set(i)
    vocabList = list(vocabSet)
    a = np.zeros((len(x),len(vocabList)))
    for index,i in enumerate(x):
        for j in i:
            a[index][vocabList.index(j)]+=1
    X = np.array(x)
    Y = np.array(y)
    P0 = a[Y == 0]
    p0Sum = np.sum(P0,axis=0)/np.sum(P0)
    p0class = len(Y[Y == 0])/float(len(Y))
    P1 = a[Y == 1]
    p1Sum = np.sum(P1,axis=0)/np.sum(P0)
    p1class = len(Y[Y == 1])/float(len(Y))

    return vocabList,p0class,p1class,p0Sum,p1Sum

def predict(wordList,vocabList,p0class,p1class,p0Sum,p1Sum):
    vocabVec = np.zeros(len(vocabList))
    for i in wordList:
        if i in vocabList:
            vocabVec[vocabList.index(i)] = 1

    wordListp0Sum = np.log(p0class)
    wordListp1Sum = np.log(p1class)

    for index,i in enumerate(vocabVec):
        if i !=0 :
            if p0Sum[index] != 0:
                wordListp0Sum += np.log(p0Sum[index])
            if p1Sum[index] != 0:
                wordListp1Sum += np.log(p1Sum[index])
    print(wordListp0Sum)
    print(wordListp1Sum)
    if wordListp0Sum <= wordListp1Sum:
        return 0
    else:
        return 1


if __name__ == '__main__':
    print('')
    vocabList, p0class, p1class, p0Sum, p1Sum = vocabVec(x,y)
    print('vocabList:',vocabList)
    print('p0class:',p0class)
    print('p1class:',p1class)
    print('p0Sum:',p0Sum)
    print('p1Sum:',p1Sum)
    print(predict(['dream','hard','sad'],vocabList, p0class, p1class, p0Sum, p1Sum))
    print(predict(['good','great','happy'],vocabList, p0class, p1class, p0Sum, p1Sum))