espnet2.legacy.nets.batch_beam_search.BatchBeamSearch

source

class espnet2.legacy.nets.batch_beam_search.BatchBeamSearch(scorers: Dict[str, ScorerInterface], weights: Dict[str, float], beam_size: int, vocab_size: int, sos: int, eos: int, token_list: List[str] = None, pre_beam_ratio: float = 1.5, pre_beam_score_key: str = None, return_hs: bool = False, hyp_primer: List[int] = None, normalize_length: bool = False)

Bases: BeamSearch

Batch beam search implementation.

Hypotheses are vectorized over the beam, and optionally over utterances as well: pass a batched encoder output (n_utt, T, D) to forward() and a whole minibatch is decoded with a single set of scorer calls, which is much faster on an accelerator than decoding one utterance at a time. Decoding a single utterance, by passing (T, D), is the n_utt == 1 case of the same search and returns a plain n-best list as before.

Initialize beam search.

• Parameters:
  • scorers (dict *[*str , ScorerInterface ]) – Dict of decoder modules e.g., Decoder, CTCPrefixScorer, LM The scorer will be ignored if it is None
  • weights (dict *[*str , float ]) – Dict of weights for each scorers The scorer will be ignored if its weight is 0
  • beam_size (int) – The number of hypotheses kept during search
  • vocab_size (int) – The number of vocabulary
  • sos (int) – Start of sequence id
  • eos (int) – End of sequence id
  • token_list (list *[*str ]) – List of tokens for debug log
  • pre_beam_score_key (str) – key of scores to perform pre-beam search
  • pre_beam_ratio (float) – beam size in the pre-beam search will be int(pre_beam_ratio * beam_size)
  • return_hs (bool) – Whether to return hidden intermediates
  • normalize_length (bool) – If true, select the best ended hypotheses based on length-normalized scores rather than the accumulated scores

batch_beam(weighted_scores: Tensor, ids: Tensor) → Tuple[Tensor, Tensor, Tensor, Tensor]

Batch-compute topk full token ids and partial token ids.

This is the single-utterance case of the pruning done in search(); it is kept as a separate method because subclasses override it.

• Parameters:
  • weighted_scores (torch.Tensor) – The weighted sum scores for each tokens. Its shape is (n_beam, self.vocab_size).
  • ids (torch.Tensor) – The partial token ids to compute topk. Its shape is (n_beam, self.pre_beam_size).
• Returns: The topk full (prev_hyp, new_token) ids and partial (prev_hyp, new_token) ids. Their shapes are all (self.beam_size,)
• Return type: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

batchfy(hyps: List[Hypothesis], n_utt: int = 1) → BatchHypothesis

Convert list to batch.

forward(x: Tensor, maxlenratio: float = 0.0, minlenratio: float = 0.0, pre_x: Tensor = None, x_lengths: Tensor | None = None, pre_x_lengths: Tensor | None = None) → List[Hypothesis] | List[List[Hypothesis]]

Perform beam search.

• Parameters:
  • x (torch.Tensor) – Encoded speech feature (T, D) for a single utterance, or (n_utt, T, D) to decode a whole minibatch with one set of scorer calls.
  • maxlenratio (float) – Input length ratio to obtain max output length. If maxlenratio=0.0 (default), it uses a end-detect function to automatically find maximum hypothesis lengths If maxlenratio<0.0, its absolute value is interpreted as a constant max output length.
  • minlenratio (float) – Input length ratio to obtain min output length. If minlenratio<0.0, its absolute value is interpreted as a constant min output length.
  • pre_x (torch.Tensor) – Encoded speech feature for sequential attn Sequential attn computes attn first on pre_x then on x, thereby attending to two sources in sequence.
  • x_lengths (torch.Tensor) – Encoder output lengths (n_utt,). Only used for a batched x; when omitted every utterance is assumed to fill the whole tensor.
  • pre_x_lengths (torch.Tensor) – Lengths of pre_x.
• Returns: N-best decoding results for a single utterance, or list[list[Hypothesis]] with one n-best list per utterance when x is batched.
• Return type: list[Hypothesis]

init_hyp(x: Tensor, x_lengths: Tensor | None = None) → BatchHypothesis

Get an initial hypothesis data.

• Parameters:
  • x (torch.Tensor) – The encoder output feature, either (T, D) for a single utterance or (n_utt, T, D) for a batch.
  • x_lengths (torch.Tensor) – Encoder output lengths (n_utt,), needed so that the CTC prefix scores ignore padded frames.
• Returns: The initial hypotheses. A single utterance starts from one hypothesis, as in BeamSearch; a batch starts from a full (n_utt, beam) grid in which only the first beam of each utterance is active, which has the same effect.
• Return type:BatchHypothesis

merge_states(states: Any, part_states: Any, part_idx: int) → Any

Merge states for new hypothesis.

• Parameters:
  • states – states of self.full_scorers
  • part_states – states of self.part_scorers
  • part_idx (int) – The new token id for part_scores
• Returns: The new score dict. : Its keys are names of self.full_scorers and self.part_scorers. Its values are states of the scorers.
• Return type: Dict[str, torch.Tensor]

post_process(i: int, maxlen: int | List[int], minlen: int | List[int], maxlenratio: float, running_hyps: BatchHypothesis, ended_hyps: List[Hypothesis] | List[List[Hypothesis]]) → BatchHypothesis

Perform post-processing of beam search iterations.

• Parameters:
  • i (int) – The length of hypothesis tokens.
  • maxlen (int or list *[*int ]) – The maximum length of tokens in beam search, per utterance when several are decoded together.
  • minlen (int or list *[*int ]) – The minimum length of tokens.
  • maxlenratio (int) – The maximum length ratio in beam search.
  • running_hyps (BatchHypothesis) – The running hypotheses in beam search.
  • ended_hyps (list) – The ended hypotheses. A flat list when a single utterance is decoded, one list per utterance otherwise. Appended to in place.
• Returns: The new running hypotheses.
• Return type:BatchHypothesis

score_full(hyp: BatchHypothesis, x: Tensor, pre_x: Tensor = None, xs_mask: Tensor | None = None, pre_xs_mask: Tensor | None = None) → Tuple[Dict[str, Tensor], Dict[str, Any]]

Score new hypothesis by self.full_scorers.

• Parameters:
  • hyp (BatchHypothesis) – Hypothesis with prefix tokens to score
  • x (torch.Tensor) – Corresponding input feature (n_batch, T, D)
  • pre_x (torch.Tensor) – Encoded speech feature for sequential attn Sequential attn computes attn first on pre_x then on x, thereby attending to two sources in sequence.
  • xs_mask (torch.Tensor) – Non-padding mask of x (n_batch, 1, T), or None when every utterance fills the whole tensor.
  • pre_xs_mask (torch.Tensor) – Non-padding mask of pre_x.
• Returns: Tuple of : score dict of hyp that has string keys of self.full_scorers and tensor score values of shape: (self.n_vocab,), and state dict that has string keys and state values of self.full_scorers
• Return type: Tuple[Dict[str, torch.Tensor], Dict[str, Any]]

score_partial(hyp: BatchHypothesis, ids: Tensor, x: Tensor, pre_x: Tensor = None) → Tuple[Dict[str, Tensor], Dict[str, Any]]

Score new hypothesis by self.part_scorers.

• Parameters:
  • hyp (BatchHypothesis) – Hypothesis with prefix tokens to score
  • ids (torch.Tensor) – 2D tensor of new partial tokens to score
  • x (torch.Tensor) – Corresponding input feature
  • pre_x (torch.Tensor) – Encoded speech feature for sequential attn (T, D) Sequential attn computes attn first on pre_x then on x, thereby attending to two sources in sequence.
• Returns: Tuple of : score dict of hyp that has string keys of self.full_scorers and tensor score values of shape: (self.n_vocab,), and state dict that has string keys and state values of self.full_scorers
• Return type: Tuple[Dict[str, torch.Tensor], Dict[str, Any]]

search(running_hyps: BatchHypothesis, x: Tensor, pre_x: Tensor = None, xs_mask: Tensor | None = None, pre_xs_mask: Tensor | None = None) → BatchHypothesis

Search new tokens for running hypotheses and encoded speech x.

• Parameters:
  • running_hyps (BatchHypothesis) – Running hypotheses on beam
  • x (torch.Tensor) – Encoded speech feature. Either (T, D) for a single utterance, which is replicated over the hypotheses, or already replicated as (n_hyp, T, D).
  • pre_x (torch.Tensor) – Encoded speech feature for sequential attention
  • xs_mask (torch.Tensor) – Non-padding mask of x
  • pre_xs_mask (torch.Tensor) – Non-padding mask of pre_x
• Returns: beam_size best hypotheses per utterance
• Return type:BatchHypothesis

static supports_xs_mask(scorer: Any) → bool

Return whether scorer.batch_score accepts an xs_mask.

unbatchfy(batch_hyps: BatchHypothesis) → List[Hypothesis]

Revert batch to list.