<jupyter_start><jupyter_text>Copyright 2020 The TensorFlow Authors.<jupyter_code>#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.<jupyter_output><empty_output><jupyter_text>Customizing MinDiffModel    View on TensorFlow.org    Run in Google Colab    View source on GitHub  Download notebook IntroductionIn most cases, using `MinDiffModel` directly as described in the ["Integrating MinDiff with MinDiffModel" guide](./integrating_min_diff_with_min_diff_model) is sufficient. However, it is possible that you will need customized behavior. The two primary reasons for this are:*   The `keras.Model` you are using has custom behavior that you want to preserve.*   You want the `MinDiffModel` to behave differently from the default.In either case, you will need to subclass `MinDiffModel` to achieve the desired results. Setup<jupyter_code>!pip install --upgrade tensorflow-model-remediation
import tensorflow as tf
tf.get_logger().setLevel('ERROR')  # Avoid TF warnings.
from tensorflow_model_remediation import min_diff
from tensorflow_model_remediation.tools.tutorials_utils import uci as tutorials_utils<jupyter_output><empty_output><jupyter_text>First, download the data. For succinctness, the input preparation logic has been factored out into helper functions as described in the [input preparation guide](./min_diff_data_preparationutility_functions_for_other_guides). You can read the full guide for details on this process.<jupyter_code># Original Dataset for training, sampled at 0.3 for reduced runtimes.
train_df = tutorials_utils.get_uci_data(split='train', sample=0.3)
train_ds = tutorials_utils.df_to_dataset(train_df, batch_size=128)

# Dataset needed to train with MinDiff.
train_with_min_diff_ds = (
    tutorials_utils.get_uci_with_min_diff_dataset(split='train', sample=0.3))<jupyter_output><empty_output><jupyter_text>Preserving Original Model Customizations`tf.keras.Model` is designed to be easily customized via subclassing as described [here](https://www.tensorflow.org/guide/keras/custom_layers_and_modelsthe_model_class). If your model has customized implementations that you wish to preserve when applying MinDiff, you will need to subclass `MinDiffModel`.  Original Custom ModelTo see how you can preserve customizations, create a custom model that sets an attribute to `True` when its custom `train_step` is called.  This is not a useful customization but will serve to illustrate behavior.<jupyter_code>class CustomModel(tf.keras.Model):

  # Customized train_step
  def train_step(self, *args, **kwargs):
    self.used_custom_train_step = True  # Marker that we can check for.
    return super(CustomModel, self).train_step(*args, **kwargs)<jupyter_output><empty_output><jupyter_text>Training such a model would look the same as a normal `Sequential` model.<jupyter_code>model = tutorials_utils.get_uci_model(model_class=CustomModel)  # Use CustomModel.

model.compile(optimizer='adam', loss='binary_crossentropy')

_ = model.fit(train_ds.take(1), epochs=1, verbose=0)

# Model has used the custom train_step.
print('Model used the custom train_step:')
print(hasattr(model, 'used_custom_train_step'))  # True<jupyter_output><empty_output><jupyter_text>Subclassing MinDiffModelIf you were to try and use `MinDiffModel` directly, the model would not use the custom `train_step`.<jupyter_code>model = tutorials_utils.get_uci_model(model_class=CustomModel)
model = min_diff.keras.MinDiffModel(model, min_diff.losses.MMDLoss())

model.compile(optimizer='adam', loss='binary_crossentropy')

_ = model.fit(train_with_min_diff_ds.take(1), epochs=1, verbose=0)

# Model has not used the custom train_step.
print('Model used the custom train_step:')
print(hasattr(model, 'used_custom_train_step'))  # False<jupyter_output><empty_output><jupyter_text>In order to use the correct `train_step` method, you need a custom class that subclasses both `MinDiffModel` and `CustomModel`.Note: Make sure to inherit from `MinDiffModel` first. This is important since you need to make sure that certain functions such as `__init__` and `call` still come from `MinDiffModel`.<jupyter_code>class CustomMinDiffModel(min_diff.keras.MinDiffModel, CustomModel):
  pass  # No need for any further implementation.<jupyter_output><empty_output><jupyter_text>Training this model will use the `train_step` from `CustomModel`.<jupyter_code>model = tutorials_utils.get_uci_model(model_class=CustomModel)

model = CustomMinDiffModel(model, min_diff.losses.MMDLoss())

model.compile(optimizer='adam', loss='binary_crossentropy')

_ = model.fit(train_with_min_diff_ds.take(1), epochs=1, verbose=0)

# Model has used the custom train_step.
print('Model used the custom train_step:')
print(hasattr(model, 'used_custom_train_step'))  # True<jupyter_output><empty_output><jupyter_text>Customizing default behaviors of `MinDiffModel`In other cases, you may want to change specific default behaviors of `MinDiffModel`. The most common use case of this is changing the default unpacking behavior to properly handle your data if you don't use `pack_min_diff_data`.When packing the data into a custom format, this might appear as follows.<jupyter_code>def _reformat_input(inputs, original_labels):
  min_diff_data = min_diff.keras.utils.unpack_min_diff_data(inputs)
  original_inputs = min_diff.keras.utils.unpack_original_inputs(inputs)

  return ({
      'min_diff_data': min_diff_data,
      'original_inputs': original_inputs}, original_labels)

customized_train_with_min_diff_ds = train_with_min_diff_ds.map(_reformat_input)<jupyter_output><empty_output><jupyter_text>The `customized_train_with_min_diff_ds` dataset returns batches composed of tuples `(x, y)` where `x` is a dict containing `min_diff_data` and `original_inputs` and `y` is the `original_labels`.<jupyter_code>for x, _ in customized_train_with_min_diff_ds.take(1):
  print('Type of x:', type(x))  # dict
  print('Keys of x:', x.keys())  # 'min_diff_data', 'original_inputs'<jupyter_output><empty_output><jupyter_text>This data format is not what `MinDiffModel` expects by default and passing `customized_train_with_min_diff_ds` to it would result in unexpected behavior. To fix this you will need to create your own subclass.<jupyter_code>class CustomUnpackingMinDiffModel(min_diff.keras.MinDiffModel):

  def unpack_min_diff_data(self, inputs):
    return inputs['min_diff_data']

  def unpack_original_inputs(self, inputs):
    return inputs['original_inputs']<jupyter_output><empty_output><jupyter_text>With this subclass, you can train as with the other examples.<jupyter_code>model = tutorials_utils.get_uci_model()
model = CustomUnpackingMinDiffModel(model, min_diff.losses.MMDLoss())

model.compile(optimizer='adam', loss='binary_crossentropy')

_ = model.fit(customized_train_with_min_diff_ds, epochs=1)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08420186505759737, 'fraction_numerical': 0.0038398244651673065, 'mean_word_length': 6.927173913043478, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 49, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '347283', 'license': 'Apache-2.0', 'n_tokens_mistral': 2248, 'n_tokens_neox': 2186, 'n_words': 660, 'path': 'docs/min_diff/guide/customizing_min_diff_model.ipynb', 'repo_name': 'sarvex/model-remediation', 'size': '7292'}