Update evo_model.py

evo_model.py

@@ -3,11 +3,23 @@ import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
 
 class EvoTransformerConfig(PretrainedConfig):
-    def __init__(self, hidden_size=384, num_layers=6, num_labels=2, **kwargs):
+    def __init__(
+        self,
+        hidden_size=384,
+        num_layers=6,
+        num_labels=2,
+        num_heads=6,
+        ffn_dim=1024,
+        use_memory=False,
+        **kwargs
+    ):
         super().__init__(**kwargs)
         self.hidden_size = hidden_size
         self.num_layers = num_layers
         self.num_labels = num_labels
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+        self.use_memory = use_memory
 
 class EvoTransformerForClassification(PreTrainedModel):
     config_class = EvoTransformerConfig
@@ -15,9 +27,20 @@ class EvoTransformerForClassification(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.config = config
+
+        # Expose architecture attributes for dashboard
+        self.num_layers = config.num_layers
+        self.num_heads = config.num_heads
+        self.ffn_dim = config.ffn_dim
+        self.use_memory = config.use_memory
+
         self.embedding = nn.Embedding(30522, config.hidden_size)  # BERT vocab size
         self.layers = nn.ModuleList([
-            nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=6, dim_feedforward=1024)
+            nn.TransformerEncoderLayer(
+                d_model=config.hidden_size,
+                nhead=config.num_heads,
+                dim_feedforward=config.ffn_dim
+            )
             for _ in range(config.num_layers)
         ])
         self.classifier = nn.Sequential(
@@ -25,6 +48,7 @@ class EvoTransformerForClassification(PreTrainedModel):
             nn.ReLU(),
             nn.Linear(256, config.num_labels)
         )
+
         self.init_weights()
 
     def forward(self, input_ids, attention_mask=None, labels=None):