23 lines
659 B
Python
23 lines
659 B
Python
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
|
|
)
|