Update init_model.py

init_model.py

@@ -1,16 +1,14 @@
-# init_model.py
-
 import os
 import torch
-from evo_model import EvoTransformerForClassification, EvoTransformerConfig
+from evo_model import EvoTransformerConfig, EvoTransformerForClassification
 
-# ✅ Make sure the folder exists
+# Ensure folder
 os.makedirs("trained_model", exist_ok=True)
 
-# ✅ Create config and model
+# Create config and model
 config = EvoTransformerConfig()
 model = EvoTransformerForClassification(config)
 
-# ✅ Save the model and config in HF format
-model.save_pretrained("trained_model")  # includes config.json + pytorch_model.bin
-print("✅ EvoTransformer initial model saved to 'trained_model/'")
+# Save config and model in Hugging Face format
+model.save_pretrained("trained_model")
+print("✅ EvoTransformer initial model + config saved to 'trained_model/'")