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

# Initialize pretrainer with the model
pretrainer = Pretrainer(model, learning_rate=1e-4, config=config)

# List of dataset paths
dataset_paths = [
    "/path/to/dataset1.parquet",
    "/path/to/dataset2.parquet"
]

# Start training with multiple datasets
pretrainer.train(
    dataset_paths=dataset_paths,
    num_epochs=10,
    batch_size=2,
    sample_size=10000
)