class_docstr: CTC Prefix Beam Search is based on the paper
  `First-Pass Large Vocabulary Continuous Speech Recognition using Bi-Directional Recurrent DNNs`
  by Awni Y. Hannun and al (https://arxiv.org/abs/1408.2873).
  
  The implementation keep tracks of the blank and non-blank probabilities.
  It also suppors n-gram scoring on words and SentencePiece tokens. The input
  is expected to be a log-probabilities tensor of shape [batch, time, vocab_size].
  
  Several heuristics are implemented to speed up the decoding process:
  - pruning of the beam : the beams are pruned if their score is lower than
      the best beam score minus the beam_prune_logp
  - pruning of the tokens : the tokens are pruned if their score is lower than
      the token_prune_min_logp
  - pruning of the history : the beams are pruned if they are the same over
      max_ngram history
  - skipping of the blank : the frame is skipped if the blank probability is
      higher than the blank_skip_threshold
  
  Note: The CTCPrefixBeamSearcher can be more unstable than the CTCBeamSearcher
  or the TorchAudioCTCPrefixBeamSearch searcher. Please, use it with caution
  and check the results carefully.
  
  Note: if the Acoustic Model is not trained, the Beam Search will
  take a lot of time. We do recommand to use Greedy Search during validation
  until the model is fully trained and ready to be evaluated on test sets.
  
  Note: This implementation does not provide the time alignment of the
  hypothesis. If you need it, please use the CTCBeamSearcher.
  
  Arguments
  ---------
  **kwargs
      see CTCBaseSearcher, arguments are directly passed.
  
  Example
  -------
  >>> import torch
  >>> from speechbrain.decoders import CTCPrefixBeamSearcher
  >>> probs = torch.tensor([[[0.2, 0.0, 0.8],
  ...                   [0.4, 0.0, 0.6]]])
  >>> log_probs = torch.log(probs)
  >>> lens = torch.tensor([1.0])
  >>> blank_index = 2
  >>> vocab_list = ['a', 'b', '-']
  >>> searcher = CTCPrefixBeamSearcher(blank_index=blank_index, vocab_list=vocab_list)
  >>> hyps = searcher(probs, lens)
class_docstr_tok: ["CTC","Prefix","Beam","Search","is","based","on","the","paper","`","First-Pass","Large","Vocabulary","Continuous","Speech","Recognition","using","Bi-Directional","Recurrent","DNNs","`","by","Awni","Y.","Hannun","and","al","(","https",":","\/\/arxiv.org\/abs\/1408.2873",")",".","The","implementation","keep","tracks","of","the","blank","and","non-blank","probabilities",".","It","also","suppors","n-gram","scoring","on","words","and","SentencePiece","tokens",".","The","input","is","expected","to","be","a","log-probabilities","tensor","of","shape","[","batch",",","time",",","vocab_size","]",".","Several","heuristics","are","implemented","to","speed","up","the","decoding","process",":","-","pruning","of","the","beam",":","the","beams","are","pruned","if","their","score","is","lower","than","the","best","beam","score","minus","the","beam_prune_logp","-","pruning","of","the","tokens",":","the","tokens","are","pruned","if","their","score","is","lower","than","the","token_prune_min_logp","-","pruning","of","the","history",":","the","beams","are","pruned","if","they","are","the","same","over","max_ngram","history","-","skipping","of","the","blank",":","the","frame","is","skipped","if","the","blank","probability","is","higher","than","the","blank_skip_threshold","Note",":","The","CTCPrefixBeamSearcher","can","be","more","unstable","than","the","CTCBeamSearcher","or","the","TorchAudioCTCPrefixBeamSearch","searcher",".","Please",",","use","it","with","caution","and","check","the","results","carefully",".","Note",":","if","the","Acoustic","Model","is","not","trained",",","the","Beam","Search","will","take","a","lot","of","time",".","We","do","recommand","to","use","Greedy","Search","during","validation","until","the","model","is","fully","trained","and","ready","to","be","evaluated","on","test","sets",".","Note",":","This","implementation","does","not","provide","the","time","alignment","of","the","hypothesis",".","If","you","need","it",",","please","use","the","CTCBeamSearcher",".","Arguments","--","--","--","--","-","*","*","kwargs","see","CTCBaseSearcher",",","arguments","are","directly","passed",".","Example","--","--","--","-",">",">",">","import","torch",">",">",">","from","speechbrain.decoders","import","CTCPrefixBeamSearcher",">",">",">","probs","=","torch.tensor","(","[","[","[","0.2",",","0.0",",","0.8","]",",","...","[","0.4",",","0.0",",","0.6","]","]","]",")",">",">",">","log_probs","=","torch.log","(","probs",")",">",">",">","lens","=","torch.tensor","(","[","1.0","]",")",">",">",">","blank_index","=","2",">",">",">","vocab_list","=","[","'","a","'",",","'","b","'",",","'-","'","]",">",">",">","searcher","=","CTCPrefixBeamSearcher","(","blank_index=blank_index",",","vocab_list=vocab_list",")",">",">",">","hyps","=","searcher","(","probs",",","lens",")"]
docstr: Score the beams with the language model if not None, and
  return the new beams.
  
  This function is modified and adapted from
  https://github.com/kensho-technologies/pyctcdecode
  
  Arguments
  ---------
  beams : list
      The list of the beams.
  cached_lm_scores : dict
      The cached language model scores.
  cached_partial_token_scores : dict
      The cached partial token scores.
  is_eos : bool (default: False)
      Whether the end of the sequence has been reached.
  
  Returns
  -------
  new_beams : list
      The list of the new beams.
docstr_tok: ["Score","the","beams","with","the","language","model","if","not","None",",","and","return","the","new","beams",".","This","function","is","modified","and","adapted","from","https",":","\/\/github.com\/kensho-technologies\/pyctcdecode","Arguments","--","--","--","--","-","beams",":","list","The","list","of","the","beams",".","cached_lm_scores",":","dict","The","cached","language","model","scores",".","cached_partial_token_scores",":","dict","The","cached","partial","token","scores",".","is_eos",":","bool","(","default",":","False",")","Whether","the","end","of","the","sequence","has","been","reached",".","Returns","--","--","--","-","new_beams",":","list","The","list","of","the","new","beams","."]
code: def get_lm_beams(
      self,
      beams: List[CTCBeam],
      cached_lm_scores: dict,
      cached_partial_token_scores: dict,
      is_eos=False,
  ) -> List[LMCTCBeam]:
      """Score the beams with the language model if not None, and
      return the new beams.
  
      This function is modified and adapted from
      https://github.com/kensho-technologies/pyctcdecode
  
      Arguments
      ---------
      beams : list
          The list of the beams.
      cached_lm_scores : dict
          The cached language model scores.
      cached_partial_token_scores : dict
          The cached partial token scores.
      is_eos : bool (default: False)
          Whether the end of the sequence has been reached.
  
      Returns
      -------
      new_beams : list
          The list of the new beams.
      """
      if self.lm is None:
          # no lm is used, lm_score is equal to score and we can return the beams
          # we have to keep track of the probabilities as well
          new_beams = []
          for beam in beams:
              new_text = self.merge_tokens(
                  beam.full_text, beam.next_word
              )
              new_beams.append(
                  LMCTCBeam(
                      text=beam.text,
                      full_text=new_text,
                      next_word="",
                      partial_word=beam.partial_word,
                      last_token=beam.last_token,
                      last_token_index=beam.last_token_index,
                      text_frames=beam.text_frames,
                      partial_frames=beam.partial_frames,
                      p=beam.p,
                      p_b=beam.p_b,
                      p_nb=beam.p_nb,
                      n_p_b=beam.n_p_b,
                      n_p_nb=beam.n_p_nb,
                      score=beam.score,
                      score_ctc=beam.score_ctc,
                      lm_score=beam.score,
                  )
              )
          return new_beams
      else:
          # lm is used, we need to compute the lm_score
          # first we compute the lm_score of the next word
          # we check if the next word is in the cache
          # if not, we compute the score and add it to the cache
          new_beams = []
          for beam in beams:
              # fast token merge
              new_text = self.merge_tokens(
                  beam.full_text, beam.next_word
              )
              cache_key = (new_text, is_eos)
              if cache_key not in cached_lm_scores:
                  prev_raw_lm_score, start_state = cached_lm_scores[
                      (beam.full_text, False)
                  ]
                  score, end_state = self.lm.score(
                      start_state, beam.next_word, is_last_word=is_eos
                  )
                  raw_lm_score = prev_raw_lm_score + score
                  cached_lm_scores[cache_key] = (
                      raw_lm_score,
                      end_state,
                  )
              lm_score, _ = cached_lm_scores[cache_key]
              word_part = beam.partial_word
  
              # we score the partial word
              if len(word_part) > 0:
                  if word_part not in cached_partial_token_scores:
                      cached_partial_token_scores[
                          word_part
                      ] = self.lm.score_partial_token(word_part)
                  lm_score += cached_partial_token_scores[word_part]
  
              new_beams.append(
                  LMCTCBeam(
                      text=beam.text,
                      full_text=new_text,
                      next_word="",
                      partial_word=beam.partial_word,
                      last_token=beam.last_token,
                      last_token_index=beam.last_token_index,
                      text_frames=beam.text_frames,
                      partial_frames=beam.partial_frames,
                      p=beam.p,
                      p_b=beam.p_b,
                      p_nb=beam.p_nb,
                      n_p_b=beam.n_p_b,
                      n_p_nb=beam.n_p_nb,
                      score=beam.score,
                      score_ctc=beam.score_ctc,
                      lm_score=beam.score + lm_score,
                  )
              )
          return new_beams
  
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 CTCPrefixBeamSearcher obj's underlying member method code with example usage: obj.get_lm_beams(beams, cached_lm_scores, cached_partial_token_scores, is_eos) and returns: new_beams, new_beams