Update configuration_tinybuddy.py

configuration_tinybuddy.py

@@ -2,16 +2,29 @@
 Configuration class for TinyBuddy-30M.
 """
 
-from dataclasses import dataclass
+from transformers import PretrainedConfig
 
-
-@dataclass
-class GPTConfig:
-    vocab_size: int = 50000
-    block_size: int = 512      # max context length
-    n_layer: int = 6
-    n_head: int = 8
-    n_embd: int = 256
-    mlp_ratio: int = 4         # hidden = mlp_ratio * n_embd
-    dropout: float = 0.0
-    tie_weights: bool = False  # False -> ~30M params; True -> ~22M
+class GPTConfig(PretrainedConfig):
+    model_type = "tinybuddy"
+    
+    def __init__(
+        self,
+        vocab_size: int = 50000,
+        block_size: int = 512,
+        n_layer: int = 6,
+        n_head: int = 8,
+        n_embd: int = 256,
+        mlp_ratio: int = 4,
+        dropout: float = 0.0,
+        tie_weights: bool = False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+        self.tie_weights = tie_weights