chore: Split dependencies

README.md

@@ -30,6 +30,7 @@ Holistic AI currently focuses on five verticals of AI trustworthiness:
 
 ```bash
 pip install holisticai  # Basic installation
+pip install holisticai[datasets]  # add datasets and plot dependencies
 pip install holisticai[bias]  # Bias mitigation support
 pip install holisticai[explainability]  # For explainability metrics and plots
 pip install holisticai[all]  # Install all packages for bias and explainability

pyproject.toml

@@ -25,19 +25,15 @@ classifiers = [
 dependencies = [
   "scikit-learn>=1.0.2",
   "pandas",
-  "matplotlib",
-  "networkx>=3.1",
-  "seaborn",
   "numpy>=1.25",
-  "pybind11>=2.12",
-  "pyarrow"
   ]
 
 [project.optional-dependencies]
+datasets = ["networkx>=3.1", "matplotlib", "seaborn",   "pybind11>=2.12",  "pyarrow", "fastparquet"]
 bias = ["networkx>=3.1", "pygad==3.3.1", "jax>=0.4.30", "jaxopt>=0.8.3", "optax>=0.2.3", "flax>=0.8.5"]
 explainability = ["lime>=0.2.0.1", "shap>=0.42.1"]
 security = ["jax>=0.4.30", "jaxopt>=0.8.3"]
-all = ["networkx>=3.1", "pygad==3.3.1", "lime>=0.2.0.1", "shap>=0.42.1",  "jax>=0.4.30", "jaxopt>=0.8.3", "optax>=0.2.3", "flax>=0.8.5"]
+all = ["networkx>=3.1", "pygad==3.3.1", "lime>=0.2.0.1", "shap>=0.42.1",  "jax>=0.4.30", "jaxopt>=0.8.3", "optax>=0.2.3", "flax>=0.8.5", "networkx>=3.1", "matplotlib", "seaborn",   "pybind11>=2.12",  "pyarrow", "fastparquet"]
 
 [project.urls]
 Documentation = "https://github.com/holistic-ai/holisticai#readme"
@@ -118,7 +114,9 @@ dependencies = [
   "jax",
   "jaxopt",
   "optax",
-  "flax"
+  "flax",
+  "pyarrow==19.0.0", 
+  "fastparquet"
 ]
 env-vars = { PYTHONPATH = "src" }
 

src/holisticai/security/metrics/_shapr.py

@@ -2,14 +2,20 @@
 
 from typing import TYPE_CHECKING, Union
 
-import jax.numpy as jnp
 import numpy as np
 import pandas as pd
 from holisticai.typing import ArrayLike
 from holisticai.utils.models.neighbors import KNeighborsClassifier
-from jax.nn import one_hot
 from sklearn.preprocessing import LabelEncoder
 
+try:
+    import jax.numpy as jnp
+    from jax.nn import one_hot
+except ImportError:
+    jnp = None
+    one_hot = None
+
+
 if TYPE_CHECKING:
     from numpy.typing import ArrayLike
 
@@ -48,13 +54,15 @@ def transform_label_to_numerical_label(labels: np.ndarray, le: LabelEncoder = No
     return np.array(labels)
 
 
-def check_and_transform_label_format(labels: np.ndarray) -> jnp.ndarray:
+def check_and_transform_label_format(labels: np.ndarray):
     """
     Check label format and transform to one-hot-encoded labels if necessary
 
     :param labels: An array of integer or string labels of shape `(nb_samples,)`, `(nb_samples, 1)` or `(nb_samples, nb_classes)`.
     :return: Labels with shape `(nb_samples, nb_classes)` (one-hot) or `(nb_samples,)` (index).
     """
+    if one_hot is None or jnp is None:
+        raise ImportError("jax or jax.nn is not installed. Please install it with `pip install jax jaxlib`.")
 
     labels = jnp.array(labels)
     return_one_hot = True
@@ -103,6 +111,9 @@ def __call__(
         train_size=1.0,
         aggregated=True,
     ) -> Union[np.ndarray, float]:
+        if one_hot is None or jnp is None:
+            raise ImportError("jax or jax.nn is not installed. Please install it with `pip install jax jaxlib`.")
+
         y_train, le = transform_label_to_numerical_label(y_train)
         y_test = transform_label_to_numerical_label(y_test, le)
         y_pred_train = transform_label_to_numerical_label(y_pred_train, le)
@@ -152,7 +163,7 @@ def shapr_score(
     y_pred_test: pd.Series,
     batch_size=500,
     train_size=1.0,
-) -> jnp.ndarray:
+):
     """
     Compute the SHAPr membership privacy risk metric [1]_ for the given classifier and training set.
 

src/holisticai/utils/models/neighbors/_classification.py

@@ -1,8 +1,16 @@
-import jax.numpy as jnp
-from jax import vmap
+try:
+    import jax.numpy as jnp
+    from jax import vmap
+except ImportError:
+    jnp = None
+    vmap = None
 
 
 class KNeighborsClassifier:
+    def __init__(self):
+        if jnp is None or vmap is None:
+            raise ImportError("jax or jax.nn is not installed. Please install it with `pip install jax jaxlib`.")
+
     def fit(self, X_train, y_train):
         self.X_train = jnp.array(X_train)
         self.y_train = jnp.array(y_train)

tests/commons/test_feature_importance_strategies.py

@@ -32,5 +32,5 @@ def test_permutation_feature_importance_call(input_data):
     fi_strategy = PermutationFeatureImportanceCalculator(random_state=RandomState(42))
     importance = fi_strategy.compute_importances(test['X'], test['y'], proxy=proxy)
     assert isinstance(importance, Importances)
-    assert np.isclose(importance['capital-gain'], 0.38461538461538486, atol=5e-2)
+    assert np.isclose(importance['capital-gain'], 0.44444444444444453, atol=5e-2)