#!/usr/bin/env python3
|
# -*- encoding: utf-8 -*-
|
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
# MIT License (https://opensource.org/licenses/MIT)
|
|
import torch
|
from typing import Dict, List, Optional, Tuple
|
|
from funasr.register import tables
|
from funasr.models.rwkv_bat.rwkv import RWKV
|
from funasr.models.transformer.layer_norm import LayerNorm
|
from funasr.models.transformer.utils.nets_utils import make_source_mask
|
from funasr.models.rwkv_bat.rwkv_subsampling import RWKVConvInput
|
|
|
@tables.register("encoder_classes", "RWKVEncoder")
|
class RWKVEncoder(torch.nn.Module):
|
"""RWKV encoder module.
|
|
Based on https://arxiv.org/pdf/2305.13048.pdf.
|
|
Args:
|
vocab_size: Vocabulary size.
|
output_size: Input/Output size.
|
context_size: Context size for WKV computation.
|
linear_size: FeedForward hidden size.
|
attention_size: SelfAttention hidden size.
|
normalization_type: Normalization layer type.
|
normalization_args: Normalization layer arguments.
|
num_blocks: Number of RWKV blocks.
|
embed_dropout_rate: Dropout rate for embedding layer.
|
att_dropout_rate: Dropout rate for the attention module.
|
ffn_dropout_rate: Dropout rate for the feed-forward module.
|
"""
|
|
def __init__(
|
self,
|
input_size: int,
|
output_size: int = 512,
|
context_size: int = 1024,
|
linear_size: Optional[int] = None,
|
attention_size: Optional[int] = None,
|
num_blocks: int = 4,
|
att_dropout_rate: float = 0.0,
|
ffn_dropout_rate: float = 0.0,
|
dropout_rate: float = 0.0,
|
subsampling_factor: int =4,
|
time_reduction_factor: int = 1,
|
kernel: int = 3,
|
**kwargs,
|
) -> None:
|
"""Construct a RWKVEncoder object."""
|
super().__init__()
|
|
self.embed = RWKVConvInput(
|
input_size,
|
[output_size//4, output_size//2, output_size],
|
subsampling_factor,
|
conv_kernel_size=kernel,
|
output_size=output_size,
|
)
|
|
self.subsampling_factor = subsampling_factor
|
|
linear_size = output_size * 4 if linear_size is None else linear_size
|
attention_size = output_size if attention_size is None else attention_size
|
|
self.rwkv_blocks = torch.nn.ModuleList(
|
[
|
RWKV(
|
output_size,
|
linear_size,
|
attention_size,
|
context_size,
|
block_id,
|
num_blocks,
|
att_dropout_rate=att_dropout_rate,
|
ffn_dropout_rate=ffn_dropout_rate,
|
dropout_rate=dropout_rate,
|
)
|
for block_id in range(num_blocks)
|
]
|
)
|
|
self.embed_norm = LayerNorm(output_size)
|
self.final_norm = LayerNorm(output_size)
|
|
self._output_size = output_size
|
self.context_size = context_size
|
|
self.num_blocks = num_blocks
|
self.time_reduction_factor = time_reduction_factor
|
|
def output_size(self) -> int:
|
return self._output_size
|
|
def forward(self, x: torch.Tensor, x_len) -> torch.Tensor:
|
"""Encode source label sequences.
|
|
Args:
|
x: Encoder input sequences. (B, L)
|
|
Returns:
|
out: Encoder output sequences. (B, U, D)
|
|
"""
|
_, length, _ = x.size()
|
|
assert (
|
length <= self.context_size * self.subsampling_factor
|
), "Context size is too short for current length: %d versus %d" % (
|
length,
|
self.context_size * self.subsampling_factor,
|
)
|
mask = make_source_mask(x_len).to(x.device)
|
x, mask = self.embed(x, mask, None)
|
x = self.embed_norm(x)
|
olens = mask.eq(0).sum(1)
|
|
if self.training:
|
for block in self.rwkv_blocks:
|
x, _ = block(x)
|
else:
|
x = self.rwkv_infer(x)
|
|
x = self.final_norm(x)
|
|
if self.time_reduction_factor > 1:
|
x = x[:,::self.time_reduction_factor,:]
|
olens = torch.floor_divide(olens-1, self.time_reduction_factor) + 1
|
|
return x, olens, None
|
|
def rwkv_infer(self, xs_pad):
|
|
batch_size = xs_pad.shape[0]
|
|
hidden_sizes = [
|
self._output_size for i in range(5)
|
]
|
|
state = [
|
torch.zeros(
|
(batch_size, 1, hidden_sizes[i], self.num_blocks),
|
dtype=torch.float32,
|
device=xs_pad.device,
|
)
|
for i in range(5)
|
]
|
|
state[4] -= 1e-30
|
|
xs_out = []
|
for t in range(xs_pad.shape[1]):
|
x_t = xs_pad[:,t,:]
|
for idx, block in enumerate(self.rwkv_blocks):
|
x_t, state = block(x_t, state=state)
|
xs_out.append(x_t)
|
xs_out = torch.cat(xs_out, dim=1)
|
return xs_out
|
