Repository: keras-team/keras-nlp

Issue:
Add a perplexity metric
Splitting this issue out from https://github.com/keras-team/keras-nlp/issues/38.

We should add a perplexity metrics as `keras_nlp.metrics.Perplexity`.


Discussion:
huggingface resource: https://huggingface.co/docs/transformers/perplexity

@abheesht17's colab: http://go/colabx-drive/1BH1lTw_qLK6671oWaoU15IrUKSQfd6oE
Looking through the colab, I think method 1 would be the correct approach? Over method 2. Perplexity seems to me to be defined on a single input sequence, so averaging over all sequences in the batch (and then all batches), seems reasonable.

I think @chenmoneygithub was going to take a look here too, to tagging for thoughts.
Hey, I can't access the notebook link (http://go/colabx-drive/1BH1lTw_qLK6671oWaoU15IrUKSQfd6oE). So, putting the link here again: https://colab.research.google.com/drive/1BH1lTw_qLK6671oWaoU15IrUKSQfd6oE?usp=sharing


Agree method1 looks correct. 

One notes about masking: checking `token_id == 0` is not a reliable way to generate mask, because users could customize their mask token id (although I do not know why people do so). The mask should be passed via `sample_weights` argument.
Ah, nice 👍🏼. I'll make the necessary changes and open a PR.
Hello, @mattdangerw, @chenmoneygithub! I've made some changes to the class. Please see this notebook: https://colab.research.google.com/drive/1XV1h5aeiy5IlHoQFjDTJ45hRC8wMSf16?usp=sharing. 

I've followed this script: https://github.com/huggingface/transformers/blob/main/examples/research_projects/codeparrot/scripts/validation_loss.py#L56-L69.

In the notebook, I've compared our results with HF's results. The perplexity scores returned by both are very close to each other!

I'll open a PR for this now :D

Resolving patch:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from keras_nlp import layers
+from keras_nlp import metrics
 from keras_nlp import tokenizers
 
 __version__ = "0.1.1"
diff --git a/keras_nlp/metrics/__init__.py b/keras_nlp/metrics/__init__.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/metrics/__init__.py
@@ -0,0 +1,15 @@
+# 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.
+
+from keras_nlp.metrics.perplexity import Perplexity
diff --git a/keras_nlp/metrics/perplexity.py b/keras_nlp/metrics/perplexity.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/metrics/perplexity.py
@@ -0,0 +1,181 @@
+# 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.
+
+"""Perplexity metric implementation based on `keras.metrics.Metric`."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class Perplexity(keras.metrics.Metric):
+    """Perplexity metric.
+
+    This class implements the perplexity metric. In short, this class calculates
+    the cross entropy loss and takes its exponent.
+    Note: This implementation is not suitable for fixed-size windows.
+
+    Args:
+        from_logits: bool. If True, `y_pred` (input to `update_state()`) should
+            be the logits as returned by the model. Otherwise, `y_pred` is a
+            tensor of probabilities.
+        mask_token_id: int. ID of the token to be masked. If provided, the mask
+            is computed for this class. Note that if this field is provided, and
+            if the `sample_weight` field in `update_state()` is also provided,
+            we will compute the final `sample_weight` as the element-wise
+            product of the mask and the `sample_weight`.
+        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.
+
+    Examples:
+
+    1. Calculate perplexity by calling update_state() and result().
+    1.1. `sample_weight`, and `mask_token_id` are not provided.
+    >>> tf.random.set_seed(42)
+    >>> perplexity = keras_nlp.metrics.Perplexity(name="perplexity")
+    >>> target = tf.random.uniform(
+    ...     shape=[2, 5],  maxval=10, dtype=tf.int32, seed=42)
+    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> perplexity.update_state(target, logits)
+    >>> perplexity.result()
+    <tf.Tensor: shape=(), dtype=float32, numpy=11.8781595>
+
+    1.2. `sample_weight` specified (masking token with ID 0).
+    >>> tf.random.set_seed(42)
+    >>> perplexity = keras_nlp.metrics.Perplexity(name="perplexity")
+    >>> target = tf.random.uniform(
+    ...     shape=[2, 5],  maxval=10, dtype=tf.int32, seed=42)
+    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> sample_weight = tf.cast(
+    ...     tf.math.logical_not(tf.equal(target, 0)), tf.float32)
+    >>> perplexity.update_state(target, logits, sample_weight)
+    >>> perplexity.result()
+    <tf.Tensor: shape=(), dtype=float32, numpy=13.1128>
+
+    2. Call perplexity directly.
+    >>> tf.random.set_seed(42)
+    >>> perplexity = keras_nlp.metrics.Perplexity(name="perplexity")
+    >>> target = tf.random.uniform(
+    ...     shape=[2, 5],  maxval=10, dtype=tf.int32, seed=42)
+    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> perplexity(target, logits)
+    <tf.Tensor: shape=(), dtype=float32, numpy=11.8781595>
+
+    3. Provide the padding token ID and let the class compute the mask on its
+       own.
+    >>> tf.random.set_seed(42)
+    >>> perplexity = keras_nlp.metrics.Perplexity(
+    ...     name="perplexity", mask_token_id=0)
+    >>> target = tf.random.uniform(
+    ...     shape=[2, 5],  maxval=10, dtype=tf.int32, seed=42)
+    >>> logits = tf.random.uniform(shape=(2, 5, 10), seed=42)
+    >>> perplexity(target, logits)
+    <tf.Tensor: shape=(), dtype=float32, numpy=13.1128>
+    """
+
+    def __init__(
+        self,
+        from_logits=False,
+        mask_token_id=None,
+        dtype=None,
+        name="perplexity",
+        **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._crossentropy = keras.losses.SparseCategoricalCrossentropy(
+            from_logits=from_logits, reduction="sum"
+        )
+
+        self.from_logits = from_logits
+        self.mask_token_id = mask_token_id
+
+        self._aggregate_crossentropy = self.add_weight(
+            name="aggregate_crossentropy",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        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):
+        # y_true shape: (batch_size, seq_len)
+        # y_pred shape: (batch_size, seq_len, vocab_size)
+        y_true = tf.cast(y_true, self.dtype)
+        y_pred = tf.cast(y_pred, self.dtype)
+
+        if sample_weight is not None:
+            sample_weight = tf.cast(sample_weight, self.dtype)
+
+        batch_size = tf.cast(tf.shape(y_true)[0], self.dtype)
+
+        if self.mask_token_id is not None:
+            mask = tf.cast(
+                tf.math.logical_not(tf.equal(y_true, self.mask_token_id)),
+                self.dtype,
+            )
+            if sample_weight is None:
+                sample_weight = mask
+            else:
+                sample_weight = tf.multiply(mask, sample_weight)
+
+        # Calculate the Cross Entropy Loss.
+        crossentropy_value = tf.cast(
+            self._crossentropy(y_true, y_pred, sample_weight=sample_weight),
+            self.dtype,
+        )  # scalar
+
+        # Divide the loss by the number of non-masked tokens
+        if sample_weight is not None:
+            crossentropy_value = crossentropy_value / tf.reduce_sum(
+                sample_weight
+            )  # scalar
+        else:
+            crossentropy_value = crossentropy_value / (
+                tf.cast(tf.shape(y_true)[0], self.dtype)
+                * tf.cast(tf.shape(y_true)[1], self.dtype)
+            )  # scalar
+
+        self._aggregate_crossentropy.assign_add(batch_size * crossentropy_value)
+        self._number_of_samples.assign_add(batch_size)
+
+    def result(self):
+        if self._number_of_samples == 0:
+            return 0.0
+        perplexity_score = tf.exp(
+            self._aggregate_crossentropy / self._number_of_samples
+        )
+        return perplexity_score
+
+    def reset_state(self):
+        self._aggregate_crossentropy.assign(0.0)
+        self._number_of_samples.assign(0.0)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "from_logits": self.from_logits,
+                "mask_token_id": self.mask_token_id,
+            }
+        )
+        return config


Test patch:
diff --git a/keras_nlp/metrics/perplexity_test.py b/keras_nlp/metrics/perplexity_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/metrics/perplexity_test.py
@@ -0,0 +1,355 @@
+# 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 Perplexity."""
+
+import tensorflow as tf
+
+from keras_nlp.metrics import Perplexity
+
+
+class PerplexityTest(tf.test.TestCase):
+    def test_vars_after_initializing_class(self):
+        perplexity = Perplexity()
+        self.assertEqual(perplexity.result().numpy(), 0.0)
+
+    def test_from_logits_without_masking(self):
+        perplexity = Perplexity(from_logits=True)
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+
+        perplexity_val = perplexity(y_true, y_pred)
+        self.assertAlmostEqual(perplexity_val.numpy(), 2.6542, delta=1e-3)
+
+    def test_from_logits_with_sample_weight(self):
+        perplexity = Perplexity(from_logits=True)
+
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+        sample_wt = tf.cast(y_true != 0, tf.int32)
+
+        perplexity_val = perplexity(y_true, y_pred, sample_wt)
+        self.assertAlmostEqual(perplexity_val.numpy(), 2.8789, delta=1e-3)
+
+    def test_from_logits_with_mask_token_id(self):
+        perplexity = Perplexity(from_logits=True, mask_token_id=0)
+
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+
+        perplexity_val = perplexity(y_true, y_pred)
+        self.assertAlmostEqual(perplexity_val.numpy(), 2.8789, delta=1e-3)
+
+    def test_from_logits_with_mask_token_id_and_sample_weight(self):
+        perplexity = Perplexity(from_logits=True, mask_token_id=0)
+
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+        sample_weight = tf.constant([[0.5, 0.1, 0.9], [1, 0.7, 0.5]])
+
+        perplexity_val = perplexity(y_true, y_pred, sample_weight)
+        self.assertAlmostEqual(perplexity_val.numpy(), 2.9442, delta=1e-3)
+
+    def test_two_inputs_from_logits(self):
+        perplexity = Perplexity(from_logits=True, mask_token_id=0)
+
+        y_true_1 = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred_1 = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+
+        perplexity_val = perplexity(y_true_1, y_pred_1)
+        self.assertAlmostEqual(perplexity_val.numpy(), 2.8789, delta=1e-3)
+
+        y_true_2 = tf.constant([[2, 0, 0], [1, 2, 3]])
+        y_pred_2 = tf.constant(
+            [
+                [
+                    [2.887, 0.885, 2.973, 2.582],
+                    [0.3838, 2.629, 1.91, 1.802],
+                    [0.2578, 1.081, 1.125, 2.773],
+                ],
+                [
+                    [1.623, 2.784, 0.2109, 2.66],
+                    [2.395, 2.01, 0.252, 1.828],
+                    [0.4482, 2.629, 0.9697, 0.998],
+                ],
+            ]
+        )
+        perplexity_val = perplexity(y_true_2, y_pred_2)
+        self.assertAlmostEqual(perplexity_val.numpy(), 3.9998, delta=1e-3)
+
+    def test_from_probs_with_sample_weight(self):
+        perplexity = Perplexity(from_logits=False)
+
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+        y_prob = tf.nn.softmax(y_pred, axis=-1)
+
+        sample_wt = tf.cast(y_true != 0, tf.int32)
+
+        perplexity_val = perplexity(y_true, y_prob, sample_wt)
+        self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
+
+    def test_from_probs_with_pad_token(self):
+        perplexity = Perplexity(from_logits=False, mask_token_id=0)
+
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+        y_prob = tf.nn.softmax(y_pred, axis=-1)
+
+        perplexity_val = perplexity(y_true, y_prob)
+        self.assertAlmostEqual(perplexity_val, 2.8789, delta=1e-3)
+
+    def test_reset_state(self):
+        y_true = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+
+        perplexity = Perplexity(from_logits=True, mask_token_id=0)
+
+        perplexity.update_state(y_true, y_pred)
+        self.assertNotEqual(perplexity.result(), 0.0)
+
+        perplexity.reset_state()
+        self.assertEqual(perplexity.result(), 0.0)
+
+    def test_update_state(self):
+        perplexity = Perplexity(from_logits=True, mask_token_id=0)
+
+        y_true_1 = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred_1 = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+
+        perplexity.update_state(y_true_1, y_pred_1)
+        perplexity_val = perplexity.result()
+        self.assertAlmostEqual(perplexity_val.numpy(), 2.8789, delta=1e-3)
+
+        y_true_2 = tf.constant([[2, 0, 0], [1, 2, 3]])
+        y_pred_2 = tf.constant(
+            [
+                [
+                    [2.887, 0.885, 2.973, 2.582],
+                    [0.3838, 2.629, 1.91, 1.802],
+                    [0.2578, 1.081, 1.125, 2.773],
+                ],
+                [
+                    [1.623, 2.784, 0.2109, 2.66],
+                    [2.395, 2.01, 0.252, 1.828],
+                    [0.4482, 2.629, 0.9697, 0.998],
+                ],
+            ]
+        )
+
+        perplexity.update_state(y_true_2, y_pred_2)
+        perplexity_val = perplexity.result()
+        self.assertAlmostEqual(perplexity_val, 3.9998, delta=1e-3)
+
+    def test_merge_state(self):
+        perplexity_1 = Perplexity(from_logits=True, mask_token_id=0)
+        perplexity_2 = Perplexity(from_logits=True, mask_token_id=0)
+
+        y_true_1 = tf.constant([[1, 3, 0], [2, 1, 3]])
+        y_pred_1 = tf.constant(
+            [
+                [
+                    [1.034, 4.797, 2.82, 1.154],
+                    [2.258, 1.591, 1.811, 1.852],
+                    [3.216, 1.037, 0.3662, 2.7],
+                ],
+                [
+                    [1.363, 1.726, 1.898, 2.582],
+                    [1.163, 1.943, 1.761, 1.497],
+                    [2.766, 1.453, 2.61, 2.805],
+                ],
+            ]
+        )
+
+        y_true_2 = tf.constant([[2, 0, 0], [1, 2, 3]])
+        y_pred_2 = tf.constant(
+            [
+                [
+                    [2.887, 0.885, 2.973, 2.582],
+                    [0.3838, 2.629, 1.91, 1.802],
+                    [0.2578, 1.081, 1.125, 2.773],
+                ],
+                [
+                    [1.623, 2.784, 0.2109, 2.66],
+                    [2.395, 2.01, 0.252, 1.828],
+                    [0.4482, 2.629, 0.9697, 0.998],
+                ],
+            ]
+        )
+
+        y_true_3 = tf.constant([[1, 3, 1], [2, 2, 1]])
+        y_pred_3 = tf.constant(
+            [
+                [
+                    [0.7383, 0.882, 0.7295, 2.64],
+                    [0.867, 1.588, 2.291, 0.967],
+                    [0.0908, 1.453, 1.5, 0.0703],
+                ],
+                [
+                    [2.783, 0.9785, 2.664, 0.507],
+                    [0.741, 1.535, 2.16, 2.531],
+                    [2.863, 1.591, 1.403, 0.885],
+                ],
+            ]
+        )
+
+        perplexity_1.update_state(y_true_1, y_pred_1)
+        perplexity_1.update_state(y_true_2, y_pred_2)
+        self.assertAlmostEqual(
+            perplexity_1.result().numpy(), 3.9998, delta=1e-3
+        )
+
+        perplexity_2.update_state(y_true_3, y_pred_3)
+        self.assertAlmostEqual(
+            perplexity_2.result().numpy(), 4.4303, delta=1e-3
+        )
+
+        merged_perplexity = Perplexity(from_logits=True, mask_token_id=0)
+        merged_perplexity.merge_state([perplexity_1, perplexity_2])
+        self.assertAlmostEqual(
+            merged_perplexity.result().numpy(), 4.1385, delta=1e-3
+        )
+
+    def test_get_config(self):
+        perplexity = Perplexity(
+            from_logits=True,
+            mask_token_id=0,
+            dtype=tf.float32,
+            name="perplexity_test",
+        )
+        config = perplexity.get_config()
+        expected_config = {
+            "from_logits": True,
+            "mask_token_id": 0,
+            "dtype": tf.float32,
+            "name": "perplexity_test",
+        }
+        self.assertEqual(config, expected_config)