Repository: keras-team/keras-nlp
Pull request #231 (2022-06-19 09:42:22)

Add Edit Distance Metric
@mattdangerw, @chenmoneygithub, here is the PR for edit distance.

I did a deep-dive on different implementations of WER/CER. I've created a document for the same:
https://docs.google.com/document/d/1IujxboocZbiv_ZAZwfxAG77Aql8O5waUazLrvRxmh-4/edit?usp=sharing

All implementations seem to use basic preprocessing/tokenisers. Please let me if I should implement WER/CER as well. If not, I'll add UTs for Edit Distance. Thanks!

`tf.vectorized_map` did not work. Hence, I've used `tf.map_fn`:

```
>>> ed = EditDistance()
>>> ed(character_splitter(["thisisfun", "bear"]), character_splitter(["yesfun", "beer"]))
...
        File "/home/abheesht/python_envs/keras_nlp/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 2133, in __init__
          self._traceback = tf_stack.extract_stack_for_node(self._c_op)
        File "/home/abheesht/python_envs/keras_nlp/lib/python3.8/site-packages/tensorflow/python/util/tf_stack.py", line 183, in extract_stack_for_node
          return _tf_stack.extract_stack_for_node(
    Node: 'loop_body/RaggedFromVariant/RaggedTensorFromVariant/pfor/TensorListConcatV2'
    PartialTensorShape: Incompatible shapes during merge: [9] vs. [4]
         [[{{node loop_body/RaggedFromVariant/RaggedTensorFromVariant/pfor/TensorListConcatV2}}]] [Op:__inference_f_819]
```



**Edit (https://github.com/keras-team/keras-nlp/pull/231/commits/8d51626d8f2a3f95ee31b380caeac225a570e4f8):**

Trial run and verifying with other open source implementations: https://colab.research.google.com/drive/1fo0jLJogVW-fooO6Nk6A8-dPPppxZxYl?usp=sharing



Patch:
diff --git a/keras_nlp/metrics/__init__.py b/keras_nlp/metrics/__init__.py
index 55ade6dc..36c05c16 100644
--- a/keras_nlp/metrics/__init__.py
+++ b/keras_nlp/metrics/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.metrics.edit_distance import EditDistance
 from keras_nlp.metrics.perplexity import Perplexity
 from keras_nlp.metrics.rouge_l import RougeL
 from keras_nlp.metrics.rouge_n import RougeN
diff --git a/keras_nlp/metrics/edit_distance.py b/keras_nlp/metrics/edit_distance.py
new file mode 100644
index 00000000..1ca5c6be
--- /dev/null
+++ b/keras_nlp/metrics/edit_distance.py
@@ -0,0 +1,224 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# 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.
+
+"""Edit Distance metric."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class EditDistance(keras.metrics.Metric):
+    """Edit Distance metric.
+
+    This class implements the edit distance metric, sometimes called
+    Levenshtein Distance, as a `keras.metrics.Metric`. Essentially, edit
+    distance is the least number of operations required to convert one string to
+    another, where an operation can be one of substitution, deletion or
+    insertion. By default, this metric will compute the normalized score, where
+    the unnormalized edit distance score is divided by the number of tokens in
+    the reference text.
+
+    This class can be used to compute character error rate (CER) and word error
+    rate (WER). You simply have to pass the appropriate tokenized text, and set
+    `normalize` to True.
+
+    Note on input shapes:
+    `y_true` and `y_pred` can either be tensors of rank 1 or ragged tensors of
+    rank 2. These tensors contain tokenized text.
+
+    Args:
+        normalize: bool. If True, the computed number of operations
+            (substitutions + deletions + insertions) across all samples is
+            divided by the aggregate number of tokens in all reference texts. If
+            False, number of operations are calculated for every sample, and
+            averaged over all the samples.
+        dtype: string or tf.dtypes.Dtype. Precision of metric computation. If
+            not specified, it defaults to tf.float32.
+        name: string. Name of the metric instance.
+        **kwargs: Other keyword arguments.
+
+    References:
+        - [Morris et al.](https://www.researchgate.net/publication/221478089)
+
+    Examples:
+
+    Various Input Types.
+
+    Single-level Python list.
+    >>> edit_distance = keras_nlp.metrics.EditDistance()
+    >>> y_true = "the tiny little cat was found under the big funny bed".split()
+    >>> y_pred = "the cat was found under the bed".split()
+    >>> edit_distance(y_true, y_pred)
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.36363637>
+
+    Rank 1 tensor.
+    >>> edit_distance = keras_nlp.metrics.EditDistance()
+    >>> y_true = tf.strings.split("the tiny little cat was found under the big funny bed")
+    >>> y_pred = tf.strings.split("the cat was found under the bed")
+    >>> edit_distance(y_true, y_pred)
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.36363637>
+
+    Nested Python list.
+    >>> edit_distance = keras_nlp.metrics.EditDistance()
+    >>> y_true = [
+    ...     "the tiny little cat was found under the big funny bed".split(),
+    ...     "it is sunny today".split(),
+    ... ]
+    >>> y_pred = [
+    ...     "the cat was found under the bed".split(),
+    ...     "it is sunny but with a hint of cloud cover".split(),
+    ... ]
+    >>> edit_distance(y_true, y_pred)
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.73333335>
+
+    Rank 2 tensor.
+    >>> edit_distance = keras_nlp.metrics.EditDistance()
+    >>> y_true = tf.strings.split([
+    ...     "the tiny little cat was found under the big funny bed",
+    ...     "it is sunny today",
+    ... ])
+    >>> y_pred = tf.strings.split([
+    ...     "the cat was found under the bed",
+    ...     "it is sunny but with a hint of cloud cover",
+    ... ])
+    >>> edit_distance(y_true, y_pred)
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.73333335>
+
+    Pass the metric to `model.compile()`.
+    >>> inputs = keras.Input(shape=(None,), dtype="string")
+    >>> outputs = tf.strings.lower(inputs)
+    >>> model = keras.Model(inputs, outputs)
+    >>> model.compile(metrics=[keras_nlp.metrics.EditDistance()])
+    >>> x = tf.strings.split(
+    ...     tf.constant(
+    ...         ["the tiny little cat was found under the big funny bed"]
+    ...     )
+    ... )
+    >>> y = tf.strings.split(["the cat was found under the bed"])
+    >>> y = tf.strings.split(["the cat was found under the bed"])
+    >>> output = model.evaluate(y, x, return_dict=True)
+    >>> output["edit_distance"]
+    0.3636363744735718
+    """
+
+    def __init__(
+        self,
+        normalize=True,
+        dtype=None,
+        name="edit_distance",
+        **kwargs,
+    ):
+        super().__init__(name=name, dtype=dtype, **kwargs)
+
+        if not tf.as_dtype(self.dtype).is_floating:
+            raise ValueError(
+                "`dtype` must be a floating point type. "
+                f"Received: dtype={dtype}"
+            )
+
+        self.normalize = normalize
+
+        self._aggregate_unnormalized_edit_distance = self.add_weight(
+            name="aggregate_unnormalized_edit_distance",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        if normalize:
+            self._aggregate_reference_length = self.add_weight(
+                name="aggregate_reference_length",
+                initializer="zeros",
+                dtype=self.dtype,
+            )
+        else:
+            self._number_of_samples = self.add_weight(
+                name="number_of_samples",
+                initializer="zeros",
+                dtype=self.dtype,
+            )
+
+    def update_state(self, y_true, y_pred, sample_weight=None):
+        def validate_and_fix_rank(inputs, tensor_name):
+            if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+                inputs = tf.ragged.constant(inputs)
+
+            if inputs.shape.rank == 1:
+                return tf.RaggedTensor.from_tensor(inputs[tf.newaxis])
+            elif inputs.shape.rank == 2:
+                return inputs
+            else:
+                raise ValueError(
+                    f"{tensor_name} must be of rank 1 or 2. "
+                    f"Found rank: {inputs.shape.rank}"
+                )
+
+        y_true = validate_and_fix_rank(y_true, "y_true")
+        y_pred = validate_and_fix_rank(y_pred, "y_pred")
+
+        if self.normalize:
+            self._aggregate_reference_length.assign_add(
+                tf.cast(tf.size(y_true.flat_values), dtype=self.dtype)
+            )
+
+        def calculate_edit_distance(args):
+            reference, hypothesis = args
+
+            reference = tf.sparse.from_dense([reference])
+            hypothesis = tf.sparse.from_dense([hypothesis])
+
+            edit_distance = tf.squeeze(
+                tf.edit_distance(
+                    hypothesis=hypothesis,
+                    truth=reference,
+                    normalize=False,
+                )
+            )
+
+            self._aggregate_unnormalized_edit_distance.assign_add(
+                tf.cast(edit_distance, dtype=self.dtype)
+            )
+            if not self.normalize:
+                self._number_of_samples.assign_add(tf.cast(1, dtype=self.dtype))
+            return 0
+
+        _ = tf.map_fn(
+            fn=calculate_edit_distance,
+            elems=(y_true, y_pred),
+            fn_output_signature=tf.int8,
+        )
+
+    def result(self):
+        if self.normalize:
+            if self._aggregate_reference_length == 0:
+                return 0.0
+            return (
+                self._aggregate_unnormalized_edit_distance
+                / self._aggregate_reference_length
+            )
+        if self._number_of_samples == 0:
+            return 0.0
+        return (
+            self._aggregate_unnormalized_edit_distance / self._number_of_samples
+        )
+
+    def reset_state(self):
+        self._aggregate_unnormalized_edit_distance.assign(0.0)
+        if self.normalize:
+            self._aggregate_reference_length.assign(0.0)
+        else:
+            self._number_of_samples.assign(0.0)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"normalize": self.normalize})
+        return config
diff --git a/keras_nlp/metrics/perplexity.py b/keras_nlp/metrics/perplexity.py
index c252d43b..416b3808 100644
--- a/keras_nlp/metrics/perplexity.py
+++ b/keras_nlp/metrics/perplexity.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Perplexity metric implementation based on `keras.metrics.Metric`."""
+"""Perplexity metric."""
 
 import tensorflow as tf
 from tensorflow import keras
diff --git a/keras_nlp/metrics/rouge_base.py b/keras_nlp/metrics/rouge_base.py
index 22d4adf3..bd9f6aa6 100644
--- a/keras_nlp/metrics/rouge_base.py
+++ b/keras_nlp/metrics/rouge_base.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""ROUGE metric implementation based on `keras.metrics.Metric`."""
+"""ROUGE metric."""
 
 
 import types
diff --git a/keras_nlp/metrics/rouge_l.py b/keras_nlp/metrics/rouge_l.py
index f6969a85..a4c9a258 100644
--- a/keras_nlp/metrics/rouge_l.py
+++ b/keras_nlp/metrics/rouge_l.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""ROUGE-L metric implementation based on `keras.metrics.Metric`."""
+"""ROUGE-L metric."""
 
 
 from keras_nlp.metrics.rouge_base import RougeBase
diff --git a/keras_nlp/metrics/rouge_n.py b/keras_nlp/metrics/rouge_n.py
index 4bfe532e..b7a1522a 100644
--- a/keras_nlp/metrics/rouge_n.py
+++ b/keras_nlp/metrics/rouge_n.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""ROUGE-N metric implementation based on `keras.metrics.Metric`."""
+"""ROUGE-N metric."""
 
 
 from keras_nlp.metrics.rouge_base import RougeBase


Test patch:
diff --git a/keras_nlp/metrics/edit_distance_test.py b/keras_nlp/metrics/edit_distance_test.py
new file mode 100644
index 00000000..80f93168
--- /dev/null
+++ b/keras_nlp/metrics/edit_distance_test.py
@@ -0,0 +1,300 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# 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.
+
+"""Tests for EditDistance."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.metrics import EditDistance
+
+
+class EditDistanceTest(tf.test.TestCase):
+    def test_initialization(self):
+        edit_distance = EditDistance()
+        result = edit_distance.result()
+
+        self.assertEqual(result, 0.0)
+
+    def test_1d_list_input_normalize(self):
+        edit_distance = EditDistance()
+        y_true = "the tiny little cat was found under the big funny bed".split()
+        y_pred = "the cat was found under the bed".split()
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.364, delta=1e-3)
+
+    def test_2d_list_input_normalize(self):
+        edit_distance = EditDistance()
+        y_true = [
+            "the tiny little cat was found under the big funny bed".split(),
+            "it is sunny today".split(),
+        ]
+        y_pred = [
+            "the cat was found under the bed".split(),
+            "it is sunny but with a hint of cloud cover".split(),
+        ]
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.733, delta=1e-3)
+
+    def test_rank_1_tensor_input_normalize(self):
+        edit_distance = EditDistance()
+        y_true = tf.strings.split(
+            "the tiny little cat was found under the big funny bed"
+        )
+        y_pred = tf.strings.split("the cat was found under the bed")
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.364, delta=1e-3)
+
+    def test_rank_2_tensor_input_normalize(self):
+        edit_distance = EditDistance()
+        y_true = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.733, delta=1e-3)
+
+    def test_rank_1_tensor_input_normalize_false(self):
+        edit_distance = EditDistance(normalize=False)
+        y_true = tf.strings.split(
+            "the tiny little cat was found under the big funny bed"
+        )
+        y_pred = tf.strings.split("the cat was found under the bed")
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val.numpy(), 4.0, delta=1e-3)
+
+    def test_rank_2_tensor_input_normalize_false(self):
+        edit_distance = EditDistance(normalize=False)
+        y_true = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        edit_distance_val = edit_distance(y_true, y_pred)
+        self.assertAlmostEqual(edit_distance_val.numpy(), 5.5, delta=1e-3)
+
+    def test_model_compile_normalize(self):
+        inputs = keras.Input(shape=(None,), dtype="string")
+        outputs = tf.strings.lower(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[EditDistance()])
+
+        x = tf.strings.split(
+            ["the tiny little cat was found under the big funny bed"]
+        )
+        y = tf.strings.split(["the cat was found under the bed"])
+
+        output = model.evaluate(y, x, return_dict=True)
+
+        self.assertAlmostEqual(output["edit_distance"], 0.364, delta=1e-3)
+
+    def test_model_compile_normalize_false(self):
+        inputs = keras.Input(shape=(None,), dtype="string")
+        outputs = tf.strings.lower(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[EditDistance(normalize=False)])
+
+        x = tf.strings.split(
+            ["the tiny little cat was found under the big funny bed"]
+        )
+        y = tf.strings.split(["the cat was found under the bed"])
+
+        output = model.evaluate(y, x, return_dict=True)
+
+        self.assertAlmostEqual(output["edit_distance"], 4.0, delta=1e-3)
+
+    def test_reset_state_normalize(self):
+        edit_distance = EditDistance()
+        y_true = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        edit_distance.update_state(y_true, y_pred)
+        edit_distance_val = edit_distance.result()
+        self.assertNotEqual(edit_distance_val.numpy(), 0.0)
+
+        edit_distance.reset_state()
+        edit_distance_val = edit_distance.result()
+        self.assertEqual(edit_distance_val, 0.0)
+
+    def test_update_state_normalize(self):
+        edit_distance = EditDistance()
+        y_true_1 = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred_1 = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        edit_distance.update_state(y_true_1, y_pred_1)
+        edit_distance_val = edit_distance.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.733, delta=1e-3)
+
+        y_true_2 = tf.strings.split(["what is your favourite show"])
+        y_pred_2 = tf.strings.split(["my favourite show is silicon valley"])
+
+        edit_distance.update_state(y_true_2, y_pred_2)
+        edit_distance_val = edit_distance.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.85, delta=1e-3)
+
+    def test_update_state_normalize_false(self):
+        edit_distance = EditDistance(normalize=False)
+        y_true_1 = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred_1 = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        edit_distance.update_state(y_true_1, y_pred_1)
+        edit_distance_val = edit_distance.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 5.5, delta=1e-3)
+
+        y_true_2 = tf.strings.split(["what is your favourite show"])
+        y_pred_2 = tf.strings.split(["my favourite show is silicon valley"])
+
+        edit_distance.update_state(y_true_2, y_pred_2)
+        edit_distance_val = edit_distance.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 5.667, delta=1e-3)
+
+    def test_merge_state_normalize(self):
+        edit_distance_1 = EditDistance()
+        edit_distance_2 = EditDistance()
+
+        y_true_1 = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred_1 = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        y_true_2 = tf.strings.split(["what is your favourite show"])
+        y_pred_2 = tf.strings.split(["my favourite show is silicon valley"])
+
+        y_true_3 = tf.strings.split(["lorem ipsum dolor sit amet"])
+        y_pred_3 = tf.strings.split(["lorem ipsum is simply dummy text"])
+
+        edit_distance_1.update_state(y_true_1, y_pred_1)
+        edit_distance_1.update_state(y_true_2, y_pred_2)
+        edit_distance_val = edit_distance_1.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.85, delta=1e-3)
+
+        edit_distance_2.update_state(y_true_3, y_pred_3)
+        edit_distance_val = edit_distance_2.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.8, delta=1e-3)
+
+        merged_edit_distance = EditDistance()
+        merged_edit_distance.merge_state([edit_distance_1, edit_distance_2])
+        edit_distance_val = merged_edit_distance.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 0.84, delta=1e-3)
+
+    def test_merge_state_normalize_false(self):
+        edit_distance_1 = EditDistance(normalize=False)
+        edit_distance_2 = EditDistance(normalize=False)
+
+        y_true_1 = tf.strings.split(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "it is sunny today",
+            ]
+        )
+        y_pred_1 = tf.strings.split(
+            [
+                "the cat was found under the bed",
+                "it is sunny but with a hint of cloud cover",
+            ]
+        )
+
+        y_true_2 = tf.strings.split(["what is your favourite show"])
+        y_pred_2 = tf.strings.split(["my favourite show is silicon valley"])
+
+        y_true_3 = tf.strings.split(["lorem ipsum dolor sit amet"])
+        y_pred_3 = tf.strings.split(["lorem ipsum is simply dummy text"])
+
+        edit_distance_1.update_state(y_true_1, y_pred_1)
+        edit_distance_1.update_state(y_true_2, y_pred_2)
+        edit_distance_val = edit_distance_1.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 5.667, delta=1e-3)
+
+        edit_distance_2.update_state(y_true_3, y_pred_3)
+        edit_distance_val = edit_distance_2.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 4.0, delta=1e-3)
+
+        merged_edit_distance = EditDistance(normalize=False)
+        merged_edit_distance.merge_state([edit_distance_1, edit_distance_2])
+        edit_distance_val = merged_edit_distance.result()
+        self.assertAlmostEqual(edit_distance_val.numpy(), 5.25, delta=1e-3)
+
+    def test_get_config(self):
+        rouge = EditDistance(
+            normalize=False,
+            dtype=tf.float32,
+            name="edit_distance_test",
+        )
+
+        config = rouge.get_config()
+        expected_config_subset = {
+            "normalize": False,
+        }
+        self.assertEqual(config, {**config, **expected_config_subset})