EvoTransformer-v2.1 / init_model.py
HemanM's picture
Create init_model.py
8e22eea verified
raw
history blame
503 Bytes
# init_model.py
import os
import torch
from evo_model import EvoTransformerForClassification, EvoTransformerConfig
# βœ… Make sure the folder exists
os.makedirs("trained_model", exist_ok=True)
# βœ… 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/'")