Sweet implementation of Viterbi in Python

An implementation of the Viterbi algorithm for HMM in Python can be as short as 10 lines of code like this:

def Decode(self, obs):
    trellis = np.zeros((self.N, len(obs)))
    backpt = np.ones((self.N, len(obs)), 'int32') * -1
    trellis[:, 0] = np.squeeze(self.initialProb * self.Obs(obs[0]))
    for t in xrange(1, len(obs)):
        trellis[:, t] = (trellis[:, t-1, None].dot(self.Obs(obs[t]).T) * self.transProb).max(0)
        backpt[:, t] = (np.tile(trellis[:, t-1, None], [1, self.N]) * self.transProb).argmax(0)
    tokens = [trellis[:, -1].argmax()]
    for i in xrange(len(obs)-1, 0, -1):
        tokens.append(backpt[tokens[-1], i])
    return tokens[::-1]

This is the implementation of Viterbi I’ve completed recently. Holy python, how sweet it is! The code is even shorter than the pseudo-code, which is taken from this book.

viterbi_algo

It takes no more than 30 lines for the complete class:

import numpy as np

class Decoder(object):
    def __init__(self, initialProb, transProb, obsProb):
        self.N = initialProb.shape[0]
        self.initialProb = initialProb
        self.transProb = transProb
        self.obsProb = obsProb
        assert self.initialProb.shape == (self.N, 1)
        assert self.transProb.shape == (self.N, self.N)
        assert self.obsProb.shape[0] == self.N

    def Obs(self, obs):
        return self.obsProb[:, obs, None]

    def Decode(self, obs):
        trellis = np.zeros((self.N, len(obs)))
        backpt = np.ones((self.N, len(obs)), 'int32') * -1

        # initialization
        trellis[:, 0] = np.squeeze(self.initialProb * self.Obs(obs[0]))

        for t in xrange(1, len(obs)):
            trellis[:, t] = (trellis[:, t-1, None].dot(self.Obs(obs[t]).T) * self.transProb).max(0)
            backpt[:, t] = (np.tile(trellis[:, t-1, None], [1, self.N]) * self.transProb).argmax(0)
        # termination
        tokens = [trellis[:, -1].argmax()]
        for i in xrange(len(obs)-1, 0, -1):
            tokens.append(backpt[tokens[-1], i])
        return tokens[::-1]

The full source code and a simple test case, as always, can be found on github.

Advertisements

3 comments

  1. I don’t know you but i love you already for this beautiful code. Although i didn’t fully understand it but replicating and running the code made my evening. Thank you. do you have any other resources that could help me understand it better or perhaps some other source code that will be great for me to lay my hands on?

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s