update rwkv_bat

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.

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

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__()

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):

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 math
 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
+from funasr.models.transformer.utils.nets_utils import pad_to_len
+
 
 class TooShortUttError(Exception):
     """Raised when the utt is too short for subsampling.