text: # Copyright 2017 Johns Hopkins University (Shinji Watanabe)
  #  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
  
  """RNN sequence-to-sequence speech recognition model (pytorch)."""
  
  import argparse
  from itertools import groupby
  import logging
  import math
  import os
  
  import chainer
  from chainer import reporter
  import editdistance
  import numpy as np
  import six
  import torch
  
  from espnet.nets.asr_interface import ASRInterface
  from espnet.nets.e2e_asr_common import label_smoothing_dist
  from espnet.nets.pytorch_backend.ctc import ctc_for
  from espnet.nets.pytorch_backend.frontends.feature_transform import (
      feature_transform_for,  # noqa: H301
  )
  from espnet.nets.pytorch_backend.frontends.frontend import frontend_for
  from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters
  from espnet.nets.pytorch_backend.initialization import set_forget_bias_to_one
  from espnet.nets.pytorch_backend.nets_utils import get_subsample
  from espnet.nets.pytorch_backend.nets_utils import pad_list
  from espnet.nets.pytorch_backend.nets_utils import to_device
  from espnet.nets.pytorch_backend.nets_utils import to_torch_tensor
  from espnet.nets.pytorch_backend.rnn.argument import (
      add_arguments_rnn_encoder_common,  # noqa: H301
      add_arguments_rnn_decoder_common,  # noqa: H301
      add_arguments_rnn_attention_common,  # noqa: H301
  )
  from espnet.nets.pytorch_backend.rnn.attentions import att_for
  from espnet.nets.pytorch_backend.rnn.decoders import decoder_for
  from espnet.nets.pytorch_backend.rnn.encoders import encoder_for
  from espnet.nets.scorers.ctc import CTCPrefixScorer
  from espnet.utils.fill_missing_args import fill_missing_args
  
  CTC_LOSS_THRESHOLD = 10000
  
  
  class Reporter(chainer.Chain):
      """A chainer reporter wrapper."""
  
      def report(self, loss_ctc, loss_att, acc, cer_ctc, cer, wer, mtl_loss):
          """Report at every step."""
          reporter.report({"loss_ctc": loss_ctc}, self)
          reporter.report({"loss_att": loss_att}, self)
          reporter.report({"acc": acc}, self)
          reporter.report({"cer_ctc": cer_ctc}, self)
          reporter.report({"cer": cer}, self)
          reporter.report({"wer": wer}, self)
          logging.info("mtl loss:" + str(mtl_loss))
          reporter.report({"loss": mtl_loss}, self)
  
  
  class E2E(ASRInterface, torch.nn.Module):
      """E2E module.
  
      :param int idim: dimension of inputs
      :param int odim: dimension of outputs
      :param Namespace args: argument Namespace containing options
  
      """
  
      @staticmethod
      def add_arguments(parser):
          """Add arguments."""
          E2E.encoder_add_arguments(parser)
          E2E.attention_add_arguments(parser)
          E2E.decoder_add_arguments(parser)
          return parser
  
      @staticmethod
      def encoder_add_arguments(parser):
          """Add arguments for the encoder."""
          group = parser.add_argument_group("E2E encoder setting")
          group = add_arguments_rnn_encoder_common(group)
          return parser
  
      @staticmethod
      def attention_add_arguments(parser):
          """Add arguments for the attention."""
          group = parser.add_argument_group("E2E attention setting")
          group = add_arguments_rnn_attention_common(group)
          return parser
  
      @staticmethod
      def decoder_add_arguments(parser):
          """Add arguments for the decoder."""
          group = parser.add_argument_group("E2E decoder setting")
          group = add_arguments_rnn_decoder_common(group)
          return parser
  
      def get_total_subsampling_factor(self):
          """Get total subsampling factor."""
          if isinstance(self.enc, torch.nn.ModuleList):
              return self.enc[0].conv_subsampling_factor * int(np.prod(self.subsample))
          else:
              return self.enc.conv_subsampling_factor * int(np.prod(self.subsample))
  
      def __init__(self, idim, odim, args):
          """Construct an E2E object.
  
          :param int idim: dimension of inputs
          :param int odim: dimension of outputs
          :param Namespace args: argument Namespace containing options
          """
          super(E2E, self).__init__()
          torch.nn.Module.__init__(self)
  
          # fill missing arguments for compatibility
          args = fill_missing_args(args, self.add_arguments)
  
          self.mtlalpha = args.mtlalpha
          assert 0.0 <= self.mtlalpha <= 1.0, "mtlalpha should be [0.0, 1.0]"
          self.etype = args.etype
          self.verbose = args.verbose
          # NOTE: for self.build method
          args.char_list = getattr(args, "char_list", None)
          self.char_list = args.char_list
          self.outdir = args.outdir
          self.space = args.sym_space
          self.blank = args.sym_blank
          self.reporter = Reporter()
  
          # below means the last number becomes eos/sos ID
          # note that sos/eos IDs are identical
          self.sos = odim - 1
          self.eos = odim - 1
  
          # subsample info
          self.subsample = get_subsample(args, mode="asr", arch="rnn")
  
          # label smoothing info
          if args.lsm_type and os.path.isfile(args.train_json):
              logging.info("Use label smoothing with " + args.lsm_type)
              labeldist = label_smoothing_dist(
                  odim, args.lsm_type, transcript=args.train_json
              )
          else:
              labeldist = None
  
          if getattr(args, "use_frontend", False):  # use getattr to keep compatibility
              self.frontend = frontend_for(args, idim)
              self.feature_transform = feature_transform_for(args, (idim - 1) * 2)
              idim = args.n_mels
          else:
              self.frontend = None
  
          # encoder
          self.enc = encoder_for(args, idim, self.subsample)
          # ctc
          self.ctc = ctc_for(args, odim)
          # attention
          self.att = att_for(args)
          # decoder
          self.dec = decoder_for(args, odim, self.sos, self.eos, self.att, labeldist)
  
          # weight initialization
          self.init_like_chainer()
  
          # options for beam search
          if args.report_cer or args.report_wer:
              recog_args = {
                  "beam_size": args.beam_size,
                  "penalty": args.penalty,
                  "ctc_weight": args.ctc_weight,
                  "maxlenratio": args.maxlenratio,
                  "minlenratio": args.minlenratio,
                  "lm_weight": args.lm_weight,
                  "rnnlm": args.rnnlm,
                  "nbest": args.nbest,
                  "space": args.sym_space,
                  "blank": args.sym_blank,
              }
  
              self.recog_args = argparse.Namespace(**recog_args)
              self.report_cer = args.report_cer
              self.report_wer = args.report_wer
          else:
              self.report_cer = False
              self.report_wer = False
          self.rnnlm = None
  
          self.logzero = -10000000000.0
          self.loss = None
          self.acc = None
  
      def init_like_chainer(self):
          """Initialize weight like chainer.
  
          chainer basically uses LeCun way: W ~ Normal(0, fan_in ** -0.5), b = 0
          pytorch basically uses W, b ~ Uniform(-fan_in**-0.5, fan_in**-0.5)
          however, there are two exceptions as far as I know.
          - EmbedID.W ~ Normal(0, 1)
          - LSTM.upward.b[forget_gate_range] = 1 (but not used in NStepLSTM)
          """
          lecun_normal_init_parameters(self)
          # exceptions
          # embed weight ~ Normal(0, 1)
          self.dec.embed.weight.data.normal_(0, 1)
          # forget-bias = 1.0
          # https://discuss.pytorch.org/t/set-forget-gate-bias-of-lstm/1745
          for i in six.moves.range(len(self.dec.decoder)):
              set_forget_bias_to_one(self.dec.decoder[i].bias_ih)
  
      def forward(self, xs_pad, ilens, ys_pad):
          """E2E forward.
  
          :param torch.Tensor xs_pad: batch of padded input sequences (B, Tmax, idim)
          :param torch.Tensor ilens: batch of lengths of input sequences (B)
          :param torch.Tensor ys_pad: batch of padded token id sequence tensor (B, Lmax)
          :return: loss value
          :rtype: torch.Tensor
          """
          # 0. Frontend
          if self.frontend is not None:
              hs_pad, hlens, mask = self.frontend(to_torch_tensor(xs_pad), ilens)
              hs_pad, hlens = self.feature_transform(hs_pad, hlens)
          else:
              hs_pad, hlens = xs_pad, ilens
  
          # 1. Encoder
          hs_pad, hlens, _ = self.enc(hs_pad, hlens)
  
          # 2. CTC loss
          if self.mtlalpha == 0:
              self.loss_ctc = None
          else:
              self.loss_ctc = self.ctc(hs_pad, hlens, ys_pad)
  
          # 3. attention loss
          if self.mtlalpha == 1:
              self.loss_att, acc = None, None
          else:
              self.loss_att, acc, _ = self.dec(hs_pad, hlens, ys_pad)
          self.acc = acc
  
          # 4. compute cer without beam search
          if self.mtlalpha == 0 or self.char_list is None:
              cer_ctc = None
          else:
              cers = []
  
              y_hats = self.ctc.argmax(hs_pad).data
              for i, y in enumerate(y_hats):
                  y_hat = [x[0] for x in groupby(y)]
                  y_true = ys_pad[i]
  
                  seq_hat = [self.char_list[int(idx)] for idx in y_hat if int(idx) != -1]
                  seq_true = [
                      self.char_list[int(idx)] for idx in y_true if int(idx) != -1
                  ]
                  seq_hat_text = "".join(seq_hat).replace(self.space, " ")
                  seq_hat_text = seq_hat_text.replace(self.blank, "")
                  seq_true_text = "".join(seq_true).replace(self.space, " ")
  
                  hyp_chars = seq_hat_text.replace(" ", "")
                  ref_chars = seq_true_text.replace(" ", "")
                  if len(ref_chars) > 0:
                      cers.append(
                          editdistance.eval(hyp_chars, ref_chars) / len(ref_chars)
                      )
  
              cer_ctc = sum(cers) / len(cers) if cers else None
  
          # 5. compute cer/wer
          if self.training or not (self.report_cer or self.report_wer):
              cer, wer = 0.0, 0.0
              # oracle_cer, oracle_wer = 0.0, 0.0
          else:
              if self.recog_args.ctc_weight > 0.0:
                  lpz = self.ctc.log_softmax(hs_pad).data
              else:
                  lpz = None
  
              word_eds, word_ref_lens, char_eds, char_ref_lens = [], [], [], []
              nbest_hyps = self.dec.recognize_beam_batch(
                  hs_pad,
                  torch.tensor(hlens),
                  lpz,
                  self.recog_args,
                  self.char_list,
                  self.rnnlm,
              )
              # remove <sos> and <eos>
              y_hats = [nbest_hyp[0]["yseq"][1:-1] for nbest_hyp in nbest_hyps]
              for i, y_hat in enumerate(y_hats):
                  y_true = ys_pad[i]
  
                  seq_hat = [self.char_list[int(idx)] for idx in y_hat if int(idx) != -1]
                  seq_true = [
                      self.char_list[int(idx)] for idx in y_true if int(idx) != -1
                  ]
                  seq_hat_text = "".join(seq_hat).replace(self.recog_args.space, " ")
                  seq_hat_text = seq_hat_text.replace(self.recog_args.blank, "")
                  seq_true_text = "".join(seq_true).replace(self.recog_args.space, " ")
  
                  hyp_words = seq_hat_text.split()
                  ref_words = seq_true_text.split()
                  word_eds.append(editdistance.eval(hyp_words, ref_words))
                  word_ref_lens.append(len(ref_words))
                  hyp_chars = seq_hat_text.replace(" ", "")
                  ref_chars = seq_true_text.replace(" ", "")
                  char_eds.append(editdistance.eval(hyp_chars, ref_chars))
                  char_ref_lens.append(len(ref_chars))
  
              wer = (
                  0.0
                  if not self.report_wer
                  else float(sum(word_eds)) / sum(word_ref_lens)
              )
              cer = (
                  0.0
                  if not self.report_cer
                  else float(sum(char_eds)) / sum(char_ref_lens)
              )
  
          alpha = self.mtlalpha
          if alpha == 0:
              self.loss = self.loss_att
              loss_att_data = float(self.loss_att)
              loss_ctc_data = None
          elif alpha == 1:
              self.loss = self.loss_ctc
              loss_att_data = None
              loss_ctc_data = float(self.loss_ctc)
          else:
              self.loss = alpha * self.loss_ctc + (1 - alpha) * self.loss_att
              loss_att_data = float(self.loss_att)
              loss_ctc_data = float(self.loss_ctc)
  
          loss_data = float(self.loss)
          if loss_data < CTC_LOSS_THRESHOLD and not math.isnan(loss_data):
              self.reporter.report(
                  loss_ctc_data, loss_att_data, acc, cer_ctc, cer, wer, loss_data
              )
          else:
              logging.warning("loss (=%f) is not correct", loss_data)
          return self.loss
  
      def scorers(self):
          """Scorers."""
          return dict(decoder=self.dec, ctc=CTCPrefixScorer(self.ctc, self.eos))
  
      def encode(self, x):
          """Encode acoustic features.
  
          :param ndarray x: input acoustic feature (T, D)
          :return: encoder outputs
          :rtype: torch.Tensor
          """
          self.eval()
          ilens = [x.shape[0]]
  
          # subsample frame
          x = x[:: self.subsample[0], :]
          p = next(self.parameters())
          h = torch.as_tensor(x, device=p.device, dtype=p.dtype)
          # make a utt list (1) to use the same interface for encoder
          hs = h.contiguous().unsqueeze(0)
  
          # 0. Frontend
          if self.frontend is not None:
              enhanced, hlens, mask = self.frontend(hs, ilens)
              hs, hlens = self.feature_transform(enhanced, hlens)
          else:
              hs, hlens = hs, ilens
  
          # 1. encoder
          hs, _, _ = self.enc(hs, hlens)
          return hs.squeeze(0)
  
      def recognize(self, x, recog_args, char_list, rnnlm=None):
          """E2E beam search.
  
          :param ndarray x: input acoustic feature (T, D)
          :param Namespace recog_args: argument Namespace containing options
          :param list char_list: list of characters
          :param torch.nn.Module rnnlm: language model module
          :return: N-best decoding results
          :rtype: list
          """
          hs = self.encode(x).unsqueeze(0)
          # calculate log P(z_t|X) for CTC scores
          if recog_args.ctc_weight > 0.0:
              lpz = self.ctc.log_softmax(hs)[0]
          else:
              lpz = None
  
          # 2. Decoder
          # decode the first utterance
          y = self.dec.recognize_beam(hs[0], lpz, recog_args, char_list, rnnlm)
          return y
  
      def recognize_batch(self, xs, recog_args, char_list, rnnlm=None):
          """E2E batch beam search.
  
          :param list xs: list of input acoustic feature arrays [(T_1, D), (T_2, D), ...]
          :param Namespace recog_args: argument Namespace containing options
          :param list char_list: list of characters
          :param torch.nn.Module rnnlm: language model module
          :return: N-best decoding results
          :rtype: list
          """
          prev = self.training
          self.eval()
          ilens = np.fromiter((xx.shape[0] for xx in xs), dtype=np.int64)
  
          # subsample frame
          xs = [xx[:: self.subsample[0], :] for xx in xs]
          xs = [to_device(self, to_torch_tensor(xx).float()) for xx in xs]
          xs_pad = pad_list(xs, 0.0)
  
          # 0. Frontend
          if self.frontend is not None:
              enhanced, hlens, mask = self.frontend(xs_pad, ilens)
              hs_pad, hlens = self.feature_transform(enhanced, hlens)
          else:
              hs_pad, hlens = xs_pad, ilens
  
          # 1. Encoder
          hs_pad, hlens, _ = self.enc(hs_pad, hlens)
  
          # calculate log P(z_t|X) for CTC scores
          if recog_args.ctc_weight > 0.0:
              lpz = self.ctc.log_softmax(hs_pad)
              normalize_score = False
          else:
              lpz = None
              normalize_score = True
  
          # 2. Decoder
          hlens = torch.tensor(list(map(int, hlens)))  # make sure hlens is tensor
          y = self.dec.recognize_beam_batch(
              hs_pad,
              hlens,
              lpz,
              recog_args,
              char_list,
              rnnlm,
              normalize_score=normalize_score,
          )
  
          if prev:
              self.train()
          return y
  
      def enhance(self, xs):
          """Forward only in the frontend stage.
  
          :param ndarray xs: input acoustic feature (T, C, F)
          :return: enhaned feature
          :rtype: torch.Tensor
          """
          if self.frontend is None:
              raise RuntimeError("Frontend does't exist")
          prev = self.training
          self.eval()
          ilens = np.fromiter((xx.shape[0] for xx in xs), dtype=np.int64)
  
          # subsample frame
          xs = [xx[:: self.subsample[0], :] for xx in xs]
          xs = [to_device(self, to_torch_tensor(xx).float()) for xx in xs]
          xs_pad = pad_list(xs, 0.0)
          enhanced, hlensm, mask = self.frontend(xs_pad, ilens)
          if prev:
              self.train()
          return enhanced.cpu().numpy(), mask.cpu().numpy(), ilens
  
      def calculate_all_attentions(self, xs_pad, ilens, ys_pad):
          """E2E attention calculation.
  
          :param torch.Tensor xs_pad: batch of padded input sequences (B, Tmax, idim)
          :param torch.Tensor ilens: batch of lengths of input sequences (B)
          :param torch.Tensor ys_pad: batch of padded token id sequence tensor (B, Lmax)
          :return: attention weights with the following shape,
              1) multi-head case => attention weights (B, H, Lmax, Tmax),
              2) other case => attention weights (B, Lmax, Tmax).
          :rtype: float ndarray
          """
          self.eval()
          with torch.no_grad():
              # 0. Frontend
              if self.frontend is not None:
                  hs_pad, hlens, mask = self.frontend(to_torch_tensor(xs_pad), ilens)
                  hs_pad, hlens = self.feature_transform(hs_pad, hlens)
              else:
                  hs_pad, hlens = xs_pad, ilens
  
              # 1. Encoder
              hpad, hlens, _ = self.enc(hs_pad, hlens)
  
              # 2. Decoder
              att_ws = self.dec.calculate_all_attentions(hpad, hlens, ys_pad)
          self.train()
          return att_ws
  
      def calculate_all_ctc_probs(self, xs_pad, ilens, ys_pad):
          """E2E CTC probability calculation.
  
          :param torch.Tensor xs_pad: batch of padded input sequences (B, Tmax)
          :param torch.Tensor ilens: batch of lengths of input sequences (B)
          :param torch.Tensor ys_pad: batch of padded token id sequence tensor (B, Lmax)
          :return: CTC probability (B, Tmax, vocab)
          :rtype: float ndarray
          """
          probs = None
          if self.mtlalpha == 0:
              return probs
  
          self.eval()
          with torch.no_grad():
              # 0. Frontend
              if self.frontend is not None:
                  hs_pad, hlens, mask = self.frontend(to_torch_tensor(xs_pad), ilens)
                  hs_pad, hlens = self.feature_transform(hs_pad, hlens)
              else:
                  hs_pad, hlens = xs_pad, ilens
  
              # 1. Encoder
              hpad, hlens, _ = self.enc(hs_pad, hlens)
  
              # 2. CTC probs
              probs = self.ctc.softmax(hpad).cpu().numpy()
          self.train()
          return probs
  
      def subsample_frames(self, x):
          """Subsample speeh frames in the encoder."""
          # subsample frame
          x = x[:: self.subsample[0], :]
          ilen = [x.shape[0]]
          h = to_device(self, torch.from_numpy(np.array(x, dtype=np.float32)))
          h.contiguous()
          return h, ilen
  