Update init_save.py

init_save.py

@@ -1,19 +1,9 @@
+from evo_model import EvoTransformerConfig, EvoTransformerForClassification
 import os
-import torch
-from evo_model import EvoTransformerForClassification, EvoTransformerConfig
 
-# ✅ Create directory if it doesn't exist
-os.makedirs("trained_model", exist_ok=True)
-
-# ✅ Create config and model
-config = EvoTransformerConfig()
-model = EvoTransformerForClassification(config)
-
-# ✅ Save like Hugging Face expects
-torch.save(model.state_dict(), "trained_model/pytorch_model.bin")
-
-# ✅ Save minimal config
-with open("trained_model/config.json", "w") as f:
-    f.write(config.to_json_string())
-
-print("✅ EvoTransformer initial model saved to 'trained_model/'")
+def initialize_and_save_model(save_path="trained_model"):
+    config = EvoTransformerConfig()
+    model = EvoTransformerForClassification(config)
+    
+    os.makedirs(save_path, exist_ok=True)
+    model.save_pretrained(save_path)