From ea16849f936d7a2e8402fd235decefe5972685ed Mon Sep 17 00:00:00 2001 From: Taisei Klasen Date: Tue, 1 Feb 2022 13:34:41 -0800 Subject: [PATCH] feat: Open LIT with a deployed model (#963) --- google/cloud/aiplatform/explain/lit.py | 275 +++++++++--- tests/unit/aiplatform/test_explain_lit.py | 503 +++++++++++++++++----- 2 files changed, 589 insertions(+), 189 deletions(-) diff --git a/google/cloud/aiplatform/explain/lit.py b/google/cloud/aiplatform/explain/lit.py index a34578abad..5032055801 100644 --- a/google/cloud/aiplatform/explain/lit.py +++ b/google/cloud/aiplatform/explain/lit.py @@ -17,6 +17,7 @@ import logging import os +from google.cloud import aiplatform from typing import Dict, List, Optional, Tuple, Union try: @@ -61,11 +62,11 @@ def __init__( ): """Construct a VertexLitDataset. Args: - dataset: - Required. A Pandas DataFrame that includes feature column names and data. - column_types: - Required. An OrderedDict of string names matching the columns of the dataset - as the key, and the associated LitType of the column. + dataset: + Required. A Pandas DataFrame that includes feature column names and data. + column_types: + Required. An OrderedDict of string names matching the columns of the dataset + as the key, and the associated LitType of the column. """ self._examples = dataset.to_dict(orient="records") self._column_types = column_types @@ -75,8 +76,109 @@ def spec(self): return dict(self._column_types) -class _VertexLitModel(lit_model.Model): - """LIT model class for the Vertex LIT integration. +class _EndpointLitModel(lit_model.Model): + """LIT model class for the Vertex LIT integration with a model deployed to an endpoint. + + This is used in the create_lit_model function. + """ + + def __init__( + self, + endpoint: Union[str, aiplatform.Endpoint], + input_types: "OrderedDict[str, lit_types.LitType]", # noqa: F821 + output_types: "OrderedDict[str, lit_types.LitType]", # noqa: F821 + model_id: Optional[str] = None, + ): + """Construct a VertexLitModel. + Args: + model: + Required. The name of the Endpoint resource. Format: + ``projects/{project}/locations/{location}/endpoints/{endpoint}`` + input_types: + Required. An OrderedDict of string names matching the features of the model + as the key, and the associated LitType of the feature. + output_types: + Required. An OrderedDict of string names matching the labels of the model + as the key, and the associated LitType of the label. + model_id: + Optional. A string of the specific model in the endpoint to create the + LIT model from. If this is not set, any usable model in the endpoint is + used to create the LIT model. + Raises: + ValueError if the model_id was not found in the endpoint. + """ + if isinstance(endpoint, str): + self._endpoint = aiplatform.Endpoint(endpoint) + else: + self._endpoint = endpoint + self._model_id = model_id + self._input_types = input_types + self._output_types = output_types + # Check if the model with the model ID has explanation enabled + if model_id: + deployed_model = next( + filter( + lambda model: model.id == model_id, self._endpoint.list_models() + ), + None, + ) + if not deployed_model: + raise ValueError( + "A model with id {model_id} was not found in the endpoint {endpoint}.".format( + model_id=model_id, endpoint=endpoint + ) + ) + self._explanation_enabled = bool(deployed_model.explanation_spec) + # Check if all models in the endpoint have explanation enabled + else: + self._explanation_enabled = all( + model.explanation_spec for model in self._endpoint.list_models() + ) + + def predict_minibatch( + self, inputs: List[lit_types.JsonDict] + ) -> List[lit_types.JsonDict]: + """Retun predictions based on a batch of inputs. + Args: + inputs: Requred. a List of instances to predict on based on the input spec. + Returns: + A list of predictions based on the output spec. + """ + instances = [] + for input in inputs: + instance = [input[feature] for feature in self._input_types] + instances.append(instance) + if self._explanation_enabled: + prediction_object = self._endpoint.explain(instances) + else: + 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 self._explanation_enabled: + for i, explanation in enumerate(prediction_object.explanations): + attributions = explanation.attributions + outputs[i]["feature_attribution"] = lit_dtypes.FeatureSalience( + attributions + ) + return outputs + + def input_spec(self) -> lit_types.Spec: + """Return a spec describing model inputs.""" + return dict(self._input_types) + + def output_spec(self) -> lit_types.Spec: + """Return a spec describing model outputs.""" + output_spec_dict = dict(self._output_types) + if self._explanation_enabled: + output_spec_dict["feature_attribution"] = lit_types.FeatureSalience( + signed=True + ) + return output_spec_dict + + +class _TensorFlowLitModel(lit_model.Model): + """LIT model class for the Vertex LIT integration with a TensorFlow saved model. This is used in the create_lit_model function. """ @@ -90,19 +192,19 @@ def __init__( ): """Construct a VertexLitModel. Args: - model: - Required. A string reference to a local TensorFlow saved model directory. - The model must have at most one input and one output tensor. - input_types: - Required. An OrderedDict of string names matching the features of the model - as the key, and the associated LitType of the feature. - output_types: - Required. An OrderedDict of string names matching the labels of the model - as the key, and the associated LitType of the label. - attribution_method: - Optional. A string to choose what attribution configuration to - set up the explainer with. Valid options are 'sampled_shapley' - or 'integrated_gradients'. + model: + Required. A string reference to a local TensorFlow saved model directory. + The model must have at most one input and one output tensor. + input_types: + Required. An OrderedDict of string names matching the features of the model + as the key, and the associated LitType of the feature. + output_types: + Required. An OrderedDict of string names matching the labels of the model + as the key, and the associated LitType of the label. + attribution_method: + Optional. A string to choose what attribution configuration to + set up the explainer with. Valid options are 'sampled_shapley' + or 'integrated_gradients'. """ self._load_model(model) self._input_types = input_types @@ -120,6 +222,12 @@ def attribution_explainer(self,) -> Optional["AttributionExplainer"]: # noqa: F def predict_minibatch( self, inputs: List[lit_types.JsonDict] ) -> List[lit_types.JsonDict]: + """Retun predictions based on a batch of inputs. + Args: + inputs: Requred. a List of instances to predict on based on the input spec. + Returns: + A list of predictions based on the output spec. + """ instances = [] for input in inputs: instance = [input[feature] for feature in self._input_types] @@ -166,7 +274,7 @@ def output_spec(self) -> lit_types.Spec: def _load_model(self, model: str): """Loads a TensorFlow saved model and populates the input and output signature attributes of the class. Args: - model: Required. A string reference to a TensorFlow saved model directory. + model: Required. A string reference to a TensorFlow saved model directory. Raises: ValueError if the model has more than one input tensor or more than one output tensor. """ @@ -188,11 +296,11 @@ def _set_up_attribution_explainer( ): """Populates the attribution explainer attribute of the class. Args: - model: Required. A string reference to a TensorFlow saved model directory. + model: Required. A string reference to a TensorFlow saved model directory. attribution_method: - Optional. A string to choose what attribution configuration to - set up the explainer with. Valid options are 'sampled_shapley' - or 'integrated_gradients'. + Optional. A string to choose what attribution configuration to + set up the explainer with. Valid options are 'sampled_shapley' + or 'integrated_gradients'. """ try: import explainable_ai_sdk @@ -228,17 +336,44 @@ def create_lit_dataset( ) -> lit_dataset.Dataset: """Creates a LIT Dataset object. Args: - dataset: - Required. A Pandas DataFrame that includes feature column names and data. - column_types: - Required. An OrderedDict of string names matching the columns of the dataset - as the key, and the associated LitType of the column. + dataset: + Required. A Pandas DataFrame that includes feature column names and data. + column_types: + Required. An OrderedDict of string names matching the columns of the dataset + as the key, and the associated LitType of the column. Returns: A LIT Dataset object that has the data from the dataset provided. """ return _VertexLitDataset(dataset, column_types) +def create_lit_model_from_endpoint( + endpoint: Union[str, aiplatform.Endpoint], + input_types: "OrderedDict[str, lit_types.LitType]", # noqa: F821 + output_types: "OrderedDict[str, lit_types.LitType]", # noqa: F821 + model_id: Optional[str] = None, +) -> lit_model.Model: + """Creates a LIT Model object. + Args: + model: + Required. The name of the Endpoint resource or an Endpoint instance. + Endpoint name format: ``projects/{project}/locations/{location}/endpoints/{endpoint}`` + input_types: + Required. An OrderedDict of string names matching the features of the model + as the key, and the associated LitType of the feature. + output_types: + Required. An OrderedDict of string names matching the labels of the model + as the key, and the associated LitType of the label. + model_id: + Optional. A string of the specific model in the endpoint to create the + LIT model from. If this is not set, any usable model in the endpoint is + used to create the LIT model. + Returns: + A LIT Model object that has the same functionality as the model provided. + """ + return _EndpointLitModel(endpoint, input_types, output_types, model_id) + + def create_lit_model( model: str, input_types: "OrderedDict[str, lit_types.LitType]", # noqa: F821 @@ -247,23 +382,23 @@ def create_lit_model( ) -> lit_model.Model: """Creates a LIT Model object. Args: - model: - Required. A string reference to a local TensorFlow saved model directory. - The model must have at most one input and one output tensor. - input_types: - Required. An OrderedDict of string names matching the features of the model - as the key, and the associated LitType of the feature. - output_types: - Required. An OrderedDict of string names matching the labels of the model - as the key, and the associated LitType of the label. - attribution_method: - Optional. A string to choose what attribution configuration to - set up the explainer with. Valid options are 'sampled_shapley' - or 'integrated_gradients'. + model: + Required. A string reference to a local TensorFlow saved model directory. + The model must have at most one input and one output tensor. + input_types: + Required. An OrderedDict of string names matching the features of the model + as the key, and the associated LitType of the feature. + output_types: + Required. An OrderedDict of string names matching the labels of the model + as the key, and the associated LitType of the label. + attribution_method: + Optional. A string to choose what attribution configuration to + set up the explainer with. Valid options are 'sampled_shapley' + or 'integrated_gradients'. Returns: A LIT Model object that has the same functionality as the model provided. """ - return _VertexLitModel(model, input_types, output_types, attribution_method) + return _TensorFlowLitModel(model, input_types, output_types, attribution_method) def open_lit( @@ -273,12 +408,12 @@ def open_lit( ): """Open LIT from the provided models and datasets. Args: - models: - Required. A list of LIT models to open LIT with. - input_types: - Required. A lit of LIT datasets to open LIT with. - open_in_new_tab: - Optional. A boolean to choose if LIT open in a new tab or not. + models: + Required. A list of LIT models to open LIT with. + input_types: + Required. A lit of LIT datasets to open LIT with. + open_in_new_tab: + Optional. A boolean to choose if LIT open in a new tab or not. Raises: ImportError if LIT is not installed. """ @@ -297,26 +432,26 @@ def set_up_and_open_lit( ) -> Tuple[lit_dataset.Dataset, lit_model.Model]: """Creates a LIT dataset and model and opens LIT. Args: - dataset: - Required. A Pandas DataFrame that includes feature column names and data. - column_types: - Required. An OrderedDict of string names matching the columns of the dataset - as the key, and the associated LitType of the column. - model: - Required. A string reference to a TensorFlow saved model directory. - The model must have at most one input and one output tensor. - input_types: - Required. An OrderedDict of string names matching the features of the model - as the key, and the associated LitType of the feature. - output_types: - Required. An OrderedDict of string names matching the labels of the model - as the key, and the associated LitType of the label. - attribution_method: - Optional. A string to choose what attribution configuration to - set up the explainer with. Valid options are 'sampled_shapley' - or 'integrated_gradients'. - open_in_new_tab: - Optional. A boolean to choose if LIT open in a new tab or not. + dataset: + Required. A Pandas DataFrame that includes feature column names and data. + column_types: + Required. An OrderedDict of string names matching the columns of the dataset + as the key, and the associated LitType of the column. + model: + Required. A string reference to a TensorFlow saved model directory. + The model must have at most one input and one output tensor. + input_types: + Required. An OrderedDict of string names matching the features of the model + as the key, and the associated LitType of the feature. + output_types: + Required. An OrderedDict of string names matching the labels of the model + as the key, and the associated LitType of the label. + attribution_method: + Optional. A string to choose what attribution configuration to + set up the explainer with. Valid options are 'sampled_shapley' + or 'integrated_gradients'. + open_in_new_tab: + Optional. A boolean to choose if LIT open in a new tab or not. Returns: A Tuple of the LIT dataset and model created. Raises: diff --git a/tests/unit/aiplatform/test_explain_lit.py b/tests/unit/aiplatform/test_explain_lit.py index 8f10193c7b..c8092b1742 100644 --- a/tests/unit/aiplatform/test_explain_lit.py +++ b/tests/unit/aiplatform/test_explain_lit.py @@ -21,16 +21,104 @@ import pytest import tensorflow as tf +from google.auth import credentials as auth_credentials +from google.cloud import aiplatform +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.compat.types import ( + endpoint as gca_endpoint, + prediction_service as gca_prediction_service, + explanation as gca_explanation, +) from google.cloud.aiplatform.explain.lit import ( create_lit_dataset, create_lit_model, + create_lit_model_from_endpoint, open_lit, set_up_and_open_lit, ) +from google.cloud.aiplatform_v1.services.endpoint_service import ( + client as endpoint_service_client, +) +from google.cloud.aiplatform_v1.services.prediction_service import ( + client as prediction_service_client, +) +from importlib import reload from lit_nlp.api import types as lit_types from lit_nlp import notebook from unittest import mock +_TEST_PROJECT = "test-project" +_TEST_LOCATION = "us-central1" +_TEST_ID = "1028944691210842416" +_TEST_ID_2 = "4366591682456584192" +_TEST_ID_3 = "5820582938582924817" +_TEST_DISPLAY_NAME = "test-display-name" +_TEST_DISPLAY_NAME_2 = "test-display-name-2" +_TEST_DISPLAY_NAME_3 = "test-display-name-3" +_TEST_ENDPOINT_NAME = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/{_TEST_ID}" +) +_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) +_TEST_EXPLANATION_METADATA = aiplatform.explain.ExplanationMetadata( + inputs={ + "features": { + "input_tensor_name": "dense_input", + "encoding": "BAG_OF_FEATURES", + "modality": "numeric", + "index_feature_mapping": ["abc", "def", "ghj"], + } + }, + outputs={"medv": {"output_tensor_name": "dense_2"}}, +) +_TEST_EXPLANATION_PARAMETERS = aiplatform.explain.ExplanationParameters( + {"sampled_shapley_attribution": {"path_count": 10}} +) +_TEST_DEPLOYED_MODELS = [ + gca_endpoint.DeployedModel(id=_TEST_ID, display_name=_TEST_DISPLAY_NAME), + gca_endpoint.DeployedModel(id=_TEST_ID_2, display_name=_TEST_DISPLAY_NAME_2), + gca_endpoint.DeployedModel(id=_TEST_ID_3, display_name=_TEST_DISPLAY_NAME_3), +] +_TEST_DEPLOYED_MODELS_WITH_EXPLANATION = [ + gca_endpoint.DeployedModel( + id=_TEST_ID, + display_name=_TEST_DISPLAY_NAME, + explanation_spec=gca_explanation.ExplanationSpec( + metadata=_TEST_EXPLANATION_METADATA, + parameters=_TEST_EXPLANATION_PARAMETERS, + ), + ), + gca_endpoint.DeployedModel( + id=_TEST_ID_2, + display_name=_TEST_DISPLAY_NAME_2, + explanation_spec=gca_explanation.ExplanationSpec( + metadata=_TEST_EXPLANATION_METADATA, + parameters=_TEST_EXPLANATION_PARAMETERS, + ), + ), + gca_endpoint.DeployedModel( + id=_TEST_ID_3, + display_name=_TEST_DISPLAY_NAME_3, + explanation_spec=gca_explanation.ExplanationSpec( + metadata=_TEST_EXPLANATION_METADATA, + parameters=_TEST_EXPLANATION_PARAMETERS, + ), + ), +] +_TEST_TRAFFIC_SPLIT = {_TEST_ID: 0, _TEST_ID_2: 100, _TEST_ID_3: 0} +_TEST_PREDICTION = [{"label": 1.0}] +_TEST_EXPLANATIONS = [gca_prediction_service.explanation.Explanation(attributions=[])] +_TEST_ATTRIBUTIONS = [ + gca_prediction_service.explanation.Attribution( + baseline_output_value=1.0, + instance_output_value=2.0, + feature_attributions={"feature_1": 3.0, "feature_2": 2.0}, + output_index=[1, 2, 3], + output_display_name="abc", + approximation_error=6.0, + output_name="xyz", + ) +] + @pytest.fixture def widget_render_mock(): @@ -57,16 +145,25 @@ def load_model_from_local_path_mock(): "feature_1": 0.01, "feature_2": 0.1, } - model_mock.explain.return_value = [ - explanation_mock - # , explanation_mock - ] + model_mock.explain.return_value = [explanation_mock] explainer_mock.return_value = model_mock yield explainer_mock @pytest.fixture -def set_up_sequential(tmpdir): +def feature_types(): + yield collections.OrderedDict( + [("feature_1", lit_types.Scalar()), ("feature_2", lit_types.Scalar())] + ) + + +@pytest.fixture +def label_types(): + yield collections.OrderedDict([("label", lit_types.RegressionScore())]) + + +@pytest.fixture +def set_up_sequential(tmpdir, feature_types, label_types): # Set up a sequential model seq_model = tf.keras.models.Sequential() seq_model.add(tf.keras.layers.Dense(32, activation="relu", input_shape=(2,))) @@ -74,10 +171,6 @@ def set_up_sequential(tmpdir): seq_model.add(tf.keras.layers.Dense(1, activation="sigmoid")) saved_model_path = str(tmpdir.mkdir("tmp")) tf.saved_model.save(seq_model, saved_model_path) - feature_types = collections.OrderedDict( - [("feature_1", lit_types.Scalar()), ("feature_2", lit_types.Scalar())] - ) - label_types = collections.OrderedDict([("label", lit_types.RegressionScore())]) yield feature_types, label_types, saved_model_path @@ -96,130 +189,302 @@ def set_up_pandas_dataframe_and_columns(): yield dataframe, columns -def test_create_lit_dataset_from_pandas_returns_dataset( - set_up_pandas_dataframe_and_columns, -): - pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns - lit_dataset = create_lit_dataset(pd_dataset, lit_columns) - expected_examples = [ - {"feature_1": 1.0, "feature_2": 3.0, "label": 1.0}, - ] +@pytest.fixture +def get_endpoint_with_models_mock(): + with mock.patch.object( + endpoint_service_client.EndpointServiceClient, "get_endpoint" + ) as get_endpoint_mock: + get_endpoint_mock.return_value = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, + name=_TEST_ENDPOINT_NAME, + deployed_models=_TEST_DEPLOYED_MODELS, + traffic_split=_TEST_TRAFFIC_SPLIT, + ) + yield get_endpoint_mock - assert lit_dataset.spec() == dict(lit_columns) - assert expected_examples == lit_dataset._examples +@pytest.fixture +def get_endpoint_with_models_with_explanation_mock(): + with mock.patch.object( + endpoint_service_client.EndpointServiceClient, "get_endpoint" + ) as get_endpoint_mock: + get_endpoint_mock.return_value = gca_endpoint.Endpoint( + display_name=_TEST_DISPLAY_NAME, + name=_TEST_ENDPOINT_NAME, + deployed_models=_TEST_DEPLOYED_MODELS_WITH_EXPLANATION, + traffic_split=_TEST_TRAFFIC_SPLIT, + ) + yield get_endpoint_mock -def test_create_lit_model_from_tensorflow_returns_model(set_up_sequential): - feature_types, label_types, saved_model_path = set_up_sequential - lit_model = create_lit_model(saved_model_path, feature_types, label_types) - 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.fixture +def predict_client_predict_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) + yield predict_mock -@mock.patch.dict(os.environ, {"LIT_PROXY_URL": "auto"}) -@pytest.mark.usefixtures( - "sampled_shapley_explainer_mock", "load_model_from_local_path_mock" -) -def test_create_lit_model_from_tensorflow_with_xai_returns_model(set_up_sequential): - feature_types, label_types, saved_model_path = set_up_sequential - lit_model = create_lit_model(saved_model_path, feature_types, label_types) - 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 - - -def test_open_lit( - set_up_sequential, set_up_pandas_dataframe_and_columns, widget_render_mock -): - pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns - lit_dataset = create_lit_dataset(pd_dataset, lit_columns) - feature_types, label_types, saved_model_path = set_up_sequential - lit_model = create_lit_model(saved_model_path, feature_types, label_types) - - open_lit({"model": lit_model}, {"dataset": lit_dataset}) - widget_render_mock.assert_called_once() - - -def test_set_up_and_open_lit( - set_up_sequential, set_up_pandas_dataframe_and_columns, widget_render_mock -): - pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns - feature_types, label_types, saved_model_path = set_up_sequential - lit_dataset, lit_model = set_up_and_open_lit( - pd_dataset, lit_columns, saved_model_path, feature_types, label_types - ) +@pytest.fixture +def predict_client_explain_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.explanations.extend(_TEST_EXPLANATIONS) + predict_mock.return_value.explanations[0].attributions.extend( + _TEST_ATTRIBUTIONS + ) + yield predict_mock + + +class TestExplainLit: + def setup_method(self): + reload(initializer) + reload(aiplatform) + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=_TEST_CREDENTIALS, + ) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_create_lit_dataset_from_pandas_returns_dataset( + self, set_up_pandas_dataframe_and_columns, + ): + pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns + lit_dataset = create_lit_dataset(pd_dataset, lit_columns) + expected_examples = [ + {"feature_1": 1.0, "feature_2": 3.0, "label": 1.0}, + ] - expected_examples = [ - {"feature_1": 1.0, "feature_2": 3.0, "label": 1.0}, - ] - test_inputs = [ - {"feature_1": 1.0, "feature_2": 2.0}, - ] - outputs = lit_model.predict_minibatch(test_inputs) + assert lit_dataset.spec() == dict(lit_columns) + assert expected_examples == lit_dataset._examples - assert lit_dataset.spec() == dict(lit_columns) - assert expected_examples == lit_dataset._examples + def test_create_lit_model_from_tensorflow_returns_model(self, set_up_sequential): + feature_types, label_types, saved_model_path = set_up_sequential + lit_model = create_lit_model(saved_model_path, feature_types, label_types) + 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 + + @mock.patch.dict(os.environ, {"LIT_PROXY_URL": "auto"}) + @pytest.mark.usefixtures( + "sampled_shapley_explainer_mock", "load_model_from_local_path_mock" + ) + def test_create_lit_model_from_tensorflow_with_xai_returns_model( + self, set_up_sequential + ): + feature_types, label_types, saved_model_path = set_up_sequential + lit_model = create_lit_model(saved_model_path, feature_types, label_types) + 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_mock", "get_endpoint_with_models_mock" + ) + @pytest.mark.parametrize("model_id", [None, _TEST_ID]) + def test_create_lit_model_from_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 + 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 - widget_render_mock.assert_called_once() + @pytest.mark.usefixtures( + "predict_client_explain_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( + 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_mock", "get_endpoint_with_models_mock" + ) + @pytest.mark.parametrize("model_id", [None, _TEST_ID]) + def test_create_lit_model_from_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 -@mock.patch.dict(os.environ, {"LIT_PROXY_URL": "auto"}) -@pytest.mark.usefixtures( - "sampled_shapley_explainer_mock", "load_model_from_local_path_mock" -) -def test_set_up_and_open_lit_with_xai( - set_up_sequential, set_up_pandas_dataframe_and_columns, widget_render_mock -): - pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns - feature_types, label_types, saved_model_path = set_up_sequential - lit_dataset, lit_model = set_up_and_open_lit( - pd_dataset, lit_columns, saved_model_path, feature_types, label_types + @pytest.mark.usefixtures( + "predict_client_explain_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( + 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 + + def test_open_lit( + self, set_up_sequential, set_up_pandas_dataframe_and_columns, widget_render_mock + ): + pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns + lit_dataset = create_lit_dataset(pd_dataset, lit_columns) + feature_types, label_types, saved_model_path = set_up_sequential + lit_model = create_lit_model(saved_model_path, feature_types, label_types) + + open_lit({"model": lit_model}, {"dataset": lit_dataset}) + widget_render_mock.assert_called_once() + + def test_set_up_and_open_lit( + self, set_up_sequential, set_up_pandas_dataframe_and_columns, widget_render_mock + ): + pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns + feature_types, label_types, saved_model_path = set_up_sequential + lit_dataset, lit_model = set_up_and_open_lit( + pd_dataset, lit_columns, saved_model_path, feature_types, label_types + ) + + expected_examples = [ + {"feature_1": 1.0, "feature_2": 3.0, "label": 1.0}, + ] + test_inputs = [ + {"feature_1": 1.0, "feature_2": 2.0}, + ] + outputs = lit_model.predict_minibatch(test_inputs) - expected_examples = [ - {"feature_1": 1.0, "feature_2": 3.0, "label": 1.0}, - ] - test_inputs = [ - {"feature_1": 1.0, "feature_2": 2.0}, - ] - outputs = lit_model.predict_minibatch(test_inputs) + assert lit_dataset.spec() == dict(lit_columns) + assert expected_examples == lit_dataset._examples - assert lit_dataset.spec() == dict(lit_columns) - assert expected_examples == lit_dataset._examples + 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 - 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 + widget_render_mock.assert_called_once() - widget_render_mock.assert_called_once() + @mock.patch.dict(os.environ, {"LIT_PROXY_URL": "auto"}) + @pytest.mark.usefixtures( + "sampled_shapley_explainer_mock", "load_model_from_local_path_mock" + ) + def test_set_up_and_open_lit_with_xai( + self, set_up_sequential, set_up_pandas_dataframe_and_columns, widget_render_mock + ): + pd_dataset, lit_columns = set_up_pandas_dataframe_and_columns + feature_types, label_types, saved_model_path = set_up_sequential + lit_dataset, lit_model = set_up_and_open_lit( + pd_dataset, lit_columns, saved_model_path, feature_types, label_types + ) + + expected_examples = [ + {"feature_1": 1.0, "feature_2": 3.0, "label": 1.0}, + ] + test_inputs = [ + {"feature_1": 1.0, "feature_2": 2.0}, + ] + outputs = lit_model.predict_minibatch(test_inputs) + + assert lit_dataset.spec() == dict(lit_columns) + assert expected_examples == lit_dataset._examples + + 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 + + widget_render_mock.assert_called_once()