updated version to support proper saving when pretraining

example.py

@@ -1,21 +1,18 @@
-data_path1 = "/root/training_data/vision-dataset/images_checkpoint.parquet"
-data_path2 = "/root/training_data/vision-dataset/vec_images_dataset.parquet"
-    
 from aiia.model import AIIABase
-from aiia.model.config import AIIAConfig
+from aiia.model import AIIAConfig
 from aiia.pretrain import Pretrainer
 
 # Create your model
-config = AIIAConfig(model_name="AIIA-Base-512x10k-small", num_hidden_layers=6, hidden_size=256)
+config = AIIAConfig(model_name="AIIA-Base-512x20k")
 model = AIIABase(config)
 
 # Initialize pretrainer with the model
-pretrainer = Pretrainer(model, learning_rate=config.learning_rate, config=config)
+pretrainer = Pretrainer(model, learning_rate=1e-4)
 
 # List of dataset paths
 dataset_paths = [
-    data_path1,
-    data_path2
+    "/path/to/dataset1.parquet",
+    "/path/to/dataset2.parquet"
 ]
 
 # Start training with multiple datasets

pyproject.toml

@@ -7,14 +7,16 @@ line-length = 88
 target-version = ['py37']
 include = '\.pyi?$'
 
+
 [project]
 name = "aiia"
-version = "0.1.1"
+version = "0.1.2"
 description = "AIIA Deep Learning Model Implementation"
 readme = "README.md"
 authors = [
-    { name="Falko Habel", email="falko.habel@gmx.de" }
+    {name = "Falko Habel", email = "falko.habel@gmx.de"}
 ]
+license = {text = "CC BY-NC 4.0"}
 dependencies = [
     "torch>=2.5.0",
     "numpy",
@@ -25,6 +27,6 @@ dependencies = [
 requires-python = ">=3.7"
 classifiers = [
     "Programming Language :: Python :: 3",
-    "License :: OSI Approved :: MIT License",
+    "License :: OSI Approved :: CC BY-NC 4.0 License",
     "Operating System :: OS Independent"
 ]

src/aiia/__init__.py

@@ -4,4 +4,4 @@ from .data.DataLoader import DataLoader
 from .pretrain.pretrainer import Pretrainer, ProjectionHead
 
 
-__version__ = "0.1.0"
+__version__ = "0.1.2"

src/aiia/pretrain/pretrainer.py

@@ -112,7 +112,7 @@ class Pretrainer:
             
         return batch_loss
 
-    def train(self, dataset_paths, column="image_bytes",  num_epochs=3, batch_size=2, sample_size=10000):
+    def train(self, dataset_paths,output_path:str="AIIA", column="image_bytes",  num_epochs=3, batch_size=2, sample_size=10000):
         """
         Train the model using multiple specified datasets.
         
@@ -186,7 +186,7 @@ class Pretrainer:
             
             if val_loss < best_val_loss:
                 best_val_loss = val_loss
-                self.model.save("AIIA-base-512")
+                self.model.save(output_path)
                 print("Best model saved!")
 
         self.save_losses('losses.csv')