class_docstr: RWKV decoder module.
  
  Based on https://arxiv.org/pdf/2305.13048.pdf.
  
  Args:
      vocab_size: Vocabulary size.
      block_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.
      rescale_every: Whether to rescale input every N blocks (inference only).
      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.
      embed_pad: Embedding padding symbol ID.
class_docstr_tok: ["RWKV","decoder","module",".","Based","on","https",":","\/\/arxiv.org\/pdf\/2305.13048.pdf",".","Args",":","vocab_size",":","Vocabulary","size",".","block_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",".","rescale_every",":","Whether","to","rescale","input","every","N","blocks","(","inference","only",")",".","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",".","embed_pad",":","Embedding","padding","symbol","ID","."]
docstr: Construct a RWKVDecoder object.
docstr_tok: ["Construct","a","RWKVDecoder","object","."]
code: def __init__(
      self,
      vocab_size: int,
      block_size: int = 512,
      context_size: int = 1024,
      linear_size: Optional[int] = None,
      attention_size: Optional[int] = None,
      normalization_type: str = "layer_norm",
      normalization_args: Dict = {},
      num_blocks: int = 4,
      rescale_every: int = 0,
      embed_dropout_rate: float = 0.0,
      att_dropout_rate: float = 0.0,
      ffn_dropout_rate: float = 0.0,
      embed_pad: int = 0,
  ) -> None:
      """Construct a RWKVDecoder object."""
      super().__init__()
  
      assert check_argument_types()
  
      norm_class, norm_args = get_normalization(
          normalization_type, **normalization_args
      )
  
      linear_size = (
          block_size * 4 if linear_size is None else linear_size
      )
      attention_size = (
          block_size if attention_size is None else attention_size
      )
  
      self.embed = torch.nn.Embedding(
          vocab_size, block_size, padding_idx=embed_pad
      )
      self.dropout_embed = torch.nn.Dropout(p=embed_dropout_rate)
  
      self.rwkv_blocks = torch.nn.ModuleList(
          [
              RWKV(
                  block_size,
                  linear_size,
                  attention_size,
                  context_size,
                  block_id,
                  num_blocks,
                  normalization_class=norm_class,
                  normalization_args=norm_args,
                  att_dropout_rate=att_dropout_rate,
                  ffn_dropout_rate=ffn_dropout_rate,
              )
              for block_id in range(num_blocks)
          ]
      )
  
      self.embed_norm = norm_class(block_size, **norm_args)
      self.final_norm = norm_class(block_size, **norm_args)
  
      self.block_size = block_size
      self.attention_size = attention_size
      self.output_size = block_size
      self.vocab_size = vocab_size
      self.context_size = context_size
  
      self.rescale_every = rescale_every
      self.rescaled_layers = False
  
      self.pad_idx = embed_pad
      self.num_blocks = num_blocks
  
      self.score_cache = {}
  
      self.device = next(self.parameters()).device
  
file_path: espnet/espnet2/asr_transducer/decoder/rwkv_decoder.py
imports: import math
  from typing import Dict, List, Optional, Tuple
  import torch
  from typeguard import check_argument_types
  from espnet2.asr_transducer.beam_search_transducer import Hypothesis
  from espnet2.asr_transducer.decoder.abs_decoder import AbsDecoder
  from espnet2.asr_transducer.decoder.blocks.rwkv import RWKV
  from espnet2.asr_transducer.normalization import get_normalization
label_desc: Use image node_id 1 to create a new RWKVDecoder object from inherited base classes: AbsDecoder with example: obj = RWKVDecoder(vocab_size, block_size, context_size, linear_size, attention_size, normalization_type, normalization_args, num_blocks, rescale_every, embed_dropout_rate, att_dropout_rate, ffn_dropout_rate, embed_pad)