example.py

@@ -6,19 +6,15 @@ from aiia.pretrain import Pretrainer
 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"
-]
+# Set checkpoint directory
+checkpoint_dir = "checkpoints/my_model"
 
-# Start training with multiple datasets
+# Start training (will automatically load checkpoint if available)
 pretrainer.train(
-    dataset_paths=dataset_paths,
-    num_epochs=10,
-    batch_size=2,
-    sample_size=10000
+    dataset_paths=["path/to/dataset1.parquet", "path/to/dataset2.parquet"],
+    output_path="trained_models/my_model",
+    checkpoint_dir=checkpoint_dir,
+    num_epochs=10
 )