fix: Fix create_lit_model_from_endpoint not accepting models that don't return a dictionary. (#1020)

Some models, like Keras squential models, don't return a dictionary for their prediction. We need to support these models as it is commonly used.

Fixes b/220167889

Author: taiseiak

google/cloud/aiplatform/explain/lit.py

@@ -18,7 +18,7 @@
 import os
 
 from google.cloud import aiplatform
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Mapping, Optional, Tuple, Union
 
 try:
     from lit_nlp.api import dataset as lit_dataset
@@ -154,7 +154,12 @@ def predict_minibatch(
             prediction_object = self._endpoint.predict(instances)
         outputs = []
         for prediction in prediction_object.predictions:
-            outputs.append({key: prediction[key] for key in self._output_types})
+            if isinstance(prediction, Mapping):
+                outputs.append({key: prediction[key] for key in self._output_types})
+            else:
+                outputs.append(
+                    {key: prediction[i] for i, key in enumerate(self._output_types)}
+                )
         if self._explanation_enabled:
             for i, explanation in enumerate(prediction_object.explanations):
                 attributions = explanation.attributions

tests/unit/aiplatform/test_explain_lit.py

@@ -105,7 +105,8 @@
     ),
 ]
 _TEST_TRAFFIC_SPLIT = {_TEST_ID: 0, _TEST_ID_2: 100, _TEST_ID_3: 0}
-_TEST_PREDICTION = [{"label": 1.0}]
+_TEST_DICT_PREDICTION = [{"label": 1.0}]
+_TEST_LIST_PREDICTION = [[1.0]]
 _TEST_EXPLANATIONS = [gca_prediction_service.explanation.Explanation(attributions=[])]
 _TEST_ATTRIBUTIONS = [
     gca_prediction_service.explanation.Attribution(
@@ -218,26 +219,54 @@ def get_endpoint_with_models_with_explanation_mock():
 
 
 @pytest.fixture
-def predict_client_predict_mock():
+def predict_client_predict_dict_mock():
     with mock.patch.object(
         prediction_service_client.PredictionServiceClient, "predict"
     ) as predict_mock:
         predict_mock.return_value = gca_prediction_service.PredictResponse(
             deployed_model_id=_TEST_ID
         )
-        predict_mock.return_value.predictions.extend(_TEST_PREDICTION)
+        predict_mock.return_value.predictions.extend(_TEST_DICT_PREDICTION)
         yield predict_mock
 
 
 @pytest.fixture
-def predict_client_explain_mock():
+def predict_client_explain_dict_mock():
     with mock.patch.object(
         prediction_service_client.PredictionServiceClient, "explain"
     ) as predict_mock:
         predict_mock.return_value = gca_prediction_service.ExplainResponse(
             deployed_model_id=_TEST_ID,
         )
-        predict_mock.return_value.predictions.extend(_TEST_PREDICTION)
+        predict_mock.return_value.predictions.extend(_TEST_DICT_PREDICTION)
+        predict_mock.return_value.explanations.extend(_TEST_EXPLANATIONS)
+        predict_mock.return_value.explanations[0].attributions.extend(
+            _TEST_ATTRIBUTIONS
+        )
+        yield predict_mock
+
+
+@pytest.fixture
+def predict_client_predict_list_mock():
+    with mock.patch.object(
+        prediction_service_client.PredictionServiceClient, "predict"
+    ) as predict_mock:
+        predict_mock.return_value = gca_prediction_service.PredictResponse(
+            deployed_model_id=_TEST_ID
+        )
+        predict_mock.return_value.predictions.extend(_TEST_LIST_PREDICTION)
+        yield predict_mock
+
+
+@pytest.fixture
+def predict_client_explain_list_mock():
+    with mock.patch.object(
+        prediction_service_client.PredictionServiceClient, "explain"
+    ) as predict_mock:
+        predict_mock.return_value = gca_prediction_service.ExplainResponse(
+            deployed_model_id=_TEST_ID,
+        )
+        predict_mock.return_value.predictions.extend(_TEST_LIST_PREDICTION)
         predict_mock.return_value.explanations.extend(_TEST_EXPLANATIONS)
         predict_mock.return_value.explanations[0].attributions.extend(
             _TEST_ATTRIBUTIONS
@@ -312,10 +341,112 @@ def test_create_lit_model_from_tensorflow_with_xai_returns_model(
             assert len(item.values()) == 2
 
     @pytest.mark.usefixtures(
-        "predict_client_predict_mock", "get_endpoint_with_models_mock"
+        "predict_client_predict_dict_mock", "get_endpoint_with_models_mock"
+    )
+    @pytest.mark.parametrize("model_id", [None, _TEST_ID])
+    def test_create_lit_model_from_dict_endpoint_returns_model(
+        self, feature_types, label_types, model_id
+    ):
+        endpoint = aiplatform.Endpoint(_TEST_ENDPOINT_NAME)
+        lit_model = create_lit_model_from_endpoint(
+            endpoint, feature_types, label_types, model_id
+        )
+        test_inputs = [
+            {"feature_1": 1.0, "feature_2": 2.0},
+        ]
+        outputs = lit_model.predict_minibatch(test_inputs)
+
+        assert lit_model.input_spec() == dict(feature_types)
+        assert lit_model.output_spec() == dict(label_types)
+        assert len(outputs) == 1
+        for item in outputs:
+            assert item.keys() == {"label"}
+            assert len(item.values()) == 1
+
+    @pytest.mark.usefixtures(
+        "predict_client_explain_dict_mock",
+        "get_endpoint_with_models_with_explanation_mock",
+    )
+    @pytest.mark.parametrize("model_id", [None, _TEST_ID])
+    def test_create_lit_model_from_dict_endpoint_with_xai_returns_model(
+        self, feature_types, label_types, model_id
+    ):
+        endpoint = aiplatform.Endpoint(_TEST_ENDPOINT_NAME)
+        lit_model = create_lit_model_from_endpoint(
+            endpoint, feature_types, label_types, model_id
+        )
+        test_inputs = [
+            {"feature_1": 1.0, "feature_2": 2.0},
+        ]
+        outputs = lit_model.predict_minibatch(test_inputs)
+
+        assert lit_model.input_spec() == dict(feature_types)
+        assert lit_model.output_spec() == dict(
+            {
+                **label_types,
+                "feature_attribution": lit_types.FeatureSalience(signed=True),
+            }
+        )
+        assert len(outputs) == 1
+        for item in outputs:
+            assert item.keys() == {"label", "feature_attribution"}
+            assert len(item.values()) == 2
+
+    @pytest.mark.usefixtures(
+        "predict_client_predict_dict_mock", "get_endpoint_with_models_mock"
+    )
+    @pytest.mark.parametrize("model_id", [None, _TEST_ID])
+    def test_create_lit_model_from_dict_endpoint_name_returns_model(
+        self, feature_types, label_types, model_id
+    ):
+        lit_model = create_lit_model_from_endpoint(
+            _TEST_ENDPOINT_NAME, feature_types, label_types, model_id
+        )
+        test_inputs = [
+            {"feature_1": 1.0, "feature_2": 2.0},
+        ]
+        outputs = lit_model.predict_minibatch(test_inputs)
+
+        assert lit_model.input_spec() == dict(feature_types)
+        assert lit_model.output_spec() == dict(label_types)
+        assert len(outputs) == 1
+        for item in outputs:
+            assert item.keys() == {"label"}
+            assert len(item.values()) == 1
+
+    @pytest.mark.usefixtures(
+        "predict_client_explain_dict_mock",
+        "get_endpoint_with_models_with_explanation_mock",
+    )
+    @pytest.mark.parametrize("model_id", [None, _TEST_ID])
+    def test_create_lit_model_from_dict_endpoint_name_with_xai_returns_model(
+        self, feature_types, label_types, model_id
+    ):
+        lit_model = create_lit_model_from_endpoint(
+            _TEST_ENDPOINT_NAME, feature_types, label_types, model_id
+        )
+        test_inputs = [
+            {"feature_1": 1.0, "feature_2": 2.0},
+        ]
+        outputs = lit_model.predict_minibatch(test_inputs)
+
+        assert lit_model.input_spec() == dict(feature_types)
+        assert lit_model.output_spec() == dict(
+            {
+                **label_types,
+                "feature_attribution": lit_types.FeatureSalience(signed=True),
+            }
+        )
+        assert len(outputs) == 1
+        for item in outputs:
+            assert item.keys() == {"label", "feature_attribution"}
+            assert len(item.values()) == 2
+
+    @pytest.mark.usefixtures(
+        "predict_client_predict_list_mock", "get_endpoint_with_models_mock"
     )
     @pytest.mark.parametrize("model_id", [None, _TEST_ID])
-    def test_create_lit_model_from_endpoint_returns_model(
+    def test_create_lit_model_from_list_endpoint_returns_model(
         self, feature_types, label_types, model_id
     ):
         endpoint = aiplatform.Endpoint(_TEST_ENDPOINT_NAME)
@@ -335,10 +466,11 @@ def test_create_lit_model_from_endpoint_returns_model(
             assert len(item.values()) == 1
 
     @pytest.mark.usefixtures(
-        "predict_client_explain_mock", "get_endpoint_with_models_with_explanation_mock"
+        "predict_client_explain_list_mock",
+        "get_endpoint_with_models_with_explanation_mock",
     )
     @pytest.mark.parametrize("model_id", [None, _TEST_ID])
-    def test_create_lit_model_from_endpoint_with_xai_returns_model(
+    def test_create_lit_model_from_list_endpoint_with_xai_returns_model(
         self, feature_types, label_types, model_id
     ):
         endpoint = aiplatform.Endpoint(_TEST_ENDPOINT_NAME)
@@ -363,10 +495,10 @@ def test_create_lit_model_from_endpoint_with_xai_returns_model(
             assert len(item.values()) == 2
 
     @pytest.mark.usefixtures(
-        "predict_client_predict_mock", "get_endpoint_with_models_mock"
+        "predict_client_predict_list_mock", "get_endpoint_with_models_mock"
     )
     @pytest.mark.parametrize("model_id", [None, _TEST_ID])
-    def test_create_lit_model_from_endpoint_name_returns_model(
+    def test_create_lit_model_from_list_endpoint_name_returns_model(
         self, feature_types, label_types, model_id
     ):
         lit_model = create_lit_model_from_endpoint(
@@ -385,10 +517,11 @@ def test_create_lit_model_from_endpoint_name_returns_model(
             assert len(item.values()) == 1
 
     @pytest.mark.usefixtures(
-        "predict_client_explain_mock", "get_endpoint_with_models_with_explanation_mock"
+        "predict_client_explain_list_mock",
+        "get_endpoint_with_models_with_explanation_mock",
     )
     @pytest.mark.parametrize("model_id", [None, _TEST_ID])
-    def test_create_lit_model_from_endpoint_name_with_xai_returns_model(
+    def test_create_lit_model_from_list_endpoint_name_with_xai_returns_model(
         self, feature_types, label_types, model_id
     ):
         lit_model = create_lit_model_from_endpoint(