Hidden Markov Model with Viterbi I have a Hidden Markov model class with basically a single method: getting the best parse of a sequence of input tokens based on Viterbi.

While I have no hardcore benchmarks, I'd love some pointers to make it even a bit faster, as it (expectedly) takes quite a long time when the number of states is over 2000.

from math import log10 as lg

class HMM(object):
    def __init__(self, states, start_probs, transition_probs, emission_probs):
        # set(str)
        self.states = states

        # map(str, float)
        self.start_probs = start_probs

        # map(str, map(str, float))
        self.transition_probs = transition_probs

        # map(str, map(str, float))
        self.emission_probs = emission_probs


    def best_parse(self, input_tokens):
        p, path = HMM.__viterbi(
            input_tokens, self.states, self.start_probs, self.transition_probs, self.emission_probs
        )
        return '{0} =&gt; {1}\t{2}'.format(' '.join(input_tokens), path, lg(p))

    @classmethod
    def __viterbi(cls, observations, states, start_probability, transitionable, emission_probability):
        trellis = [{}]

        path = {}

        for state in states:
            trellis[0][state] = \
                start_probability.get(state, 0.0)
            path[state] = [state]

        for i, e in enumerate(observations[0:-1], 1):
            trellis.append({})
            path_for_current_state = {}

            for state in states:
                (probability, candidate_state) = \
                    max((trellis[i - 1][s] *
                        transitionable.get(s, {}).get(state, 0.0) *
                        emission_probability.get(state, {}).get(observations[i], float('-inf')), s)
                    for s in states)

                trellis[i][state] = probability
                path_for_current_state[state] = path[candidate_state] + [state]

            path = path_for_current_state


        # trellis[n-1][state] is the most likely parse for state.
        (probability, state) = max(
            (trellis[len(observations) - 1][state], state)
            for state in states
        )

        return (probability, ' '.join(t for (t, _) in path[state]))
1. Bugs


No use is made of the last observation. For example, if we have this model:

A = 'A', None
B = 'B', None
states = [A, B]
initial = {A:.75, B:.25}
transition = {}
emission = {A:dict(a=1), B:dict(b=1)}


And then we give it the single observation b:

&gt;&gt;&gt; print(HMM(states, initial, transition, emission).best_parse('b'))
b =&gt; A  -0.12493873660829995


The conclusion is that the most likely path was to start in state A. But this is impossible, because state A can't emit the observation b!

If you look at the Wikipedia description of the Viterbi algorithm, you'll see that the initial state probabilities have to take the first observation into account. So where the code in the post has:

trellis[0][state] = start_probability.get(state, 0.0)


it should instead be:

trellis[0][state] = (start_probability.get(state, 0.0)
    * emission_probability.get(state, {}).get(observations[0], 0.0))


and then the loop over observations[0:-1] should be over observations[1:].
If the emission table has no entry for an observation, then the probability defaults to \$−∞\$:

get(observations[i], float('-inf'))


But probabilities must be between 0 and 1, so \$−∞\$ is not a probability. This leads to wrong answers because the multiplication \$0 × −∞\$ is likely to turn up somewhere in the computation, and \$0 × −∞\$ is not a number:

&gt;&gt;&gt; 0 * float('-inf')
nan


(Probably you were thinking "log-likelihood" when you wrote float('-inf'). See §2.5 below.)


2. Review


There are no docstrings. What do these methods do? What should I pass? What do they return?
When you have a "class with basically a single method" then what you are looking for is a function. Not everything needs to be a class!
The result of the best_parse method is returned as a formatted string. This is not likely to be useful for anything other than development and debugging. It would be more useful to return the highest probability path as a list.
It seems very limiting that states are required to be pairs whose first element is a string.
The code computes probabilities by multiplying. But over the course of a lengthy sequence of observations this is likely to underflow the range of double-precision floating-point numbers. To avoid this, it would be better to work with the logarithms of probabilities, and add them instead of multiplying.
Just as you only need to keep two copies of the path dictionary (one for the previous observation, one for the current observation), you only need to keep two copies of the trellis dictionary.
The paths are extended by concatenatation: path[candidate_state] + [state]. But this requires copying out the whole of path[candidate_state], and as the paths get longer this takes more and more time. It would be better to represent the paths as linked lists, so that they can be constructed in constant time. Then the best path can be reconstructed by following the links.
The code requires observations to be a sequence so that it can be sliced. This means that all observations have to be acquired before you can start running the Viterbi algorithm. But since observations may take time to acquire, it would be nice if the Viterbi algorithm could be interleaved with the acquisition of the observations. This would be easy to do in Python by iterating over observations instead of slicing it.
Every time around the inner loop, this code:

max((trellis[i - 1][s] *
    transitionable.get(s, {}).get(state, 0.0) *
    emission_probability.get(state, {}).get(observations[i], float('-inf')), s)
for s in states)


has to construct two empty dictionaries (that will most likely not be needed and just get discarded), do four dictionary lookups, and build a tuple. All of this work could be avoided if you pre-processed the data so that (i) states are represented by consecutive integers starting at zero; (ii) observables are also represented by consecutive integers starting at zero; (iii) transitionable and emission_probability are represented by NumPy arrays of logarithms of probabilities (instead of dictionaries of dictionaries of probabilities).

Now, the multiplications become additions, and using NumPy we can do all the additions simultaneously like this:

logprob[:, np.newaxis] + transition + emission[:, event]



3. Revised code

import numpy as np

def likeliest_path(initial, transition, emission, events):
    """Find the likeliest path in a hidden Markov Model resulting in the
    given events.

    Arguments:
    initial: arraylike(n) --- probability of starting in each state
    transition: arraylike(n, n) -- probability of transition between states
    emission: arraylike(n, e) -- probability of emitting each event in
        each state
    events -- iterable of events

    Returns:
    path: list(int) -- list of states in the most probable path
    p: float -- log-likelihood of that path

    """
    # Use log-likelihoods to avoid floating-point underflow. Note that
    # we want -inf for the log of zero, so suppress warnings here.
    with np.errstate(divide='ignore'):
        initial = np.log10(initial)
        transition = np.log10(transition)
        emission = np.log10(emission)

    # List of arrays giving most likely previous state for each state.
    prev = []

    events = iter(events)
    logprob = initial + emission[:, next(events)]
    for event in events:
        # p[i, j] is log-likelihood of being in state j, having come from i.
        p = logprob[:, np.newaxis] + transition + emission[:, event]
        prev.append(np.argmax(p, axis=0))
        logprob = np.max(p, axis=0)

    # Most likely final state.
    best_state = np.argmax(logprob)

    # Reconstruct path by following links and then reversing.
    state = best_state
    path = [state]
    for p in reversed(prev):
        state = p[state]
        path.append(state)
    return path[::-1], logprob[best_state]