Repository: keras-team/keras-nlp
Pull request #699 (2023-01-27 18:42:36)

Add OPT Backbone and Tokenizer
I'm pretty sure this is the rough model structure. Sadly this is not quite 1-1 with the BART decoder as claimed, OPT will project from a separate `embedding_dim` to a `hidden_dim`, somewhat similar to ALBERT.

Will leave this as a draft until I have finished checkpoint conversion, in case a few more changes are needed!

mattdangerw: > Great start! Do we not need a preprocessor or is that just WIP?

Preprocessor may not be a useful as the raw tokenizer for the type of pretraining we are interested in doing, so I was punting on it for now. We should add it as a follow up, but not in this PR!
mattdangerw: OK! Few updates on this.

First, the 350m checkpoint is an outlier. It is the only checkpoint with a projection from the token embedding to the hidden dim (and back), it applies layernorm in a different place then other checkpoints, and it has no final layernorm. Rather than handle these discrepancies, and I will just ignore the 350m checkpoint.

Second, I will update this to stay in sync with https://github.com/keras-team/keras-nlp/pull/704

With that, this should be now ready for review.
mattdangerw: > Looks great, but I'm confused why we wouldn't need preprocessors for our use case (esp. if the tokenization does not support truncation or padding). Seems like a followup is required before we can pre-train.

I'm down to add the preprocessor as a follow up! We should model small, incremental PRs. So throwing it all together here would be bad.

I am also worried there is a technical point that keeps getting lost though--the real preprocessing for OPT and GPT style pretraining is **not representable with a preprocessing layer**! See https://github.com/keras-team/keras-nlp/issues/701. This isn't something to be fixed up at some point, it's a fundamental incompatibility with the assumptions in the `dataset.map(preprocessing_layer)` assumption, namely that a raw sample is 1-1 mappable with an output sample.

OPT pre-training preprocessing will look something like

```
ds = load(...)
ds.map(tokenizer).map(add_eos)
ds.rebatch(sequence_length).batch(batch_size)
```

This is quite readable and natural looking a tokenizer. This couldn't possibly work with preprocessing layers unless we wanted preprocessors to output raggeds, which seems way out of scope for this PR.

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 8be1c247..db70d43b 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -40,6 +40,8 @@ from keras_nlp.models.gpt2.gpt2_causal_lm_preprocessor import (
 )
 from keras_nlp.models.gpt2.gpt2_preprocessor import GPT2Preprocessor
 from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
+from keras_nlp.models.opt.opt_backbone import OPTBackbone
+from keras_nlp.models.opt.opt_tokenizer import OPTTokenizer
 from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
 from keras_nlp.models.roberta.roberta_masked_lm import RobertaMaskedLM
diff --git a/keras_nlp/models/opt/__init__.py b/keras_nlp/models/opt/__init__.py
new file mode 100644
index 00000000..ba0c2545
--- /dev/null
+++ b/keras_nlp/models/opt/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2023 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.
diff --git a/keras_nlp/models/opt/opt_backbone.py b/keras_nlp/models/opt/opt_backbone.py
new file mode 100644
index 00000000..c88f208b
--- /dev/null
+++ b/keras_nlp/models/opt/opt_backbone.py
@@ -0,0 +1,161 @@
+# 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.
+
+"""OPT backbone model."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.token_and_position_embedding import (
+    TokenAndPositionEmbedding,
+)
+from keras_nlp.layers.transformer_decoder import TransformerDecoder
+from keras_nlp.models.backbone import Backbone
+
+
+def opt_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class OPTBackbone(Backbone):
+    """OPT encoder-decoder network.
+
+    This class implements a Transformer-based decoder model as described in
+    ["OPT: Open Pre-trained Transformer Language Models"](https://arxiv.org/abs/2205.01068).
+    The default constructor gives a fully customizable, randomly initialized OPT
+    model with any number of layers, heads, and embedding dimensions. To load
+    preset architectures and weights, use the `from_preset` constructor.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind. The underlying model is provided by a
+    third party and subject to a separate license, available
+    [here](https://github.com/facebookresearch/fairseq/).
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer encoder layers and
+            transformer decoder layers.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: int. The hidden size of the transformer decoder layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer decoder layer.
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+
+    Examples:
+    ```python
+    input_data = {
+        "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+    }
+    # Randomly initialized OPT decoder model with a custom config
+    model = keras_nlp.models.OPTBackbone(
+        vocabulary_size=50265,
+        num_layers=6,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+    )
+    output = model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        dropout=0.1,
+        max_sequence_length=2048,
+        **kwargs,
+    ):
+        # Decoder inputs.
+        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
+        padding_mask = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+
+        # Embed tokens and positions.
+        x = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            sequence_length=max_sequence_length,
+            embedding_dim=hidden_dim,
+            embeddings_initializer=opt_kernel_initializer(),
+            name="embeddings",
+        )(token_ids)
+
+        # Apply successive transformer decoder blocks.
+        for i in range(num_layers):
+            x = TransformerDecoder(
+                intermediate_dim=intermediate_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                activation="relu",
+                layer_norm_epsilon=1e-5,
+                normalize_first=True,
+                kernel_initializer=opt_kernel_initializer(),
+                name=f"transformer_layer_{i}",
+            )(x, decoder_padding_mask=padding_mask)
+
+        # Add a final layer norm.
+        x = keras.layers.LayerNormalization(
+            name="layer_norm",
+            axis=-1,
+            epsilon=1e-5,
+            dtype=tf.float32,
+        )(x)
+
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs={
+                "token_ids": token_ids,
+                "padding_mask": padding_mask,
+            },
+            outputs=x,
+            **kwargs,
+        )
+
+        # All references to `self` below this line
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.max_sequence_length = max_sequence_length
+
+    def get_config(self):
+        return {
+            "vocabulary_size": self.vocabulary_size,
+            "num_layers": self.num_layers,
+            "num_heads": self.num_heads,
+            "hidden_dim": self.hidden_dim,
+            "intermediate_dim": self.intermediate_dim,
+            "dropout": self.dropout,
+            "max_sequence_length": self.max_sequence_length,
+        }
+
+    @property
+    def token_embedding(self):
+        return self.get_layer("embeddings").token_embedding
diff --git a/keras_nlp/models/opt/opt_tokenizer.py b/keras_nlp/models/opt/opt_tokenizer.py
new file mode 100644
index 00000000..b7f6d466
--- /dev/null
+++ b/keras_nlp/models/opt/opt_tokenizer.py
@@ -0,0 +1,110 @@
+# Copyright 2023 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.
+
+"""OPT tokenizer."""
+
+from tensorflow import keras
+
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class OPTTokenizer(BytePairTokenizer):
+    """An OPT tokenizer using Byte-Pair Encoding subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.BytePairTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by OPT
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a OPT preset.
+
+    This tokenizer does not provide truncation or padding of inputs.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space.
+
+    Examples:
+
+    Batched inputs.
+    >>> vocab = {"<pad>": 1, "
": 2, "a": 3, "Ġquick": 4, "Ġfox": 5}
+    >>> merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
+    >>> merges += ["Ġ f", "o x", "Ġf ox"]
+    >>> tokenizer = keras_nlp.models.OPTTokenizer(
+    ...     vocabulary=vocab,
+    ...     merges=merges,
+    ... )
+    >>> tokenizer(["a quick fox", "a fox quick"])
+    <tf.RaggedTensor [[3, 4, 5], [3, 5, 4]]>
+
+    Unbatched input.
+    >>> vocab = {"<pad>": 1, "
": 2, "a": 3, "Ġquick": 4, "Ġfox": 5}
+    >>> merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
+    >>> merges += ["Ġ f", "o x", "Ġf ox"]
+    >>> tokenizer = keras_nlp.models.OPTTokenizer(
+    ...     vocabulary=vocab,
+    ...     merges=merges,
+    ... )
+    >>> tokenizer("a quick fox")
+    <tf.Tensor: shape=(4,), dtype=int32, numpy=array([3, 4, 5], dtype=int32)>
+
+    Detokenization.
+    >>> vocab = {"<pad>": 1, "
": 2, "Ġquick": 4, "Ġfox": 5}
+    >>> merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
+    >>> merges += ["Ġ f", "o x", "Ġf ox"]
+    >>> tokenizer = keras_nlp.models.OPTTokenizer(
+    ...     vocabulary=vocab,
+    ...     merges=merges,
+    ... )
+    >>> tokenizer.detokenize(tokenizer(" quick fox")).numpy().decode('utf-8')
+    ' quick fox'
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        merges,
+        **kwargs,
+    ):
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            **kwargs,
+        )
+
+        # We use `"
"` as both a start and end token, as OPT was only
+        # pre-trained with `"
"` marking document boundaries.
+        start_token = "
"
+        pad_token = "<pad>"
+        end_token = "
"
+        for token in [start_token, pad_token, end_token]:
+            if token not in self.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        self.start_token_id = self.token_to_id(start_token)
+        self.pad_token_id = self.token_to_id(pad_token)
+        self.end_token_id = self.token_to_id(end_token)


Test patch:
diff --git a/keras_nlp/models/opt/opt_backbone_test.py b/keras_nlp/models/opt/opt_backbone_test.py
new file mode 100644
index 00000000..4283dc13
--- /dev/null
+++ b/keras_nlp/models/opt/opt_backbone_test.py
@@ -0,0 +1,96 @@
+# Copyright 2023 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.
+"""Test for OPT backbone models."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.opt.opt_backbone import OPTBackbone
+
+
+class OPTTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.model = OPTBackbone(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=32,
+            intermediate_dim=128,
+            max_sequence_length=128,
+        )
+        self.batch_size = 8
+        self.input_batch = {
+            "token_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+            "padding_mask": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+        }
+
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_valid_call_opt(self):
+        self.model(self.input_batch)
+
+        # Check default name passed through
+        self.assertRegexpMatches(self.model.name, "opt_backbone")
+
+    def test_variable_sequence_length_call_opt(self):
+        for seq_length in (25, 50, 75):
+            input_data = {
+                "token_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "padding_mask": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+            }
+            self.model(input_data)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_opt_compile(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_opt_compile_batched_ds(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        model_output = self.model(self.input_batch)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, OPTBackbone)
+
+        # Check that output matches.
+        restored_output = restored_model(self.input_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/opt/opt_tokenizer_test.py b/keras_nlp/models/opt/opt_tokenizer_test.py
new file mode 100644
index 00000000..2e0eb461
--- /dev/null
+++ b/keras_nlp/models/opt/opt_tokenizer_test.py
@@ -0,0 +1,87 @@
+# Copyright 2023 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 OPT tokenizer layer."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.opt.opt_tokenizer import OPTTokenizer
+
+
+class OPTTokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.vocab = {
+            "
": 0,
+            "<pad>": 1,
+            "
": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "Ġa i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "Ġai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["pla ne"]
+        self.merges = merges
+
+        self.tokenizer = OPTTokenizer(vocabulary=self.vocab, merges=self.merges)
+
+    def test_tokenize(self):
+        input_data = " airplane at airport"
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [3, 4, 5, 3, 6])
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant([" airplane at airport", " kohli is the best"])
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [[3, 4, 5, 3, 6], [7, 8, 9, 10, 11]])
+
+    def test_detokenize(self):
+        input_tokens = [3, 4, 5, 3, 6]
+        output = self.tokenizer.detokenize(input_tokens)
+        self.assertEqual(output, " airplane at airport")
+
+    def test_vocabulary_size(self):
+        self.assertEqual(self.tokenizer.vocabulary_size(), 12)
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        input_data = tf.constant([" airplane at airport"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )