Viterbi algorithm on Markov Model

viterbi.py

"""
Viterbi algorithm on Markov Model.
Copyright (c) 2012 by Naoaki Okazaki
"""

import json
import math
import operator
import sys

DEFAULT = -100.     # The score for unseen parameters.

class Node:
    def __init__(self, prev=None, score=None):
        self.prev = prev
        self.score = score

def viterbi(S, T, seq):
    tbl = [{} for x in seq]

    for label in S.iterkeys():
        tbl[0][label] = Node(None, S[label].get(seq[0], DEFAULT))

    for t in range(1, len(seq)):
        for cur in S.iterkeys():
            node = Node()
            for prev in S.iterkeys():
                score = tbl[t-1][prev].score + T[prev].get(cur, DEFAULT)
                if node.score is None or node.score < score:
                    node = Node(prev, score)
            node.score += S[cur].get(seq[t], DEFAULT)
            tbl[t][cur] = node

    t = len(seq)-1
    L = ['' for i in seq]
    L[t] = max(tbl[t].iteritems(), key=lambda x: x[1].score)[0]
    while 0 < t:
        L[t-1] = tbl[t][L[t]].prev
        t -= 1
    return L

if __name__ == '__main__':
    model = json.load(open(sys.argv[1]))

    for line in sys.stdin:
        tokens = line.strip('\n').split(' ')
        labels = viterbi(model['S'], model['T'], tokens)
        print ' '.join(['%s/%s' % item for item in zip(tokens, labels)])