Refactor TrainingConfig to improve initialization and cache directory handling

Author: DefinetlyNotAI

v4n1_Generator.py

@@ -43,7 +43,7 @@
 dataset_ranges = [10, 100, 1000, 5000, 10000, 17500, 25000]
 
 for dr in dataset_ranges:
-    dataset_path = f"{cfg.DATA_CACHE_DIR}/dataset_{dr}.pt"
+    dataset_path = f"{cfg.DATASET_CACHE_DIR}/dataset_{dr}.pt"
 
     # Skip if already exists
     if os.path.exists(dataset_path):
@@ -57,7 +57,7 @@
     smaller_existing.sort(reverse=True)
 
     for sr in smaller_existing:
-        candidate_path = f"{cfg.DATA_CACHE_DIR}/dataset_{sr}.pt"
+        candidate_path = f"{cfg.DATASET_CACHE_DIR}/dataset_{sr}.pt"
         if os.path.exists(candidate_path):
             data = torch.load(candidate_path, map_location="cpu")
             base_texts, base_labels = data["texts"], data["labels"]

v4n1_Trainer.py

@@ -38,7 +38,7 @@ def train(config: TrainingConfig, resources: dict):
     generate = DataGen(cfg=config)
 
     # Generate dataset
-    dataset_path = f"{config.DATA_CACHE_DIR}/dataset_{config.DATASET_SIZE}.pt"
+    dataset_path = f"{config.DATASET_CACHE_DIR}/dataset_{config.DATASET_SIZE}.pt"
     if os.path.exists(dataset_path):
         data = torch.load(dataset_path)
         texts, labels = data["texts"], data["labels"]
@@ -86,9 +86,6 @@ def train(config: TrainingConfig, resources: dict):
     # ---------------- CONFIG ----------------
     cfg = TrainingConfig()
     cfg.update({
-        # Model / caching / logging
-        "MODEL_NAME": "Model_Sense.4n1",  # Name of the model for identification and caching
-
         # Training parameters
         "BATCH_SIZE": 32,  # Number of samples per training batch
         "MAX_EPOCHS": 35,  # Maximum number of training epochs
@@ -101,8 +98,6 @@ def train(config: TrainingConfig, resources: dict):
         "LR_DECAY": 0.9,  # Factor to multiply learning rate after decay
         "AUTO_CONTINUE": False,  # Whether to automatically continue training and ignore EARLY_STOPPING_PATIENCE
 
-        # Dataset / data generation
-        "DATASET_SIZE": 25000,
         # Number of samples to generate for training (not the same as for the training rounds themselves)
         "TEXT_MAX_LEN": 128,  # Maximum length of generated text samples
         "TEXT_MAX_LEN_JUMP_RANGE": 10,  # Range for random variation in text length

vulnscan/config.py

@@ -4,21 +4,25 @@
 
 
 class TrainingConfig:
-    def __init__(self, model_name: str = "Model_Sense.4n1"):
+    def __init__(self):
+        """
+        Configuration class for training settings and hyperparameters.
+
+        You must call the update method and set MODEL_NAME.
+        """
+
         # Model / caching / logging
-        self.MODEL_NAME = model_name
+        self.MODEL_NAME = None
+        self.writer = None
+        self.LOG_FILE = None
+        self.EMBED_CACHE_DIR = None
+
         self.CACHE_DIR = os.path.join(os.getcwd(), "cache")
+        self.DATASET_CACHE_DIR = f"{self.CACHE_DIR}/dataset"
 
         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"
-        self.EMBED_CACHE_DIR = f"{self.CACHE_DIR}/{self.MODEL_NAME}/round_{self.MODEL_ROUND}/embeddings"
-        self.DATA_CACHE_DIR = f"{self.CACHE_DIR}/dataset"
-
-        # TensorBoard
-        self.writer = SummaryWriter(log_dir=f"{self.CACHE_DIR}/{self.MODEL_NAME}/round_{self.MODEL_ROUND}/tensorboard_logs")
-
         # Training parameters
         self.BATCH_SIZE: int = 16
         self.MAX_EPOCHS: int = 35
@@ -53,11 +57,6 @@ def __init__(self, model_name: str = "Model_Sense.4n1"):
         self.DEVICE: str = "cuda" if torch.cuda.is_available() else "cpu"
         self.RAM_THRESHOLD: float = 0.85
 
-        # Create necessary folders
-        os.makedirs(self.CACHE_DIR, exist_ok=True)
-        os.makedirs(self.EMBED_CACHE_DIR, exist_ok=True)
-        os.makedirs(self.DATA_CACHE_DIR, exist_ok=True)
-
     @staticmethod
     def __get_existing_rounds(cache_dir: str) -> list[int]:
         """
@@ -95,8 +94,9 @@ def update(self, updates):
         if 'MODEL_NAME' in dict(items) or 'CACHE_DIR' in dict(items) or 'MODEL_ROUND' in dict(items):
             self.LOG_FILE = f"{self.CACHE_DIR}/{self.MODEL_NAME}/training.log"
             self.EMBED_CACHE_DIR = f"{self.CACHE_DIR}/{self.MODEL_NAME}/round_{self.MODEL_ROUND}/embeddings"
-            self.writer = SummaryWriter(log_dir=f"{self.CACHE_DIR}/{self.MODEL_NAME}/round_{self.MODEL_ROUND}/tensorboard_logs")
+            self.writer = SummaryWriter(
+                log_dir=f"{self.CACHE_DIR}/{self.MODEL_NAME}/round_{self.MODEL_ROUND}/tensorboard_logs")
             os.makedirs(self.EMBED_CACHE_DIR, exist_ok=True)
-        if 'CACHE_DIR' in dict(items):
-            self.DATA_CACHE_DIR = f"{self.CACHE_DIR}/dataset"
-            os.makedirs(self.DATA_CACHE_DIR, exist_ok=True)
+        if 'DATASET_CACHE_DIR' in dict(items):
+            self.DATASET_CACHE_DIR = f"{self.CACHE_DIR}/dataset"
+            os.makedirs(self.DATASET_CACHE_DIR, exist_ok=True)

vulnscan/genData.py

@@ -84,7 +84,7 @@ def dataset(self, gpt_tokenizer: PreTrainedTokenizerFast, gpt_model: PreTrainedM
                 labels.append(int(sensitive))
             except KeyboardInterrupt:
                 sys.exit(f"\nDataset generation interrupted by user early. Premature dataset exit.")
-        torch.save({"texts": dataset, "labels": labels}, f"{self.cfg.DATA_CACHE_DIR}/dataset_{self.cfg.DATASET_SIZE}.pt")
+        torch.save({"texts": dataset, "labels": labels}, f"{self.cfg.DATASET_CACHE_DIR}/dataset_{self.cfg.DATASET_SIZE}.pt")
         return dataset, labels
 
     # ---------------- EMBEDDINGS ----------------