text: #!/usr/bin/env python3
  
  # Copyright (c)  2020  Xiaomi Corporation (authors: Daniel Povey, Haowen Qiu, Mingshuang Luo)
  #                                           2021 Pingfeng Luo
  # Apache 2.0
  
  import k2
  import logging
  import numpy as np
  import os
  import torch
  from k2 import Fsa, SymbolTable
  from kaldialign import edit_distance
  from pathlib import Path
  from typing import List
  from typing import Union
  
  from lhotse import CutSet
  from lhotse.dataset import K2SpeechRecognitionDataset, SingleCutSampler
  from snowfall.common import average_checkpoint, store_transcripts
  from snowfall.common import find_first_disambig_symbol
  from snowfall.common import get_texts
  from snowfall.common import load_checkpoint
  from snowfall.common import setup_logger
  from snowfall.common import write_error_stats
  from snowfall.decoding.graph import compile_HLG
  from snowfall.lexicon import Lexicon
  from snowfall.models import AcousticModel
  from snowfall.models.tdnn_lstm import TdnnLstm1b
  from snowfall.training.ctc_graph import build_ctc_topo
  from snowfall.training.mmi_graph import create_bigram_phone_lm
  from snowfall.training.mmi_graph import get_phone_symbols
  
  
  @torch.no_grad()
  def decode(dataloader: torch.utils.data.DataLoader, model: AcousticModel,
             device: Union[str, torch.device], HLG: Fsa, symbols: SymbolTable):
      num_cuts = 0
      results = []  # a list of pair (ref_words, hyp_words)
      for batch_idx, batch in enumerate(dataloader):
          feature = batch['inputs']
          supervisions = batch['supervisions']
          supervision_segments = torch.stack(
              (supervisions['sequence_idx'],
               torch.floor_divide(supervisions['start_frame'],
                                  model.subsampling_factor),
               torch.floor_divide(supervisions['num_frames'],
                                  model.subsampling_factor)), 1).to(torch.int32)
          supervision_segments = torch.clamp(supervision_segments, min=0)
          indices = torch.argsort(supervision_segments[:, 2], descending=True)
          supervision_segments = supervision_segments[indices]
          texts = supervisions['text']
          assert feature.ndim == 3
  
          feature = feature.to(device)
          # at entry, feature is [N, T, C]
          feature = feature.permute(0, 2, 1)  # now feature is [N, C, T]
          with torch.no_grad():
              nnet_output = model(feature)
          # nnet_output is [N, C, T]
          nnet_output = nnet_output.permute(0, 2,
                                            1)  # now nnet_output is [N, T, C]
  
          #  blank_bias = -3.0
          #  nnet_output[:, :, 0] += blank_bias
  
          dense_fsa_vec = k2.DenseFsaVec(nnet_output, supervision_segments)
          # assert HLG.is_cuda()
          assert HLG.device == nnet_output.device, \
              f"Check failed: HLG.device ({HLG.device}) == nnet_output.device ({nnet_output.device})"
          # TODO(haowen): with a small `beam`, we may get empty `target_graph`,
          # thus `tot_scores` will be `inf`. Definitely we need to handle this later.
          lattices = k2.intersect_dense_pruned(HLG, dense_fsa_vec, 20.0, 7.0, 30,
                                               10000)
  
          # lattices = k2.intersect_dense(HLG, dense_fsa_vec, 10.0)
          best_paths = k2.shortest_path(lattices, use_double_scores=True)
          assert best_paths.shape[0] == len(texts)
          hyps = get_texts(best_paths, indices)
          assert len(hyps) == len(texts)
  
          for i in range(len(texts)):
              hyp_words = [symbols.get(x) for x in hyps[i]]
              ref_words = texts[i].split(' ')
              results.append((ref_words, hyp_words))
  
          if batch_idx % 10 == 0:
              logging.info(
                  'batch {}, cuts processed until now is {}'.format(
                      batch_idx, num_cuts))
  
          num_cuts += len(texts)
  
      return results
  
  def get_parser():
      import argparse
      parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
      parser.add_argument('--epoch', default=9, type=int)
      return parser
  
  
  def main():
      args = get_parser().parse_args()
      exp_dir = Path('exp-lstm-adam-mmi-bigram-musan')
      setup_logger('{}/log/log-decode'.format(exp_dir), log_level='debug')
  
      # load L, G, symbol_table
      lang_dir = Path('data/lang_nosp')
      lexicon = Lexicon(lang_dir)
  
      phone_ids = lexicon.phone_symbols()
  
      phone_ids_with_blank = [0] + phone_ids
      ctc_topo = k2.arc_sort(build_ctc_topo(phone_ids_with_blank))
  
      logging.debug("About to load model")
      # Note: Use "export CUDA_VISIBLE_DEVICES=N" to setup device id to N
      # device = torch.device('cuda', 1)
      device = torch.device('cuda')
      model = TdnnLstm1b(num_features=40,
                         num_classes=len(phone_ids) + 1,  # +1 for the blank symbol
                         subsampling_factor=3)
  
      checkpoint = os.path.join(exp_dir, f'epoch-{args.epoch}.pt')
      load_checkpoint(checkpoint, model)
      model.to(device)
      model.eval()
  
      if not os.path.exists(lang_dir / 'HLG.pt'):
          logging.debug("Loading L_disambig.fst.txt")
          with open(lang_dir / 'L_disambig.fst.txt') as f:
              L = k2.Fsa.from_openfst(f.read(), acceptor=False)
          logging.debug("Loading G.fst.txt")
          with open(lang_dir / 'G.fst.txt') as f:
              G = k2.Fsa.from_openfst(f.read(), acceptor=False)
          first_phone_disambig_id = find_first_disambig_symbol(lexicon.phones)
          first_word_disambig_id = find_first_disambig_symbol(lexicon.words)
          HLG = compile_HLG(L=L,
                           G=G,
                           H=ctc_topo,
                           labels_disambig_id_start=first_phone_disambig_id,
                           aux_labels_disambig_id_start=first_word_disambig_id)
          torch.save(HLG.as_dict(), lang_dir / 'HLG.pt')
      else:
          logging.debug("Loading pre-compiled LG")
          d = torch.load(lang_dir / 'HLG.pt')
          HLG = k2.Fsa.from_dict(d)
  
      # load dataset
      feature_dir = Path('exp/data')
      logging.debug("About to get test cuts")
      cuts_test = CutSet.from_json(feature_dir / 'cuts_test.json.gz')
  
      logging.info("About to create test dataset")
      test = K2SpeechRecognitionDataset(cuts_test)
      sampler = SingleCutSampler(cuts_test, max_frames=40000)
      logging.info("About to create test dataloader")
      test_dl = torch.utils.data.DataLoader(test, batch_size=None, sampler=sampler, num_workers=1)
  
      #  if not torch.cuda.is_available():
      #  logging.error('No GPU detected!')
      #  sys.exit(-1)
  
      logging.debug("convert HLG to device")
      HLG = HLG.to(device)
      HLG.aux_labels = k2.ragged.remove_values_eq(HLG.aux_labels, 0)
      HLG.requires_grad_(False)
  
      logging.debug("About to decode")
      results = decode(dataloader=test_dl,
                       model=model,
                       device=device,
                       HLG=HLG,
                       symbols=lexicon.words)
  
      s = ''
      results2 = []
      for ref, hyp in results:
          s += f'ref={ref}\n'
          s += f'hyp={hyp}\n'
          results2.append((list(''.join(ref)), list(''.join(hyp))))
      logging.info(s)
      # compute WER
      dists = [edit_distance(r, h) for r, h in results]
      dists2 = [edit_distance(r, h) for r, h in results2]
      errors = {
          key: sum(dist[key] for dist in dists)
          for key in ['sub', 'ins', 'del', 'total']
      }
      errors2 = {
          key: sum(dist[key] for dist in dists2)
          for key in ['sub', 'ins', 'del', 'total']
      }
      total_words = sum(len(ref) for ref, _ in results)
      total_chars = sum(len(ref) for ref, _ in results2)
      # Print Kaldi-like message:
      # %WER 8.20 [ 4459 / 54402, 695 ins, 427 del, 3337 sub ]
      logging.info(
          f'%WER {errors["total"] / total_words:.2%} '
          f'[{errors["total"]} / {total_words}, {errors["ins"]} ins, {errors["del"]} del, {errors["sub"]} sub ]'
      )
      logging.info(
          f'%CER {errors2["total"] / total_chars:.2%} '
          f'[{errors2["total"]} / {total_chars}, {errors2["ins"]} ins, {errors2["del"]} del, {errors2["sub"]} sub ]'
      )
  
  torch.set_num_threads(1)
  torch.set_num_interop_threads(1)
  
  if __name__ == '__main__':
      main()
  
  