From 7d24de1f7e5cacb1ed7bc0fd049053103f24ca56 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 16 Apr 2025 22:59:13 +0200 Subject: [PATCH] updated example usage to feature the checkpoint handling --- example.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/example.py b/example.py index 8dbb67e..7be0e5a 100644 --- a/example.py +++ b/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 -) \ No newline at end of file + dataset_paths=["path/to/dataset1.parquet", "path/to/dataset2.parquet"], + output_path="trained_models/my_model", + checkpoint_dir=checkpoint_dir, + num_epochs=10 +)