viterbi.py

def get_transition_probs(labels, tag_transformer):
    tags = [tag for tag, idx in tag_transformer.items()]
    n = len(tags)
    start_cnts, transition_cnts = np.zeros(n), np.zeros((n, n))
    
    for label in labels:
        for idx in range(len(label)):
            if idx == 0:
                start_cnts[tag_transformer[label[idx]]] += 1
            else:
                transition_cnts[tag_transformer[label[idx-1]], tag_transformer[label[idx]]] += 1
                
    start_cnts = start_cnts/np.sum(start_cnts)
    transition_cnts = (transition_cnts.T/np.sum(transition_cnts, axis=1)).T
    
    return start_cnts, transition_cnts

def viterbi_decoding(emission_probs, transition_probs, start_probs, tag_inverse_transformer):
    n, m = emission_probs.shape
    viterbi_state = np.zeros((n, m))
    
    for i in range(n):
        if i == 0:
            viterbi_state[i] = emission_probs[i] * start_probs
        else:
            viterbi_state[i] = np.max(np.multiply.outer(viterbi_state[i-1], emission_probs[i]) * transition_probs, axis=0)
    
    output_states = np.zeros(n)
    
    for i in reversed(range(n)):
        if i == n-1:
            output_states[i] = np.argmax(viterbi_state[i])
        else:
            nxt_state = int(output_states[i+1])
            output_states[i] = np.argmax(viterbi_state[i] * emission_probs[i+1,nxt_state] * transition_probs[:,nxt_state])
            
    output_states = [tag_inverse_transformer[int(x)] for x in output_states]
    
    return output_states