24 lines
569 B
Python
24 lines
569 B
Python
from aiia.model import AIIABase
|
|
from aiia.model import AIIAConfig
|
|
from aiia.pretrain import Pretrainer
|
|
|
|
# Create your model
|
|
config = AIIAConfig(model_name="AIIA-Base-512x20k")
|
|
model = AIIABase(config)
|
|
|
|
# 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
|
|
) |