added SVM to tests

ombhojane · Sep 26, 2024 · 16c7906 · 16c7906
1 parent 709d835
commit 16c7906
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/explainableai/model_selection.py b/explainableai/model_selection.py
@@ -9,14 +9,18 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-def compare_models(X_train, y_train, X_test, y_test):
-    models = {
+def get_default_models():
+    return {
         'Logistic Regression': LogisticRegression(max_iter=1000),
         'Random Forest': RandomForestClassifier(n_estimators=100, random_state=42),
         'SVM': SVC(probability=True, random_state=42),
         'XGBoost': XGBClassifier(n_estimators=100, random_state=42),
         'Neural Network': MLPClassifier(hidden_layer_sizes=(100, 50), max_iter=1000, random_state=42)
     }
+
+def compare_models(X_train, y_train, X_test, y_test, models=None):
+    if models is None:
+        models = get_default_models()
 
     results = {}
     for name, model in models.items():

diff --git a/tests/__pycache__/test_xai_wrapper.cpython-311-pytest-8.0.2.pyc b/tests/__pycache__/test_xai_wrapper.cpython-311-pytest-8.0.2.pyc
diff --git a/tests/test_xai_wrapper.py b/tests/test_xai_wrapper.py
@@ -5,6 +5,7 @@
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.linear_model import LogisticRegression
 from xgboost import XGBClassifier
+from sklearn.svm import SVC
 from sklearn.neural_network import MLPClassifier
 from explainableai import XAIWrapper
 import os
@@ -23,7 +24,8 @@ def sample_models():
         'Random Forest': RandomForestClassifier(n_estimators=10, random_state=42),
         'Logistic Regression': LogisticRegression(max_iter=1000),
         'XGBoost': XGBClassifier(n_estimators=10, random_state=42),
-        'Neural Network': MLPClassifier(hidden_layer_sizes=(10,), max_iter=1000, random_state=42)
+        'Neural Network': MLPClassifier(hidden_layer_sizes=(10,), max_iter=1000, random_state=42),
+        'SVM': SVC(probability=True, random_state=42)
     }
 
 def test_xai_wrapper_initialization(sample_data, sample_models):
@@ -44,7 +46,7 @@ def test_xai_wrapper_fit(sample_data, sample_models):
     assert hasattr(xai.model, 'predict')
     assert hasattr(xai.model, 'predict_proba')
 
-@pytest.mark.parametrize("model_name", ['Random Forest', 'Logistic Regression', 'XGBoost', 'Neural Network'])
+@pytest.mark.parametrize("model_name", ['Random Forest', 'Logistic Regression', 'XGBoost', 'Neural Network', 'SVM'])
 def test_xai_wrapper_analyze_with_different_models(sample_data, sample_models, model_name):
     X, y = sample_data
     models = {model_name: sample_models[model_name]}