example.py

@@ -8,15 +8,19 @@ 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)
 
-# Set checkpoint directory
-checkpoint_dir = "checkpoints/my_model"
+# List of dataset paths
+dataset_paths = [
+    "/path/to/dataset1.parquet",
+    "/path/to/dataset2.parquet"
+]
 
-# Start training (will automatically load checkpoint if available)
+# Start training with multiple datasets
 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
-)
+    dataset_paths=dataset_paths,
+    num_epochs=10,
+    batch_size=2,
+    sample_size=10000
+)