Viterbi decoder for top_k best paths

lucas_alvarez_lacasa · October 7, 2022, 10:10am

Hello team!.
I was wondering if there’s a similar function to the viterbi_decode but instead of returning the best path, it returns the top_k best paths.

I couldn’t find anything on the internet.

Let me know. Best.

Bhack · October 8, 2022, 1:04pm

No but I think that you can quite easily adapt this with the same TF ops:

gist.github.com

https://gist.github.com/PetrochukM/afaa3613a99a8e7213d2efdd02ae4762

top_k_viterbi.py

import torch

# Credits to AllenNLP for the base implementation and base tests:
# https://github.com/allenai/allennlp/blob/master/allennlp/nn/util.py#L174

# Modified AllenNLP `viterbi_decode` to support `top_k` sequences efficiently.
def viterbi_decode(tag_sequence: torch.Tensor, transition_matrix: torch.Tensor, top_k: int=5):
    """
    Perform Viterbi decoding in log space over a sequence given a transition matrix
    specifying pairwise (transition) potentials between tags and a matrix of shape

This file has been truncated. show original

/cc @XiaoquanKong

lucas_alvarez_lacasa · October 8, 2022, 1:40pm

Thank you Bhack! This was the same implementation I found yesterday and actually did what you suggested, translating the Torch functions to Tensorflow.

As soon as I have the working code, I’ll post it here so anyone can use it in the future.

Bhack · October 8, 2022, 1:46pm

You could try to propose it to keras-nlp if they are interested: