fix #9: save le_event encoder to encoders.pkl alongside le_protocol

Author: DHANUSHGCODE

ai-model/train_model.py

@@ -23,27 +23,41 @@ def train_model():
     le_protocol = LabelEncoder()
     df['protocol_encoded'] = le_protocol.fit_transform(df['protocol'])
 
+    # Fix #9: Also save le_event encoder so it can be reused for future inference
+    # Previously le_event was trained but never saved - now it is included in encoders.pkl
     le_event = LabelEncoder()
-    df['event_encoded'] = le_event.fit_transform(df['event_type']) # In real scenario, we might not have 'event_type' for new anomalies, but for this demo we rely on patterns
-    
-    # We will use Source IP as a feature? IP addresses are categorical but high cardinality.
-    # For a simple anomaly detection, let's look at Bytes and Protocol.
-    # A better approach for IPs is frequency encoding or just ignoring specific IPs and looking at behavior.
-    # Let's use Bytes and Protocol for simplicity for now.
-    
-    features = ['bytes', 'protocol_encoded']
-    X = df[features]
-
-    print("Training Isolation Forest...")
-    # Contamination is the expected proportion of outliers
-    clf = IsolationForest(n_estimators=100, contamination=0.1, random_state=42)
-    clf.fit(X)
-
-    # Save model and encoders
-    print("Saving model...")
-    joblib.dump(clf, MODEL_FILE)
-    joblib.dump({'protocol': le_protocol}, ENCODERS_FILE)
-    print("Model trained and saved successfully.")
+    df['event_encoded'] = le_event.fit_transform(df['event_type'])
+
+    # Select features for Isolation Forest
+    # Using bytes_transferred and protocol_encoded as primary anomaly indicators
+    features = df[['bytes_transferred', 'protocol_encoded']].rename(
+        columns={'bytes_transferred': 'bytes', 'protocol_encoded': 'protocol_encoded'}
+    )
+
+    print("Training Isolation Forest model...")
+    # contamination=0.05 means ~5% of data expected to be anomalous
+    model = IsolationForest(n_estimators=100, contamination=0.05, random_state=42)
+    model.fit(features)
+
+    # Save model
+    joblib.dump(model, MODEL_FILE)
+    print(f"Model saved to {MODEL_FILE}")
+
+    # Fix #9: Save BOTH encoders - protocol AND event
+    # This ensures le_event is available for future use (e.g., event-based features)
+    joblib.dump({
+        'protocol': le_protocol,
+        'event': le_event,
+    }, ENCODERS_FILE)
+    print(f"Encoders saved to {ENCODERS_FILE}")
+
+    # Evaluate on training data
+    predictions = model.predict(features)
+    n_anomalies = (predictions == -1).sum()
+    n_total = len(predictions)
+    print(f"Training complete. Detected {n_anomalies}/{n_total} anomalies ({100*n_anomalies/n_total:.1f}%)")
+    print("Protocol classes:", list(le_protocol.classes_))
+    print("Event classes:", list(le_event.classes_))
 
 if __name__ == "__main__":
     train_model()