Source code for torchaudio.models.rnnt_decoder
from typing import Callable, Dict, List, Optional, Tuple
import torch
from torchaudio.models import RNNT
__all__ = ["Hypothesis", "RNNTBeamSearch"]
Hypothesis = Tuple[List[int], torch.Tensor, List[List[torch.Tensor]], float]
Hypothesis.__doc__ = """Hypothesis generated by RNN-T beam search decoder,
represented as tuple of (tokens, prediction network output, prediction network state, score).
"""
def _get_hypo_tokens(hypo: Hypothesis) -> List[int]:
return hypo[0]
def _get_hypo_predictor_out(hypo: Hypothesis) -> torch.Tensor:
return hypo[1]
def _get_hypo_state(hypo: Hypothesis) -> List[List[torch.Tensor]]:
return hypo[2]
def _get_hypo_score(hypo: Hypothesis) -> float:
return hypo[3]
def _get_hypo_key(hypo: Hypothesis) -> str:
return str(hypo[0])
def _batch_state(hypos: List[Hypothesis]) -> List[List[torch.Tensor]]:
states: List[List[torch.Tensor]] = []
for i in range(len(_get_hypo_state(hypos[0]))):
batched_state_components: List[torch.Tensor] = []
for j in range(len(_get_hypo_state(hypos[0])[i])):
batched_state_components.append(torch.cat([_get_hypo_state(hypo)[i][j] for hypo in hypos]))
states.append(batched_state_components)
return states
def _slice_state(states: List[List[torch.Tensor]], idx: int, device: torch.device) -> List[List[torch.Tensor]]:
idx_tensor = torch.tensor([idx], device=device)
return [[state.index_select(0, idx_tensor) for state in state_tuple] for state_tuple in states]
def _default_hypo_sort_key(hypo: Hypothesis) -> float:
return _get_hypo_score(hypo) / (len(_get_hypo_tokens(hypo)) + 1)
def _compute_updated_scores(
hypos: List[Hypothesis],
next_token_probs: torch.Tensor,
beam_width: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
hypo_scores = torch.tensor([_get_hypo_score(h) for h in hypos]).unsqueeze(1)
nonblank_scores = hypo_scores + next_token_probs[:, :-1] # [beam_width, num_tokens - 1]
nonblank_nbest_scores, nonblank_nbest_idx = nonblank_scores.reshape(-1).topk(beam_width)
nonblank_nbest_hypo_idx = nonblank_nbest_idx.div(nonblank_scores.shape[1], rounding_mode="trunc")
nonblank_nbest_token = nonblank_nbest_idx % nonblank_scores.shape[1]
return nonblank_nbest_scores, nonblank_nbest_hypo_idx, nonblank_nbest_token
def _remove_hypo(hypo: Hypothesis, hypo_list: List[Hypothesis]) -> None:
for i, elem in enumerate(hypo_list):
if _get_hypo_key(hypo) == _get_hypo_key(elem):
del hypo_list[i]
break
[docs]class RNNTBeamSearch(torch.nn.Module):
r"""Beam search decoder for RNN-T model.
Args:
model (RNNT): RNN-T model to use.
blank (int): index of blank token in vocabulary.
temperature (float, optional): temperature to apply to joint network output.
Larger values yield more uniform samples. (Default: 1.0)
hypo_sort_key (Callable[[Hypothesis], float] or None, optional): callable that computes a score
for a given hypothesis to rank hypotheses by. If ``None``, defaults to callable that returns
hypothesis score normalized by token sequence length. (Default: None)
step_max_tokens (int, optional): maximum number of tokens to emit per input time step. (Default: 100)
"""
def __init__(
self,
model: RNNT,
blank: int,
temperature: float = 1.0,
hypo_sort_key: Optional[Callable[[Hypothesis], float]] = None,
step_max_tokens: int = 100,
) -> None:
super().__init__()
self.model = model
self.blank = blank
self.temperature = temperature
if hypo_sort_key is None:
self.hypo_sort_key = _default_hypo_sort_key
else:
self.hypo_sort_key = hypo_sort_key
self.step_max_tokens = step_max_tokens
def _init_b_hypos(self, hypo: Optional[Hypothesis], device: torch.device) -> List[Hypothesis]:
if hypo is not None:
token = _get_hypo_tokens(hypo)[-1]
state = _get_hypo_state(hypo)
else:
token = self.blank
state = None
one_tensor = torch.tensor([1], device=device)
pred_out, _, pred_state = self.model.predict(torch.tensor([[token]], device=device), one_tensor, state)
init_hypo = (
[token],
pred_out[0].detach(),
pred_state,
0.0,
)
return [init_hypo]
def _gen_next_token_probs(
self, enc_out: torch.Tensor, hypos: List[Hypothesis], device: torch.device
) -> torch.Tensor:
one_tensor = torch.tensor([1], device=device)
predictor_out = torch.stack([_get_hypo_predictor_out(h) for h in hypos], dim=0)
joined_out, _, _ = self.model.join(
enc_out,
one_tensor,
predictor_out,
torch.tensor([1] * len(hypos), device=device),
) # [beam_width, 1, 1, num_tokens]
joined_out = torch.nn.functional.log_softmax(joined_out / self.temperature, dim=3)
return joined_out[:, 0, 0]
def _gen_b_hypos(
self,
b_hypos: List[Hypothesis],
a_hypos: List[Hypothesis],
next_token_probs: torch.Tensor,
key_to_b_hypo: Dict[str, Hypothesis],
) -> List[Hypothesis]:
for i in range(len(a_hypos)):
h_a = a_hypos[i]
append_blank_score = _get_hypo_score(h_a) + next_token_probs[i, -1]
if _get_hypo_key(h_a) in key_to_b_hypo:
h_b = key_to_b_hypo[_get_hypo_key(h_a)]
_remove_hypo(h_b, b_hypos)
score = float(torch.tensor(_get_hypo_score(h_b)).logaddexp(append_blank_score))
else:
score = float(append_blank_score)
h_b = (
_get_hypo_tokens(h_a),
_get_hypo_predictor_out(h_a),
_get_hypo_state(h_a),
score,
)
b_hypos.append(h_b)
key_to_b_hypo[_get_hypo_key(h_b)] = h_b
_, sorted_idx = torch.tensor([_get_hypo_score(hypo) for hypo in b_hypos]).sort()
return [b_hypos[idx] for idx in sorted_idx]
def _gen_a_hypos(
self,
a_hypos: List[Hypothesis],
b_hypos: List[Hypothesis],
next_token_probs: torch.Tensor,
t: int,
beam_width: int,
device: torch.device,
) -> List[Hypothesis]:
(
nonblank_nbest_scores,
nonblank_nbest_hypo_idx,
nonblank_nbest_token,
) = _compute_updated_scores(a_hypos, next_token_probs, beam_width)
if len(b_hypos) < beam_width:
b_nbest_score = -float("inf")
else:
b_nbest_score = _get_hypo_score(b_hypos[-beam_width])
base_hypos: List[Hypothesis] = []
new_tokens: List[int] = []
new_scores: List[float] = []
for i in range(beam_width):
score = float(nonblank_nbest_scores[i])
if score > b_nbest_score:
a_hypo_idx = int(nonblank_nbest_hypo_idx[i])
base_hypos.append(a_hypos[a_hypo_idx])
new_tokens.append(int(nonblank_nbest_token[i]))
new_scores.append(score)
if base_hypos:
new_hypos = self._gen_new_hypos(base_hypos, new_tokens, new_scores, t, device)
else:
new_hypos: List[Hypothesis] = []
return new_hypos
def _gen_new_hypos(
self,
base_hypos: List[Hypothesis],
tokens: List[int],
scores: List[float],
t: int,
device: torch.device,
) -> List[Hypothesis]:
tgt_tokens = torch.tensor([[token] for token in tokens], device=device)
states = _batch_state(base_hypos)
pred_out, _, pred_states = self.model.predict(
tgt_tokens,
torch.tensor([1] * len(base_hypos), device=device),
states,
)
new_hypos: List[Hypothesis] = []
for i, h_a in enumerate(base_hypos):
new_tokens = _get_hypo_tokens(h_a) + [tokens[i]]
new_hypos.append((new_tokens, pred_out[i].detach(), _slice_state(pred_states, i, device), scores[i]))
return new_hypos
def _search(
self,
enc_out: torch.Tensor,
hypo: Optional[Hypothesis],
beam_width: int,
) -> List[Hypothesis]:
n_time_steps = enc_out.shape[1]
device = enc_out.device
a_hypos: List[Hypothesis] = []
b_hypos = self._init_b_hypos(hypo, device)
for t in range(n_time_steps):
a_hypos = b_hypos
b_hypos = torch.jit.annotate(List[Hypothesis], [])
key_to_b_hypo: Dict[str, Hypothesis] = {}
symbols_current_t = 0
while a_hypos:
next_token_probs = self._gen_next_token_probs(enc_out[:, t : t + 1], a_hypos, device)
next_token_probs = next_token_probs.cpu()
b_hypos = self._gen_b_hypos(b_hypos, a_hypos, next_token_probs, key_to_b_hypo)
if symbols_current_t == self.step_max_tokens:
break
a_hypos = self._gen_a_hypos(
a_hypos,
b_hypos,
next_token_probs,
t,
beam_width,
device,
)
if a_hypos:
symbols_current_t += 1
_, sorted_idx = torch.tensor([self.hypo_sort_key(hypo) for hypo in b_hypos]).topk(beam_width)
b_hypos = [b_hypos[idx] for idx in sorted_idx]
return b_hypos
[docs] def forward(self, input: torch.Tensor, length: torch.Tensor, beam_width: int) -> List[Hypothesis]:
r"""Performs beam search for the given input sequence.
T: number of frames;
D: feature dimension of each frame.
Args:
input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D).
length (torch.Tensor): number of valid frames in input
sequence, with shape () or (1,).
beam_width (int): beam size to use during search.
Returns:
List[Hypothesis]: top-``beam_width`` hypotheses found by beam search.
"""
assert input.dim() == 2 or (
input.dim() == 3 and input.shape[0] == 1
), "input must be of shape (T, D) or (1, T, D)"
if input.dim() == 2:
input = input.unsqueeze(0)
assert length.shape == () or length.shape == (1,), "length must be of shape () or (1,)"
if input.dim() == 0:
input = input.unsqueeze(0)
enc_out, _ = self.model.transcribe(input, length)
return self._search(enc_out, None, beam_width)
[docs] @torch.jit.export
def infer(
self,
input: torch.Tensor,
length: torch.Tensor,
beam_width: int,
state: Optional[List[List[torch.Tensor]]] = None,
hypothesis: Optional[Hypothesis] = None,
) -> Tuple[List[Hypothesis], List[List[torch.Tensor]]]:
r"""Performs beam search for the given input sequence in streaming mode.
T: number of frames;
D: feature dimension of each frame.
Args:
input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D).
length (torch.Tensor): number of valid frames in input
sequence, with shape () or (1,).
beam_width (int): beam size to use during search.
state (List[List[torch.Tensor]] or None, optional): list of lists of tensors
representing transcription network internal state generated in preceding
invocation. (Default: ``None``)
hypothesis (Hypothesis or None): hypothesis from preceding invocation to seed
search with. (Default: ``None``)
Returns:
(List[Hypothesis], List[List[torch.Tensor]]):
List[Hypothesis]
top-``beam_width`` hypotheses found by beam search.
List[List[torch.Tensor]]
list of lists of tensors representing transcription network
internal state generated in current invocation.
"""
assert input.dim() == 2 or (
input.dim() == 3 and input.shape[0] == 1
), "input must be of shape (T, D) or (1, T, D)"
if input.dim() == 2:
input = input.unsqueeze(0)
assert length.shape == () or length.shape == (1,), "length must be of shape () or (1,)"
if length.dim() == 0:
length = length.unsqueeze(0)
enc_out, _, state = self.model.transcribe_streaming(input, length, state)
return self._search(enc_out, hypothesis, beam_width), state
