from src.aiia.model import AIIAmoe
from src.aiia.model import AIIAConfig
from src.aiia.pretrain import Pretrainer

# Create your model
config = AIIAConfig(num_experts=5)
model = AIIAmoe(config)
model.save_pretrained("test")
model = AIIAmoe.from_pretrained("test")

pretrainer = Pretrainer(model, learning_rate=1e-4, config=config)

# Set checkpoint directory
checkpoint_dir = "checkpoints/my_model"

# Start training (will automatically load checkpoint if available)
pretrainer.train(
    dataset_paths=["path/to/dataset1.parquet", "path/to/dataset2.parquet"],
    output_path="trained_models/my_model",
    checkpoint_dir=checkpoint_dir,
    num_epochs=10
)