more detailed training process to help development of run.py

Author: ollycassidy13

nids/__init__.py

@@ -1,7 +1,6 @@
 # nids/__init__.py
 
+# Import necessary modules and functions
 from .data_preprocessing import load_and_preprocess_data
-from .model import Net, train_model
+from .model import train_model
 from .logging import setup_logging, log_prediction
-from .prediction import run_prediction
-from .retraining import retrain

nids/data_preprocessing.py

@@ -5,6 +5,7 @@
 from sklearn.preprocessing import StandardScaler
 from sklearn.model_selection import train_test_split
 import numpy as np
+import pickle
 
 def load_and_preprocess_data(csv_files_path):
     # Load all CSV files
@@ -35,10 +36,13 @@ def load_and_preprocess_data(csv_files_path):
         raise ValueError("The target label column is not found in the dataset.")
 
     # Encode categorical variables
-    data[label_column] = data[label_column].astype('category').cat.codes
+    data[label_column] = data[label_column].astype('category')
+    class_mapping = dict(enumerate(data[label_column].cat.categories))
+    data[label_column] = data[label_column].cat.codes
 
     # Print unique values of the target labels
     print(f"Unique target labels: {data[label_column].unique()}")
+    print(f"Class mapping: {class_mapping}")
 
     # Replace infinite values with NaN
     data.replace([np.inf, -np.inf], np.nan, inplace=True)
@@ -61,4 +65,16 @@ def load_and_preprocess_data(csv_files_path):
     # Split data into training and testing sets
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
-    return X_train, X_test, y_train, y_test, scaler
+    # Save the class mapping, number of features, and feature names
+    metadata = {
+        'num_features': X_train.shape[1],
+        'num_classes': len(class_mapping),
+        'class_mapping': class_mapping,
+        'feature_names': list(X.columns)
+    }
+    with open('nids/model_metadata.pkl', 'wb') as f:
+        pickle.dump(metadata, f)
+
+    print("Metadata (number of features, classes, class mapping and faeture names) saved. ")
+
+    return X_train, X_test, y_train, y_test, scaler

nids/model.pth

@@ -1,4 +1,4 @@
 pandas
 scikit-learn
 torch
-kafka-python
+scapy

nids/model_metadata.pkl

@@ -19,13 +19,5 @@
     torch.save(model.state_dict(), 'nids/model.pth')
     with open('nids/scaler.pkl', 'wb') as f:
         pickle.dump(scaler, f)
-    
-    # Save number of features and classes
-    with open('nids/model_metadata.pkl', 'wb') as f:
-        metadata = {
-            'num_features': X_train.shape[1],
-            'num_classes': num_classes
-        }
-        pickle.dump(metadata, f)
 
-    print("Model, scaler, and metadata (number of features and classes) saved.")
+    print("Model and scaler saved.")