chore: import upstream snapshot with attribution

nemo/collections/asr/modules/flashlight_decoder.py

@@ -0,0 +1,295 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import math
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.core.classes import NeuralModule
+
+
+class _TokensWrapper:
+    def __init__(self, vocabulary: List[str], tokenizer: TokenizerSpec):
+        self.vocabulary = vocabulary
+        self.tokenizer = tokenizer
+
+        if tokenizer is None:
+            self.reverse_map = {self.vocabulary[i]: i for i in range(len(self.vocabulary))}
+
+        self.vocab_len = len(self.vocabulary)
+
+        if (self.tokenizer is not None) and hasattr(self.tokenizer, 'unk_id') and self.tokenizer.unk_id is not None:
+            self.unknown_id = self.tokenizer.unk_id
+        elif ' ' in self.vocabulary:
+            self.unknown_id = self.token_to_id(' ')
+        elif '<unk>' in self.vocabulary:
+            self.unknown_id = self.token_to_id('<unk>')
+        else:
+            self.unknown_id = -1
+
+    @property
+    def blank(self):
+        return self.vocab_len
+
+    @property
+    def unk_id(self):
+        return self.unknown_id
+
+    @property
+    def vocab(self):
+        return self.vocabulary
+
+    @property
+    def vocab_size(self):
+        # the +1 is because we add the blank id
+        return self.vocab_len + 1
+
+    def token_to_id(self, token: str):
+        if token == self.blank:
+            return -1
+
+        if self.tokenizer is not None:
+            return self.tokenizer.token_to_id(token)
+        else:
+            return self.reverse_map[token]
+
+    def text_to_tokens(self, text: str):
+        if self.tokenizer is not None:
+            return self.tokenizer.text_to_tokens(text)
+        else:
+            return list(text)
+
+
+class FlashLightKenLMBeamSearchDecoder(NeuralModule):
+    '''
+    @property
+    def input_types(self):
+        """Returns definitions of module input ports.
+        """
+        return {
+            "log_probs": NeuralType(('B', 'T', 'D'), LogprobsType()),
+        }
+
+    @property
+    def output_types(self):
+        """Returns definitions of module output ports.
+        """
+        return {"hypos": NeuralType(('B'), PredictionsType())}
+    '''
+
+    def __init__(
+        self,
+        lm_path: str,
+        vocabulary: List[str],
+        tokenizer: Optional[TokenizerSpec] = None,
+        lexicon_path: Optional[str] = None,
+        boost_path: Optional[str] = None,
+        beam_size: int = 32,
+        beam_size_token: int = 32,
+        beam_threshold: float = 25.0,
+        lm_weight: float = 2.0,
+        word_score: float = -1.0,
+        unk_weight: float = -math.inf,
+        sil_weight: float = 0.0,
+    ):
+
+        try:
+            from flashlight.lib.text.decoder import (
+                CriterionType,
+                KenLM,
+                LexiconDecoder,
+                LexiconDecoderOptions,
+                SmearingMode,
+                Trie,
+            )
+            from flashlight.lib.text.dictionary import create_word_dict, load_words
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "FlashLightKenLMBeamSearchDecoder requires the installation of flashlight python bindings "
+                "from https://github.com/flashlight/text. Please follow the build instructions there."
+            )
+
+        super().__init__()
+
+        self.criterion_type = CriterionType.CTC
+        self.tokenizer_wrapper = _TokensWrapper(vocabulary, tokenizer)
+        self.vocab_size = self.tokenizer_wrapper.vocab_size
+        self.blank = self.tokenizer_wrapper.blank
+        self.silence = self.tokenizer_wrapper.unk_id
+
+        if lexicon_path is not None:
+            self.lexicon = load_words(lexicon_path)
+            self.word_dict = create_word_dict(self.lexicon)
+            self.unk_word = self.word_dict.get_index("<unk>")
+
+            # loads in the boosted words if given via a file
+            if boost_path is not None:
+                with open(boost_path, 'r', encoding='utf_8') as fr:
+                    boost_words = [line.strip().split('\t') for line in fr]
+                    boost_words = {w[0]: w[1] for w in boost_words}
+            else:
+                boost_words = {}
+
+            # add OOV boosted words to word_dict so it gets picked up in LM obj creation
+            for word in boost_words.keys():
+                if word not in self.lexicon:
+                    self.word_dict.add_entry(word)
+
+            # loads in the kenlm binary and combines in with the dictionary object from the lexicon
+            # this gives a mapping between each entry in the kenlm binary and its mapping to whatever
+            # numeraire is used by the AM, which is explicitly mapped via the lexicon
+            # this information is ued to build a vocabulary trie for decoding
+            self.lm = KenLM(lm_path, self.word_dict)
+            self.trie = Trie(self.vocab_size, self.silence)
+
+            start_state = self.lm.start(False)
+            for i, (word, spellings) in enumerate(self.lexicon.items()):
+                word_idx = self.word_dict.get_index(word)
+                _, score = self.lm.score(start_state, word_idx)
+                for spelling in spellings:
+                    spelling_idxs = [self.tokenizer_wrapper.token_to_id(token) for token in spelling]
+                    if self.tokenizer_wrapper.unk_id in spelling_idxs:
+                        print(f'tokenizer has unknown id for word[ {word} ] {spelling} {spelling_idxs}', flush=True)
+                        continue
+                    self.trie.insert(
+                        spelling_idxs, word_idx, score if word not in boost_words else float(boost_words[word])
+                    )
+            # handle OOV boosted words
+            for word, boost in boost_words.items():
+                if word not in self.lexicon:
+                    word_idx = self.word_dict.get_index(word)
+                    spelling = self.tokenizer_wrapper.text_to_tokens(word)
+                    spelling_idxs = [self.tokenizer_wrapper.token_to_id(token) for token in spelling]
+                    if self.tokenizer_wrapper.unk_id in spelling_idxs:
+                        print(f'tokenizer has unknown id for word[ {word} ] {spelling} {spelling_idxs}', flush=True)
+                        continue
+                    self.trie.insert(spelling_idxs, word_idx, float(boost))
+            self.trie.smear(SmearingMode.MAX)
+
+            self.decoder_opts = LexiconDecoderOptions(
+                beam_size=beam_size,
+                beam_size_token=int(beam_size_token),
+                beam_threshold=beam_threshold,
+                lm_weight=lm_weight,
+                word_score=word_score,
+                unk_score=unk_weight,
+                sil_score=sil_weight,
+                log_add=False,
+                criterion_type=self.criterion_type,
+            )
+
+            self.decoder = LexiconDecoder(
+                self.decoder_opts,
+                self.trie,
+                self.lm,
+                self.silence,
+                self.blank,
+                self.unk_word,
+                [],
+                False,
+            )
+        else:
+            from flashlight.lib.text.decoder import LexiconFreeDecoder, LexiconFreeDecoderOptions
+
+            d = {
+                w: [[w]]
+                for w in self.tokenizer_wrapper.vocab + ([] if '<unk>' in self.tokenizer_wrapper.vocab else ['<unk>'])
+            }
+            self.word_dict = create_word_dict(d)
+            self.lm = KenLM(lm_path, self.word_dict)
+            self.decoder_opts = LexiconFreeDecoderOptions(
+                beam_size=beam_size,
+                beam_size_token=int(beam_size_token),
+                beam_threshold=beam_threshold,
+                lm_weight=lm_weight,
+                sil_score=sil_weight,
+                log_add=False,
+                criterion_type=self.criterion_type,
+            )
+            self.decoder = LexiconFreeDecoder(self.decoder_opts, self.lm, self.silence, self.blank, [])
+
+    def _get_tokens(self, idxs: List[int]):
+        """Normalize tokens by handling CTC blank, ASG replabels, etc."""
+
+        idxs = (g[0] for g in itertools.groupby(idxs))
+        if self.silence < 0:
+            idxs = filter(lambda x: x != self.blank and x != self.silence, idxs)
+        else:
+            idxs = filter(lambda x: x != self.blank, idxs)
+        idxs = list(idxs)
+        if idxs[0] == self.silence:
+            idxs = idxs[1:]
+        if idxs[-1] == self.silence:
+            idxs = idxs[:-1]
+
+        return torch.LongTensor(idxs)
+
+    def _get_timesteps(self, token_idxs: List[int]):
+        """Returns frame numbers corresponding to every non-blank token.
+        Parameters
+        ----------
+        token_idxs : List[int]
+            IDs of decoded tokens.
+        Returns
+        -------
+        List[int]
+            Frame numbers corresponding to every non-blank token.
+        """
+
+        timesteps = []
+        for i, token_idx in enumerate(token_idxs):
+            if token_idx == self.blank:
+                continue
+            if i == 0 or token_idx != token_idxs[i - 1]:
+                timesteps.append(i)
+
+        return timesteps
+
+    # @typecheck(ignore_collections=True)
+    @torch.no_grad()
+    def forward(self, log_probs: Union[np.ndarray, torch.Tensor]):
+        if isinstance(log_probs, np.ndarray):
+            log_probs = torch.from_numpy(log_probs).float()
+        if log_probs.dim() == 2:
+            log_probs = log_probs.unsqueeze(0)
+
+        emissions = log_probs.cpu().contiguous()
+
+        B, T, N = emissions.size()
+        hypos = []
+        # we iterate over the batch dimension of our input tensor log probabilities
+        for b in range(B):
+            # the flashlight C++ expects a C style pointer, so the memory address
+            # which is what we obtain here. Then we pass it to pybinding method which
+            # is bound to the underlying C++ code
+            emissions_ptr = emissions.data_ptr() + 4 * b * emissions.stride(0)
+            results = self.decoder.decode(emissions_ptr, T, N)
+
+            hypos.append(
+                [
+                    {
+                        "tokens": self._get_tokens(result.tokens),
+                        "score": result.score,
+                        "timesteps": self._get_timesteps(result.tokens),
+                        "words": [self.word_dict.get_entry(x) for x in result.words if x >= 0],
+                    }
+                    for result in results
+                ]
+            )
+
+        return hypos