We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Call `super.config()` in `BartBackbone`'s `get_config()`
We should call `super().config()` here: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bart/bart_backbone.py#L238, update the config with `BARTBackbone`-specific keys and return the config. Check this for reference: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bert/bert_backbone.py#L204.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `keras_nlp/models/bart/bart_backbone.py`
Content:
```
1 # Copyright 2022 The KerasNLP Authors
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """BART backbone model."""
16 
17 import copy
18 
19 import tensorflow as tf
20 from tensorflow import keras
21 
22 from keras_nlp.layers.position_embedding import PositionEmbedding
23 from keras_nlp.layers.transformer_decoder import TransformerDecoder
24 from keras_nlp.layers.transformer_encoder import TransformerEncoder
25 from keras_nlp.models.backbone import Backbone
26 from keras_nlp.models.bart.bart_presets import backbone_presets
27 from keras_nlp.utils.python_utils import classproperty
28 
29 
30 def bart_kernel_initializer(stddev=0.02):
31     return keras.initializers.TruncatedNormal(stddev=stddev)
32 
33 
34 @keras.utils.register_keras_serializable(package="keras_nlp")
35 class BartBackbone(Backbone):
36     """BART encoder-decoder network.
37 
38     This class implements a Transformer-based encoder-decoder model as
39     described in
40     ["BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension"](https://arxiv.org/abs/1910.13461).
41 
42     The default constructor gives a fully customizable, randomly initialized BART
43     model with any number of layers, heads, and embedding dimensions. To load
44     preset architectures and weights, use the `from_preset` constructor.
45 
46     Disclaimer: Pre-trained models are provided on an "as is" basis, without
47     warranties or conditions of any kind. The underlying model is provided by a
48     third party and subject to a separate license, available
49     [here](https://github.com/facebookresearch/fairseq/).
50 
51     Args:
52         vocabulary_size: int. The size of the token vocabulary.
53         num_layers: int. The number of transformer encoder layers and
54             transformer decoder layers.
55         num_heads: int. The number of attention heads for each transformer.
56             The hidden size must be divisible by the number of attention heads.
57         hidden_dim: int. The size of the transformer encoding and pooler layers.
58         intermediate_dim: int. The output dimension of the first Dense layer in
59             a two-layer feedforward network for each transformer.
60         dropout: float. Dropout probability for the Transformer encoder.
61         max_sequence_length: int. The maximum sequence length that this encoder
62             can consume. If None, `max_sequence_length` uses the value from
63             sequence length. This determines the variable shape for positional
64             embeddings.
65 
66     Examples:
67     ```python
68     input_data = {
69         "encoder_token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
70         "encoder_padding_mask": tf.constant(
71             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
72         ),
73         "decoder_token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
74         "decoder_padding_mask": tf.constant(
75             [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0], shape=(1, 12)
76         ),
77     }
78 
79     # Randomly initialized BART encoder-decoder model with a custom config
80     model = keras_nlp.models.BartBackbone(
81         vocabulary_size=50265,
82         num_layers=6,
83         num_heads=12,
84         hidden_dim=768,
85         intermediate_dim=3072,
86         max_sequence_length=12,
87     )
88     output = model(input_data)
89     ```
90     """
91 
92     def __init__(
93         self,
94         vocabulary_size,
95         num_layers,
96         num_heads,
97         hidden_dim,
98         intermediate_dim,
99         dropout=0.1,
100         max_sequence_length=1024,
101         **kwargs,
102     ):
103         # Encoder inputs
104         encoder_token_id_input = keras.Input(
105             shape=(None,), dtype="int32", name="encoder_token_ids"
106         )
107         encoder_padding_mask = keras.Input(
108             shape=(None,), dtype="int32", name="encoder_padding_mask"
109         )
110 
111         # Decoder inputs.
112         decoder_token_id_input = keras.Input(
113             shape=(None,), dtype="int32", name="decoder_token_ids"
114         )
115         decoder_padding_mask = keras.Input(
116             shape=(None,), dtype="int32", name="decoder_padding_mask"
117         )
118 
119         # Token embedding layer. This layer is shared by encoder and decoder.
120         token_embedding_layer = keras.layers.Embedding(
121             input_dim=vocabulary_size,
122             output_dim=hidden_dim,
123             embeddings_initializer=bart_kernel_initializer(),
124             name="token_embedding",
125         )
126 
127         # ===== Encoder =====
128 
129         # Embed tokens and positions.
130         token_embedding = token_embedding_layer(encoder_token_id_input)
131         # Position embedding parameters are not shared by encode and decoder.
132         position_embedding = PositionEmbedding(
133             initializer=bart_kernel_initializer(),
134             sequence_length=max_sequence_length,
135             name="encoder_position_embedding",
136         )(token_embedding)
137 
138         # Sum, normalize and apply dropout to embeddings.
139         x = keras.layers.Add()((token_embedding, position_embedding))
140         x = keras.layers.LayerNormalization(
141             name="encoder_embeddings_layer_norm",
142             axis=-1,
143             epsilon=1e-5,
144             dtype=tf.float32,
145         )(x)
146         x = keras.layers.Dropout(
147             dropout,
148             name="encoder_embeddings_dropout",
149         )(x)
150 
151         # Apply successive transformer encoder blocks.
152         for i in range(num_layers):
153             x = TransformerEncoder(
154                 num_heads=num_heads,
155                 intermediate_dim=intermediate_dim,
156                 activation=lambda x: keras.activations.gelu(
157                     x, approximate=False
158                 ),
159                 dropout=dropout,
160                 layer_norm_epsilon=1e-5,
161                 kernel_initializer=bart_kernel_initializer(),
162                 name=f"transformer_encoder_layer_{i}",
163             )(x, padding_mask=encoder_padding_mask)
164 
165         encoder_output = x
166 
167         # ===== Decoder =====
168 
169         # Embed tokens and positions.
170         token_embedding = token_embedding_layer(decoder_token_id_input)
171         # Position embedding parameters are not shared by encode and decoder.
172         position_embedding = PositionEmbedding(
173             initializer=bart_kernel_initializer(),
174             sequence_length=max_sequence_length,
175             name="decoder_position_embedding",
176         )(token_embedding)
177 
178         # Sum, normalize and apply dropout to embeddings.
179         x = keras.layers.Add()((token_embedding, position_embedding))
180         x = keras.layers.LayerNormalization(
181             name="decoder_embeddings_layer_norm",
182             axis=-1,
183             epsilon=1e-5,
184             dtype=tf.float32,
185         )(x)
186         x = keras.layers.Dropout(
187             dropout,
188             name="decoder_embeddings_dropout",
189         )(x)
190 
191         # Apply successive transformer decoder blocks.
192         for i in range(num_layers):
193             transformer_decoder_layer = TransformerDecoder(
194                 intermediate_dim=intermediate_dim,
195                 num_heads=num_heads,
196                 dropout=dropout,
197                 activation=lambda x: keras.activations.gelu(
198                     x, approximate=False
199                 ),
200                 layer_norm_epsilon=1e-5,
201                 kernel_initializer=bart_kernel_initializer(),
202                 name=f"transformer_decoder_layer_{i}",
203                 has_cross_attention=True,
204             )
205             x = transformer_decoder_layer(
206                 decoder_sequence=x,
207                 encoder_sequence=encoder_output,
208                 decoder_padding_mask=decoder_padding_mask,
209                 encoder_padding_mask=encoder_padding_mask,
210             )
211 
212         decoder_output = x
213 
214         # Instantiate using Functional API Model constructor
215         super().__init__(
216             inputs={
217                 "encoder_token_ids": encoder_token_id_input,
218                 "encoder_padding_mask": encoder_padding_mask,
219                 "decoder_token_ids": decoder_token_id_input,
220                 "decoder_padding_mask": decoder_padding_mask,
221             },
222             outputs={
223                 "encoder_sequence_output": encoder_output,
224                 "decoder_sequence_output": decoder_output,
225             },
226             **kwargs,
227         )
228 
229         # All references to `self` below this line
230         self.vocabulary_size = vocabulary_size
231         self.num_layers = num_layers
232         self.num_heads = num_heads
233         self.hidden_dim = hidden_dim
234         self.intermediate_dim = intermediate_dim
235         self.dropout = dropout
236         self.max_sequence_length = max_sequence_length
237 
238     def get_config(self):
239         return {
240             "vocabulary_size": self.vocabulary_size,
241             "num_layers": self.num_layers,
242             "num_heads": self.num_heads,
243             "hidden_dim": self.hidden_dim,
244             "intermediate_dim": self.intermediate_dim,
245             "dropout": self.dropout,
246             "max_sequence_length": self.max_sequence_length,
247             "name": self.name,
248             "trainable": self.trainable,
249         }
250 
251     @property
252     def token_embedding(self):
253         return self.get_layer("token_embedding")
254 
255     @classproperty
256     def presets(cls):
257         return copy.deepcopy(backbone_presets)
258
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
--- a/keras_nlp/models/bart/bart_backbone.py
+++ b/keras_nlp/models/bart/bart_backbone.py
@@ -236,17 +236,20 @@
         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,
-            "name": self.name,
-            "trainable": self.trainable,
-        }
+        config = super().get_config()
+        config.update(
+            {
+                "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,
+            }
+        )
+
+        return config
 
     @property
     def token_embedding(self):

{"golden_diff": "diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py\n--- a/keras_nlp/models/bart/bart_backbone.py\n+++ b/keras_nlp/models/bart/bart_backbone.py\n@@ -236,17 +236,20 @@\n         self.max_sequence_length = max_sequence_length\n \n     def get_config(self):\n-        return {\n-            \"vocabulary_size\": self.vocabulary_size,\n-            \"num_layers\": self.num_layers,\n-            \"num_heads\": self.num_heads,\n-            \"hidden_dim\": self.hidden_dim,\n-            \"intermediate_dim\": self.intermediate_dim,\n-            \"dropout\": self.dropout,\n-            \"max_sequence_length\": self.max_sequence_length,\n-            \"name\": self.name,\n-            \"trainable\": self.trainable,\n-        }\n+        config = super().get_config()\n+        config.update(\n+            {\n+                \"vocabulary_size\": self.vocabulary_size,\n+                \"num_layers\": self.num_layers,\n+                \"num_heads\": self.num_heads,\n+                \"hidden_dim\": self.hidden_dim,\n+                \"intermediate_dim\": self.intermediate_dim,\n+                \"dropout\": self.dropout,\n+                \"max_sequence_length\": self.max_sequence_length,\n+            }\n+        )\n+\n+        return config\n \n     @property\n     def token_embedding(self):\n", "issue": "Call `super.config()` in `BartBackbone`'s `get_config()`\nWe should call `super().config()` here: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bart/bart_backbone.py#L238, update the config with `BARTBackbone`-specific keys and return the config. Check this for reference: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bert/bert_backbone.py#L204.\n", "before_files": [{"content": "# Copyright 2022 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"BART backbone model.\"\"\"\n\nimport copy\n\nimport tensorflow as tf\nfrom tensorflow import keras\n\nfrom keras_nlp.layers.position_embedding import PositionEmbedding\nfrom keras_nlp.layers.transformer_decoder import TransformerDecoder\nfrom keras_nlp.layers.transformer_encoder import TransformerEncoder\nfrom keras_nlp.models.backbone import Backbone\nfrom keras_nlp.models.bart.bart_presets import backbone_presets\nfrom keras_nlp.utils.python_utils import classproperty\n\n\ndef bart_kernel_initializer(stddev=0.02):\n    return keras.initializers.TruncatedNormal(stddev=stddev)\n\n\n@keras.utils.register_keras_serializable(package=\"keras_nlp\")\nclass BartBackbone(Backbone):\n    \"\"\"BART encoder-decoder network.\n\n    This class implements a Transformer-based encoder-decoder model as\n    described in\n    [\"BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension\"](https://arxiv.org/abs/1910.13461).\n\n    The default constructor gives a fully customizable, randomly initialized BART\n    model with any number of layers, heads, and embedding dimensions. To load\n    preset architectures and weights, use the `from_preset` constructor.\n\n    Disclaimer: Pre-trained models are provided on an \"as is\" basis, without\n    warranties or conditions of any kind. The underlying model is provided by a\n    third party and subject to a separate license, available\n    [here](https://github.com/facebookresearch/fairseq/).\n\n    Args:\n        vocabulary_size: int. The size of the token vocabulary.\n        num_layers: int. The number of transformer encoder layers and\n            transformer decoder layers.\n        num_heads: int. The number of attention heads for each transformer.\n            The hidden size must be divisible by the number of attention heads.\n        hidden_dim: int. The size of the transformer encoding and pooler layers.\n        intermediate_dim: int. The output dimension of the first Dense layer in\n            a two-layer feedforward network for each transformer.\n        dropout: float. Dropout probability for the Transformer encoder.\n        max_sequence_length: int. The maximum sequence length that this encoder\n            can consume. If None, `max_sequence_length` uses the value from\n            sequence length. This determines the variable shape for positional\n            embeddings.\n\n    Examples:\n    ```python\n    input_data = {\n        \"encoder_token_ids\": tf.ones(shape=(1, 12), dtype=tf.int64),\n        \"encoder_padding_mask\": tf.constant(\n            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)\n        ),\n        \"decoder_token_ids\": tf.ones(shape=(1, 12), dtype=tf.int64),\n        \"decoder_padding_mask\": tf.constant(\n            [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0], shape=(1, 12)\n        ),\n    }\n\n    # Randomly initialized BART encoder-decoder model with a custom config\n    model = keras_nlp.models.BartBackbone(\n        vocabulary_size=50265,\n        num_layers=6,\n        num_heads=12,\n        hidden_dim=768,\n        intermediate_dim=3072,\n        max_sequence_length=12,\n    )\n    output = model(input_data)\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        vocabulary_size,\n        num_layers,\n        num_heads,\n        hidden_dim,\n        intermediate_dim,\n        dropout=0.1,\n        max_sequence_length=1024,\n        **kwargs,\n    ):\n        # Encoder inputs\n        encoder_token_id_input = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"encoder_token_ids\"\n        )\n        encoder_padding_mask = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"encoder_padding_mask\"\n        )\n\n        # Decoder inputs.\n        decoder_token_id_input = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"decoder_token_ids\"\n        )\n        decoder_padding_mask = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"decoder_padding_mask\"\n        )\n\n        # Token embedding layer. This layer is shared by encoder and decoder.\n        token_embedding_layer = keras.layers.Embedding(\n            input_dim=vocabulary_size,\n            output_dim=hidden_dim,\n            embeddings_initializer=bart_kernel_initializer(),\n            name=\"token_embedding\",\n        )\n\n        # ===== Encoder =====\n\n        # Embed tokens and positions.\n        token_embedding = token_embedding_layer(encoder_token_id_input)\n        # Position embedding parameters are not shared by encode and decoder.\n        position_embedding = PositionEmbedding(\n            initializer=bart_kernel_initializer(),\n            sequence_length=max_sequence_length,\n            name=\"encoder_position_embedding\",\n        )(token_embedding)\n\n        # Sum, normalize and apply dropout to embeddings.\n        x = keras.layers.Add()((token_embedding, position_embedding))\n        x = keras.layers.LayerNormalization(\n            name=\"encoder_embeddings_layer_norm\",\n            axis=-1,\n            epsilon=1e-5,\n            dtype=tf.float32,\n        )(x)\n        x = keras.layers.Dropout(\n            dropout,\n            name=\"encoder_embeddings_dropout\",\n        )(x)\n\n        # Apply successive transformer encoder blocks.\n        for i in range(num_layers):\n            x = TransformerEncoder(\n                num_heads=num_heads,\n                intermediate_dim=intermediate_dim,\n                activation=lambda x: keras.activations.gelu(\n                    x, approximate=False\n                ),\n                dropout=dropout,\n                layer_norm_epsilon=1e-5,\n                kernel_initializer=bart_kernel_initializer(),\n                name=f\"transformer_encoder_layer_{i}\",\n            )(x, padding_mask=encoder_padding_mask)\n\n        encoder_output = x\n\n        # ===== Decoder =====\n\n        # Embed tokens and positions.\n        token_embedding = token_embedding_layer(decoder_token_id_input)\n        # Position embedding parameters are not shared by encode and decoder.\n        position_embedding = PositionEmbedding(\n            initializer=bart_kernel_initializer(),\n            sequence_length=max_sequence_length,\n            name=\"decoder_position_embedding\",\n        )(token_embedding)\n\n        # Sum, normalize and apply dropout to embeddings.\n        x = keras.layers.Add()((token_embedding, position_embedding))\n        x = keras.layers.LayerNormalization(\n            name=\"decoder_embeddings_layer_norm\",\n            axis=-1,\n            epsilon=1e-5,\n            dtype=tf.float32,\n        )(x)\n        x = keras.layers.Dropout(\n            dropout,\n            name=\"decoder_embeddings_dropout\",\n        )(x)\n\n        # Apply successive transformer decoder blocks.\n        for i in range(num_layers):\n            transformer_decoder_layer = TransformerDecoder(\n                intermediate_dim=intermediate_dim,\n                num_heads=num_heads,\n                dropout=dropout,\n                activation=lambda x: keras.activations.gelu(\n                    x, approximate=False\n                ),\n                layer_norm_epsilon=1e-5,\n                kernel_initializer=bart_kernel_initializer(),\n                name=f\"transformer_decoder_layer_{i}\",\n                has_cross_attention=True,\n            )\n            x = transformer_decoder_layer(\n                decoder_sequence=x,\n                encoder_sequence=encoder_output,\n                decoder_padding_mask=decoder_padding_mask,\n                encoder_padding_mask=encoder_padding_mask,\n            )\n\n        decoder_output = x\n\n        # Instantiate using Functional API Model constructor\n        super().__init__(\n            inputs={\n                \"encoder_token_ids\": encoder_token_id_input,\n                \"encoder_padding_mask\": encoder_padding_mask,\n                \"decoder_token_ids\": decoder_token_id_input,\n                \"decoder_padding_mask\": decoder_padding_mask,\n            },\n            outputs={\n                \"encoder_sequence_output\": encoder_output,\n                \"decoder_sequence_output\": decoder_output,\n            },\n            **kwargs,\n        )\n\n        # All references to `self` below this line\n        self.vocabulary_size = vocabulary_size\n        self.num_layers = num_layers\n        self.num_heads = num_heads\n        self.hidden_dim = hidden_dim\n        self.intermediate_dim = intermediate_dim\n        self.dropout = dropout\n        self.max_sequence_length = max_sequence_length\n\n    def get_config(self):\n        return {\n            \"vocabulary_size\": self.vocabulary_size,\n            \"num_layers\": self.num_layers,\n            \"num_heads\": self.num_heads,\n            \"hidden_dim\": self.hidden_dim,\n            \"intermediate_dim\": self.intermediate_dim,\n            \"dropout\": self.dropout,\n            \"max_sequence_length\": self.max_sequence_length,\n            \"name\": self.name,\n            \"trainable\": self.trainable,\n        }\n\n    @property\n    def token_embedding(self):\n        return self.get_layer(\"token_embedding\")\n\n    @classproperty\n    def presets(cls):\n        return copy.deepcopy(backbone_presets)\n", "path": "keras_nlp/models/bart/bart_backbone.py"}], "after_files": [{"content": "# Copyright 2022 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"BART backbone model.\"\"\"\n\nimport copy\n\nimport tensorflow as tf\nfrom tensorflow import keras\n\nfrom keras_nlp.layers.position_embedding import PositionEmbedding\nfrom keras_nlp.layers.transformer_decoder import TransformerDecoder\nfrom keras_nlp.layers.transformer_encoder import TransformerEncoder\nfrom keras_nlp.models.backbone import Backbone\nfrom keras_nlp.models.bart.bart_presets import backbone_presets\nfrom keras_nlp.utils.python_utils import classproperty\n\n\ndef bart_kernel_initializer(stddev=0.02):\n    return keras.initializers.TruncatedNormal(stddev=stddev)\n\n\n@keras.utils.register_keras_serializable(package=\"keras_nlp\")\nclass BartBackbone(Backbone):\n    \"\"\"BART encoder-decoder network.\n\n    This class implements a Transformer-based encoder-decoder model as\n    described in\n    [\"BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension\"](https://arxiv.org/abs/1910.13461).\n\n    The default constructor gives a fully customizable, randomly initialized BART\n    model with any number of layers, heads, and embedding dimensions. To load\n    preset architectures and weights, use the `from_preset` constructor.\n\n    Disclaimer: Pre-trained models are provided on an \"as is\" basis, without\n    warranties or conditions of any kind. The underlying model is provided by a\n    third party and subject to a separate license, available\n    [here](https://github.com/facebookresearch/fairseq/).\n\n    Args:\n        vocabulary_size: int. The size of the token vocabulary.\n        num_layers: int. The number of transformer encoder layers and\n            transformer decoder layers.\n        num_heads: int. The number of attention heads for each transformer.\n            The hidden size must be divisible by the number of attention heads.\n        hidden_dim: int. The size of the transformer encoding and pooler layers.\n        intermediate_dim: int. The output dimension of the first Dense layer in\n            a two-layer feedforward network for each transformer.\n        dropout: float. Dropout probability for the Transformer encoder.\n        max_sequence_length: int. The maximum sequence length that this encoder\n            can consume. If None, `max_sequence_length` uses the value from\n            sequence length. This determines the variable shape for positional\n            embeddings.\n\n    Examples:\n    ```python\n    input_data = {\n        \"encoder_token_ids\": tf.ones(shape=(1, 12), dtype=tf.int64),\n        \"encoder_padding_mask\": tf.constant(\n            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)\n        ),\n        \"decoder_token_ids\": tf.ones(shape=(1, 12), dtype=tf.int64),\n        \"decoder_padding_mask\": tf.constant(\n            [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0], shape=(1, 12)\n        ),\n    }\n\n    # Randomly initialized BART encoder-decoder model with a custom config\n    model = keras_nlp.models.BartBackbone(\n        vocabulary_size=50265,\n        num_layers=6,\n        num_heads=12,\n        hidden_dim=768,\n        intermediate_dim=3072,\n        max_sequence_length=12,\n    )\n    output = model(input_data)\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        vocabulary_size,\n        num_layers,\n        num_heads,\n        hidden_dim,\n        intermediate_dim,\n        dropout=0.1,\n        max_sequence_length=1024,\n        **kwargs,\n    ):\n        # Encoder inputs\n        encoder_token_id_input = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"encoder_token_ids\"\n        )\n        encoder_padding_mask = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"encoder_padding_mask\"\n        )\n\n        # Decoder inputs.\n        decoder_token_id_input = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"decoder_token_ids\"\n        )\n        decoder_padding_mask = keras.Input(\n            shape=(None,), dtype=\"int32\", name=\"decoder_padding_mask\"\n        )\n\n        # Token embedding layer. This layer is shared by encoder and decoder.\n        token_embedding_layer = keras.layers.Embedding(\n            input_dim=vocabulary_size,\n            output_dim=hidden_dim,\n            embeddings_initializer=bart_kernel_initializer(),\n            name=\"token_embedding\",\n        )\n\n        # ===== Encoder =====\n\n        # Embed tokens and positions.\n        token_embedding = token_embedding_layer(encoder_token_id_input)\n        # Position embedding parameters are not shared by encode and decoder.\n        position_embedding = PositionEmbedding(\n            initializer=bart_kernel_initializer(),\n            sequence_length=max_sequence_length,\n            name=\"encoder_position_embedding\",\n        )(token_embedding)\n\n        # Sum, normalize and apply dropout to embeddings.\n        x = keras.layers.Add()((token_embedding, position_embedding))\n        x = keras.layers.LayerNormalization(\n            name=\"encoder_embeddings_layer_norm\",\n            axis=-1,\n            epsilon=1e-5,\n            dtype=tf.float32,\n        )(x)\n        x = keras.layers.Dropout(\n            dropout,\n            name=\"encoder_embeddings_dropout\",\n        )(x)\n\n        # Apply successive transformer encoder blocks.\n        for i in range(num_layers):\n            x = TransformerEncoder(\n                num_heads=num_heads,\n                intermediate_dim=intermediate_dim,\n                activation=lambda x: keras.activations.gelu(\n                    x, approximate=False\n                ),\n                dropout=dropout,\n                layer_norm_epsilon=1e-5,\n                kernel_initializer=bart_kernel_initializer(),\n                name=f\"transformer_encoder_layer_{i}\",\n            )(x, padding_mask=encoder_padding_mask)\n\n        encoder_output = x\n\n        # ===== Decoder =====\n\n        # Embed tokens and positions.\n        token_embedding = token_embedding_layer(decoder_token_id_input)\n        # Position embedding parameters are not shared by encode and decoder.\n        position_embedding = PositionEmbedding(\n            initializer=bart_kernel_initializer(),\n            sequence_length=max_sequence_length,\n            name=\"decoder_position_embedding\",\n        )(token_embedding)\n\n        # Sum, normalize and apply dropout to embeddings.\n        x = keras.layers.Add()((token_embedding, position_embedding))\n        x = keras.layers.LayerNormalization(\n            name=\"decoder_embeddings_layer_norm\",\n            axis=-1,\n            epsilon=1e-5,\n            dtype=tf.float32,\n        )(x)\n        x = keras.layers.Dropout(\n            dropout,\n            name=\"decoder_embeddings_dropout\",\n        )(x)\n\n        # Apply successive transformer decoder blocks.\n        for i in range(num_layers):\n            transformer_decoder_layer = TransformerDecoder(\n                intermediate_dim=intermediate_dim,\n                num_heads=num_heads,\n                dropout=dropout,\n                activation=lambda x: keras.activations.gelu(\n                    x, approximate=False\n                ),\n                layer_norm_epsilon=1e-5,\n                kernel_initializer=bart_kernel_initializer(),\n                name=f\"transformer_decoder_layer_{i}\",\n                has_cross_attention=True,\n            )\n            x = transformer_decoder_layer(\n                decoder_sequence=x,\n                encoder_sequence=encoder_output,\n                decoder_padding_mask=decoder_padding_mask,\n                encoder_padding_mask=encoder_padding_mask,\n            )\n\n        decoder_output = x\n\n        # Instantiate using Functional API Model constructor\n        super().__init__(\n            inputs={\n                \"encoder_token_ids\": encoder_token_id_input,\n                \"encoder_padding_mask\": encoder_padding_mask,\n                \"decoder_token_ids\": decoder_token_id_input,\n                \"decoder_padding_mask\": decoder_padding_mask,\n            },\n            outputs={\n                \"encoder_sequence_output\": encoder_output,\n                \"decoder_sequence_output\": decoder_output,\n            },\n            **kwargs,\n        )\n\n        # All references to `self` below this line\n        self.vocabulary_size = vocabulary_size\n        self.num_layers = num_layers\n        self.num_heads = num_heads\n        self.hidden_dim = hidden_dim\n        self.intermediate_dim = intermediate_dim\n        self.dropout = dropout\n        self.max_sequence_length = max_sequence_length\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"vocabulary_size\": self.vocabulary_size,\n                \"num_layers\": self.num_layers,\n                \"num_heads\": self.num_heads,\n                \"hidden_dim\": self.hidden_dim,\n                \"intermediate_dim\": self.intermediate_dim,\n                \"dropout\": self.dropout,\n                \"max_sequence_length\": self.max_sequence_length,\n            }\n        )\n\n        return config\n\n    @property\n    def token_embedding(self):\n        return self.get_layer(\"token_embedding\")\n\n    @classproperty\n    def presets(cls):\n        return copy.deepcopy(backbone_presets)\n", "path": "keras_nlp/models/bart/bart_backbone.py"}]}