diff --git a/funasr/models/rwkv_bat/model.py b/funasr/models/rwkv_bat/model.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/funasr/models/rwkv_bat/rwkv.py b/funasr/models/rwkv_bat/rwkv.py index 422e1c8fe..bd218a282 100644 --- a/funasr/models/rwkv_bat/rwkv.py +++ b/funasr/models/rwkv_bat/rwkv.py @@ -1,16 +1,15 @@ -"""Receptance Weighted Key Value (RWKV) block definition. - -Based/modified from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/src/model.py - -""" - -from typing import Dict, Optional, Tuple +#!/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, Optional, Tuple -from funasr.models.rwkv_bat.rwkv_attention import EncoderSelfAttention, DecoderSelfAttention -from funasr.models.rwkv_bat.rwkv_feed_forward import FeedForward from funasr.models.transformer.layer_norm import LayerNorm +from funasr.models.rwkv_bat.rwkv_feed_forward import FeedForward +from funasr.models.rwkv_bat.rwkv_attention import EncoderSelfAttention, DecoderSelfAttention + class RWKV(torch.nn.Module): """RWKV module. diff --git a/funasr/models/rwkv_bat/rwkv_attention.py b/funasr/models/rwkv_bat/rwkv_attention.py index 5384fb9ca..c085874e4 100644 --- a/funasr/models/rwkv_bat/rwkv_attention.py +++ b/funasr/models/rwkv_bat/rwkv_attention.py @@ -1,17 +1,14 @@ -"""Attention (time mixing) modules for RWKV block. - -Based/Modified from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/src/model.py. - -Some variables are renamed according to https://github.com/huggingface/transformers/blob/main/src/transformers/models/rwkv/modeling_rwkv.py. - -""" # noqa +#!/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 math -from importlib.util import find_spec +import torch from pathlib import Path +from importlib.util import find_spec from typing import List, Optional, Tuple, Union -import torch wkv_kernel_encoder = None wkv_kernel_decoder = None diff --git a/funasr/models/rwkv_bat/rwkv_encoder.py b/funasr/models/rwkv_bat/rwkv_encoder.py index af702e91b..c0e5f4255 100644 --- a/funasr/models/rwkv_bat/rwkv_encoder.py +++ b/funasr/models/rwkv_bat/rwkv_encoder.py @@ -1,17 +1,20 @@ -"""RWKV encoder definition for Transducer models.""" - -import math -from typing import Dict, List, Optional, Tuple +#!/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.models.encoder.abs_encoder import AbsEncoder +from funasr.register import tables from funasr.models.rwkv_bat.rwkv import RWKV from funasr.models.transformer.layer_norm import LayerNorm -from funasr.models.rwkv_bat.rwkv_subsampling import RWKVConvInput from funasr.models.transformer.utils.nets_utils import make_source_mask +from funasr.models.rwkv_bat.rwkv_subsampling import RWKVConvInput -class RWKVEncoder(AbsEncoder): + +@tables.register("encoder_classes", "RWKVEncoder") +class RWKVEncoder(torch.nn.Module): """RWKV encoder module. Based on https://arxiv.org/pdf/2305.13048.pdf. @@ -44,6 +47,7 @@ class RWKVEncoder(AbsEncoder): subsampling_factor: int =4, time_reduction_factor: int = 1, kernel: int = 3, + **kwargs, ) -> None: """Construct a RWKVEncoder object.""" super().__init__() diff --git a/funasr/models/rwkv_bat/rwkv_feed_forward.py b/funasr/models/rwkv_bat/rwkv_feed_forward.py index ddb42859e..32949ab02 100644 --- a/funasr/models/rwkv_bat/rwkv_feed_forward.py +++ b/funasr/models/rwkv_bat/rwkv_feed_forward.py @@ -1,14 +1,10 @@ -"""Feed-forward (channel mixing) module for RWKV block. - -Based/Modified from https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4/src/model.py - -Some variables are renamed according to https://github.com/huggingface/transformers/blob/main/src/transformers/models/rwkv/modeling_rwkv.py. - -""" # noqa - -from typing import List, Optional, Tuple +#!/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 List, Optional, Tuple class FeedForward(torch.nn.Module): diff --git a/funasr/models/rwkv_bat/rwkv_subsampling.py b/funasr/models/rwkv_bat/rwkv_subsampling.py index 54ad1f5ad..a688acaff 100644 --- a/funasr/models/rwkv_bat/rwkv_subsampling.py +++ b/funasr/models/rwkv_bat/rwkv_subsampling.py @@ -1,19 +1,13 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- +# -*- encoding: utf-8 -*- +# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved. +# MIT License (https://opensource.org/licenses/MIT) -# Copyright 2019 Shigeki Karita -# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) - -"""Subsampling layer definition.""" -import numpy as np -import torch -import torch.nn.functional as F -from funasr.models.transformer.embedding import PositionalEncoding -import logging -from funasr.models.scama.utils import sequence_mask -from funasr.models.transformer.utils.nets_utils import sub_factor_to_params, pad_to_len -from typing import Optional, Tuple, Union import math +import torch +from typing import Optional, Tuple, Union +from funasr.models.transformer.utils.nets_utils import pad_to_len + class TooShortUttError(Exception): """Raised when the utt is too short for subsampling.