Repository: keras-team/keras-nlp
Pull request #88 (2022-04-05 16:54:34)

Add a layer for packing inputs for BERT-likes
Fixes https://github.com/keras-team/keras-nlp/issues/56


Patch:
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index 4c2fd2cf..a4cc2508 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -15,6 +15,7 @@
 from keras_nlp.layers.fnet_encoder import FNetEncoder
 from keras_nlp.layers.mlm_head import MLMHead
 from keras_nlp.layers.mlm_mask_generator import MLMMaskGenerator
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
 from keras_nlp.layers.start_end_packer import StartEndPacker
diff --git a/keras_nlp/layers/multi_segment_packer.py b/keras_nlp/layers/multi_segment_packer.py
new file mode 100644
index 00000000..a6cb001d
--- /dev/null
+++ b/keras_nlp/layers/multi_segment_packer.py
@@ -0,0 +1,223 @@
+# 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.
+
+"""BERT token packing layer."""
+
+import tensorflow as tf
+import tensorflow_text as tf_text
+from tensorflow import keras
+
+
+class MultiSegmentPacker(keras.layers.Layer):
+    """Packs multiple sequences into a single fixed width model input.
+
+    This layer packs multiple input sequences into a single fixed width sequence
+    containing start and end delimeters, forming an dense input suitable for a
+    classification task for BERT and BERT-like models.
+
+    Takes as input a list or tuple of token segments. The layer will process
+    inputs as follows:
+     - Truncate all input segments to fit within `sequence_length` according to
+       the `truncator` strategy.
+     - Concatenate all input segments, adding a single `start_value` at the
+       start of the entire sequence, and multiple `end_value`s at the end of
+       each segment.
+     - Pad the resulting sequence to `sequence_length` using `pad_tokens`.
+     - Calculate a separate tensor of "segment ids", with integer type and the
+       same shape as the packed token output, where each integer index of the
+       segment the token originated from. The segment id of the `start_value`
+       is always 0, and the segment id of each `end_value` is the segment that
+       precedes it.
+
+    If inputs are batched, inputs should be `tf.RaggedTensor`s with shape
+    `[batch_size, None]` and will be packed and converted to a dense tensor with
+    shape `[batch_size, sequence_length]`.
+
+    If inputs are unbatched, inputs should be dense rank-1 tensors of any shape,
+    and will be packed to shape `[sequence_length]`.
+
+    Args:
+        sequence_length: The desired output length.
+        start_value: The id or token that is to be placed at the start of each
+            sequence (called "[CLS]" for BERT). The dtype must mach the dtype of
+            the input tensors to the layer.
+        end_value: The id or token that is to be placed at the end of each
+            input segment (called "[SEP]" for BERT). The dtype much mach the
+            dtype of the input tensors to the layer.
+        pad_value: The id or token that is to be placed into the unused
+            positions after the last segment in the sequence
+            (called "[PAD]" for BERT).
+        truncator: The algorithm to truncate a list of batched segments to fit a
+            per-example length limit. The value can be either `round_robin` or
+            `waterfall`:
+                - `"round_robin"`: Available space is assigned one token at a
+                    time in a round-robin fashion to the inputs that still need
+                    some, until the limit is reached.
+                - `"waterfall"`: The allocation of the budget is done using a
+                    "waterfall" algorithm that allocates quota in a
+                    left-to-right manner and fills up the buckets until we run
+                    out of budget. It support arbitrary number of segments.
+
+    Returns:
+        A tuple with two elements. The first is the dense, packed token
+        sequence. The second is an integer tensor of the same shape, containing
+        the segment ids.
+
+    Examples:
+
+    *Pack a single input for classification.*
+    >>> seq1 = tf.constant([1, 2, 3, 4])
+    >>> packer = keras_nlp.layers.MultiSegmentPacker(
+    ...     8, start_value=101, end_value=102)
+    >>> packer(seq1)
+    (<tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([101, 1, 2, 3, 4, 102, 0, 0], dtype=int32)>,
+     <tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)>)
+
+    *Pack multiple inputs for classification.*
+    >>> seq1 = tf.constant([1, 2, 3, 4])
+    >>> seq2 = tf.constant([11, 12, 13, 14])
+    >>> packer = keras_nlp.layers.MultiSegmentPacker(
+    ...     8, start_value=101, end_value=102)
+    >>> packer((seq1, seq2))
+    (<tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([101,   1,   2,   3, 102,  11,  12, 102], dtype=int32)>,
+     <tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([0, 0, 0, 0, 0, 1, 1, 1], dtype=int32)>)
+
+    Reference:
+        [Devlin et al., 2018](https://arxiv.org/abs/1810.04805).
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        start_value,
+        end_value,
+        pad_value=None,
+        truncator="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.sequence_length = sequence_length
+        if truncator not in ("round_robin", "waterfall"):
+            raise ValueError(
+                "Only 'round_robin' and 'waterfall' algorithms are "
+                "supported. Received %s" % truncator
+            )
+        self.truncator = truncator
+        self.start_value = start_value
+        self.end_value = end_value
+        self.pad_value = pad_value
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "start_value": self.start_value,
+                "end_value": self.end_value,
+                "pad_value": self.pad_value,
+                "truncator": self.truncator,
+            }
+        )
+        return config
+
+    def _sanitize_inputs(self, inputs):
+        """Force inputs to a list of rank 2 ragged tensors."""
+        # Sanitize inputs.
+        if not isinstance(inputs, (list, tuple)):
+            inputs = [inputs]
+        if not inputs:
+            raise ValueError("At least one input is required for packing")
+        input_ranks = [x.shape.rank for x in inputs]
+        if not all(0 < rank < 3 for rank in input_ranks):
+            raise ValueError(
+                "All inputs for packing must have rank 1 or 2. "
+                f"Received input ranks: {input_ranks}"
+            )
+        if None in input_ranks or len(set(input_ranks)) > 1:
+            raise ValueError(
+                "All inputs for packing must have the same rank. "
+                f"Received input ranks: {input_ranks}"
+            )
+        return inputs
+
+    def _trim_inputs(self, inputs):
+        """Trim inputs to desired length."""
+        num_special_tokens = len(inputs) + 1
+        if self.truncator == "round_robin":
+            return tf_text.RoundRobinTrimmer(
+                self.sequence_length - num_special_tokens
+            ).trim(inputs)
+        elif self.truncator == "waterfall":
+            return tf_text.WaterfallTrimmer(
+                self.sequence_length - num_special_tokens
+            ).trim(inputs)
+        else:
+            raise ValueError("Unsupported truncator: %s" % self.truncator)
+
+    def _combine_inputs(self, segments):
+        """Combine inputs with start and end values added."""
+        dtype = segments[0].dtype
+        start_value = tf.convert_to_tensor(self.start_value, dtype=dtype)
+        end_value = tf.convert_to_tensor(self.end_value, dtype=dtype)
+
+        start_column = tf.tile([start_value], [segments[0].nrows()])
+        start_column = tf.expand_dims(start_column, 1)
+        end_column = tf.tile([end_value], [segments[0].nrows()])
+        end_column = tf.expand_dims(end_column, 1)
+        ones_column = tf.ones_like(start_column, dtype=tf.int32)
+
+        segments_to_combine = [start_column]
+        segment_ids_to_combine = [ones_column * 0]
+        for i, seg in enumerate(segments):
+            # Combine all segments adding end tokens.
+            segments_to_combine.append(seg)
+            segments_to_combine.append(end_column)
+
+            # Combine segment ids accounting for end tokens.
+            segment_ids_to_combine.append(tf.ones_like(seg, dtype=tf.int32) * i)
+            segment_ids_to_combine.append(ones_column * i)
+
+        token_ids = tf.concat(segments_to_combine, 1)
+        segment_ids = tf.concat(segment_ids_to_combine, 1)
+        return token_ids, segment_ids
+
+    def call(self, inputs):
+        inputs = self._sanitize_inputs(inputs)
+
+        # If rank 1, add a batch dim and convert to ragged.
+        rank_1 = inputs[0].shape.rank == 1
+        if rank_1:
+            inputs = [tf.expand_dims(x, 0) for x in inputs]
+            inputs = [tf.RaggedTensor.from_tensor(x) for x in inputs]
+
+        segments = self._trim_inputs(inputs)
+        token_ids, segment_ids = self._combine_inputs(segments)
+
+        # Pad to dense tensor output.
+        shape = tf.cast([-1, self.sequence_length], "int64")
+        token_ids = token_ids.to_tensor(
+            shape=shape, default_value=self.pad_value
+        )
+        segment_ids = segment_ids.to_tensor(shape=shape)
+
+        # Remove the batch dim if added.
+        if rank_1:
+            token_ids = tf.squeeze(token_ids, 0)
+            segment_ids = tf.squeeze(segment_ids, 0)
+
+        return (token_ids, segment_ids)


Test patch:
diff --git a/keras_nlp/layers/multi_segment_packer_test.py b/keras_nlp/layers/multi_segment_packer_test.py
new file mode 100644
index 00000000..a03b20cd
--- /dev/null
+++ b/keras_nlp/layers/multi_segment_packer_test.py
@@ -0,0 +1,178 @@
+# 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 Transformer Decoder."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
+
+
+class MultiSegmentPackerTest(tf.test.TestCase):
+    def test_trim_single_input_ints(self):
+        input_data = tf.range(3, 10)
+        packer = MultiSegmentPacker(8, start_value=1, end_value=2)
+        output = packer(input_data)
+        self.assertAllEqual(
+            output, ([1, 3, 4, 5, 6, 7, 8, 2], [0, 0, 0, 0, 0, 0, 0, 0])
+        )
+
+    def test_trim_single_input_strings(self):
+        input_data = tf.constant(["a", "b", "c", "d"])
+        packer = MultiSegmentPacker(5, start_value="[CLS]", end_value="[SEP]")
+        output = packer(input_data)
+        self.assertAllEqual(
+            output, (["[CLS]", "a", "b", "c", "[SEP]"], [0, 0, 0, 0, 0])
+        )
+
+    def test_trim_multiple_inputs_round_robin(self):
+        seq1 = tf.constant(["a", "b", "c"])
+        seq2 = tf.constant(["x", "y", "z"])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="round_robin"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                [0, 0, 0, 0, 1, 1, 1],
+            ),
+        )
+
+    def test_trim_multiple_inputs_waterfall(self):
+        seq1 = tf.constant(["a", "b", "c"])
+        seq2 = tf.constant(["x", "y", "z"])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                ["[CLS]", "a", "b", "c", "[SEP]", "x", "[SEP]"],
+                [0, 0, 0, 0, 0, 1, 1],
+            ),
+        )
+
+    def test_trim_batched_inputs_round_robin(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b", "c"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="round_robin"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                [
+                    ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                    ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                ],
+                [
+                    [0, 0, 0, 0, 1, 1, 1],
+                    [0, 0, 0, 0, 1, 1, 1],
+                ],
+            ),
+        )
+
+    def test_trim_batched_inputs_waterfall(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                [
+                    ["[CLS]", "a", "b", "c", "[SEP]", "x", "[SEP]"],
+                    ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                ],
+                [
+                    [0, 0, 0, 0, 0, 1, 1],
+                    [0, 0, 0, 0, 1, 1, 1],
+                ],
+            ),
+        )
+
+    def test_pad_inputs(self):
+        seq1 = tf.constant(["a"])
+        seq2 = tf.constant(["x"])
+        packer = MultiSegmentPacker(
+            6, start_value="[CLS]", end_value="[SEP]", pad_value="[PAD]"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                ["[CLS]", "a", "[SEP]", "x", "[SEP]", "[PAD]"],
+                [0, 0, 0, 1, 1, 0],
+            ),
+        )
+
+    def test_pad_batched_inputs(self):
+        seq1 = tf.ragged.constant([["a"], ["a"]])
+        seq2 = tf.ragged.constant([["x"], ["x", "y"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", pad_value="[PAD]"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                [
+                    ["[CLS]", "a", "[SEP]", "x", "[SEP]", "[PAD]", "[PAD]"],
+                    ["[CLS]", "a", "[SEP]", "x", "y", "[SEP]", "[PAD]"],
+                ],
+                [
+                    [0, 0, 0, 1, 1, 0, 0],
+                    [0, 0, 0, 1, 1, 1, 0],
+                ],
+            ),
+        )
+
+    def test_config(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        original_packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        cloned_packer = MultiSegmentPacker.from_config(
+            original_packer.get_config()
+        )
+        self.assertAllEqual(
+            original_packer([seq1, seq2]),
+            cloned_packer([seq1, seq2]),
+        )
+
+    def test_saving(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        inputs = (
+            keras.Input(dtype="string", ragged=True, shape=(None,)),
+            keras.Input(dtype="string", ragged=True, shape=(None,)),
+        )
+        outputs = packer(inputs)
+        model = keras.Model(inputs, outputs)
+        model.save(self.get_temp_dir())
+        restored_model = keras.models.load_model(self.get_temp_dir())
+        self.assertAllEqual(
+            model((seq1, seq2)),
+            restored_model((seq1, seq2)),
+        )