Repository: tensorflow/addons

Issue:
Unable to restore model using crf_decode
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Cataline 10.15.3
- TensorFlow version and how it was installed (source or binary): 2.2.0 (from pypi)
- TensorFlow-Addons version and how it was installed (source or binary): 0.7.0 (from pypi)
- Python version: 3.7.4
- Is GPU used? (yes/no): no

**Describe the bug**
When trying to load a model with `tf.keras.models.load_model` which uses `tfa.text.crf.crf_decode` it fails on restoring the custom object. This seems to occur as it can't build the object because the init method takes extra arguments. Below there is a simple (standalone) example that fails and the traceback of the error.

**Code to reproduce the issue**

    import numpy as np
    import tensorflow as tf
    import tensorflow_addons as tfa


    class CRFLoss:
        def __init__(self, transition):
            self.transition = transition
            self.__name__ = "CRFLoss"

        def __call__(self, true, pred):
            log_likelihood, _ = tfa.text.crf.crf_log_likelihood(
                inputs=pred,
                tag_indices=true,
                sequence_lengths=tf.ones(tf.shape(true)[0]) * 10,
                transition_params=self.transition,
            )
            return tf.reduce_mean(-log_likelihood)

        def get_config(self):
            """Return helper config."""
            config = {"transition": self.transition.numpy().tolist()}
            return config

        @classmethod
        def from_config(cls, config):
            config["transition"] = tf.constant(config["transition"], dtype=tf.float32)
            return cls(**config)


    def create_model():
        input_tensor = tf.keras.Input(shape=(10, 3), dtype=tf.float32, name="input_tensor")
        seq_len = tf.keras.Input(shape=(), dtype=tf.int32, name="seq_len")
        transition = tf.constant([[1, 1, 0], [0, 1, 1], [1, 0, 1]], dtype=tf.float32)

        output = tf.multiply(input_tensor, tf.constant(1.0))
        decoded, _ = tfa.text.crf.crf_decode(input_tensor, transition, seq_len)

        loss = CRFLoss(transition)

        model = tf.keras.Model(inputs=[input_tensor, seq_len], outputs=[output, decoded], name="example_model")
        model.compile(optimizer="Adam", loss={"tf_op_layer_Mul": loss})
        return model

    x_data = {
        "input_tensor": np.random.random_sample((5, 10, 3)).astype(dtype=np.float32),
        "seq_len": np.array([10] * 5, dtype=np.int32),
    }
    y_data = {"tf_op_layer_Mul": np.random.randint(0, 3, (5, 10))}

    model = create_model()
    model.fit(x_data, y_data)
    model.predict({"input_tensor": tf.expand_dims(x_data["input_tensor"][0], 0), "seq_len": np.array([10])})

    tf.saved_model.save(model, "./example_model/")

    tf.keras.backend.clear_session()
    model = tf.keras.models.load_model(
        "./example_model/",
        custom_objects={"CrfDecodeForwardRnnCell": tfa.text.crf.CrfDecodeForwardRnnCell, "CRFLoss": CRFLoss}
    )
    model.fit(x_data, y_data)
    model.predict({"input_tensor": tf.expand_dims(x_data["input_tensor"][0], 0), "seq_len": np.array([10])})


**Other info / logs**
```
Traceback (most recent call last):
  File "addon_issue_snippet.py", line 60, in <module>
    custom_objects={"CrfDecodeForwardRnnCell": tfa.text.crf.CrfDecodeForwardRnnCell, "CRFLoss": CRFLoss}
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 190, in load_model
    return saved_model_load.load(filepath, compile)
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 116, in load
    model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
  File "venv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 604, in load_internal
    export_dir)
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 188, in __init__
    super(KerasObjectLoader, self).__init__(*args, **kwargs)
  File "venv/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 123, in __init__
    self._load_all()
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 209, in _load_all
    self._layer_nodes = self._load_layers()
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 309, in _load_layers
    layers[node_id] = self._load_layer(proto.user_object, node_id)
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 337, in _load_layer
    obj, setter = revive_custom_object(proto.identifier, metadata)
  File "venv/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 782, in revive_custom_object
    .format(identifier))
ValueError: Unable to restore custom object of type _tf_keras_rnn_layer currently. Please make sure that the layer implements `get_config`and `from_config` when saving. In addition, please use the `custom_objects` arg when calling `load_model()`.
```

**Possible solution**
I've tried to debug this a bit and the solution I've found to work is to implement `get_config` and `from_config` on `tfa.text.crf.CrfDecodeForwardRnnCell`.
For now I'm patching those methods in the object but it isn't very clean. If the proposed solution is good I'd gladly add the methods and create a PR.


Discussion:
Hi @AntPeixe, your solution is correct IMO. Feel free to open an PR and request my review. Thank you!

Resolving patch:
diff --git a/tensorflow_addons/text/crf.py b/tensorflow_addons/text/crf.py
--- a/tensorflow_addons/text/crf.py
+++ b/tensorflow_addons/text/crf.py
@@ -422,6 +422,20 @@ def call(self, inputs, state):
         backpointers = tf.cast(backpointers, dtype=tf.int32)
         return backpointers, new_state
 
+    def get_config(self) -> dict:
+        config = {
+            "transition_params": tf.squeeze(self._transition_params, 0).numpy().tolist()
+        }
+        base_config = super(CrfDecodeForwardRnnCell, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    @classmethod
+    def from_config(cls, config: dict) -> "CrfDecodeForwardRnnCell":
+        config["transition_params"] = np.array(
+            config["transition_params"], dtype=np.float32
+        )
+        return cls(**config)
+
 
 def crf_decode_forward(
     inputs: TensorLike,


Test patch:
diff --git a/tensorflow_addons/text/tests/crf_test.py b/tensorflow_addons/text/tests/crf_test.py
--- a/tensorflow_addons/text/tests/crf_test.py
+++ b/tensorflow_addons/text/tests/crf_test.py
@@ -365,3 +365,48 @@ def test_tf_function():
     transition_params = tf.random.uniform([num_tags, num_tags])
     sequence_length = tf.ones([batch_size], dtype=tf.int32)
     crf_decode(potentials, transition_params, sequence_length)
+
+
+def test_crf_decode_save_load(tmpdir):
+    tf.keras.backend.clear_session()
+    input_tensor = tf.keras.Input(shape=(10, 3), dtype=tf.float32, name="input_tensor")
+    seq_len = tf.keras.Input(shape=(), dtype=tf.int32, name="seq_len")
+    transition = tf.constant([[1, 1, 0], [0, 1, 1], [1, 0, 1]], dtype=tf.float32)
+
+    output = tf.multiply(input_tensor, tf.constant(1.0))
+    decoded, _ = text.crf_decode(input_tensor, transition, seq_len)
+
+    model = tf.keras.Model(
+        inputs=[input_tensor, seq_len], outputs=[output, decoded], name="example_model"
+    )
+    model.compile(optimizer="Adam")
+
+    x_data = {
+        "input_tensor": np.random.random_sample((5, 10, 3)).astype(dtype=np.float32),
+        "seq_len": np.array([10] * 5, dtype=np.int32),
+    }
+    y_data = {"tf_op_layer_Mul": np.random.randint(0, 3, (5, 10))}
+
+    model.fit(x_data, y_data)
+    model.predict(
+        {
+            "input_tensor": tf.expand_dims(x_data["input_tensor"][0], 0),
+            "seq_len": np.array([10]),
+        }
+    )
+
+    temp_dir = str(tmpdir.mkdir("model"))
+    tf.saved_model.save(model, temp_dir)
+
+    tf.keras.backend.clear_session()
+    model = tf.keras.models.load_model(
+        temp_dir,
+        custom_objects={"CrfDecodeForwardRnnCell": text.crf.CrfDecodeForwardRnnCell},
+    )
+    model.fit(x_data, y_data)
+    model.predict(
+        {
+            "input_tensor": tf.expand_dims(x_data["input_tensor"][0], 0),
+            "seq_len": np.array([10]),
+        }
+    )