File size: 382 Bytes
984ce24
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from transformers import PretrainedConfig


class RabbitConfig(PretrainedConfig):
    model_type = "rabbit_ssm"

    def __init__(
        self,
        vocab_size: int = 50280,
        d_model: int = 2560,
        n_layers: int = 64,
        **kwargs,
    ):
        self.d_model = d_model
        self.n_layers = n_layers
        super().__init__(vocab_size=vocab_size, **kwargs)