Maximum Likelihood Estimation (MLE) for Hidden Markov Model (HMM)

hmmtrain.py

"""
Maximum Likelihood Estimation (MLE) for Hidden Markov Model (HMM).
Copyright (c) 2012 by Naoaki Okazaki
"""

import collections
import json
import math
import sys

def logprob(V):
    n = sum(V.itervalues())
    for x, f in V.iteritems():
        V[x] = math.log(f / n)

def train(D):
    S = collections.defaultdict(lambda: collections.defaultdict(float))
    T = collections.defaultdict(lambda: collections.defaultdict(float))

    for seq in D:
        prev = None
        for token, label in seq:
            S[label][token] += 1
            if prev is not None:
                T[prev][label] += 1
            prev = label

    map(logprob, S.itervalues())
    map(logprob, T.itervalues())
    return S, T

def readiter(fi):
    seq = []
    for line in fi:
        line = line.strip('\n')
        if not line:
            yield seq
            seq = []
        else:
            seq.append(line.split('\t'))

if __name__ == '__main__':
    S, T = train(readiter(sys.stdin))
    json.dump({'S': S, 'T': T}, sys.stdout, indent=2)