Update evo_model.py

evo_model.py

@@ -1,54 +1,57 @@
 import torch
-from torch import nn
+import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
-from model import EvoTransformer  # assumes your core model is in model.py
 
 class EvoTransformerConfig(PretrainedConfig):
-    model_type = "evo-transformer"
-
-    def __init__(
-        self,
-        vocab_size=30522,
-        d_model=256,
-        nhead=4,
-        dim_feedforward=512,
-        num_hidden_layers=4,
-        **kwargs
-    ):
+    def __init__(self, hidden_size=384, num_layers=6, num_labels=2, **kwargs):
         super().__init__(**kwargs)
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.nhead = nhead
-        self.dim_feedforward = dim_feedforward
-        self.num_hidden_layers = num_hidden_layers
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_labels = num_labels
 
 class EvoTransformerForClassification(PreTrainedModel):
     config_class = EvoTransformerConfig
 
     def __init__(self, config):
         super().__init__(config)
-        self.model = EvoTransformer(
-            vocab_size=config.vocab_size,
-            d_model=config.d_model,
-            nhead=config.nhead,
-            dim_feedforward=config.dim_feedforward,
-            num_layers=config.num_hidden_layers
+        self.config = config
+        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)
+            for _ in range(config.num_layers)
+        ])
+        self.classifier = nn.Sequential(
+            nn.Linear(config.hidden_size, 256),
+            nn.ReLU(),
+            nn.Linear(256, config.num_labels)
         )
-        self.classifier = nn.Linear(config.d_model, 2)  # 2-way classification
+        self.init_weights()
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        x = self.embedding(input_ids)  # [batch, seq_len, hidden_size]
+        x = x.transpose(0, 1)  # Transformer expects [seq_len, batch, hidden_size]
+        for layer in self.layers:
+            x = layer(x, src_key_padding_mask=(attention_mask == 0) if attention_mask is not None else None)
+        x = x.mean(dim=0)  # mean pooling over seq_len
+        logits = self.classifier(x)
 
-    def forward(self, input_ids):
-        hidden = self.model(input_ids)  # (batch_size, seq_len, d_model)
-        pooled = hidden[:, 0, :]        # Use the first token as a summary
-        logits = self.classifier(pooled)  # (batch_size, 2)
+        if labels is not None:
+            loss = nn.functional.cross_entropy(logits, labels)
+            return loss, logits
         return logits
 
     def save_pretrained(self, save_directory):
-        torch.save(self.model.state_dict(), f"{save_directory}/pytorch_model.bin")
-        self.config.save_pretrained(save_directory)
+        import os, json
+        os.makedirs(save_directory, exist_ok=True)
+        torch.save(self.state_dict(), f"{save_directory}/pytorch_model.bin")
+        with open(f"{save_directory}/config.json", "w") as f:
+            f.write(self.config.to_json_string())
 
     @classmethod
     def from_pretrained(cls, load_directory):
-        config = EvoTransformerConfig.from_pretrained(load_directory)
+        config_path = f"{load_directory}/config.json"
+        model_path = f"{load_directory}/pytorch_model.bin"
+        config = EvoTransformerConfig.from_json_file(config_path)
         model = cls(config)
-        model.model.load_state_dict(torch.load(f"{load_directory}/pytorch_model.bin"))
+        model.load_state_dict(torch.load(model_path, map_location="cpu"))
         return model