Update model.py

model.py

@@ -1,36 +1,57 @@
-# ✅ model.py
 import torch
 import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
 
-class EvoTransformerBlock(nn.Module):
-    def __init__(self, d_model, nhead, dim_feedforward):
-        super().__init__()
-        self.layer = nn.TransformerEncoderLayer(
-            d_model=d_model,
-            nhead=nhead,
-            dim_feedforward=dim_feedforward,
-            batch_first=True
-        )
+class EvoTransformerConfig(PretrainedConfig):
+    def __init__(self, hidden_size=384, num_layers=6, num_labels=2, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_labels = num_labels
 
-    def forward(self, x):
-        return self.layer(x)
+class EvoTransformerForClassification(PreTrainedModel):
+    config_class = EvoTransformerConfig
 
-class EvoTransformer(nn.Module):
-    def __init__(self, vocab_size, d_model=256, nhead=4, dim_feedforward=512, num_layers=4):
-        super().__init__()
-        self.embedding = nn.Embedding(vocab_size, d_model)
-        self.encoder = nn.Sequential(*[
-            EvoTransformerBlock(d_model, nhead, dim_feedforward) for _ in range(num_layers)
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.embedding = nn.Embedding(30522, config.hidden_size)
+        self.layers = nn.ModuleList([
+            nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=6, dim_feedforward=1024)
+            for _ in range(config.num_layers)
         ])
-        self.pooler = nn.AdaptiveAvgPool1d(1)
         self.classifier = nn.Sequential(
-            nn.Linear(d_model, d_model // 2),
+            nn.Linear(config.hidden_size, 256),
             nn.ReLU(),
-            nn.Linear(d_model // 2, 2)
+            nn.Linear(256, config.num_labels)
         )
+        self.init_weights()
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        x = self.embedding(input_ids)
+        x = x.transpose(0, 1)
+        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)
+        logits = self.classifier(x)
+
+        if labels is not None:
+            loss = nn.functional.cross_entropy(logits, labels)
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}
+
+    def save_pretrained(self, save_directory):
+        import os
+        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())
 
-    def forward(self, x):
-        x = self.embedding(x)
-        x = self.encoder(x)
-        x = self.pooler(x.transpose(1, 2)).squeeze(-1)
-        return self.classifier(x)
+    @classmethod
+    def from_pretrained(cls, 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.load_state_dict(torch.load(model_path, map_location="cpu"))
+        return model