Refactor config and trainer scripts to expose get_existing_rounds method and streamline model round initialization

Author: DefinetlyNotAI

v4n1_Trainer.py

@@ -133,12 +133,13 @@ def train(config: TrainingConfig, resources: dict):
             name = "Sense"
         else:
             name = "SenseMacro"
-
+        existing_rounds = cfg.get_existing_rounds(cfg.CACHE_DIR)  # Auto-increment round based on existing folders
+        model_round = max(existing_rounds) + 1 if existing_rounds else 1
         cfg.update({
             # Model / caching / logging
             "MODEL_NAME": f"Model_{name}.4n1",  # Name of the model for identification and caching
-            "DATASET_SIZE": dataset,
-            # Number of samples to generate for training (not the same as for the training rounds themselves)
+            "DATASET_SIZE": dataset,            # Number of samples to generate for training (not the same as for the training rounds themselves)
+            "MODEL_ROUND": model_round          # Current training round (auto-incremented)
         })
-        log(message=f"Training {name} with {dataset} dataset...", cfg=cfg)
+        log(message=f"Training 'Model_{name}.4n1/round_{model_round}/' with {dataset} dataset...", cfg=cfg)
         train(config=cfg, resources=train_init)

vulnscan/config.py

@@ -9,7 +9,7 @@ def __init__(self, model_name: str = "Model_Sense.4n1"):
         self.MODEL_NAME = model_name
         self.CACHE_DIR = os.path.join(os.getcwd(), "cache")
 
-        existing_rounds = self.__get_existing_rounds(self.CACHE_DIR)  # Auto-increment round based on existing folders
+        existing_rounds = self.get_existing_rounds(self.CACHE_DIR)  # Auto-increment round based on existing folders
         self.MODEL_ROUND = max(existing_rounds) + 1 if existing_rounds else 1
 
         self.LOG_FILE = f"{self.CACHE_DIR}/{self.MODEL_NAME}/training.log"
@@ -59,7 +59,7 @@ def __init__(self, model_name: str = "Model_Sense.4n1"):
         os.makedirs(self.DATA_CACHE_DIR, exist_ok=True)
 
     @staticmethod
-    def __get_existing_rounds(cache_dir: str) -> list[int]:
+    def get_existing_rounds(cache_dir: str) -> list[int]:
         """
         Returns a list of round numbers based on existing folders in the cache directory.
         """