class_docstr: This class implements the CTC prefix score of Algorithm 2 in
  reference: https://www.merl.com/publications/docs/TR2017-190.pdf.
  Official implementation: https://github.com/espnet/espnet/blob/master/espnet/nets/ctc_prefix_score.py
  
  Arguments
  ---------
  x : torch.Tensor
      The encoder states.
  enc_lens : torch.Tensor
      The actual length of each enc_states sequence.
  batch_size : int
      The size of the batch.
  beam_size : int
      The width of beam.
  blank_index : int
      The index of the blank token.
  eos_index : int
      The index of the end-of-sequence (eos) token.
  ctc_window_size: int
      Compute the ctc scores over the time frames using windowing based on attention peaks.
      If 0, no windowing applied.
class_docstr_tok: ["This","class","implements","the","CTC","prefix","score","of","Algorithm","2","in","reference",":","https",":","\/\/www.merl.com\/publications\/docs\/TR2017-190.pdf",".","Official","implementation",":","https",":","\/\/github.com\/espnet\/espnet\/blob\/master\/espnet\/nets\/ctc_prefix_score.py","Arguments","--","--","--","--","-","x",":","torch.Tensor","The","encoder","states",".","enc_lens",":","torch.Tensor","The","actual","length","of","each","enc_states","sequence",".","batch_size",":","int","The","size","of","the","batch",".","beam_size",":","int","The","width","of","beam",".","blank_index",":","int","The","index","of","the","blank","token",".","eos_index",":","int","The","index","of","the","end-of-sequence","(","eos",")","token",".","ctc_window_size",":","int","Compute","the","ctc","scores","over","the","time","frames","using","windowing","based","on","attention","peaks",".","If","0",",","no","windowing","applied","."]
docstr: This method if one step of forwarding operation
  for the prefix ctc scorer.
  
  Arguments
  ---------
  inp_tokens : torch.Tensor
      The last chars of prefix label sequences g, where h = g + c.
  states : tuple
      Previous ctc states.
  candidates : torch.Tensor
      (batch_size * beam_size, ctc_beam_size), The topk candidates for rescoring.
      If given, performing partial ctc scoring.
  attn : torch.Tensor
      (batch_size * beam_size, max_enc_len), The attention weights.
docstr_tok: ["This","method","if","one","step","of","forwarding","operation","for","the","prefix","ctc","scorer",".","Arguments","--","--","--","--","-","inp_tokens",":","torch.Tensor","The","last","chars","of","prefix","label","sequences","g",",","where","h","=","g","+","c.","states",":","tuple","Previous","ctc","states",".","candidates",":","torch.Tensor","(","batch_size","*","beam_size",",","ctc_beam_size",")",",","The","topk","candidates","for","rescoring",".","If","given",",","performing","partial","ctc","scoring",".","attn",":","torch.Tensor","(","batch_size","*","beam_size",",","max_enc_len",")",",","The","attention","weights","."]
code: def forward_step(
      self, inp_tokens, states, candidates=None, attn=None
  ):
      """This method if one step of forwarding operation
      for the prefix ctc scorer.
  
      Arguments
      ---------
      inp_tokens : torch.Tensor
          The last chars of prefix label sequences g, where h = g + c.
      states : tuple
          Previous ctc states.
      candidates : torch.Tensor
          (batch_size * beam_size, ctc_beam_size), The topk candidates for rescoring.
          If given, performing partial ctc scoring.
      attn : torch.Tensor
          (batch_size * beam_size, max_enc_len), The attention weights.
      """
  
      n_bh = inp_tokens.size(0)
      beam_size = n_bh // self.batch_size
      last_char = inp_tokens
      self.prefix_length += 1
      self.num_candidates = (
          self.vocab_size if candidates is None else candidates.size(-1)
      )
      if states is None:
          # r_prev: (L, 2, batch_size * beam_size)
          r_prev = torch.full(
              (self.max_enc_len, 2, self.batch_size, beam_size),
              self.minus_inf,
              device=self.device,
          )
  
          # Accumulate blank posteriors at each step
          r_prev[:, 1] = torch.cumsum(
              self.x[0, :, :, self.blank_index], 0
          ).unsqueeze(2)
          r_prev = r_prev.view(-1, 2, n_bh)
          psi_prev = torch.full(
              (n_bh, self.vocab_size),
              0.0,
              device=self.device,
          )
      else:
          r_prev, psi_prev = states
  
      # for partial search
      if candidates is not None:
          # The first index of each candidate.
          cand_offset = self.batch_index * self.vocab_size
          scoring_table = torch.full(
              (n_bh, self.vocab_size),
              -1,
              dtype=torch.long,
              device=self.device,
          )
          # Assign indices of candidates to their positions in the table
          col_index = torch.arange(n_bh, device=self.device).unsqueeze(
              1
          )
          scoring_table[col_index, candidates] = torch.arange(
              self.num_candidates, device=self.device
          )
          # Select candidates indices for scoring
          scoring_index = (
              candidates
              + cand_offset.unsqueeze(1)
              .repeat(1, beam_size)
              .view(-1, 1)
          ).view(-1)
          x_inflate = torch.index_select(
              self.x.view(2, -1, self.batch_size * self.vocab_size),
              2,
              scoring_index,
          ).view(2, -1, n_bh, self.num_candidates)
      # for full search
      else:
          scoring_table = None
          # Inflate x to (2, -1, batch_size * beam_size, num_candidates)
          # It is used to compute forward probs in a batched way
          x_inflate = (
              self.x.unsqueeze(3)
              .repeat(1, 1, 1, beam_size, 1)
              .view(2, -1, n_bh, self.num_candidates)
          )
  
      # Prepare forward probs
      r = torch.full(
          (
              self.max_enc_len,
              2,
              n_bh,
              self.num_candidates,
          ),
          self.minus_inf,
          device=self.device,
      )
      r.fill_(self.minus_inf)
  
      # (Alg.2-6)
      if self.prefix_length == 0:
          r[0, 0] = x_inflate[0, 0]
      # (Alg.2-10): phi = prev_nonblank + prev_blank = r_t-1^nb(g) + r_t-1^b(g)
      r_sum = torch.logsumexp(r_prev, 1)
      phi = r_sum.unsqueeze(2).repeat(1, 1, self.num_candidates)
  
      # (Alg.2-10): if last token of prefix g in candidates, phi = prev_b + 0
      if candidates is not None:
          for i in range(n_bh):
              pos = scoring_table[i, last_char[i]]
              if pos != -1:
                  phi[:, i, pos] = r_prev[:, 1, i]
      else:
          for i in range(n_bh):
              phi[:, i, last_char[i]] = r_prev[:, 1, i]
  
      # Start, end frames for scoring (|g| < |h|).
      # Scoring based on attn peak if ctc_window_size > 0
      if self.ctc_window_size == 0 or attn is None:
          start = max(1, self.prefix_length)
          end = self.max_enc_len
      else:
          _, attn_peak = torch.max(attn, dim=1)
          max_frame = torch.max(attn_peak).item() + self.ctc_window_size
          min_frame = torch.min(attn_peak).item() - self.ctc_window_size
          start = max(max(1, self.prefix_length), int(min_frame))
          end = min(self.max_enc_len, int(max_frame))
  
      # Compute forward prob log(r_t^nb(h)) and log(r_t^b(h)):
      for t in range(start, end):
          # (Alg.2-11): dim=0, p(h|cur step is nonblank) = [p(prev step=y) + phi] * p(c)
          rnb_prev = r[t - 1, 0]
          # (Alg.2-12): dim=1, p(h|cur step is blank) = [p(prev step is blank) + p(prev step is nonblank)] * p(blank)
          rb_prev = r[t - 1, 1]
          r_ = torch.stack(
              [rnb_prev, phi[t - 1], rnb_prev, rb_prev]
          ).view(2, 2, n_bh, self.num_candidates)
          r[t] = torch.logsumexp(r_, 1) + x_inflate[:, t]
  
      # Compute the predix prob, psi
      psi_init = r[start - 1, 0].unsqueeze(0)
      # phi is prob at t-1 step, shift one frame and add it to the current prob p(c)
      phix = (
          torch.cat((phi[0].unsqueeze(0), phi[:-1]), dim=0)
          + x_inflate[0]
      )
      # (Alg.2-13): psi = psi + phi * p(c)
      if candidates is not None:
          psi = torch.full(
              (n_bh, self.vocab_size),
              self.minus_inf,
              device=self.device,
          )
          psi_ = torch.logsumexp(
              torch.cat((phix[start:end], psi_init), dim=0), dim=0
          )
          # only assign prob to candidates
          for i in range(n_bh):
              psi[i, candidates[i]] = psi_[i]
      else:
          psi = torch.logsumexp(
              torch.cat((phix[start:end], psi_init), dim=0), dim=0
          )
  
      # (Alg.2-3): if c = <eos>, psi = log(r_T^n(g) + r_T^b(g)), where T is the length of max frames
      for i in range(n_bh):
          psi[i, self.eos_index] = r_sum[
              self.last_frame_index[i // beam_size], i
          ]
  
      # Exclude blank probs for joint scoring
      psi[:, self.blank_index] = self.minus_inf
  
      return psi - psi_prev, (r, psi, scoring_table)
  
imports: from itertools import groupby
  from speechbrain.dataio.dataio import length_to_mask
  import math
  import dataclasses
  import numpy
  import heapq
  import logging
  import torch
  import warnings
  from typing import Dict, List, Optional, Union, Any, Tuple
label_desc: Use image node_id 2 for calling the CTCPrefixScore obj's underlying member method code with example usage: obj.forward_step(inp_tokens, states, candidates, attn) and returns: unknown,  r,  psi,  scoring_table