Merge pull request #42 from ruivieira/FAI-710

Move simple and counterfactual predictions to the model package

Author: ruivieira

tests/test_counterfactualexplainer.py

@@ -7,9 +7,8 @@
 from java.util import Random
 
 from trustyai.explainers import CounterfactualExplainer
-from trustyai.local.counterfactual import counterfactual_prediction
 from trustyai.model import (
-    FeatureFactory,
+    counterfactual_prediction,
     output, Model, feature,
 )
 from trustyai.utils import TestUtils

tests/test_limeexplainer.py

@@ -6,9 +6,8 @@
 import pytest
 
 from trustyai.explainers import LimeExplainer
-from trustyai.local.counterfactual import simple_prediction
 from trustyai.utils import TestUtils
-from trustyai.model import feature
+from trustyai.model import feature, simple_prediction
 
 from org.kie.kogito.explainability.local import (
     LocalExplanationException,

tests/test_shap.py

@@ -6,8 +6,7 @@
 import pytest
 
 from trustyai.explainers import SHAPExplainer
-from trustyai.local.counterfactual import simple_prediction
-from trustyai.model import feature, PredictionInput
+from trustyai.model import feature, PredictionInput, simple_prediction
 from trustyai.utils import TestUtils
 
 

trustyai/local/counterfactual.py

@@ -1,11 +1,14 @@
 # pylint: disable = import-error, too-few-public-methods, invalid-name, duplicate-code
 """General model classes"""
-from typing import List
+from typing import List, Optional, Tuple
+import uuid as _uuid
 
+from java.lang import Long
 from java.util.concurrent import CompletableFuture, ForkJoinPool
 from jpype import JImplements, JOverride, _jcustomizer, _jclass
 from org.kie.kogito.explainability.model import (
     CounterfactualPrediction as _CounterfactualPrediction,
+    DataDistribution,
     DataDomain as _DataDomain,
     Feature,
     FeatureFactory as _FeatureFactory,
@@ -22,6 +25,8 @@
     CounterfactualEntity,
 )
 
+from trustyai.model.domain import feature_domain
+
 CounterfactualPrediction = _CounterfactualPrediction
 DataDomain = _DataDomain
 FeatureFactory = _FeatureFactory
@@ -272,3 +277,45 @@ def feature(name: str, dtype: str, value=None) -> Feature:
     else:
         _feature = FeatureFactory.newObjectFeature(name, value)
     return _feature
+
+
+def simple_prediction(
+    input_features: List[Feature],
+    outputs: List[Output],
+) -> SimplePrediction:
+    """Helper to build SimplePrediction"""
+    return SimplePrediction(PredictionInput(input_features), PredictionOutput(outputs))
+
+
+# pylint: disable=too-many-arguments
+def counterfactual_prediction(
+    input_features: List[Feature],
+    outputs: List[Output],
+    domains: List[Optional[Tuple]],
+    constraints: Optional[List[bool]] = None,
+    data_distribution: Optional[DataDistribution] = None,
+    uuid: Optional[_uuid.UUID] = None,
+    timeout: Optional[float] = None,
+) -> CounterfactualPrediction:
+    """Helper to build CounterfactualPrediction"""
+    if not uuid:
+        uuid = _uuid.uuid4()
+    if timeout:
+        timeout = Long(timeout)
+    if not constraints:
+        constraints = [False] * len(input_features)
+
+    # build the feature domains from the Python tuples
+    java_domains = _jclass.JClass("java.util.Arrays").asList(
+        [feature_domain(domain) for domain in domains]
+    )
+
+    return CounterfactualPrediction(
+        PredictionInput(input_features),
+        PredictionOutput(outputs),
+        PredictionFeatureDomain(java_domains),
+        constraints,
+        data_distribution,
+        uuid,
+        timeout,
+    )