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/'")