hmm.py

# coding=utf-8

import sys
import json
from math import log
from collections import defaultdict as dict

from common import read_sequence_data
from common import plot_sequence_data

def argmax(d):
    k = max(d, key = lambda x: d[x])
    return k, d[k]

class HiddenMarkovModel:
    def __init__(self):

        self.prob_s3 = dict(float)
        self.prob_xs = dict(float)
        self.D = dict()
        self.V = dict(set)
        
    def train(self, X, S, smooth = 0.5):
        x_count  = dict(int)
        xs_count = dict(int)
        
        unigram = dict(int)
        bigram  = dict(int)
        trigram = dict(int)

        # maximum likelihood estimate
        for x_list, s_list in zip(X, S):
            if len(x_list) == 0 or len(x_list) != len(s_list):
                print >> sys.stderr, 'ERROR: x_list =', x_list, ', s_list =', s_list
                continue
            bigram[('*', '*')] += 1
            bigram[('*', s_list[0])] += 1
            trigram[('*', '*', s_list[0])] += 1       
            if len(s_list) > 1:
                trigram[('*', s_list[0], s_list[1])] += 1
            
            for p in range(len(x_list)):
                x, s = x_list[p], s_list[p]
                
                self.V[x].add(s)
                x_count[x] += 1
                xs_count[(x, s)] += 1
                unigram[s] += 1

                if p < 1: continue
                s_i = s_list[p - 1]
                bigram[(s_i, s)] += 1

                if p < 2: continue
                s_i = s_list[p - 2]
                s_j = s_list[p - 1]
                trigram[(s_i, s_j, s)] += 1

            bigram[(s_list[-1], 'STOP')] += 1
        
            if len(s_list) < 2: continue
            trigram[(s_list[-2], s_list[-1], 'STOP')] += 1

        self.D = unigram

        # smoothing
        for s in self.D:
            trigram[('*', '*', s)] += smooth
        bigram[('*', '*')] += smooth * len(self.D)
        
        for s_i in self.D:
            for s_j in self.D:
                trigram[('*', s_i, s_j)] += smooth
            bigram[('*', s_i)] += smooth * len(self.D)

        for s_i in self.D:
            for s_j in self.D:
                trigram[(s_i, s_j, 'STOP')] += smooth
                bigram[(s_i, s_j)] += smooth

        # normalize
        for s_i in self.D:
            for s_j in self.D:
                for s_k in self.D:
                    trigram[(s_i, s_j, s_k)] += smooth
                bigram[(s_i, s_j)] += smooth * len(self.D)

        for x, s in xs_count:
            self.prob_xs[(x, s)] = 1.0 * xs_count[(x, s)] / unigram[s]

        for s_i, s_j, s_k in trigram:
            self.prob_s3[(s_i, s_j, s_k)] = 1.0 * trigram[(s_i, s_j, s_k)] / bigram[(s_i, s_j)]


    def viterbi(self, x_list):

        n = len(x_list)
        x_list = [''] + x_list        
        s_list = [''] * (n + 1)

        T = dict()        
        T[-1] = T[0] = ['*']
        for i in range(1, n + 1):
            if x_list[i] in self.V:
                T[i] = self.V[x_list[i]]
            else:
                T[i] = set(self.D)

        pi = dict(float)
        bp = dict(str)
        
        pi[(0, '*', '*')] = 1.0
        for k in range(1, n + 1):
            for u in T[k - 1]:
                for v in T[k]:
                    K = dict()
                    for w in T[k - 2]:
                        if (x_list[k], v) in self.prob_xs:
                            p = self.prob_xs[(x_list[k], v)]
                        else:
                            p = 1.0
                        K[w] = pi[(k - 1, w, u)] * self.prob_s3[(w, u, v)] * p

                    bp[(k, u, v)], pi[(k, u, v)] = argmax(K)

        K = dict()
        for u in T[n - 1]:
            for v in T[n]:
                K[(u, v)] = pi[(n, u, v)] * self.prob_s3[(u, v, 'STOP')]
        
        (s_list[n - 1], s_list[n]), val = argmax(K)

        for k in range(n - 2, 0, -1):    
            s_list[k] = bp[(k + 2, s_list[k + 1], s_list[k + 2])]
            
        return s_list[1 : ]
  
    
    def baseline(self, X, S):

        m = n = c = 0

        k, v = argmax(self.D)

        for x_list, s_list in zip(X, S):
            s_list_pred = [''] * len(s_list)

            for i in range(len(x_list)):
                x = x_list[i]
                if not x in self.V:
                    s_list_pred[i] = k  
                else:
                    d = {(x, s) : self.prob_xs[(x, s)] for s in self.V[x] if (x, s) in self.prob_xs}
                    (x_max, s_max), v_max = argmax(d)
                    s_list_pred[i] = s_max

            for p, q in zip(s_list, s_list_pred):
                if p == q: m += 1
        
            n += len(s_list)
            c += 1

            # print >> sys.stderr, 'Ground truth:'
            # plot_sequence_data(x_list, s_list)            
            # print >> sys.stderr, 'Tagging result:'
            # plot_sequence_data(x_list, s_list_pred)            

        # print >> sys.stderr, 'Baseline accuracy : %lf%% (%d/%d)' % (100.0 * m / n, m, n)

        return 1.0 * m / n


 
    def test(self, X, S):

        m = n = c = 0
        for x_list, s_list in zip(X, S):
            s_list_pred = self.viterbi(x_list)
            assert(len(s_list) == len(s_list_pred))
        
            for p, q in zip(s_list, s_list_pred):
                if p == q: m += 1
        
            n += len(s_list)
            c += 1

            # print >> sys.stderr, 'Ground truth:'
            # plot_sequence_data(x_list, s_list)            
            # print >> sys.stderr, 'Tagging result:'
            # plot_sequence_data(x_list, s_list_pred)            
        
        # print >> sys.stderr, 'Accuracy for HMM POS-tagger : %lf%% (%d/%d)' % (100.0 * m / n, m, n)

        return 1.0 * m / n


if __name__ == '__main__':

    train_path = 'data/pos_tagging.train'
    test_path = 'data/pos_tagging.test'
    
    # train_path = 'data/pos-tagging/ictrain'
    # test_path = 'data/pos-tagging/ictest'

    X_train, S_train = read_sequence_data(open(train_path))
    X_test, S_test = read_sequence_data(open(test_path))

    tagger = HiddenMarkovModel()

    tagger.train(X_train, S_train)
    baseline_acc = tagger.baseline(X_test, S_test)
    viterbi_acc = tagger.test(X_test, S_test)

    print >> sys.stderr, 'Baseline accuracy %lf%%' % (100.0 * baseline_acc)
    print >> sys.stderr, 'Accuracy for HMM POS-tagger %lf%%' % (100.0 * viterbi_acc)