Update init_model.py

init_model.py

@@ -3,8 +3,15 @@ from transformers import BertTokenizer
 from evo_model import EvoTransformerConfig, EvoTransformerForClassification
 
 def initialize_and_save_model():
-    # Step 1: Initialize configuration
-    config = EvoTransformerConfig()
+    # Step 1: Initialize configuration with architecture info
+    config = EvoTransformerConfig(
+        hidden_size=384,
+        num_layers=6,
+        num_labels=2,
+        num_heads=6,
+        ffn_dim=1024,
+        use_memory=False
+    )
 
     # Step 2: Initialize model
     model = EvoTransformerForClassification(config)
@@ -19,10 +26,5 @@ def initialize_and_save_model():
 
     print("✅ EvoTransformer and tokenizer initialized and saved to 'trained_model/'")
 
-def load_model():
-    model = EvoTransformerForClassification.from_pretrained("trained_model")
-    return model
-
-# Optional: reinitialize if run directly
 if __name__ == "__main__":
     initialize_and_save_model()