Spaces:
Runtime error
Runtime error
| import pandas as pd | |
| import numpy as np | |
| import os | |
| import joblib | |
| from sklearn.ensemble import RandomForestClassifier | |
| from sklearn.preprocessing import StandardScaler, LabelEncoder | |
| from sklearn.model_selection import train_test_split | |
| BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| DATASET_PATH = os.path.join(BASE_DIR, "datasets", "dataset_2.csv") | |
| MODEL_DIR = os.path.join(BASE_DIR, "ml_models") | |
| os.makedirs(MODEL_DIR, exist_ok=True) | |
| MODEL_PATH = os.path.join(MODEL_DIR, "device_classifier.pkl") | |
| SCALER_PATH = os.path.join(BASE_DIR, "datasets", "scaler.pkl") | |
| LE_PATH = os.path.join(BASE_DIR, "datasets", "label_encoder.pkl") | |
| def train(): | |
| print("Loading dataset...") | |
| df = pd.read_csv(DATASET_PATH) | |
| print("Columns available:", list(df.columns)) | |
| print("Preprocessing data...") | |
| # Drop NaNs | |
| df.dropna(subset=['device_type', 'packet_size', 'protocol_type', 'bandwidth_usage', 'transmission_time'], inplace=True) | |
| # Encode protocol (TCP/UDP mapping) | |
| print("Encoding protocol...") | |
| df['protocol'] = df['protocol_type'].apply(lambda x: 0 if str(x).upper() == 'TCP' else (1 if str(x).upper() == 'UDP' else 2)) | |
| # Map features | |
| df['flow_duration'] = df['transmission_time'] | |
| df['packet_rate'] = df['bandwidth_usage'] # Simple mapping | |
| features = ['packet_size', 'protocol', 'packet_rate', 'flow_duration'] | |
| X = df[features] | |
| y = df['device_type'] | |
| print("Encoding Labels...") | |
| le = LabelEncoder() | |
| y_encoded = le.fit_transform(y) | |
| print("Scaling Features...") | |
| scaler = StandardScaler() | |
| X_scaled = scaler.fit_transform(X) | |
| print("Splitting...") | |
| X_train, X_test, y_train, y_test = train_test_split(X_scaled, y_encoded, test_size=0.2, random_state=42) | |
| print("Training Model...") | |
| clf = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1) | |
| clf.fit(X_train, y_train) | |
| print(f"Accuracy: {clf.score(X_test, y_test):.4f}") | |
| print("Saving artifacts...") | |
| pd.DataFrame(X_test, columns=features).to_csv(os.path.join(BASE_DIR, "datasets", "X_test.csv"), index=False) | |
| pd.DataFrame(y_test, columns=['device_type']).to_csv(os.path.join(BASE_DIR, "datasets", "y_test.csv"), index=False) | |
| joblib.dump(clf, MODEL_PATH) | |
| joblib.dump(scaler, SCALER_PATH) | |
| joblib.dump(le, LE_PATH) | |
| print(f"Saved model to {MODEL_PATH}") | |
| print(f"Saved scaler to {SCALER_PATH}") | |
| print(f"Saved label encoder to {LE_PATH}") | |
| print("Done!") | |
| if __name__ == "__main__": | |
| train() | |