update rwkv_bat

shixian.shi

2024-02-21 9a6c6ab5ea25bf2d6e874010ae101e29bde2a217

update rwkv_bat

 funasr/models/rwkv_bat/model.py

 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 @@
        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 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.