Update rita_configuration.py

rita_configuration.py

@@ -8,8 +8,8 @@ class RITAConfig(PretrainedConfig):
 
     def __init__(
         self,
-        in_vocab_size=128,
-        out_vocab_size=32,
+        vocab_size=26,
+        vocab_size=26,
         d_model=768,
         num_layers=12,
         max_seq_len=1024,
@@ -20,8 +20,7 @@ class RITAConfig(PretrainedConfig):
         **kwargs,
     ):
         super().__init__(eos_token_id=eos_token_id, **kwargs)
-        self.in_vocab_size = in_vocab_size
-        self.out_vocab_size = out_vocab_size
+        self.vocab_size = vocab_size
         self.d_model = d_model
         self.num_heads = num_heads
         self.d_feedforward = d_model*ff_ratio