wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
1555 lines
72 KiB
Python
1555 lines
72 KiB
Python
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# 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
|
|
import random
|
|
from collections import OrderedDict
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional, Set, Tuple
|
|
|
|
import torch
|
|
import torch.distributed
|
|
from omegaconf import DictConfig, ListConfig, open_dict
|
|
from torch import nn
|
|
|
|
from nemo.collections.asr.models.configs import CacheAwareStreamingConfig
|
|
from nemo.collections.asr.parts.mixins.streaming import StreamingEncoder
|
|
from nemo.collections.asr.parts.submodules.causal_convs import CausalConv1D
|
|
from nemo.collections.asr.parts.submodules.conformer_modules import ConformerLayer
|
|
from nemo.collections.asr.parts.submodules.multi_head_attention import (
|
|
LocalAttRelPositionalEncoding,
|
|
MultiHeadAttention,
|
|
PositionalEncoding,
|
|
RelPositionalEncoding,
|
|
RelPositionMultiHeadAttention,
|
|
RelPositionMultiHeadAttentionLongformer,
|
|
RoPEMultiHeadAttention,
|
|
RotaryPositionalEncoding,
|
|
)
|
|
from nemo.collections.asr.parts.submodules.subsampling import (
|
|
ConvSubsampling,
|
|
StackingSubsampling,
|
|
SubsamplingReductionModule,
|
|
)
|
|
from nemo.collections.asr.parts.utils import adapter_utils
|
|
from nemo.collections.asr.parts.utils.regularization_utils import compute_stochastic_depth_drop_probs
|
|
from nemo.core.classes.common import typecheck
|
|
from nemo.core.classes.exportable import Exportable
|
|
from nemo.core.classes.mixins import AccessMixin, adapter_mixins
|
|
from nemo.core.classes.module import NeuralModule
|
|
from nemo.core.neural_types import (
|
|
AcousticEncodedRepresentation,
|
|
BoolType,
|
|
ChannelType,
|
|
LengthsType,
|
|
NeuralType,
|
|
SpectrogramType,
|
|
)
|
|
from nemo.utils import logging
|
|
|
|
__all__ = ['ConformerEncoder', 'ConformerMultiLayerFeatureExtractor']
|
|
|
|
|
|
class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin):
|
|
"""
|
|
The encoder for ASR model of Conformer.
|
|
Based on this paper:
|
|
'Conformer: Convolution-augmented Transformer for Speech Recognition' by Anmol Gulati et al.
|
|
https://arxiv.org/abs/2005.08100
|
|
|
|
Args:
|
|
feat_in (int): the size of feature channels
|
|
n_layers (int): number of layers of ConformerBlock
|
|
d_model (int): the hidden size of the model
|
|
feat_out (int): the size of the output features
|
|
Defaults to -1 (means feat_out is d_model)
|
|
subsampling (str): the method of subsampling:
|
|
choices = ['vggnet', 'striding', 'dw-striding', 'stacking', 'stacking_norm']
|
|
Defaults to striding.
|
|
subsampling_factor (int): the subsampling factor which should be power of 2
|
|
Defaults to 4.
|
|
subsampling_conv_chunking_factor(int): optionally, force chunk inputs (helpful for large inputs)
|
|
Should be power of 2, 1 (auto-chunking, default), or -1 (no chunking)
|
|
subsampling_conv_channels (int): the size of the convolutions in the subsampling module
|
|
Defaults to -1 which would set it to d_model.
|
|
reduction (str, Optional): the method of reduction, choices=['pooling', 'striding']. If no value
|
|
is passed, then no reduction is performed and the models runs with the original 4x subsampling.
|
|
reduction_position (int, Optional): the index of the layer to apply reduction. If -1, apply reduction
|
|
at the end.
|
|
reduction_factor (int): the reduction factor which should be either 1 or a power of 2
|
|
Defaults to 1.
|
|
ff_expansion_factor (int): the expansion factor in feed forward layers
|
|
Defaults to 4.
|
|
self_attention_model (str): the type of the attention layer and positional encoding.
|
|
|
|
'rel_pos':
|
|
relative positional embedding and Transformer-XL
|
|
'rel_pos_local_attn':
|
|
relative positional embedding and Transformer-XL with local attention using
|
|
overlapping chunks. Attention context is determined by att_context_size parameter.
|
|
'abs_pos':
|
|
absolute positional embedding and Transformer
|
|
'rope':
|
|
rotary position embedding
|
|
|
|
Default is rel_pos.
|
|
pos_emb_max_len (int): the maximum length of positional embeddings
|
|
Defaults to 5000
|
|
rope_base (float): theta base for the rotary position embedding.
|
|
Defaults to 10000.
|
|
rotary_fraction (float): fraction of the per-head dim to rotate.
|
|
Defaults to 1.0.
|
|
n_heads (int): number of heads in multi-headed attention layers
|
|
Defaults to 4.
|
|
att_context_size (List[Union[List[int],int]]): specifies the context sizes on each side.
|
|
Each context size should be a list of two integers like `[100, 100]`.
|
|
A list of context sizes like `[[100,100]`, `[100,50]]` can also be passed. -1 means unlimited context.
|
|
Defaults to `[-1, -1]`
|
|
att_context_probs (List[float]): a list of probabilities of each one of the att_context_size
|
|
when a list of them is passed. If not specified, uniform distribution is being used.
|
|
Defaults to None
|
|
att_chunk_context_size (List[List[int]]): specifies the context sizes for unified (offline/streaming) ASR training.
|
|
It defines the range of Left, Middle, and Right context sizes for the attention mechanism.
|
|
At each streaming step, the context size is sampled from the range of Left, Middle, and Right context sizes.
|
|
Example: att_chunk_context_size=[[70],[1,2,7,13],[0,1,3,7,13]] -> sampling -> [70, 2, 3] -> attention mask generation
|
|
att_context_style (str): 'regular', 'chunked_limited', or 'chunked_limited_with_rc'.
|
|
Defaults to 'regular'
|
|
xscaling (bool): enables scaling the inputs to the multi-headed attention layers by `sqrt(d_model)`.
|
|
Defaults to True.
|
|
untie_biases (bool): whether to not share (untie) the bias weights between layers of Transformer-XL
|
|
Defaults to True.
|
|
conv_kernel_size (int): the size of the convolutions in the convolutional modules
|
|
Defaults to 31.
|
|
conv_norm_type (str): the type of the normalization in the convolutional modules
|
|
Defaults to 'batch_norm'.
|
|
conv_context_size (list): it can be"causal" or a list of two integers
|
|
while `conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size`.
|
|
`None` means `[(conv_kernel_size-1)//2`, `(conv_kernel_size-1)//2]`, and 'causal' means
|
|
`[(conv_kernel_size-1), 0]`.
|
|
Defaults to None.
|
|
conv_context_style (str): 'regular' or 'dcc'
|
|
DCC - Dynamic Chunked Convolution that is used for unified ASR training.
|
|
Defaults to 'regular'.
|
|
conv_dual_mode (bool): specifies if convolution should be dual mode when dual_offline mode is being used.
|
|
When enables, the left half of the convolution kernel would get masked in streaming cases.
|
|
Defaults to False.
|
|
use_bias (bool): Use bias in all Linear and Conv1d layers from each ConformerLayer to improve
|
|
activation flow and stabilize training of huge models.
|
|
Defaults to True.
|
|
dropout (float): the dropout rate used in all layers except the attention layers
|
|
Defaults to 0.1.
|
|
dropout_pre_encoder (float): the dropout rate used before the encoder
|
|
Defaults to 0.1.
|
|
dropout_emb (float): the dropout rate used for the positional embeddings
|
|
Defaults to 0.1.
|
|
dropout_att (float): the dropout rate used for the attention layer
|
|
Defaults to 0.0.
|
|
stochastic_depth_drop_prob (float): if non-zero, will randomly drop
|
|
layers during training. The higher this value, the more often layers
|
|
are dropped. Defaults to 0.0.
|
|
stochastic_depth_mode (str): can be either "linear" or "uniform". If
|
|
set to "uniform", all layers have the same probability of drop. If
|
|
set to "linear", the drop probability grows linearly from 0 for the
|
|
first layer to the desired value for the final layer. Defaults to
|
|
"linear".
|
|
stochastic_depth_start_layer (int): starting layer for stochastic depth.
|
|
All layers before this will never be dropped. Note that drop
|
|
probability will be adjusted accordingly if mode is "linear" when
|
|
start layer is > 1. Defaults to 1.
|
|
global_tokens (int): number of tokens to be used for global attention.
|
|
Only relevant if self_attention_model is 'rel_pos_local_attn'.
|
|
Defaults to 0.
|
|
global_tokens_spacing (int): how far apart the global tokens are
|
|
Defaults to 1.
|
|
global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate.
|
|
Defaults to False.
|
|
use_pytorch_sdpa (bool): use torch sdpa instead of manual attention.
|
|
Defaults to False.
|
|
use_pytorch_sdpa_backends (list[str]): list of backend names to use in sdpa.
|
|
None or empty list means all backends. e.g. ["MATH"]
|
|
Defaults to None.
|
|
bypass_pre_encode: if True, skip the pre-encoder module and the `audio_signal` should be pre-encoded
|
|
embeddings. The `audio_signal` input supports two formats depending on the `bypass_pre_encode`
|
|
boolean flag. This determines the required format of the input variable `audio_signal`.
|
|
Defaults to `bypass_pre_encode=False`. `bypass_pre_encode=True` is used for the cases
|
|
where frame-level, context-independent embeddings are needed to be saved or reused.
|
|
(e.g., speaker cache in streaming speaker diarization)
|
|
sync_max_audio_length (bool): when true, performs NCCL all_reduce to allocate the same amount of memory for
|
|
positional encoding buffers on all GPUs. Disabling this setting may help with deadlocks in certain
|
|
scenarios such as model parallelism, or generally when this module is not being ran on some GPUs
|
|
as a part of the training step.
|
|
"""
|
|
|
|
def input_example(self, max_batch=1, max_dim=256):
|
|
"""
|
|
Generates input examples for tracing etc.
|
|
Returns:
|
|
A tuple of input examples.
|
|
"""
|
|
dev = next(self.parameters()).device
|
|
if self.export_cache_support:
|
|
window_size = max_dim
|
|
if self.streaming_cfg is not None:
|
|
if isinstance(self.streaming_cfg.chunk_size, list):
|
|
chunk_size = self.streaming_cfg.chunk_size[1]
|
|
else:
|
|
chunk_size = self.streaming_cfg.chunk_size
|
|
if isinstance(self.streaming_cfg.pre_encode_cache_size, list):
|
|
pre_encode_cache_size = self.streaming_cfg.pre_encode_cache_size[1]
|
|
else:
|
|
pre_encode_cache_size = self.streaming_cfg.pre_encode_cache_size
|
|
window_size = chunk_size + pre_encode_cache_size
|
|
input_example = torch.randn(max_batch, self._feat_in, window_size, device=dev)
|
|
input_example_length = torch.randint(
|
|
window_size // 4, window_size, (max_batch,), device=dev, dtype=torch.int64
|
|
)
|
|
cache_last_channel, cache_last_time, cache_last_channel_len = self.get_initial_cache_state(
|
|
batch_size=max_batch, device=dev, max_dim=max_dim
|
|
)
|
|
all_input_example = tuple(
|
|
[
|
|
input_example,
|
|
input_example_length,
|
|
cache_last_channel.transpose(0, 1),
|
|
cache_last_time.transpose(0, 1),
|
|
cache_last_channel_len,
|
|
]
|
|
)
|
|
else:
|
|
input_example = torch.randn(max_batch, self._feat_in, max_dim, device=dev)
|
|
input_example_length = torch.randint(max_dim // 4, max_dim, (max_batch,), device=dev, dtype=torch.int64)
|
|
all_input_example = tuple([input_example, input_example_length])
|
|
|
|
return all_input_example
|
|
|
|
@property
|
|
def input_types(self):
|
|
"""Returns definitions of module input ports."""
|
|
return OrderedDict(
|
|
{
|
|
"audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
|
|
"length": NeuralType(tuple('B'), LengthsType()),
|
|
"cache_last_channel": NeuralType(('D', 'B', 'T', 'D'), ChannelType(), optional=True),
|
|
"cache_last_time": NeuralType(('D', 'B', 'D', 'T'), ChannelType(), optional=True),
|
|
"cache_last_channel_len": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
"bypass_pre_encode": NeuralType(tuple(), BoolType(), optional=True),
|
|
}
|
|
)
|
|
|
|
@property
|
|
def input_types_for_export(self):
|
|
"""Returns definitions of module input ports."""
|
|
return OrderedDict(
|
|
{
|
|
"audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
|
|
"length": NeuralType(tuple('B'), LengthsType()),
|
|
"cache_last_channel": NeuralType(('B', 'D', 'T', 'D'), ChannelType(), optional=True),
|
|
"cache_last_time": NeuralType(('B', 'D', 'D', 'T'), ChannelType(), optional=True),
|
|
"cache_last_channel_len": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
"bypass_pre_encode": NeuralType(tuple(), BoolType(), optional=True),
|
|
}
|
|
)
|
|
|
|
@property
|
|
def output_types(self):
|
|
"""Returns definitions of module output ports."""
|
|
return OrderedDict(
|
|
{
|
|
"outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
|
|
"encoded_lengths": NeuralType(tuple('B'), LengthsType()),
|
|
"cache_last_channel_next": NeuralType(('D', 'B', 'T', 'D'), ChannelType(), optional=True),
|
|
"cache_last_time_next": NeuralType(('D', 'B', 'D', 'T'), ChannelType(), optional=True),
|
|
"cache_last_channel_next_len": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
}
|
|
)
|
|
|
|
@property
|
|
def output_types_for_export(self):
|
|
"""Returns definitions of module output ports."""
|
|
return OrderedDict(
|
|
{
|
|
"outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
|
|
"encoded_lengths": NeuralType(tuple('B'), LengthsType()),
|
|
"cache_last_channel_next": NeuralType(('B', 'D', 'T', 'D'), ChannelType(), optional=True),
|
|
"cache_last_time_next": NeuralType(('B', 'D', 'D', 'T'), ChannelType(), optional=True),
|
|
"cache_last_channel_next_len": NeuralType(tuple('B'), LengthsType(), optional=True),
|
|
}
|
|
)
|
|
|
|
@property
|
|
def disabled_deployment_input_names(self):
|
|
if not self.export_cache_support:
|
|
return set(["cache_last_channel", "cache_last_time", "cache_last_channel_len"])
|
|
else:
|
|
return set()
|
|
|
|
@property
|
|
def disabled_deployment_output_names(self):
|
|
if not self.export_cache_support:
|
|
return set(["cache_last_channel_next", "cache_last_time_next", "cache_last_channel_next_len"])
|
|
else:
|
|
return set()
|
|
|
|
def __init__(
|
|
self,
|
|
feat_in,
|
|
n_layers,
|
|
d_model,
|
|
feat_out=-1,
|
|
causal_downsampling=False,
|
|
subsampling='striding',
|
|
subsampling_factor=4,
|
|
subsampling_conv_chunking_factor=1,
|
|
subsampling_conv_channels=-1,
|
|
reduction=None,
|
|
reduction_position=None,
|
|
reduction_factor=1,
|
|
ff_expansion_factor=4,
|
|
self_attention_model='rel_pos',
|
|
n_heads=4,
|
|
att_context_size=None,
|
|
att_context_probs=None,
|
|
att_chunk_context_size=None,
|
|
att_context_style='regular',
|
|
xscaling=True,
|
|
untie_biases=True,
|
|
pos_emb_max_len=5000,
|
|
conv_kernel_size=31,
|
|
conv_norm_type='batch_norm',
|
|
conv_context_size=None,
|
|
conv_context_style='regular',
|
|
use_bias=True,
|
|
dropout=0.1,
|
|
dropout_pre_encoder=0.1,
|
|
dropout_emb=0.1,
|
|
dropout_att=0.0,
|
|
stochastic_depth_drop_prob: float = 0.0,
|
|
stochastic_depth_mode: str = "linear",
|
|
stochastic_depth_start_layer: int = 1,
|
|
global_tokens: int = 0,
|
|
global_tokens_spacing: int = 1,
|
|
global_attn_separate: bool = False,
|
|
use_pytorch_sdpa: bool = False,
|
|
use_pytorch_sdpa_backends=None,
|
|
sync_max_audio_length: bool = True,
|
|
rope_base: float = 10000.0,
|
|
rotary_fraction: float = 1.0,
|
|
):
|
|
super().__init__()
|
|
d_ff = d_model * ff_expansion_factor
|
|
self.d_model = d_model
|
|
self.n_layers = n_layers
|
|
self._feat_in = feat_in
|
|
self.att_context_style = att_context_style
|
|
self.subsampling_factor = subsampling_factor
|
|
self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
|
|
|
|
self.self_attention_model = self_attention_model
|
|
self.global_tokens = global_tokens
|
|
self.global_attn_separate = global_attn_separate
|
|
self.global_tokens_spacing = global_tokens_spacing
|
|
self.use_pytorch_sdpa = use_pytorch_sdpa
|
|
if use_pytorch_sdpa_backends is None:
|
|
use_pytorch_sdpa_backends = []
|
|
self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
|
|
self.sync_max_audio_length = sync_max_audio_length
|
|
|
|
assert conv_context_style in ["regular", "dcc"], f"Invalid conv_context_style: {conv_context_style}!"
|
|
self.conv_context_style = conv_context_style
|
|
self.conv_kernel_size = conv_kernel_size
|
|
|
|
# Setting up the att_chunk_context_size
|
|
if att_chunk_context_size is not None:
|
|
assert (
|
|
att_context_style == "chunked_limited_with_rc"
|
|
), "att_chunk_context_size is only supported for chunked_limited_with_rc attention style!"
|
|
assert (
|
|
len(att_chunk_context_size) == 3
|
|
), "att_chunk_context_size must have 3 elements: [left_context, chunk_size, right_context]"
|
|
self.att_chunk_context_size = att_chunk_context_size
|
|
else:
|
|
self.att_chunk_context_size = None
|
|
|
|
# Setting up the att_context_size
|
|
(
|
|
self.att_context_size_all,
|
|
self.att_context_size,
|
|
self.att_context_probs,
|
|
self.conv_context_size,
|
|
) = self._calc_context_sizes(
|
|
att_context_style=att_context_style,
|
|
att_context_size=att_context_size,
|
|
att_context_probs=att_context_probs,
|
|
conv_context_size=conv_context_size,
|
|
conv_kernel_size=conv_kernel_size,
|
|
)
|
|
|
|
if xscaling:
|
|
self.xscale = math.sqrt(d_model)
|
|
else:
|
|
self.xscale = None
|
|
|
|
# Subsampling
|
|
if subsampling_conv_channels == -1:
|
|
subsampling_conv_channels = d_model
|
|
if subsampling and subsampling_factor > 1:
|
|
if subsampling in ['stacking', 'stacking_norm']:
|
|
# stacking_norm has an extra layer norm after stacking comparing to stacking
|
|
self.pre_encode = StackingSubsampling(
|
|
subsampling_factor=subsampling_factor,
|
|
feat_in=feat_in,
|
|
feat_out=d_model,
|
|
norm=True if subsampling == 'stacking_norm' else False,
|
|
)
|
|
else:
|
|
self.pre_encode = ConvSubsampling(
|
|
subsampling=subsampling,
|
|
subsampling_factor=subsampling_factor,
|
|
feat_in=feat_in,
|
|
feat_out=d_model,
|
|
conv_channels=subsampling_conv_channels,
|
|
subsampling_conv_chunking_factor=subsampling_conv_chunking_factor,
|
|
activation=nn.ReLU(True),
|
|
is_causal=causal_downsampling,
|
|
)
|
|
else:
|
|
self.pre_encode = nn.Linear(feat_in, d_model)
|
|
|
|
# Reduction
|
|
if reduction and reduction_factor > 1:
|
|
assert reduction_position >= -1 and reduction_position < n_layers
|
|
self.reduction_subsampling = SubsamplingReductionModule(
|
|
reduction=reduction,
|
|
d_model=d_model,
|
|
reduction_factor=reduction_factor,
|
|
)
|
|
self.reduction_position = reduction_position
|
|
else:
|
|
self.reduction_subsampling = None
|
|
self.reduction_position = None
|
|
|
|
self._feat_out = d_model
|
|
|
|
# Biases for relative positional encoding
|
|
if not untie_biases and self_attention_model == "rel_pos":
|
|
d_head = d_model // n_heads
|
|
pos_bias_u = nn.Parameter(torch.Tensor(n_heads, d_head))
|
|
pos_bias_v = nn.Parameter(torch.Tensor(n_heads, d_head))
|
|
nn.init.zeros_(pos_bias_u)
|
|
nn.init.zeros_(pos_bias_v)
|
|
else:
|
|
pos_bias_u = None
|
|
pos_bias_v = None
|
|
|
|
# Positional encodings
|
|
self.pos_emb_max_len = pos_emb_max_len
|
|
if self_attention_model == "rel_pos":
|
|
self.pos_enc = RelPositionalEncoding(
|
|
d_model=d_model,
|
|
dropout_rate=dropout_pre_encoder,
|
|
max_len=pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
dropout_rate_emb=dropout_emb,
|
|
)
|
|
elif self_attention_model == 'rel_pos_local_attn':
|
|
if max(att_context_size) <= 0:
|
|
raise ValueError("When using local attention, context size must be set > 0")
|
|
self.pos_enc = LocalAttRelPositionalEncoding(
|
|
att_context_size=att_context_size,
|
|
d_model=d_model,
|
|
dropout_rate=dropout,
|
|
max_len=pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
dropout_rate_emb=dropout_emb,
|
|
)
|
|
elif self_attention_model == "abs_pos":
|
|
pos_bias_u = None
|
|
pos_bias_v = None
|
|
self.pos_enc = PositionalEncoding(
|
|
d_model=d_model, dropout_rate=dropout_pre_encoder, max_len=pos_emb_max_len, xscale=self.xscale
|
|
)
|
|
elif self_attention_model == "rope":
|
|
self.dropout_pre_encoder = torch.nn.Dropout(dropout_pre_encoder)
|
|
self.pos_enc = RotaryPositionalEncoding(
|
|
d_k=d_model // n_heads,
|
|
rotary_fraction=rotary_fraction,
|
|
rope_base=rope_base,
|
|
max_len=pos_emb_max_len,
|
|
)
|
|
else:
|
|
raise ValueError(f"Not valid self_attention_model: '{self_attention_model}'!")
|
|
|
|
layer_pos_enc = self.pos_enc if self_attention_model == 'rope' else None
|
|
self.layers = nn.ModuleList()
|
|
for i in range(n_layers):
|
|
layer = ConformerLayer(
|
|
d_model=d_model,
|
|
d_ff=d_ff,
|
|
self_attention_model=self_attention_model,
|
|
global_tokens=global_tokens,
|
|
global_tokens_spacing=global_tokens_spacing,
|
|
global_attn_separate=global_attn_separate,
|
|
n_heads=n_heads,
|
|
conv_kernel_size=conv_kernel_size,
|
|
conv_norm_type=conv_norm_type,
|
|
conv_context_size=self.conv_context_size,
|
|
dropout=dropout,
|
|
dropout_att=dropout_att,
|
|
pos_bias_u=pos_bias_u,
|
|
pos_bias_v=pos_bias_v,
|
|
att_context_size=self.att_context_size,
|
|
use_bias=use_bias,
|
|
use_pytorch_sdpa=self.use_pytorch_sdpa,
|
|
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
|
|
pos_enc=layer_pos_enc,
|
|
)
|
|
self.layers.append(layer)
|
|
|
|
if feat_out > 0 and feat_out != self._feat_out:
|
|
self.out_proj = nn.Linear(self._feat_out, feat_out)
|
|
self._feat_out = feat_out
|
|
else:
|
|
self.out_proj = None
|
|
self._feat_out = d_model
|
|
self.set_max_audio_length(self.pos_emb_max_len)
|
|
self.use_pad_mask = True
|
|
|
|
self.setup_streaming_params()
|
|
self.export_cache_support = False
|
|
|
|
self.layer_drop_probs = compute_stochastic_depth_drop_probs(
|
|
len(self.layers), stochastic_depth_drop_prob, stochastic_depth_mode, stochastic_depth_start_layer
|
|
)
|
|
# will be set in self.forward() if defined in AccessMixin config
|
|
self.interctc_capture_at_layers = None
|
|
|
|
def forward_for_export(
|
|
self, audio_signal, length, cache_last_channel=None, cache_last_time=None, cache_last_channel_len=None
|
|
):
|
|
"""
|
|
Forward function for model export. Please see `forward()` for more details.
|
|
"""
|
|
if cache_last_channel is not None:
|
|
cache_last_channel = cache_last_channel.transpose(0, 1)
|
|
cache_last_time = cache_last_time.transpose(0, 1)
|
|
|
|
rets = self.forward_internal(
|
|
audio_signal,
|
|
length,
|
|
cache_last_channel=cache_last_channel,
|
|
cache_last_time=cache_last_time,
|
|
cache_last_channel_len=cache_last_channel_len,
|
|
)
|
|
rets = self.streaming_post_process(rets, keep_all_outputs=False)
|
|
if len(rets) == 2:
|
|
return rets
|
|
elif rets[2] is None and rets[3] is None and rets[4] is None:
|
|
return (rets[0], rets[1])
|
|
else:
|
|
return (
|
|
rets[0],
|
|
rets[1],
|
|
rets[2].transpose(0, 1),
|
|
rets[3].transpose(0, 1),
|
|
rets[4],
|
|
)
|
|
|
|
def streaming_post_process(self, rets, keep_all_outputs=True):
|
|
"""
|
|
Post-process the output of the forward function for streaming.
|
|
|
|
Args:
|
|
rets: The output of the forward function.
|
|
keep_all_outputs: Whether to keep all outputs.
|
|
"""
|
|
if len(rets) == 2:
|
|
return rets[0], rets[1], None, None, None
|
|
|
|
(encoded, encoded_len, cache_last_channel_next, cache_last_time_next, cache_last_channel_next_len) = rets
|
|
|
|
if cache_last_channel_next is not None and self.streaming_cfg.last_channel_cache_size >= 0:
|
|
if self.streaming_cfg.last_channel_cache_size > 0:
|
|
cache_last_channel_next = cache_last_channel_next[
|
|
:, :, -self.streaming_cfg.last_channel_cache_size :, :
|
|
]
|
|
|
|
if self.streaming_cfg.valid_out_len > 0 and (not keep_all_outputs or self.att_context_style == "regular"):
|
|
encoded = encoded[:, :, : self.streaming_cfg.valid_out_len]
|
|
encoded_len = torch.clamp(encoded_len, max=self.streaming_cfg.valid_out_len)
|
|
|
|
return (encoded, encoded_len, cache_last_channel_next, cache_last_time_next, cache_last_channel_next_len)
|
|
|
|
@typecheck()
|
|
def forward(
|
|
self,
|
|
audio_signal,
|
|
length,
|
|
cache_last_channel=None,
|
|
cache_last_time=None,
|
|
cache_last_channel_len=None,
|
|
bypass_pre_encode=False,
|
|
):
|
|
"""
|
|
Forward function for the ConformerEncoder accepting an audio signal and its corresponding length.
|
|
The ``audio_signal`` input supports two formats depending on ``bypass_pre_encode``:
|
|
|
|
- ``bypass_pre_encode=False`` (default): ``audio_signal`` must be a tensor
|
|
containing audio features. Shape: ``(batch, feat_in, n_frames)``.
|
|
- ``bypass_pre_encode=True``: ``audio_signal`` must be a tensor containing
|
|
pre-encoded embeddings. Shape: ``(batch, n_frame, d_model)``.
|
|
"""
|
|
if not bypass_pre_encode and audio_signal.shape[-2] != self._feat_in:
|
|
raise ValueError(
|
|
f"If bypass_pre_encode is False, audio_signal should have shape "
|
|
f"(batch, {self._feat_in}, n_frame) but got last dimension {audio_signal.shape[-2]}."
|
|
)
|
|
if bypass_pre_encode and audio_signal.shape[-1] != self.d_model:
|
|
raise ValueError(
|
|
f"If bypass_pre_encode is True, audio_signal should have shape "
|
|
f"(batch, n_frame, {self.d_model}) but got last dimension {audio_signal.shape[-1]}."
|
|
)
|
|
|
|
if bypass_pre_encode:
|
|
self.update_max_seq_length(seq_length=audio_signal.size(1), device=audio_signal.device)
|
|
else:
|
|
self.update_max_seq_length(seq_length=audio_signal.size(2), device=audio_signal.device)
|
|
return self.forward_internal(
|
|
audio_signal,
|
|
length,
|
|
cache_last_channel=cache_last_channel,
|
|
cache_last_time=cache_last_time,
|
|
cache_last_channel_len=cache_last_channel_len,
|
|
bypass_pre_encode=bypass_pre_encode,
|
|
)
|
|
|
|
def forward_internal(
|
|
self,
|
|
audio_signal,
|
|
length,
|
|
cache_last_channel=None,
|
|
cache_last_time=None,
|
|
cache_last_channel_len=None,
|
|
bypass_pre_encode=False,
|
|
):
|
|
"""
|
|
The ``audio_signal`` input supports two formats depending on ``bypass_pre_encode``:
|
|
|
|
- ``bypass_pre_encode=False`` (default): ``audio_signal`` must be a tensor
|
|
containing audio features. Shape: ``(batch, feat_in, n_frames)``.
|
|
- ``bypass_pre_encode=True``: ``audio_signal`` must be a tensor containing
|
|
pre-encoded embeddings. Shape: ``(batch, n_frame, d_model)``.
|
|
|
|
``bypass_pre_encode=True`` is used in cases where frame-level, context-independent embeddings are
|
|
needed to be saved or reused (e.g., speaker cache in streaming speaker diarization).
|
|
"""
|
|
if length is None:
|
|
length = audio_signal.new_full(
|
|
(audio_signal.size(0),), audio_signal.size(-1), dtype=torch.int64, device=audio_signal.device
|
|
)
|
|
|
|
# select a random att_context_size with the distribution specified by att_context_probs during training
|
|
# for non-validation cases like test, validation or inference, it uses the first mode in self.att_context_size
|
|
if self.training and len(self.att_context_size_all) > 1:
|
|
cur_att_context_size = random.choices(self.att_context_size_all, weights=self.att_context_probs)[0]
|
|
else:
|
|
cur_att_context_size = self.att_context_size
|
|
|
|
if not bypass_pre_encode:
|
|
audio_signal = torch.transpose(audio_signal, 1, 2)
|
|
|
|
if isinstance(self.pre_encode, nn.Linear):
|
|
audio_signal = self.pre_encode(audio_signal)
|
|
else:
|
|
audio_signal, length = self.pre_encode(x=audio_signal, lengths=length)
|
|
length = length.to(torch.int64)
|
|
# `self.streaming_cfg` is set by setup_streaming_cfg(), called in the init
|
|
if self.streaming_cfg.drop_extra_pre_encoded > 0 and cache_last_channel is not None:
|
|
audio_signal = audio_signal[:, self.streaming_cfg.drop_extra_pre_encoded :, :]
|
|
length = (length - self.streaming_cfg.drop_extra_pre_encoded).clamp(min=0)
|
|
|
|
if self.reduction_position is not None and cache_last_channel is not None:
|
|
raise ValueError("Caching with reduction feature is not supported yet!")
|
|
|
|
max_audio_length = audio_signal.size(1)
|
|
if cache_last_channel is not None:
|
|
cache_len = self.streaming_cfg.last_channel_cache_size
|
|
cache_keep_size = max_audio_length - self.streaming_cfg.cache_drop_size
|
|
max_audio_length = max_audio_length + cache_len
|
|
padding_length = length + cache_len
|
|
offset = torch.neg(cache_last_channel_len) + cache_len
|
|
else:
|
|
padding_length = length
|
|
cache_last_channel_next = None
|
|
cache_len = 0
|
|
offset = None
|
|
|
|
if self.self_attention_model == 'rope':
|
|
if self.xscale:
|
|
audio_signal = audio_signal * self.xscale
|
|
audio_signal = self.dropout_pre_encoder(audio_signal)
|
|
pos_emb = None
|
|
else:
|
|
audio_signal, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len)
|
|
|
|
# Create the self-attention and padding masks
|
|
pad_mask, att_mask = self._create_masks(
|
|
att_context_size=cur_att_context_size,
|
|
padding_length=padding_length,
|
|
max_audio_length=max_audio_length,
|
|
offset=offset,
|
|
device=audio_signal.device,
|
|
)
|
|
|
|
if cache_last_channel is not None:
|
|
pad_mask = pad_mask[:, cache_len:]
|
|
if att_mask is not None:
|
|
att_mask = att_mask[:, cache_len:]
|
|
# Convert caches from the tensor to list
|
|
cache_last_time_next = []
|
|
cache_last_channel_next = []
|
|
|
|
for lth, (drop_prob, layer) in enumerate(zip(self.layer_drop_probs, self.layers)):
|
|
original_signal = audio_signal
|
|
if cache_last_channel is not None:
|
|
cache_last_channel_cur = cache_last_channel[lth]
|
|
cache_last_time_cur = cache_last_time[lth]
|
|
else:
|
|
cache_last_channel_cur = None
|
|
cache_last_time_cur = None
|
|
audio_signal = layer(
|
|
x=audio_signal,
|
|
att_mask=att_mask,
|
|
pos_emb=pos_emb,
|
|
pad_mask=pad_mask,
|
|
cache_last_channel=cache_last_channel_cur,
|
|
cache_last_time=cache_last_time_cur,
|
|
)
|
|
|
|
if cache_last_channel_cur is not None:
|
|
(audio_signal, cache_last_channel_cur, cache_last_time_cur) = audio_signal
|
|
cache_last_channel_next.append(cache_last_channel_cur)
|
|
cache_last_time_next.append(cache_last_time_cur)
|
|
|
|
# applying stochastic depth logic from https://arxiv.org/abs/2102.03216
|
|
if self.training and drop_prob > 0.0:
|
|
should_drop = torch.rand(1) < drop_prob
|
|
# adjusting to match expectation
|
|
if should_drop:
|
|
# that's not efficient, but it's hard to implement distributed
|
|
# version of dropping layers without deadlock or random seed meddling
|
|
# so multiplying the signal by 0 to ensure all weights get gradients
|
|
audio_signal = audio_signal * 0.0 + original_signal
|
|
else:
|
|
# not doing this operation if drop prob is 0 as it's identity in that case
|
|
audio_signal = (audio_signal - original_signal) / (1.0 - drop_prob) + original_signal
|
|
|
|
if self.reduction_position == lth:
|
|
audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length)
|
|
max_audio_length = audio_signal.size(1)
|
|
# Don't update the audio_signal here because then it will again scale the audio_signal
|
|
# and cause an increase in the WER
|
|
if self.self_attention_model != 'rope':
|
|
_, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len)
|
|
pad_mask, att_mask = self._create_masks(
|
|
att_context_size=cur_att_context_size,
|
|
padding_length=length,
|
|
max_audio_length=max_audio_length,
|
|
offset=offset,
|
|
device=audio_signal.device,
|
|
)
|
|
# saving tensors if required for interctc loss
|
|
if self.is_access_enabled(getattr(self, "model_guid", None)):
|
|
if self.interctc_capture_at_layers is None:
|
|
self.interctc_capture_at_layers = self.access_cfg.get('interctc', {}).get('capture_layers', [])
|
|
if lth in self.interctc_capture_at_layers:
|
|
lth_audio_signal = audio_signal
|
|
if self.out_proj is not None:
|
|
lth_audio_signal = self.out_proj(audio_signal)
|
|
# shape is the same as the shape of audio_signal output, i.e. [B, D, T]
|
|
self.register_accessible_tensor(
|
|
name=f'interctc/layer_output_{lth}', tensor=torch.transpose(lth_audio_signal, 1, 2)
|
|
)
|
|
self.register_accessible_tensor(name=f'interctc/layer_length_{lth}', tensor=length)
|
|
|
|
if self.out_proj is not None:
|
|
audio_signal = self.out_proj(audio_signal)
|
|
|
|
# Reduction
|
|
if self.reduction_position == -1:
|
|
audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length)
|
|
|
|
audio_signal = torch.transpose(audio_signal, 1, 2)
|
|
length = length.to(dtype=torch.int64)
|
|
|
|
if cache_last_channel is not None:
|
|
cache_last_channel_next = torch.stack(cache_last_channel_next, dim=0)
|
|
cache_last_time_next = torch.stack(cache_last_time_next, dim=0)
|
|
return (
|
|
audio_signal,
|
|
length,
|
|
cache_last_channel_next,
|
|
cache_last_time_next,
|
|
torch.clamp(cache_last_channel_len + cache_keep_size, max=cache_len),
|
|
)
|
|
else:
|
|
return audio_signal, length
|
|
|
|
def update_max_seq_length(self, seq_length: int, device):
|
|
"""
|
|
Updates the maximum sequence length for the model.
|
|
|
|
Args:
|
|
seq_length (int): New maximum sequence length.
|
|
device (torch.device): Device to use for computations.
|
|
"""
|
|
# Find global max audio length across all nodes
|
|
if self.sync_max_audio_length and torch.distributed.is_initialized():
|
|
global_max_len = torch.tensor([seq_length], dtype=torch.float32, device=device)
|
|
|
|
# Update across all ranks in the distributed system
|
|
torch.distributed.all_reduce(global_max_len, op=torch.distributed.ReduceOp.MAX)
|
|
|
|
seq_length = global_max_len.int().item()
|
|
|
|
if seq_length > self.max_audio_length:
|
|
self.set_max_audio_length(seq_length)
|
|
|
|
def set_max_audio_length(self, max_audio_length):
|
|
"""
|
|
Sets maximum input length.
|
|
Pre-calculates internal seq_range mask.
|
|
|
|
Args:
|
|
max_audio_length (int): New maximum sequence length.
|
|
"""
|
|
self.max_audio_length = max_audio_length
|
|
device = next(self.parameters()).device
|
|
dtype = next(self.parameters()).dtype
|
|
self.pos_enc.extend_pe(max_audio_length, device, dtype)
|
|
|
|
def _create_masks(self, att_context_size, padding_length, max_audio_length, offset, device):
|
|
if self.self_attention_model != "rel_pos_local_attn":
|
|
att_mask = torch.ones(1, max_audio_length, max_audio_length, dtype=torch.bool, device=device)
|
|
|
|
if self.att_context_style == "regular":
|
|
if att_context_size[0] >= 0:
|
|
att_mask = att_mask.triu(diagonal=-att_context_size[0])
|
|
if att_context_size[1] >= 0:
|
|
att_mask = att_mask.tril(diagonal=att_context_size[1])
|
|
elif self.att_context_style == "chunked_limited":
|
|
# When right context is unlimited, just the left side of the masking need to get updated
|
|
if att_context_size[1] == -1:
|
|
if att_context_size[0] >= 0:
|
|
att_mask = att_mask.triu(diagonal=-att_context_size[0])
|
|
else:
|
|
chunk_size = att_context_size[1] + 1
|
|
# left_chunks_num specifies the number of chunks to be visible by each chunk on the left side
|
|
if att_context_size[0] >= 0:
|
|
left_chunks_num = att_context_size[0] // chunk_size
|
|
else:
|
|
left_chunks_num = 10000
|
|
|
|
chunk_idx = torch.arange(0, max_audio_length, dtype=torch.int, device=att_mask.device)
|
|
chunk_idx = torch.div(chunk_idx, chunk_size, rounding_mode="trunc")
|
|
diff_chunks = chunk_idx.unsqueeze(1) - chunk_idx.unsqueeze(0)
|
|
chunked_limited_mask = torch.logical_and(
|
|
torch.le(diff_chunks, left_chunks_num), torch.ge(diff_chunks, 0)
|
|
)
|
|
att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0))
|
|
elif self.att_context_style == "chunked_limited_with_rc" and sum(att_context_size) != -3:
|
|
assert (
|
|
len(att_context_size) == 3
|
|
), "att_context_size must have 3 elements: [left_context, chunk_size, right_context]"
|
|
|
|
left_context_frames = att_context_size[0]
|
|
chunk_size_frames = att_context_size[1]
|
|
right_context_frames = att_context_size[2]
|
|
assert chunk_size_frames >= 1, "chunk_size_frames must be greater than 0!"
|
|
# Calculate chunk index for each frame (which processing group it belongs to)
|
|
frame_idx = torch.arange(0, max_audio_length, dtype=torch.int, device=att_mask.device)
|
|
chunk_idx = torch.div(frame_idx, chunk_size_frames, rounding_mode="trunc")
|
|
|
|
window_start = chunk_idx * chunk_size_frames - left_context_frames
|
|
window_start = torch.maximum(window_start, torch.zeros_like(window_start))
|
|
window_end = chunk_idx * chunk_size_frames + chunk_size_frames - 1 + right_context_frames
|
|
|
|
window_end = torch.minimum(window_end, torch.full_like(window_end, max_audio_length - 1))
|
|
# Create the mask: frame i can see frame j if window_start[i] <= j <= window_end[i]
|
|
j_indices = frame_idx.unsqueeze(0) # [1, T]
|
|
window_start_expanded = window_start.unsqueeze(1) # [T, 1]
|
|
window_end_expanded = window_end.unsqueeze(1) # [T, 1]
|
|
|
|
chunked_limited_mask = torch.logical_and(
|
|
j_indices >= window_start_expanded, j_indices <= window_end_expanded
|
|
)
|
|
att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0))
|
|
else:
|
|
att_mask = None
|
|
|
|
# pad_mask is the masking to be used to ignore paddings
|
|
pad_mask = torch.arange(0, max_audio_length, device=device).expand(
|
|
padding_length.size(0), -1
|
|
) < padding_length.unsqueeze(-1)
|
|
|
|
if offset is not None:
|
|
pad_mask_off = torch.arange(0, max_audio_length, device=device).expand(
|
|
padding_length.size(0), -1
|
|
) >= offset.unsqueeze(-1)
|
|
pad_mask = pad_mask_off.logical_and(pad_mask)
|
|
|
|
if att_mask is not None:
|
|
# pad_mask_for_att_mask is the mask which helps to ignore paddings
|
|
pad_mask_for_att_mask = pad_mask.unsqueeze(1).repeat([1, max_audio_length, 1])
|
|
pad_mask_for_att_mask = torch.logical_and(pad_mask_for_att_mask, pad_mask_for_att_mask.transpose(1, 2))
|
|
# att_mask is the masking to be used by the MHA layers to ignore the tokens not supposed to be visible
|
|
att_mask = att_mask[:, :max_audio_length, :max_audio_length]
|
|
# paddings should also get ignored, so pad_mask_for_att_mask is used to ignore their corresponding scores
|
|
att_mask = torch.logical_and(pad_mask_for_att_mask, att_mask.to(pad_mask_for_att_mask.device))
|
|
att_mask = ~att_mask
|
|
|
|
pad_mask = ~pad_mask
|
|
return pad_mask, att_mask
|
|
|
|
def enable_pad_mask(self, on=True):
|
|
"""
|
|
Enables or disables the pad mask and assign the boolean state `on`.
|
|
|
|
Returns:
|
|
mask (bool): The current state of the pad mask.
|
|
"""
|
|
# On inference, user may choose to disable pad mask
|
|
mask = self.use_pad_mask
|
|
self.use_pad_mask = on
|
|
return mask
|
|
|
|
def _calc_context_sizes(
|
|
self, att_context_size, att_context_probs, att_context_style, conv_context_size, conv_kernel_size
|
|
):
|
|
# convert att_context_size to a standard list of lists
|
|
if att_context_size:
|
|
att_context_size_all = list(att_context_size)
|
|
if isinstance(att_context_size_all[0], int):
|
|
att_context_size_all = [att_context_size_all]
|
|
for i, att_cs in enumerate(att_context_size_all):
|
|
if isinstance(att_cs, ListConfig):
|
|
att_context_size_all[i] = list(att_cs)
|
|
if att_context_style == "chunked_limited":
|
|
if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0:
|
|
raise ValueError(f"att_context_size[{i}][0] % (att_context_size[{i}][1] + 1) should be zero!")
|
|
if att_cs[1] < 0 and len(att_context_size_all) <= 1:
|
|
raise ValueError(
|
|
f"Right context (att_context_size[{i}][1]) can not be unlimited for chunked_limited style!"
|
|
)
|
|
else:
|
|
att_context_size_all = [[-1, -1]]
|
|
|
|
if att_context_style == "chunked_limited_with_rc":
|
|
att_context_size_all = [[-1, -1, -1]]
|
|
|
|
if att_context_probs:
|
|
if len(att_context_probs) != len(att_context_size_all):
|
|
raise ValueError("The size of the att_context_probs should be the same as att_context_size.")
|
|
att_context_probs = list(att_context_probs)
|
|
if sum(att_context_probs) != 1:
|
|
raise ValueError(
|
|
"The sum of numbers in att_context_probs should be equal to one to be a distribution."
|
|
)
|
|
else:
|
|
att_context_probs = [1.0 / len(att_context_size_all)] * len(att_context_size_all)
|
|
|
|
if conv_context_size is not None:
|
|
if isinstance(conv_context_size, ListConfig):
|
|
conv_context_size = list(conv_context_size)
|
|
if not isinstance(conv_context_size, list) and not isinstance(conv_context_size, str):
|
|
raise ValueError(
|
|
"Invalid conv_context_size! It should be the string 'causal' or a list of two integers."
|
|
)
|
|
if conv_context_size == "causal":
|
|
conv_context_size = [conv_kernel_size - 1, 0]
|
|
else:
|
|
if conv_context_size[0] + conv_context_size[1] + 1 != conv_kernel_size:
|
|
raise ValueError(f"Invalid conv_context_size: {self.conv_context_size}!")
|
|
else:
|
|
conv_context_size = [(conv_kernel_size - 1) // 2, (conv_kernel_size - 1) // 2]
|
|
return att_context_size_all, att_context_size_all[0], att_context_probs, conv_context_size
|
|
|
|
def set_default_att_context_size(self, att_context_size):
|
|
"""
|
|
Sets the default attention context size from `att_context_size` argument.
|
|
|
|
Args:
|
|
att_context_size (list): The attention context size to be set.
|
|
"""
|
|
if att_context_size not in self.att_context_size_all:
|
|
logging.warning(
|
|
f"att_context_size={att_context_size} is not among the list of the supported "
|
|
f"look-aheads: {self.att_context_size_all}"
|
|
)
|
|
if att_context_size is not None:
|
|
self.att_context_size = att_context_size
|
|
|
|
self.setup_streaming_params()
|
|
|
|
def setup_streaming_params(
|
|
self,
|
|
chunk_size: int = None,
|
|
shift_size: int = None,
|
|
left_chunks: int = None,
|
|
att_context_size: list = None,
|
|
max_context: int = 10000,
|
|
):
|
|
"""
|
|
This function sets the needed values and parameters to perform streaming.
|
|
The configuration would be stored in self.streaming_cfg.
|
|
The streaming configuration is needed to simulate streaming inference.
|
|
|
|
Args:
|
|
chunk_size (int): overrides the chunk size
|
|
shift_size (int): overrides the shift size for chunks
|
|
left_chunks (int): overrides the number of left chunks visible to each chunk
|
|
max_context (int): the value used for the cache size of last_channel layers
|
|
if left context is set to infinity (-1)
|
|
Defaults to -1 (means feat_out is d_model)
|
|
"""
|
|
streaming_cfg = CacheAwareStreamingConfig()
|
|
|
|
# When att_context_size is not specified, it uses the default_att_context_size
|
|
if att_context_size is None:
|
|
att_context_size = self.att_context_size
|
|
|
|
if chunk_size is not None:
|
|
if chunk_size < 1:
|
|
raise ValueError("chunk_size needs to be a number larger or equal to one.")
|
|
lookahead_steps = chunk_size - 1
|
|
streaming_cfg.cache_drop_size = chunk_size - shift_size
|
|
elif self.att_context_style == "chunked_limited":
|
|
lookahead_steps = att_context_size[1]
|
|
streaming_cfg.cache_drop_size = 0
|
|
elif self.att_context_style == "chunked_limited_with_rc":
|
|
lookahead_steps = att_context_size[2] * self.n_layers + self.conv_context_size[1] * self.n_layers
|
|
streaming_cfg.cache_drop_size = 0
|
|
elif self.att_context_style == "regular":
|
|
lookahead_steps = att_context_size[1] * self.n_layers + self.conv_context_size[1] * self.n_layers
|
|
streaming_cfg.cache_drop_size = lookahead_steps
|
|
else:
|
|
streaming_cfg.cache_drop_size = 0
|
|
lookahead_steps = None
|
|
|
|
if chunk_size is None:
|
|
streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context
|
|
else:
|
|
if left_chunks is None:
|
|
streaming_cfg.last_channel_cache_size = (
|
|
att_context_size[0] if att_context_size[0] >= 0 else max_context
|
|
)
|
|
logging.warning(
|
|
f"left_chunks is not set. Setting it to default: {streaming_cfg.last_channel_cache_size}."
|
|
)
|
|
else:
|
|
streaming_cfg.last_channel_cache_size = left_chunks * chunk_size
|
|
|
|
if hasattr(self.pre_encode, "get_sampling_frames"):
|
|
sampling_frames = self.pre_encode.get_sampling_frames()
|
|
else:
|
|
sampling_frames = 0
|
|
|
|
if isinstance(sampling_frames, list):
|
|
streaming_cfg.chunk_size = [
|
|
sampling_frames[0] + self.subsampling_factor * lookahead_steps,
|
|
sampling_frames[1] + self.subsampling_factor * lookahead_steps,
|
|
]
|
|
else:
|
|
streaming_cfg.chunk_size = sampling_frames * (1 + lookahead_steps)
|
|
|
|
if isinstance(sampling_frames, list):
|
|
streaming_cfg.shift_size = [
|
|
sampling_frames[0] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size),
|
|
sampling_frames[1] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size),
|
|
]
|
|
else:
|
|
streaming_cfg.shift_size = sampling_frames * (1 + lookahead_steps - streaming_cfg.cache_drop_size)
|
|
|
|
if isinstance(streaming_cfg.shift_size, list):
|
|
streaming_cfg.valid_out_len = (
|
|
streaming_cfg.shift_size[1] - sampling_frames[1]
|
|
) // self.subsampling_factor + 1
|
|
else:
|
|
streaming_cfg.valid_out_len = streaming_cfg.shift_size // self.subsampling_factor
|
|
|
|
if hasattr(self.pre_encode, "get_streaming_cache_size"):
|
|
streaming_cfg.pre_encode_cache_size = self.pre_encode.get_streaming_cache_size()
|
|
else:
|
|
streaming_cfg.pre_encode_cache_size = 0
|
|
|
|
if isinstance(streaming_cfg.pre_encode_cache_size, list):
|
|
if streaming_cfg.pre_encode_cache_size[1] >= 1:
|
|
streaming_cfg.drop_extra_pre_encoded = (
|
|
1 + (streaming_cfg.pre_encode_cache_size[1] - 1) // self.subsampling_factor
|
|
)
|
|
else:
|
|
streaming_cfg.drop_extra_pre_encoded = 0
|
|
else:
|
|
streaming_cfg.drop_extra_pre_encoded = streaming_cfg.pre_encode_cache_size // self.subsampling_factor
|
|
|
|
for m in self.layers.modules():
|
|
if hasattr(m, "_max_cache_len"):
|
|
if isinstance(m, MultiHeadAttention):
|
|
m.cache_drop_size = streaming_cfg.cache_drop_size
|
|
if isinstance(m, CausalConv1D):
|
|
m.cache_drop_size = streaming_cfg.cache_drop_size
|
|
|
|
self.streaming_cfg = streaming_cfg
|
|
|
|
def get_initial_cache_state(self, batch_size=1, dtype=torch.float32, device=None, max_dim=0):
|
|
if device is None:
|
|
device = next(self.parameters()).device
|
|
if max_dim > 0:
|
|
create_tensor = torch.randn
|
|
else:
|
|
create_tensor = torch.zeros
|
|
last_time_cache_size = self.conv_context_size[0]
|
|
cache_last_channel = create_tensor(
|
|
(
|
|
len(self.layers),
|
|
batch_size,
|
|
self.streaming_cfg.last_channel_cache_size,
|
|
self.d_model,
|
|
),
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
cache_last_time = create_tensor(
|
|
(len(self.layers), batch_size, self.d_model, last_time_cache_size),
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
if max_dim > 0:
|
|
cache_last_channel_len = torch.randint(
|
|
0,
|
|
min(max_dim, self.streaming_cfg.last_channel_cache_size),
|
|
(batch_size,),
|
|
device=device,
|
|
dtype=torch.int64,
|
|
)
|
|
for i in range(batch_size):
|
|
cache_last_channel[:, i, cache_last_channel_len[i] :, :] = 0
|
|
# what is the right rule to zero out cache_last_time?
|
|
if cache_last_channel_len[i] == 0:
|
|
cache_last_time[:, i, :, :] = 0
|
|
else:
|
|
cache_last_channel_len = torch.zeros(batch_size, device=device, dtype=torch.int64)
|
|
return cache_last_channel, cache_last_time, cache_last_channel_len
|
|
|
|
def change_attention_model(
|
|
self,
|
|
self_attention_model: str = None,
|
|
att_context_size: List[int] = None,
|
|
update_config: bool = True,
|
|
device: torch.device = None,
|
|
rope_base: float = None,
|
|
rotary_fraction: float = None,
|
|
):
|
|
"""
|
|
Update the self_attention_model which changes the positional encoding and attention layers.
|
|
|
|
Args:
|
|
self_attention_model (str): type of the attention layer and positional encoding
|
|
|
|
'rel_pos':
|
|
relative positional embedding and Transformer-XL
|
|
|
|
'rel_pos_local_attn':
|
|
relative positional embedding and Transformer-XL with local attention using
|
|
overlapping windows. Attention context is determined by att_context_size parameter.
|
|
|
|
'abs_pos':
|
|
absolute positional embedding and Transformer
|
|
|
|
'rope':
|
|
rotary position embedding
|
|
|
|
If None is provided, the self_attention_model isn't changed. Defaults to None.
|
|
att_context_size (List[int]): List of 2 ints corresponding to left and right attention context sizes,
|
|
or None to keep as it is. Defaults to None.
|
|
update_config (bool): Whether to update the config or not with the new attention model.
|
|
Defaults to True.
|
|
device (torch.device): If provided, new layers will be moved to the device.
|
|
Defaults to None.
|
|
rope_base (float): Theta base for rotary position embedding. Only used when
|
|
``self_attention_model='rope'``. If None, the stored config value is kept.
|
|
rotary_fraction (float): Fraction of the per-head dim to rotate. Only used when
|
|
``self_attention_model='rope'``. If None, the stored config value is kept.
|
|
"""
|
|
|
|
if att_context_size:
|
|
att_context_size = list(att_context_size)
|
|
else:
|
|
att_context_size = self.att_context_size
|
|
|
|
if self_attention_model is None:
|
|
self_attention_model = self.self_attention_model
|
|
|
|
if rope_base is None:
|
|
rope_base = getattr(self._cfg, 'rope_base', 10000.0)
|
|
if rotary_fraction is None:
|
|
rotary_fraction = getattr(self._cfg, 'rotary_fraction', 1.0)
|
|
|
|
if self_attention_model == 'rel_pos_local_attn' and max(att_context_size) <= 0:
|
|
raise ValueError("When using local attention, context size must be set > 0")
|
|
|
|
if self_attention_model == "rel_pos":
|
|
new_pos_enc = RelPositionalEncoding(
|
|
d_model=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout,
|
|
max_len=self._cfg.pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
dropout_rate_emb=self._cfg.dropout_emb,
|
|
)
|
|
elif self_attention_model == 'rel_pos_local_attn':
|
|
new_pos_enc = LocalAttRelPositionalEncoding(
|
|
att_context_size=att_context_size,
|
|
d_model=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout,
|
|
max_len=self._cfg.pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
dropout_rate_emb=self._cfg.dropout_emb,
|
|
)
|
|
elif self_attention_model == "abs_pos":
|
|
new_pos_enc = PositionalEncoding(
|
|
d_model=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout,
|
|
max_len=self._cfg.pos_emb_max_len,
|
|
xscale=self.xscale,
|
|
)
|
|
elif self_attention_model == "rope":
|
|
self.dropout_pre_encoder = torch.nn.Dropout(getattr(self._cfg, 'dropout_pre_encoder', 0.1))
|
|
new_pos_enc = RotaryPositionalEncoding(
|
|
d_k=self._cfg.d_model // self._cfg.n_heads,
|
|
rotary_fraction=rotary_fraction,
|
|
rope_base=rope_base,
|
|
max_len=self._cfg.pos_emb_max_len,
|
|
)
|
|
else:
|
|
raise ValueError(f"Not valid self_attention_model: '{self_attention_model}'!")
|
|
|
|
if device is not None:
|
|
new_pos_enc = new_pos_enc.to(device=device)
|
|
del self.pos_enc
|
|
self.pos_enc = new_pos_enc
|
|
self.self_attention_model = self_attention_model
|
|
self.att_context_size = att_context_size
|
|
self.set_max_audio_length(self.pos_emb_max_len)
|
|
|
|
for _, m in self.named_modules():
|
|
if type(m) == ConformerLayer:
|
|
if self_attention_model == 'rel_pos':
|
|
new_attn = RelPositionMultiHeadAttention(
|
|
n_head=self._cfg.n_heads,
|
|
n_feat=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout_att,
|
|
max_cache_len=att_context_size[0],
|
|
pos_bias_u=None,
|
|
pos_bias_v=None,
|
|
use_bias=getattr(self._cfg, 'use_bias', True),
|
|
use_pytorch_sdpa=self.use_pytorch_sdpa,
|
|
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
|
|
)
|
|
elif self_attention_model == 'rel_pos_local_attn':
|
|
new_attn = RelPositionMultiHeadAttentionLongformer(
|
|
n_head=self._cfg.n_heads,
|
|
n_feat=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout_att,
|
|
max_cache_len=att_context_size[0],
|
|
att_context_size=att_context_size,
|
|
pos_bias_u=None,
|
|
pos_bias_v=None,
|
|
use_bias=getattr(self._cfg, 'use_bias', True),
|
|
use_pytorch_sdpa=self.use_pytorch_sdpa,
|
|
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
|
|
)
|
|
elif self_attention_model == 'abs_pos':
|
|
new_attn = MultiHeadAttention(
|
|
n_head=self._cfg.n_heads,
|
|
n_feat=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout_att,
|
|
max_cache_len=att_context_size[0],
|
|
use_bias=getattr(self._cfg, 'use_bias', True),
|
|
use_pytorch_sdpa=self.use_pytorch_sdpa,
|
|
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
|
|
)
|
|
elif self_attention_model == 'rope':
|
|
new_attn = RoPEMultiHeadAttention(
|
|
n_head=self._cfg.n_heads,
|
|
n_feat=self._cfg.d_model,
|
|
dropout_rate=self._cfg.dropout_att,
|
|
pos_enc=new_pos_enc,
|
|
max_cache_len=att_context_size[0],
|
|
use_bias=getattr(self._cfg, 'use_bias', True),
|
|
use_pytorch_sdpa=self.use_pytorch_sdpa,
|
|
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"'{self_attention_model}' is not not a valid value for 'self_attention_model', "
|
|
f"valid values can be from ['rel_pos', 'rel_pos_local_attn', 'abs_pos', 'rope']"
|
|
)
|
|
if device is not None:
|
|
new_attn = new_attn.to(device=device)
|
|
new_attn.load_state_dict(m.self_attn.state_dict(), strict=False)
|
|
del m.self_attn
|
|
m.self_attn = new_attn
|
|
m.self_attention_model = self_attention_model
|
|
|
|
if update_config:
|
|
with open_dict(self._cfg):
|
|
self._cfg.self_attention_model = self_attention_model
|
|
self._cfg.att_context_size = att_context_size
|
|
if self_attention_model == 'rope':
|
|
self._cfg.rope_base = rope_base
|
|
self._cfg.rotary_fraction = rotary_fraction
|
|
|
|
def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int):
|
|
"""
|
|
Update the conv_chunking_factor (int)
|
|
Default is 1 (auto)
|
|
Set it to -1 (disabled) or to a specific value (power of 2) if you OOM in the conv subsampling layers
|
|
|
|
|
|
Args:
|
|
subsampling_conv_chunking_factor (int)
|
|
"""
|
|
|
|
if not hasattr(self.pre_encode, "change_subsampling_conv_chunking_factor"):
|
|
logging.info("Model pre_encoder doesn't have a change_subsampling_conv_chunking_factor method ")
|
|
return
|
|
|
|
self.pre_encode.change_subsampling_conv_chunking_factor(
|
|
subsampling_conv_chunking_factor=subsampling_conv_chunking_factor
|
|
)
|
|
|
|
|
|
class ConformerEncoderAdapter(ConformerEncoder, adapter_mixins.AdapterModuleMixin):
|
|
"""This class inherits from ConformerEncoder and wraps the adapter mixin class."""
|
|
|
|
# Higher level forwarding
|
|
def add_adapter(self, name: str, cfg: dict):
|
|
cfg = self._update_adapter_cfg_input_dim(cfg)
|
|
for conformer_layer in self.layers: # type: adapter_mixins.AdapterModuleMixin
|
|
conformer_layer.add_adapter(name, cfg)
|
|
|
|
def is_adapter_available(self) -> bool:
|
|
return any([conformer_layer.is_adapter_available() for conformer_layer in self.layers])
|
|
|
|
def set_enabled_adapters(self, name: Optional[str] = None, enabled: bool = True):
|
|
for conformer_layer in self.layers: # type: adapter_mixins.AdapterModuleMixin
|
|
conformer_layer.set_enabled_adapters(name=name, enabled=enabled)
|
|
|
|
def get_enabled_adapters(self) -> List[str]:
|
|
names = set([])
|
|
for conformer_layer in self.layers: # type: adapter_mixins.AdapterModuleMixin
|
|
names.update(conformer_layer.get_enabled_adapters())
|
|
|
|
names = sorted(list(names))
|
|
return names
|
|
|
|
def _update_adapter_cfg_input_dim(self, cfg: DictConfig):
|
|
cfg = adapter_utils.update_adapter_cfg_input_dim(self, cfg, module_dim=self.d_model)
|
|
return cfg
|
|
|
|
def get_accepted_adapter_types(
|
|
self,
|
|
) -> Set[type]:
|
|
types = super().get_accepted_adapter_types()
|
|
|
|
if len(types) == 0:
|
|
self.set_accepted_adapter_types(
|
|
[
|
|
adapter_utils.LINEAR_ADAPTER_CLASSPATH,
|
|
adapter_utils.MHA_ADAPTER_CLASSPATH,
|
|
adapter_utils.RELMHA_ADAPTER_CLASSPATH,
|
|
]
|
|
)
|
|
types = self.get_accepted_adapter_types()
|
|
return types
|
|
|
|
|
|
class ConformerMultiLayerFeatureExtractor(NeuralModule, Exportable, AccessMixin):
|
|
"""
|
|
A wrapper module that extracts features from multiple layers of a ConformerEncoder,
|
|
by reusing existing mechanisim for interctc loss.
|
|
To use it, set `layer_idx_list` to specify the indices of layers to extract from.
|
|
Also, you can specify an `aggretator` module to aggregate the features from different layers,
|
|
default not aggregating.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
encoder: ConformerEncoder,
|
|
layer_idx_list: Optional[List[int]] = None,
|
|
aggregator: Optional[NeuralModule] = None,
|
|
detach: bool = False,
|
|
convert_to_cpu: bool = False,
|
|
include_final_output: bool = True,
|
|
):
|
|
"""
|
|
This class is used to extract features from different layers of the ConformerEncoder.
|
|
Args:
|
|
encoder: ConformerEncoder instance.
|
|
layer_idx_list: List of layer indices to extract features from. If None, all layers are extracted.
|
|
Negative indices follow standard Python convention (``-1`` == ``num_layers - 1``).
|
|
aggregator: Aggregator instance. If None, the features are returned as a list.
|
|
detach: If True, the features are detached from the graph.
|
|
convert_to_cpu: If True, the features are converted to CPU.
|
|
include_final_output: If True (default), the post-encoder final output (the value returned
|
|
by ``encoder.forward()`` after its out_proj + optional reduction) is appended to the
|
|
feature list, after all intermediate-layer captures. This is distinct from capturing
|
|
``num_layers - 1`` via ``layer_idx_list`` — that captures the raw last-layer activation,
|
|
whereas this captures the post-projection final output. Set False to return only the
|
|
intermediate captures.
|
|
"""
|
|
super().__init__()
|
|
self.encoder = encoder
|
|
self.num_layers = len(encoder.layers)
|
|
self.layer_idx_list = []
|
|
self.include_final_output = include_final_output
|
|
if not layer_idx_list:
|
|
layer_idx_list = list(range(self.num_layers))
|
|
for lid in layer_idx_list:
|
|
if lid < -self.num_layers or lid >= self.num_layers:
|
|
raise ValueError(f"Invalid layer index {lid} for ConformerEncoder with {self.num_layers} layers.")
|
|
if lid < 0:
|
|
lid = self.num_layers + lid
|
|
self.layer_idx_list.append(lid)
|
|
self.layer_idx_list.sort()
|
|
logging.info(
|
|
f"Extracting ConformerEncoder features from layers: {self.layer_idx_list}"
|
|
+ (" (+ final encoder output)" if self.include_final_output else "")
|
|
)
|
|
# Layers past the last captured intermediate index contribute no gradient to any
|
|
# downstream loss (their output is not consumed), so Adam would never allocate optimizer
|
|
# state for them and DCP resume would fail with "Missing key in checkpoint state_dict".
|
|
# Freeze them so the optimizer never adopts them in the first place. Skip when
|
|
# ``include_final_output`` is set, since that path backprops through every layer.
|
|
if not self.include_final_output and self.layer_idx_list:
|
|
max_used_layer = self.layer_idx_list[-1]
|
|
for layer in encoder.layers[max_used_layer + 1 :]:
|
|
for p in layer.parameters():
|
|
p.requires_grad_(False)
|
|
self.enc_access_cfg = {
|
|
"interctc": {
|
|
"capture_layers": self.layer_idx_list,
|
|
},
|
|
"detach": detach,
|
|
"convert_to_cpu": convert_to_cpu,
|
|
}
|
|
self.aggregator = aggregator
|
|
|
|
def forward(
|
|
self, audio_signal, length, cache_last_channel=None, cache_last_time=None, cache_last_channel_len=None
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Args:
|
|
same interface as ConformerEncoder.forward()
|
|
Returns:
|
|
- Tuple[List[Tensor[B,D,T]], List[Tensor[B]]] if aggregator is None
|
|
- Tuple[Tensor[B,H,T], Tensor[B]] if aggregator is not None, where H is the hidden size of the aggregator
|
|
"""
|
|
old_access_flag = self.is_access_enabled(guid=getattr(self, "model_guid", None))
|
|
self.update_access_cfg(self.enc_access_cfg, guid=getattr(self, "model_guid", None))
|
|
self.set_access_enabled(access_enabled=True, guid=getattr(self, "model_guid", None))
|
|
|
|
encoder_ret = self.encoder(
|
|
audio_signal=audio_signal,
|
|
length=length,
|
|
cache_last_channel=cache_last_channel,
|
|
cache_last_time=cache_last_time,
|
|
cache_last_channel_len=cache_last_channel_len,
|
|
)
|
|
# ConformerEncoder.forward_internal returns (audio_signal, length) when caches are unused
|
|
# and a 5-tuple (+ cache tensors) otherwise. First two elements are always the final
|
|
# [B, D, T] output and the [B] length.
|
|
encoder_out, encoder_out_len = encoder_ret[0], encoder_ret[1]
|
|
|
|
# Chunk of code adapted from ConformerEncoder.forward_internal()
|
|
total_registry = {}
|
|
for module_registry in self.get_module_registry(self.encoder).values():
|
|
for key in module_registry:
|
|
if key.startswith("interctc/") and key in total_registry:
|
|
raise RuntimeError(f"layer {key} has been logged multiple times!")
|
|
total_registry.update(module_registry)
|
|
|
|
encoded_list = []
|
|
encoded_len_list = []
|
|
for layer_idx in self.layer_idx_list:
|
|
try:
|
|
layer_outputs = total_registry[f"interctc/layer_output_{layer_idx}"]
|
|
layer_lengths = total_registry[f"interctc/layer_length_{layer_idx}"]
|
|
except KeyError:
|
|
raise RuntimeError(
|
|
f"Intermediate layer {layer_idx} was not captured! "
|
|
"Check the layer index and the number of ConformerEncoder layers."
|
|
)
|
|
if len(layer_outputs) > 1 or len(layer_lengths) > 1:
|
|
raise RuntimeError("Make sure encoder.forward is called exactly one time")
|
|
encoded_list.append(layer_outputs[0]) # [B, D, T]
|
|
encoded_len_list.append(layer_lengths[0]) # [B]
|
|
|
|
if self.include_final_output:
|
|
# encoder_out is already [B, D, T] (ConformerEncoder.forward_internal transposes it
|
|
# on line ~765), matching the layout of the captured interctc tensors above.
|
|
encoded_list.append(encoder_out)
|
|
encoded_len_list.append(encoder_out_len)
|
|
|
|
self.encoder.reset_registry()
|
|
self.set_access_enabled(access_enabled=old_access_flag, guid=getattr(self, "model_guid", None))
|
|
# End of the adapted chunk
|
|
|
|
if self.aggregator is not None:
|
|
return self.aggregator(encoded_list, encoded_len_list) # Tensor[B,H,T], Tensor[B]
|
|
else:
|
|
return encoded_list, encoded_len_list # List[Tensor[B,D,T]], List[Tensor[B]]
|
|
|
|
|
|
# Register any additional information
|
|
if adapter_mixins.get_registered_adapter(ConformerEncoder) is None:
|
|
adapter_mixins.register_adapter(base_class=ConformerEncoder, adapter_class=ConformerEncoderAdapter)
|
|
|
|
|
|
@dataclass
|
|
class ConformerChangeConfig:
|
|
"""
|
|
Change self_attention_model for Conformer.
|
|
|
|
Options:
|
|
'rel_pos': relative positional embedding and Transformer-XL
|
|
'rel_pos_local_attn': relative positional embedding and Transformer-XL with local attention using
|
|
overlapping chunks. Attention context is determined by att_context_size parameter.
|
|
'abs_pos': absolute positional embedding and Transformer
|
|
'rope': rotary position embedding
|
|
"""
|
|
|
|
# If None is provided, self_attention_model is not changed.
|
|
self_attention_model: Optional[str] = None
|
|
|
|
# Change the attention context size by providing 2 integers,
|
|
# corresponding to left and right context, or -1 for full context.
|
|
# If None is provided, the attention context size isn't changed.
|
|
att_context_size: Optional[List[int]] = None
|
|
|
|
# Rotary position embedding parameters; only used when self_attention_model is
|
|
# being set to (or already is) 'rope'. If None, the values from the stored
|
|
# config are kept.
|
|
rope_base: Optional[float] = None
|
|
rotary_fraction: Optional[float] = None
|