ctc_decoder.py

from ctcdecode import CTCBeamDecoder
import torch.nn.functional as F
import torch
import numpy as np
import editdistance


class Decoder:
    def __init__(self, labels, lm_path=None, alpha=1, beta=1.5, cutoff_top_n=40, cutoff_prob=0.99, beam_width=200, num_processes=24, blank_id=0):
        self.vocab_list = ['_'] + labels # NOTE: blank symbol
        self._decoder = CTCBeamDecoder(['_@'] + labels[1:], lm_path, alpha, beta, cutoff_top_n, cutoff_prob, beam_width, num_processes, blank_id)
        # NOTE: the whitespace symbol is replaced with an @ symbol for explicit modeling in char-based LMs

    def convert_to_string(self, tokens, seq_len=None):
        if not seq_len:
            seq_len = tokens.size(0)
        out = []
        for i in range(seq_len):
            if len(out) == 0:
                if tokens[i] != 0:
                    out.append(tokens[i])
            else:
                if tokens[i] != 0 and tokens[i] != tokens[i - 1]:
                    out.append(tokens[i])
        return ''.join(self.vocab_list[i] for i in out)
    
    def decode_beam(self, logits, seq_lens):
        decoded = []
        tlogits = logits.transpose(0, 1)
        beam_result, beam_scores, timesteps, out_seq_len = self._decoder.decode(tlogits.softmax(-1), seq_lens)
        for i in range(tlogits.size(0)):
            output_str = ''.join(map(lambda x: self.vocab_list[x], beam_result[i][0][:out_seq_len[i][0]]))
            decoded.append(output_str)
        return decoded

    def decode_greedy(self, logits, seq_lens):
        decoded = []
        tlogits = logits.transpose(0, 1)
        _, tokens = torch.max(tlogits, 2)
        for i in range(tlogits.size(0)):
            output_str = self.convert_to_string(tokens[i], seq_lens[i])
            decoded.append(output_str)
        return decoded
    
    def get_mean(self, decoded, gt, individual_length, func):
        total_norm  = 0.0
        length      = len(decoded)
        for i in range(0, length):
            val         = float(func(decoded[i], gt[i]))
            total_norm += val / individual_length
        return total_norm / length

    def wer(self, r, h):
        # initialisation
        d = np.zeros((len(r)+1)*(len(h)+1), dtype=np.uint8)
        d = d.reshape((len(r)+1, len(h)+1))
        for i in range(len(r)+1):
            for j in range(len(h)+1):
                if i == 0:
                    d[0][j] = j
                elif j == 0:
                    d[i][0] = i

        # computation
        for i in range(1, len(r)+1):
            for j in range(1, len(h)+1):
                if r[i-1] == h[j-1]:
                    d[i][j] = d[i-1][j-1]
                else:
                    substitution = d[i-1][j-1] + 1
                    insertion    = d[i][j-1] + 1
                    deletion     = d[i-1][j] + 1
                    d[i][j] = min(substitution, insertion, deletion)

        return d[len(r)][len(h)]

    def wer_sentence(self, r, h):
        return self.wer(r.split(), h.split())
    
    def cer_batch(self, decoded, gt):
        assert len(decoded) == len(gt), 'batch size mismatch: {}!={}'.format(len(decoded), len(gt))
        mean_indiv_len = np.mean([len(s) for s in gt])
        
        return self.get_mean(decoded, gt, mean_indiv_len, editdistance.eval)
        
    def wer_batch(self, decoded, gt):
        assert len(decoded) == len(gt), 'batch size mismatch: {}!={}'.format(len(decoded), len(gt))
        mean_indiv_len = np.mean([len(s.split()) for s in gt])
        
        return self.get_mean(decoded, gt, mean_indiv_len, self.wer_sentence)