File size: 3,555 Bytes
071812d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)
#
# Copied from https://github.com/k2-fsa/sherpa/blob/master/sherpa/bin/conformer_rnnt/decode.py
#
# See LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
from typing import List
import torch
from sherpa import RnntConformerModel, greedy_search, modified_beam_search
from torch.nn.utils.rnn import pad_sequence
LOG_EPS = math.log(1e-10)
@torch.no_grad()
def run_model_and_do_greedy_search(
model: RnntConformerModel,
features: List[torch.Tensor],
) -> List[List[int]]:
"""Run RNN-T model with the given features and use greedy search
to decode the output of the model.
Args:
model:
The RNN-T model.
features:
A list of 2-D tensors. Each entry is of shape
(num_frames, feature_dim).
Returns:
Return a list-of-list containing the decoding token IDs.
"""
features_length = torch.tensor(
[f.size(0) for f in features],
dtype=torch.int64,
)
features = pad_sequence(
features,
batch_first=True,
padding_value=LOG_EPS,
)
device = model.device
features = features.to(device)
features_length = features_length.to(device)
encoder_out, encoder_out_length = model.encoder(
features=features,
features_length=features_length,
)
hyp_tokens = greedy_search(
model=model,
encoder_out=encoder_out,
encoder_out_length=encoder_out_length.cpu(),
)
return hyp_tokens
@torch.no_grad()
def run_model_and_do_modified_beam_search(
model: RnntConformerModel,
features: List[torch.Tensor],
num_active_paths: int,
) -> List[List[int]]:
"""Run RNN-T model with the given features and use greedy search
to decode the output of the model.
Args:
model:
The RNN-T model.
features:
A list of 2-D tensors. Each entry is of shape
(num_frames, feature_dim).
num_active_paths:
Used only when decoding_method is modified_beam_search.
It specifies number of active paths for each utterance. Due to
merging paths with identical token sequences, the actual number
may be less than "num_active_paths".
Returns:
Return a list-of-list containing the decoding token IDs.
"""
features_length = torch.tensor(
[f.size(0) for f in features],
dtype=torch.int64,
)
features = pad_sequence(
features,
batch_first=True,
padding_value=LOG_EPS,
)
device = model.device
features = features.to(device)
features_length = features_length.to(device)
encoder_out, encoder_out_length = model.encoder(
features=features,
features_length=features_length,
)
hyp_tokens = modified_beam_search(
model=model,
encoder_out=encoder_out,
encoder_out_length=encoder_out_length.cpu(),
num_active_paths=num_active_paths,
)
return hyp_tokens
|