Update configuration_lightning.py

configuration_lightning.py

@@ -1,27 +1,31 @@
 from transformers import PretrainedConfig
 
 
-class LightningConfig(PretrainedConfig):
-    model_type = "lightning"
-
-    def __init__(
-        self,
-        vocab_size=50000,
-        d_model=256,
-        nhead=4,
-        num_layers=4,
-        dropout=0.1,
-        max_seq_len=160,
-        **kwargs
-    ):
-        super().__init__(
-            tie_word_embeddings=False,
-            **kwargs
-        )
-
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.nhead = nhead
-        self.num_layers = num_layers
-        self.dropout = dropout
+lass LightningConfig(PretrainedConfig):
+
+    model_type = "lightning"
+
+    def __init__(
+        self,
+        vocab_size=50000,
+        d_model=256,
+        nhead=4,
+        num_layers=4,
+        dropout=0.1,
+        max_seq_len=160,
+        **kwargs
+    ):
+        super().__init__(
+            tie_word_embeddings=False,
+            **kwargs
+        )
+
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.nhead = nhead
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.num_hidden_layers = num_layers
+        self.num_attention_heads = nhead
+        self.hidden_size = d_model
         self.max_seq_len = max_seq_len