Update evo_model.py

evo_model.py

@@ -0,0 +1,52 @@
+# ✅ evo_model.py – HF-compatible wrapper for EvoTransformer
+
+import torch
+from torch import 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
+    ):
+        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
+
+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
+        )
+
+    def forward(self, input_ids):
+        return self.model(input_ids)
+
+    def save_pretrained(self, save_directory):
+        torch.save(self.model.state_dict(), f"{save_directory}/pytorch_model.bin")
+        self.config.save_pretrained(save_directory)
+
+    @classmethod
+    def from_pretrained(cls, load_directory):
+        config = EvoTransformerConfig.from_pretrained(load_directory)
+        model = cls(config)
+        model.model.load_state_dict(torch.load(f"{load_directory}/pytorch_model.bin"))
+        return model