26 lines
632 B
Python
26 lines
632 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")
|
|
|
|
# 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
|
|
) |