kaldi_decoder.py

#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from concurrent.futures import ThreadPoolExecutor
import logging
from omegaconf import MISSING
import os
import torch
from typing import Optional
import warnings


from dataclasses import dataclass
from fairseq.dataclass import FairseqDataclass
from .kaldi_initializer import KaldiInitializerConfig, initalize_kaldi


logger = logging.getLogger(__name__)


@dataclass
class KaldiDecoderConfig(FairseqDataclass):
    hlg_graph_path: Optional[str] = None
    output_dict: str = MISSING

    kaldi_initializer_config: Optional[KaldiInitializerConfig] = None

    acoustic_scale: float = 0.5
    max_active: int = 10000
    beam_delta: float = 0.5
    hash_ratio: float = 2.0

    is_lattice: bool = False
    lattice_beam: float = 10.0
    prune_interval: int = 25
    determinize_lattice: bool = True
    prune_scale: float = 0.1
    max_mem: int = 0
    phone_determinize: bool = True
    word_determinize: bool = True
    minimize: bool = True

    num_threads: int = 1


class KaldiDecoder(object):
    def __init__(
        self,
        cfg: KaldiDecoderConfig,
        beam: int,
        nbest: int = 1,
    ):
        try:
            from kaldi.asr import FasterRecognizer, LatticeFasterRecognizer
            from kaldi.base import set_verbose_level
            from kaldi.decoder import (
                FasterDecoder,
                FasterDecoderOptions,
                LatticeFasterDecoder,
                LatticeFasterDecoderOptions,
            )
            from kaldi.lat.functions import DeterminizeLatticePhonePrunedOptions
            from kaldi.fstext import read_fst_kaldi, SymbolTable
        except:
            warnings.warn(
                "pykaldi is required for this functionality. Please install from https://github.com/pykaldi/pykaldi"
            )

        # set_verbose_level(2)

        self.acoustic_scale = cfg.acoustic_scale
        self.nbest = nbest

        if cfg.hlg_graph_path is None:
            assert (
                cfg.kaldi_initializer_config is not None
            ), "Must provide hlg graph path or kaldi initializer config"
            cfg.hlg_graph_path = initalize_kaldi(cfg.kaldi_initializer_config)

        assert os.path.exists(cfg.hlg_graph_path), cfg.hlg_graph_path

        if cfg.is_lattice:
            self.dec_cls = LatticeFasterDecoder
            opt_cls = LatticeFasterDecoderOptions
            self.rec_cls = LatticeFasterRecognizer
        else:
            assert self.nbest == 1, "nbest > 1 requires lattice decoder"
            self.dec_cls = FasterDecoder
            opt_cls = FasterDecoderOptions
            self.rec_cls = FasterRecognizer

        self.decoder_options = opt_cls()
        self.decoder_options.beam = beam
        self.decoder_options.max_active = cfg.max_active
        self.decoder_options.beam_delta = cfg.beam_delta
        self.decoder_options.hash_ratio = cfg.hash_ratio

        if cfg.is_lattice:
            self.decoder_options.lattice_beam = cfg.lattice_beam
            self.decoder_options.prune_interval = cfg.prune_interval
            self.decoder_options.determinize_lattice = cfg.determinize_lattice
            self.decoder_options.prune_scale = cfg.prune_scale
            det_opts = DeterminizeLatticePhonePrunedOptions()
            det_opts.max_mem = cfg.max_mem
            det_opts.phone_determinize = cfg.phone_determinize
            det_opts.word_determinize = cfg.word_determinize
            det_opts.minimize = cfg.minimize
            self.decoder_options.det_opts = det_opts

        self.output_symbols = {}
        with open(cfg.output_dict, "r") as f:
            for line in f:
                items = line.rstrip().split()
                assert len(items) == 2
                self.output_symbols[int(items[1])] = items[0]

        logger.info(f"Loading FST from {cfg.hlg_graph_path}")
        self.fst = read_fst_kaldi(cfg.hlg_graph_path)
        self.symbol_table = SymbolTable.read_text(cfg.output_dict)

        self.executor = ThreadPoolExecutor(max_workers=cfg.num_threads)

    def generate(self, models, sample, **unused):
        """Generate a batch of inferences."""
        # model.forward normally channels prev_output_tokens into the decoder
        # separately, but SequenceGenerator directly calls model.encoder
        encoder_input = {
            k: v for k, v in sample["net_input"].items() if k != "prev_output_tokens"
        }
        emissions, padding = self.get_emissions(models, encoder_input)
        return self.decode(emissions, padding)

    def get_emissions(self, models, encoder_input):
        """Run encoder and normalize emissions"""
        model = models[0]

        all_encoder_out = [m(**encoder_input) for m in models]

        if len(all_encoder_out) > 1:

            if "encoder_out" in all_encoder_out[0]:
                encoder_out = {
                    "encoder_out": sum(e["encoder_out"] for e in all_encoder_out)
                    / len(all_encoder_out),
                    "encoder_padding_mask": all_encoder_out[0]["encoder_padding_mask"],
                }
                padding = encoder_out["encoder_padding_mask"]
            else:
                encoder_out = {
                    "logits": sum(e["logits"] for e in all_encoder_out)
                    / len(all_encoder_out),
                    "padding_mask": all_encoder_out[0]["padding_mask"],
                }
                padding = encoder_out["padding_mask"]
        else:
            encoder_out = all_encoder_out[0]
            padding = (
                encoder_out["padding_mask"]
                if "padding_mask" in encoder_out
                else encoder_out["encoder_padding_mask"]
            )

        if hasattr(model, "get_logits"):
            emissions = model.get_logits(encoder_out, normalize=True)
        else:
            emissions = model.get_normalized_probs(encoder_out, log_probs=True)

        return (
            emissions.cpu().float().transpose(0, 1),
            padding.cpu() if padding is not None and padding.any() else None,
        )

    def decode_one(self, logits, padding):
        from kaldi.matrix import Matrix

        decoder = self.dec_cls(self.fst, self.decoder_options)
        asr = self.rec_cls(
            decoder, self.symbol_table, acoustic_scale=self.acoustic_scale
        )

        if padding is not None:
            logits = logits[~padding]

        mat = Matrix(logits.numpy())

        out = asr.decode(mat)

        if self.nbest > 1:
            from kaldi.fstext import shortestpath
            from kaldi.fstext.utils import (
                convert_compact_lattice_to_lattice,
                convert_lattice_to_std,
                convert_nbest_to_list,
                get_linear_symbol_sequence,
            )

            lat = out["lattice"]

            sp = shortestpath(lat, nshortest=self.nbest)

            sp = convert_compact_lattice_to_lattice(sp)
            sp = convert_lattice_to_std(sp)
            seq = convert_nbest_to_list(sp)

            results = []
            for s in seq:
                _, o, w = get_linear_symbol_sequence(s)
                words = list(self.output_symbols[z] for z in o)
                results.append(
                    {
                        "tokens": words,
                        "words": words,
                        "score": w.value,
                        "emissions": logits,
                    }
                )
            return results
        else:
            words = out["text"].split()
            return [
                {
                    "tokens": words,
                    "words": words,
                    "score": out["likelihood"],
                    "emissions": logits,
                }
            ]

    def decode(self, emissions, padding):
        if padding is None:
            padding = [None] * len(emissions)

        ret = list(
            map(
                lambda e, p: self.executor.submit(self.decode_one, e, p),
                emissions,
                padding,
            )
        )
        return ret