We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Use KerasTuner to hyper-parameter search for the BERT finetuning script
From the [BERT paper](https://arxiv.org/pdf/1810.04805.pdf)...

```
For fine-tuning, most model hyperparameters are
the same as in pre-training, with the exception of
the batch size, learning rate, and number of training
epochs. The dropout probability was always
kept at 0.1. The optimal hyperparameter values
are task-specific, but we found the following range
of possible values to work well across all tasks:

• Batch size: 16, 32
• Learning rate (Adam): 5e-5, 3e-5, 2e-5
• Number of epochs: 2, 3, 4
```

We should allow our [BERT finetuning script](https://github.com/keras-team/keras-nlp/blob/master/examples/bert/run_glue_finetuning.py) to do this search automatically. [KerasTuner](https://keras.io/keras_tuner/) is a good fit for this.

Steps:
 - [ ] Add an setup.py `examples` dependency on keras-tuner.
 - [ ] Remove epochs, batch size and learning rage arguments from run_glue_finetuning.py.
 - [ ] Use keras tuner to hyperparemeter search on the above value ranges with the validation set.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 # Copyright 2021 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 """Setup script."""
16 
17 import pathlib
18 
19 from setuptools import find_packages
20 from setuptools import setup
21 
22 HERE = pathlib.Path(__file__).parent
23 README = (HERE / "README.md").read_text()
24 
25 setup(
26     name="keras-nlp",
27     description=(
28         "Industry-strength Natural Language Processing extensions for Keras."
29     ),
30     long_description=README,
31     long_description_content_type="text/markdown",
32     version="0.2.0-dev.1",
33     url="https://github.com/keras-team/keras-nlp",
34     author="Keras team",
35     author_email="keras-nlp@google.com",
36     license="Apache License 2.0",
37     install_requires=[
38         "absl-py",
39         "numpy",
40         "packaging",
41         "tensorflow",
42         "tensorflow_text",
43     ],
44     extras_require={
45         "tests": [
46             "black",
47             "flake8",
48             "isort",
49             "pytest",
50             "pytest-cov",
51         ],
52         "examples": [
53             "datasets",  # For GLUE in BERT example.
54             "nltk",
55             "wikiextractor",
56         ],
57     },
58     classifiers=[
59         "Programming Language :: Python",
60         "Programming Language :: Python :: 3.7",
61         "Operating System :: Unix",
62         "Operating System :: Microsoft :: Windows",
63         "Operating System :: MacOS",
64         "Intended Audience :: Science/Research",
65         "Topic :: Scientific/Engineering",
66         "Topic :: Software Development",
67     ],
68     packages=find_packages(exclude=("*_test.py",)),
69 )
70
```

Path: `examples/bert/run_glue_finetuning.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 """Run finetuning on a GLUE task."""
15 
16 import json
17 
18 import datasets
19 import tensorflow as tf
20 import tensorflow_text as tftext
21 from absl import app
22 from absl import flags
23 from tensorflow import keras
24 
25 FLAGS = flags.FLAGS
26 
27 flags.DEFINE_string(
28     "bert_config_file",
29     None,
30     "The json config file for the bert model parameters.",
31 )
32 
33 flags.DEFINE_string(
34     "vocab_file",
35     None,
36     "The vocabulary file that the BERT model was trained on.",
37 )
38 
39 flags.DEFINE_string(
40     "saved_model_input",
41     None,
42     "The directory containing the input pretrained model to finetune.",
43 )
44 
45 flags.DEFINE_string(
46     "saved_model_output", None, "The directory to save the finetuned model in."
47 )
48 
49 
50 flags.DEFINE_string(
51     "task_name", "mrpc", "The name of the GLUE task to finetune on."
52 )
53 
54 flags.DEFINE_bool(
55     "do_lower_case",
56     True,
57     "Whether to lower case the input text. Should be True for uncased "
58     "models and False for cased models.",
59 )
60 
61 flags.DEFINE_bool(
62     "do_evaluation",
63     True,
64     "Whether to run evaluation on the validation set for a given task.",
65 )
66 
67 flags.DEFINE_integer("batch_size", 32, "The batch size.")
68 
69 flags.DEFINE_integer("epochs", 3, "The number of training epochs.")
70 
71 flags.DEFINE_float("learning_rate", 2e-5, "The initial learning rate for Adam.")
72 
73 flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
74 
75 
76 def pack_inputs(
77     inputs,
78     seq_length,
79     start_of_sequence_id,
80     end_of_segment_id,
81     padding_id,
82 ):
83     # In case inputs weren't truncated (as they should have been),
84     # fall back to some ad-hoc truncation.
85     trimmed_segments = tftext.RoundRobinTrimmer(
86         seq_length - len(inputs) - 1
87     ).trim(inputs)
88     # Combine segments.
89     segments_combined, segment_ids = tftext.combine_segments(
90         trimmed_segments,
91         start_of_sequence_id=start_of_sequence_id,
92         end_of_segment_id=end_of_segment_id,
93     )
94     # Pad to dense Tensors.
95     input_word_ids, _ = tftext.pad_model_inputs(
96         segments_combined, seq_length, pad_value=padding_id
97     )
98     input_type_ids, input_mask = tftext.pad_model_inputs(
99         segment_ids, seq_length, pad_value=0
100     )
101     # Assemble nest of input tensors as expected by BERT model.
102     return {
103         "input_ids": input_word_ids,
104         "input_mask": input_mask,
105         "segment_ids": input_type_ids,
106     }
107 
108 
109 def load_data(task_name):
110     if task_name in ("cola", "sst2"):
111         feature_names = ("sentence",)
112     elif task_name in ("mrpc", "stsb", "rte", "wnli"):
113         feature_names = ("sentence1", "sentence2")
114     elif task_name in ("mnli", "mnli_matched", "mnli_mismatched"):
115         feature_names = ("premise", "hypothesis")
116     elif task_name in "qnli":
117         feature_names = ("question", "sentence")
118     elif task_name in "qqp":
119         feature_names = ("question1", "question2")
120     else:
121         raise ValueError(f"Unkown task_name {task_name}.")
122 
123     test_suffix = ""
124     if task_name in ("mnli", "mnli_matched"):
125         # For "mnli", just run default to "mnli_matched".
126         task_name = "mnli"
127         test_suffix = "_matched"
128     elif task_name in ("mnli_mismatched",):
129         task_name = "mnli"
130         test_suffix = "_mismatched"
131 
132     def to_tf_dataset(split):
133         # Format each sample as a tuple of string features and an int label.
134         features = tuple([split[f] for f in feature_names])
135         label = tf.cast(split["label"], tf.int32)
136         return tf.data.Dataset.from_tensor_slices((features, label))
137 
138     data = datasets.load_dataset("glue", task_name)
139     data.set_format(type="tensorflow")
140     train_ds = to_tf_dataset(data["train"])
141     test_ds = to_tf_dataset(data["test" + test_suffix])
142     validation_ds = to_tf_dataset(data["validation" + test_suffix])
143     return train_ds, test_ds, validation_ds
144 
145 
146 class BertClassificationFinetuner(keras.Model):
147     """Adds a classification head to a pre-trained BERT model for finetuning"""
148 
149     def __init__(self, bert_model, hidden_size, num_classes, **kwargs):
150         super().__init__(**kwargs)
151         self.bert_model = bert_model
152         self._pooler_layer = keras.layers.Dense(
153             hidden_size,
154             activation="tanh",
155             name="pooler",
156         )
157         self._logit_layer = tf.keras.layers.Dense(
158             num_classes,
159             name="logits",
160         )
161 
162     def call(self, inputs):
163         outputs = self.bert_model(inputs)
164         # Get the first [CLS] token from each output.
165         outputs = outputs[:, 0, :]
166         outputs = self._pooler_layer(outputs)
167         return self._logit_layer(outputs)
168 
169 
170 def main(_):
171     print(f"Reading input model from {FLAGS.saved_model_input}")
172     model = keras.models.load_model(FLAGS.saved_model_input)
173 
174     vocab = []
175     with open(FLAGS.vocab_file, "r") as vocab_file:
176         for line in vocab_file:
177             vocab.append(line.strip())
178     tokenizer = tftext.BertTokenizer(
179         FLAGS.vocab_file,
180         lower_case=FLAGS.do_lower_case,
181         token_out_type=tf.int32,
182     )
183     start_id = vocab.index("[CLS]")
184     end_id = vocab.index("[SEP]")
185     pad_id = vocab.index("[PAD]")
186 
187     with open(FLAGS.bert_config_file, "r") as bert_config_file:
188         bert_config = json.loads(bert_config_file.read())
189 
190     def preprocess_data(inputs, labels):
191         inputs = [tokenizer.tokenize(x).merge_dims(1, -1) for x in inputs]
192         inputs = pack_inputs(
193             inputs,
194             FLAGS.max_seq_length,
195             start_of_sequence_id=start_id,
196             end_of_segment_id=end_id,
197             padding_id=pad_id,
198         )
199         return inputs, labels
200 
201     # Read and preprocess GLUE task data.
202     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
203     train_ds = train_ds.batch(FLAGS.batch_size).map(
204         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
205     )
206     validation_ds = validation_ds.batch(FLAGS.batch_size).map(
207         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
208     )
209     test_ds = test_ds.batch(FLAGS.batch_size).map(
210         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
211     )
212 
213     finetuning_model = BertClassificationFinetuner(
214         bert_model=model,
215         hidden_size=bert_config["hidden_size"],
216         num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
217     )
218     finetuning_model.compile(
219         optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate),
220         loss="sparse_categorical_crossentropy",
221         metrics=["accuracy"],
222     )
223     finetuning_model.fit(
224         train_ds, epochs=FLAGS.epochs, validation_data=validation_ds
225     )
226 
227     if FLAGS.do_evaluation:
228         print("Evaluating on test set.")
229         finetuning_model.evaluate(test_ds)
230 
231     # TODO(mattdangerw): After incorporating keras_nlp tokenization, save an
232     # end-to-end model includeing preprocessing that operates on raw strings.
233     if FLAGS.saved_model_output:
234         print(f"Saving to {FLAGS.saved_model_output}")
235         finetuning_model.save(FLAGS.saved_model_output)
236 
237 
238 if __name__ == "__main__":
239     flags.mark_flag_as_required("vocab_file")
240     flags.mark_flag_as_required("bert_config_file")
241     flags.mark_flag_as_required("saved_model_input")
242     app.run(main)
243
```
--- 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/examples/bert/run_glue_finetuning.py b/examples/bert/run_glue_finetuning.py
--- a/examples/bert/run_glue_finetuning.py
+++ b/examples/bert/run_glue_finetuning.py
@@ -16,6 +16,7 @@
 import json
 
 import datasets
+import keras_tuner
 import tensorflow as tf
 import tensorflow_text as tftext
 from absl import app
@@ -68,8 +69,6 @@
 
 flags.DEFINE_integer("epochs", 3, "The number of training epochs.")
 
-flags.DEFINE_float("learning_rate", 2e-5, "The initial learning rate for Adam.")
-
 flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
 
 
@@ -167,9 +166,32 @@
         return self._logit_layer(outputs)
 
 
+class BertHyperModel(keras_tuner.HyperModel):
+    """Creates a hypermodel to help with the search space for finetuning."""
+
+    def __init__(self, bert_config):
+        self.bert_config = bert_config
+
+    def build(self, hp):
+        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
+        bert_config = self.bert_config
+        finetuning_model = BertClassificationFinetuner(
+            bert_model=model,
+            hidden_size=bert_config["hidden_size"],
+            num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
+        )
+        finetuning_model.compile(
+            optimizer=keras.optimizers.Adam(
+                learning_rate=hp.Choice("lr", [5e-5, 4e-5, 3e-5, 2e-5])
+            ),
+            loss="sparse_categorical_crossentropy",
+            metrics=["accuracy"],
+        )
+        return finetuning_model
+
+
 def main(_):
     print(f"Reading input model from {FLAGS.saved_model_input}")
-    model = keras.models.load_model(FLAGS.saved_model_input)
 
     vocab = []
     with open(FLAGS.vocab_file, "r") as vocab_file:
@@ -200,6 +222,7 @@
 
     # Read and preprocess GLUE task data.
     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
+
     train_ds = train_ds.batch(FLAGS.batch_size).map(
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
@@ -210,18 +233,27 @@
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
 
-    finetuning_model = BertClassificationFinetuner(
-        bert_model=model,
-        hidden_size=bert_config["hidden_size"],
-        num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
-    )
-    finetuning_model.compile(
-        optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate),
-        loss="sparse_categorical_crossentropy",
-        metrics=["accuracy"],
+    # Create a hypermodel object for a RandomSearch.
+    hypermodel = BertHyperModel(bert_config)
+
+    # Initialize the random search over the 4 learning rate parameters, for 4
+    # trials and 3 epochs for each trial.
+    tuner = keras_tuner.RandomSearch(
+        hypermodel=hypermodel,
+        objective=keras_tuner.Objective("val_loss", direction="min"),
+        max_trials=4,
+        overwrite=True,
+        project_name="hyperparameter_tuner_results",
     )
-    finetuning_model.fit(
-        train_ds, epochs=FLAGS.epochs, validation_data=validation_ds
+
+    tuner.search(train_ds, epochs=FLAGS.epochs, validation_data=validation_ds)
+
+    # Extract the best hyperparameters after the search.
+    best_hp = tuner.get_best_hyperparameters()[0]
+    finetuning_model = tuner.get_best_models()[0]
+
+    print(
+        f"The best hyperparameters found are:\nLearning Rate: {best_hp['lr']}"
     )
 
     if FLAGS.do_evaluation:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -53,6 +53,7 @@
             "datasets",  # For GLUE in BERT example.
             "nltk",
             "wikiextractor",
+            "keras-tuner",
         ],
     },
     classifiers=[

{"golden_diff": "diff --git a/examples/bert/run_glue_finetuning.py b/examples/bert/run_glue_finetuning.py\n--- a/examples/bert/run_glue_finetuning.py\n+++ b/examples/bert/run_glue_finetuning.py\n@@ -16,6 +16,7 @@\n import json\n \n import datasets\n+import keras_tuner\n import tensorflow as tf\n import tensorflow_text as tftext\n from absl import app\n@@ -68,8 +69,6 @@\n \n flags.DEFINE_integer(\"epochs\", 3, \"The number of training epochs.\")\n \n-flags.DEFINE_float(\"learning_rate\", 2e-5, \"The initial learning rate for Adam.\")\n-\n flags.DEFINE_integer(\"max_seq_length\", 128, \"Maximum sequence length.\")\n \n \n@@ -167,9 +166,32 @@\n         return self._logit_layer(outputs)\n \n \n+class BertHyperModel(keras_tuner.HyperModel):\n+    \"\"\"Creates a hypermodel to help with the search space for finetuning.\"\"\"\n+\n+    def __init__(self, bert_config):\n+        self.bert_config = bert_config\n+\n+    def build(self, hp):\n+        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)\n+        bert_config = self.bert_config\n+        finetuning_model = BertClassificationFinetuner(\n+            bert_model=model,\n+            hidden_size=bert_config[\"hidden_size\"],\n+            num_classes=3 if FLAGS.task_name in (\"mnli\", \"ax\") else 2,\n+        )\n+        finetuning_model.compile(\n+            optimizer=keras.optimizers.Adam(\n+                learning_rate=hp.Choice(\"lr\", [5e-5, 4e-5, 3e-5, 2e-5])\n+            ),\n+            loss=\"sparse_categorical_crossentropy\",\n+            metrics=[\"accuracy\"],\n+        )\n+        return finetuning_model\n+\n+\n def main(_):\n     print(f\"Reading input model from {FLAGS.saved_model_input}\")\n-    model = keras.models.load_model(FLAGS.saved_model_input)\n \n     vocab = []\n     with open(FLAGS.vocab_file, \"r\") as vocab_file:\n@@ -200,6 +222,7 @@\n \n     # Read and preprocess GLUE task data.\n     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)\n+\n     train_ds = train_ds.batch(FLAGS.batch_size).map(\n         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n     )\n@@ -210,18 +233,27 @@\n         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n     )\n \n-    finetuning_model = BertClassificationFinetuner(\n-        bert_model=model,\n-        hidden_size=bert_config[\"hidden_size\"],\n-        num_classes=3 if FLAGS.task_name in (\"mnli\", \"ax\") else 2,\n-    )\n-    finetuning_model.compile(\n-        optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate),\n-        loss=\"sparse_categorical_crossentropy\",\n-        metrics=[\"accuracy\"],\n+    # Create a hypermodel object for a RandomSearch.\n+    hypermodel = BertHyperModel(bert_config)\n+\n+    # Initialize the random search over the 4 learning rate parameters, for 4\n+    # trials and 3 epochs for each trial.\n+    tuner = keras_tuner.RandomSearch(\n+        hypermodel=hypermodel,\n+        objective=keras_tuner.Objective(\"val_loss\", direction=\"min\"),\n+        max_trials=4,\n+        overwrite=True,\n+        project_name=\"hyperparameter_tuner_results\",\n     )\n-    finetuning_model.fit(\n-        train_ds, epochs=FLAGS.epochs, validation_data=validation_ds\n+\n+    tuner.search(train_ds, epochs=FLAGS.epochs, validation_data=validation_ds)\n+\n+    # Extract the best hyperparameters after the search.\n+    best_hp = tuner.get_best_hyperparameters()[0]\n+    finetuning_model = tuner.get_best_models()[0]\n+\n+    print(\n+        f\"The best hyperparameters found are:\\nLearning Rate: {best_hp['lr']}\"\n     )\n \n     if FLAGS.do_evaluation:\ndiff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -53,6 +53,7 @@\n             \"datasets\",  # For GLUE in BERT example.\n             \"nltk\",\n             \"wikiextractor\",\n+            \"keras-tuner\",\n         ],\n     },\n     classifiers=[\n", "issue": "Use KerasTuner to hyper-parameter search for the BERT finetuning script\nFrom the [BERT paper](https://arxiv.org/pdf/1810.04805.pdf)...\r\n\r\n```\r\nFor fine-tuning, most model hyperparameters are\r\nthe same as in pre-training, with the exception of\r\nthe batch size, learning rate, and number of training\r\nepochs. The dropout probability was always\r\nkept at 0.1. The optimal hyperparameter values\r\nare task-specific, but we found the following range\r\nof possible values to work well across all tasks:\r\n\r\n\u2022 Batch size: 16, 32\r\n\u2022 Learning rate (Adam): 5e-5, 3e-5, 2e-5\r\n\u2022 Number of epochs: 2, 3, 4\r\n```\r\n\r\nWe should allow our [BERT finetuning script](https://github.com/keras-team/keras-nlp/blob/master/examples/bert/run_glue_finetuning.py) to do this search automatically. [KerasTuner](https://keras.io/keras_tuner/) is a good fit for this.\r\n\r\nSteps:\r\n - [ ] Add an setup.py `examples` dependency on keras-tuner.\r\n - [ ] Remove epochs, batch size and learning rage arguments from run_glue_finetuning.py.\r\n - [ ] Use keras tuner to hyperparemeter search on the above value ranges with the validation set.\n", "before_files": [{"content": "# Copyright 2021 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\"\"\"Setup script.\"\"\"\n\nimport pathlib\n\nfrom setuptools import find_packages\nfrom setuptools import setup\n\nHERE = pathlib.Path(__file__).parent\nREADME = (HERE / \"README.md\").read_text()\n\nsetup(\n    name=\"keras-nlp\",\n    description=(\n        \"Industry-strength Natural Language Processing extensions for Keras.\"\n    ),\n    long_description=README,\n    long_description_content_type=\"text/markdown\",\n    version=\"0.2.0-dev.1\",\n    url=\"https://github.com/keras-team/keras-nlp\",\n    author=\"Keras team\",\n    author_email=\"keras-nlp@google.com\",\n    license=\"Apache License 2.0\",\n    install_requires=[\n        \"absl-py\",\n        \"numpy\",\n        \"packaging\",\n        \"tensorflow\",\n        \"tensorflow_text\",\n    ],\n    extras_require={\n        \"tests\": [\n            \"black\",\n            \"flake8\",\n            \"isort\",\n            \"pytest\",\n            \"pytest-cov\",\n        ],\n        \"examples\": [\n            \"datasets\",  # For GLUE in BERT example.\n            \"nltk\",\n            \"wikiextractor\",\n        ],\n    },\n    classifiers=[\n        \"Programming Language :: Python\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Operating System :: Unix\",\n        \"Operating System :: Microsoft :: Windows\",\n        \"Operating System :: MacOS\",\n        \"Intended Audience :: Science/Research\",\n        \"Topic :: Scientific/Engineering\",\n        \"Topic :: Software Development\",\n    ],\n    packages=find_packages(exclude=(\"*_test.py\",)),\n)\n", "path": "setup.py"}, {"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\"\"\"Run finetuning on a GLUE task.\"\"\"\n\nimport json\n\nimport datasets\nimport tensorflow as tf\nimport tensorflow_text as tftext\nfrom absl import app\nfrom absl import flags\nfrom tensorflow import keras\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    \"bert_config_file\",\n    None,\n    \"The json config file for the bert model parameters.\",\n)\n\nflags.DEFINE_string(\n    \"vocab_file\",\n    None,\n    \"The vocabulary file that the BERT model was trained on.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_input\",\n    None,\n    \"The directory containing the input pretrained model to finetune.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_output\", None, \"The directory to save the finetuned model in.\"\n)\n\n\nflags.DEFINE_string(\n    \"task_name\", \"mrpc\", \"The name of the GLUE task to finetune on.\"\n)\n\nflags.DEFINE_bool(\n    \"do_lower_case\",\n    True,\n    \"Whether to lower case the input text. Should be True for uncased \"\n    \"models and False for cased models.\",\n)\n\nflags.DEFINE_bool(\n    \"do_evaluation\",\n    True,\n    \"Whether to run evaluation on the validation set for a given task.\",\n)\n\nflags.DEFINE_integer(\"batch_size\", 32, \"The batch size.\")\n\nflags.DEFINE_integer(\"epochs\", 3, \"The number of training epochs.\")\n\nflags.DEFINE_float(\"learning_rate\", 2e-5, \"The initial learning rate for Adam.\")\n\nflags.DEFINE_integer(\"max_seq_length\", 128, \"Maximum sequence length.\")\n\n\ndef pack_inputs(\n    inputs,\n    seq_length,\n    start_of_sequence_id,\n    end_of_segment_id,\n    padding_id,\n):\n    # In case inputs weren't truncated (as they should have been),\n    # fall back to some ad-hoc truncation.\n    trimmed_segments = tftext.RoundRobinTrimmer(\n        seq_length - len(inputs) - 1\n    ).trim(inputs)\n    # Combine segments.\n    segments_combined, segment_ids = tftext.combine_segments(\n        trimmed_segments,\n        start_of_sequence_id=start_of_sequence_id,\n        end_of_segment_id=end_of_segment_id,\n    )\n    # Pad to dense Tensors.\n    input_word_ids, _ = tftext.pad_model_inputs(\n        segments_combined, seq_length, pad_value=padding_id\n    )\n    input_type_ids, input_mask = tftext.pad_model_inputs(\n        segment_ids, seq_length, pad_value=0\n    )\n    # Assemble nest of input tensors as expected by BERT model.\n    return {\n        \"input_ids\": input_word_ids,\n        \"input_mask\": input_mask,\n        \"segment_ids\": input_type_ids,\n    }\n\n\ndef load_data(task_name):\n    if task_name in (\"cola\", \"sst2\"):\n        feature_names = (\"sentence\",)\n    elif task_name in (\"mrpc\", \"stsb\", \"rte\", \"wnli\"):\n        feature_names = (\"sentence1\", \"sentence2\")\n    elif task_name in (\"mnli\", \"mnli_matched\", \"mnli_mismatched\"):\n        feature_names = (\"premise\", \"hypothesis\")\n    elif task_name in \"qnli\":\n        feature_names = (\"question\", \"sentence\")\n    elif task_name in \"qqp\":\n        feature_names = (\"question1\", \"question2\")\n    else:\n        raise ValueError(f\"Unkown task_name {task_name}.\")\n\n    test_suffix = \"\"\n    if task_name in (\"mnli\", \"mnli_matched\"):\n        # For \"mnli\", just run default to \"mnli_matched\".\n        task_name = \"mnli\"\n        test_suffix = \"_matched\"\n    elif task_name in (\"mnli_mismatched\",):\n        task_name = \"mnli\"\n        test_suffix = \"_mismatched\"\n\n    def to_tf_dataset(split):\n        # Format each sample as a tuple of string features and an int label.\n        features = tuple([split[f] for f in feature_names])\n        label = tf.cast(split[\"label\"], tf.int32)\n        return tf.data.Dataset.from_tensor_slices((features, label))\n\n    data = datasets.load_dataset(\"glue\", task_name)\n    data.set_format(type=\"tensorflow\")\n    train_ds = to_tf_dataset(data[\"train\"])\n    test_ds = to_tf_dataset(data[\"test\" + test_suffix])\n    validation_ds = to_tf_dataset(data[\"validation\" + test_suffix])\n    return train_ds, test_ds, validation_ds\n\n\nclass BertClassificationFinetuner(keras.Model):\n    \"\"\"Adds a classification head to a pre-trained BERT model for finetuning\"\"\"\n\n    def __init__(self, bert_model, hidden_size, num_classes, **kwargs):\n        super().__init__(**kwargs)\n        self.bert_model = bert_model\n        self._pooler_layer = keras.layers.Dense(\n            hidden_size,\n            activation=\"tanh\",\n            name=\"pooler\",\n        )\n        self._logit_layer = tf.keras.layers.Dense(\n            num_classes,\n            name=\"logits\",\n        )\n\n    def call(self, inputs):\n        outputs = self.bert_model(inputs)\n        # Get the first [CLS] token from each output.\n        outputs = outputs[:, 0, :]\n        outputs = self._pooler_layer(outputs)\n        return self._logit_layer(outputs)\n\n\ndef main(_):\n    print(f\"Reading input model from {FLAGS.saved_model_input}\")\n    model = keras.models.load_model(FLAGS.saved_model_input)\n\n    vocab = []\n    with open(FLAGS.vocab_file, \"r\") as vocab_file:\n        for line in vocab_file:\n            vocab.append(line.strip())\n    tokenizer = tftext.BertTokenizer(\n        FLAGS.vocab_file,\n        lower_case=FLAGS.do_lower_case,\n        token_out_type=tf.int32,\n    )\n    start_id = vocab.index(\"[CLS]\")\n    end_id = vocab.index(\"[SEP]\")\n    pad_id = vocab.index(\"[PAD]\")\n\n    with open(FLAGS.bert_config_file, \"r\") as bert_config_file:\n        bert_config = json.loads(bert_config_file.read())\n\n    def preprocess_data(inputs, labels):\n        inputs = [tokenizer.tokenize(x).merge_dims(1, -1) for x in inputs]\n        inputs = pack_inputs(\n            inputs,\n            FLAGS.max_seq_length,\n            start_of_sequence_id=start_id,\n            end_of_segment_id=end_id,\n            padding_id=pad_id,\n        )\n        return inputs, labels\n\n    # Read and preprocess GLUE task data.\n    train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)\n    train_ds = train_ds.batch(FLAGS.batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    validation_ds = validation_ds.batch(FLAGS.batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    test_ds = test_ds.batch(FLAGS.batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n\n    finetuning_model = BertClassificationFinetuner(\n        bert_model=model,\n        hidden_size=bert_config[\"hidden_size\"],\n        num_classes=3 if FLAGS.task_name in (\"mnli\", \"ax\") else 2,\n    )\n    finetuning_model.compile(\n        optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate),\n        loss=\"sparse_categorical_crossentropy\",\n        metrics=[\"accuracy\"],\n    )\n    finetuning_model.fit(\n        train_ds, epochs=FLAGS.epochs, validation_data=validation_ds\n    )\n\n    if FLAGS.do_evaluation:\n        print(\"Evaluating on test set.\")\n        finetuning_model.evaluate(test_ds)\n\n    # TODO(mattdangerw): After incorporating keras_nlp tokenization, save an\n    # end-to-end model includeing preprocessing that operates on raw strings.\n    if FLAGS.saved_model_output:\n        print(f\"Saving to {FLAGS.saved_model_output}\")\n        finetuning_model.save(FLAGS.saved_model_output)\n\n\nif __name__ == \"__main__\":\n    flags.mark_flag_as_required(\"vocab_file\")\n    flags.mark_flag_as_required(\"bert_config_file\")\n    flags.mark_flag_as_required(\"saved_model_input\")\n    app.run(main)\n", "path": "examples/bert/run_glue_finetuning.py"}], "after_files": [{"content": "# Copyright 2021 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\"\"\"Setup script.\"\"\"\n\nimport pathlib\n\nfrom setuptools import find_packages\nfrom setuptools import setup\n\nHERE = pathlib.Path(__file__).parent\nREADME = (HERE / \"README.md\").read_text()\n\nsetup(\n    name=\"keras-nlp\",\n    description=(\n        \"Industry-strength Natural Language Processing extensions for Keras.\"\n    ),\n    long_description=README,\n    long_description_content_type=\"text/markdown\",\n    version=\"0.2.0-dev.1\",\n    url=\"https://github.com/keras-team/keras-nlp\",\n    author=\"Keras team\",\n    author_email=\"keras-nlp@google.com\",\n    license=\"Apache License 2.0\",\n    install_requires=[\n        \"absl-py\",\n        \"numpy\",\n        \"packaging\",\n        \"tensorflow\",\n        \"tensorflow_text\",\n    ],\n    extras_require={\n        \"tests\": [\n            \"black\",\n            \"flake8\",\n            \"isort\",\n            \"pytest\",\n            \"pytest-cov\",\n        ],\n        \"examples\": [\n            \"datasets\",  # For GLUE in BERT example.\n            \"nltk\",\n            \"wikiextractor\",\n            \"keras-tuner\",\n        ],\n    },\n    classifiers=[\n        \"Programming Language :: Python\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Operating System :: Unix\",\n        \"Operating System :: Microsoft :: Windows\",\n        \"Operating System :: MacOS\",\n        \"Intended Audience :: Science/Research\",\n        \"Topic :: Scientific/Engineering\",\n        \"Topic :: Software Development\",\n    ],\n    packages=find_packages(exclude=(\"*_test.py\",)),\n)\n", "path": "setup.py"}, {"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\"\"\"Run finetuning on a GLUE task.\"\"\"\n\nimport json\n\nimport datasets\nimport keras_tuner\nimport tensorflow as tf\nimport tensorflow_text as tftext\nfrom absl import app\nfrom absl import flags\nfrom tensorflow import keras\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    \"bert_config_file\",\n    None,\n    \"The json config file for the bert model parameters.\",\n)\n\nflags.DEFINE_string(\n    \"vocab_file\",\n    None,\n    \"The vocabulary file that the BERT model was trained on.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_input\",\n    None,\n    \"The directory containing the input pretrained model to finetune.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_output\", None, \"The directory to save the finetuned model in.\"\n)\n\n\nflags.DEFINE_string(\n    \"task_name\", \"mrpc\", \"The name of the GLUE task to finetune on.\"\n)\n\nflags.DEFINE_bool(\n    \"do_lower_case\",\n    True,\n    \"Whether to lower case the input text. Should be True for uncased \"\n    \"models and False for cased models.\",\n)\n\nflags.DEFINE_bool(\n    \"do_evaluation\",\n    True,\n    \"Whether to run evaluation on the validation set for a given task.\",\n)\n\nflags.DEFINE_integer(\"batch_size\", 32, \"The batch size.\")\n\nflags.DEFINE_integer(\"epochs\", 3, \"The number of training epochs.\")\n\nflags.DEFINE_integer(\"max_seq_length\", 128, \"Maximum sequence length.\")\n\n\ndef pack_inputs(\n    inputs,\n    seq_length,\n    start_of_sequence_id,\n    end_of_segment_id,\n    padding_id,\n):\n    # In case inputs weren't truncated (as they should have been),\n    # fall back to some ad-hoc truncation.\n    trimmed_segments = tftext.RoundRobinTrimmer(\n        seq_length - len(inputs) - 1\n    ).trim(inputs)\n    # Combine segments.\n    segments_combined, segment_ids = tftext.combine_segments(\n        trimmed_segments,\n        start_of_sequence_id=start_of_sequence_id,\n        end_of_segment_id=end_of_segment_id,\n    )\n    # Pad to dense Tensors.\n    input_word_ids, _ = tftext.pad_model_inputs(\n        segments_combined, seq_length, pad_value=padding_id\n    )\n    input_type_ids, input_mask = tftext.pad_model_inputs(\n        segment_ids, seq_length, pad_value=0\n    )\n    # Assemble nest of input tensors as expected by BERT model.\n    return {\n        \"input_ids\": input_word_ids,\n        \"input_mask\": input_mask,\n        \"segment_ids\": input_type_ids,\n    }\n\n\ndef load_data(task_name):\n    if task_name in (\"cola\", \"sst2\"):\n        feature_names = (\"sentence\",)\n    elif task_name in (\"mrpc\", \"stsb\", \"rte\", \"wnli\"):\n        feature_names = (\"sentence1\", \"sentence2\")\n    elif task_name in (\"mnli\", \"mnli_matched\", \"mnli_mismatched\"):\n        feature_names = (\"premise\", \"hypothesis\")\n    elif task_name in \"qnli\":\n        feature_names = (\"question\", \"sentence\")\n    elif task_name in \"qqp\":\n        feature_names = (\"question1\", \"question2\")\n    else:\n        raise ValueError(f\"Unkown task_name {task_name}.\")\n\n    test_suffix = \"\"\n    if task_name in (\"mnli\", \"mnli_matched\"):\n        # For \"mnli\", just run default to \"mnli_matched\".\n        task_name = \"mnli\"\n        test_suffix = \"_matched\"\n    elif task_name in (\"mnli_mismatched\",):\n        task_name = \"mnli\"\n        test_suffix = \"_mismatched\"\n\n    def to_tf_dataset(split):\n        # Format each sample as a tuple of string features and an int label.\n        features = tuple([split[f] for f in feature_names])\n        label = tf.cast(split[\"label\"], tf.int32)\n        return tf.data.Dataset.from_tensor_slices((features, label))\n\n    data = datasets.load_dataset(\"glue\", task_name)\n    data.set_format(type=\"tensorflow\")\n    train_ds = to_tf_dataset(data[\"train\"])\n    test_ds = to_tf_dataset(data[\"test\" + test_suffix])\n    validation_ds = to_tf_dataset(data[\"validation\" + test_suffix])\n    return train_ds, test_ds, validation_ds\n\n\nclass BertClassificationFinetuner(keras.Model):\n    \"\"\"Adds a classification head to a pre-trained BERT model for finetuning\"\"\"\n\n    def __init__(self, bert_model, hidden_size, num_classes, **kwargs):\n        super().__init__(**kwargs)\n        self.bert_model = bert_model\n        self._pooler_layer = keras.layers.Dense(\n            hidden_size,\n            activation=\"tanh\",\n            name=\"pooler\",\n        )\n        self._logit_layer = tf.keras.layers.Dense(\n            num_classes,\n            name=\"logits\",\n        )\n\n    def call(self, inputs):\n        outputs = self.bert_model(inputs)\n        # Get the first [CLS] token from each output.\n        outputs = outputs[:, 0, :]\n        outputs = self._pooler_layer(outputs)\n        return self._logit_layer(outputs)\n\n\nclass BertHyperModel(keras_tuner.HyperModel):\n    \"\"\"Creates a hypermodel to help with the search space for finetuning.\"\"\"\n\n    def __init__(self, bert_config):\n        self.bert_config = bert_config\n\n    def build(self, hp):\n        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)\n        bert_config = self.bert_config\n        finetuning_model = BertClassificationFinetuner(\n            bert_model=model,\n            hidden_size=bert_config[\"hidden_size\"],\n            num_classes=3 if FLAGS.task_name in (\"mnli\", \"ax\") else 2,\n        )\n        finetuning_model.compile(\n            optimizer=keras.optimizers.Adam(\n                learning_rate=hp.Choice(\"lr\", [5e-5, 4e-5, 3e-5, 2e-5])\n            ),\n            loss=\"sparse_categorical_crossentropy\",\n            metrics=[\"accuracy\"],\n        )\n        return finetuning_model\n\n\ndef main(_):\n    print(f\"Reading input model from {FLAGS.saved_model_input}\")\n\n    vocab = []\n    with open(FLAGS.vocab_file, \"r\") as vocab_file:\n        for line in vocab_file:\n            vocab.append(line.strip())\n    tokenizer = tftext.BertTokenizer(\n        FLAGS.vocab_file,\n        lower_case=FLAGS.do_lower_case,\n        token_out_type=tf.int32,\n    )\n    start_id = vocab.index(\"[CLS]\")\n    end_id = vocab.index(\"[SEP]\")\n    pad_id = vocab.index(\"[PAD]\")\n\n    with open(FLAGS.bert_config_file, \"r\") as bert_config_file:\n        bert_config = json.loads(bert_config_file.read())\n\n    def preprocess_data(inputs, labels):\n        inputs = [tokenizer.tokenize(x).merge_dims(1, -1) for x in inputs]\n        inputs = pack_inputs(\n            inputs,\n            FLAGS.max_seq_length,\n            start_of_sequence_id=start_id,\n            end_of_segment_id=end_id,\n            padding_id=pad_id,\n        )\n        return inputs, labels\n\n    # Read and preprocess GLUE task data.\n    train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)\n\n    train_ds = train_ds.batch(FLAGS.batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    validation_ds = validation_ds.batch(FLAGS.batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    test_ds = test_ds.batch(FLAGS.batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n\n    # Create a hypermodel object for a RandomSearch.\n    hypermodel = BertHyperModel(bert_config)\n\n    # Initialize the random search over the 4 learning rate parameters, for 4\n    # trials and 3 epochs for each trial.\n    tuner = keras_tuner.RandomSearch(\n        hypermodel=hypermodel,\n        objective=keras_tuner.Objective(\"val_loss\", direction=\"min\"),\n        max_trials=4,\n        overwrite=True,\n        project_name=\"hyperparameter_tuner_results\",\n    )\n\n    tuner.search(train_ds, epochs=FLAGS.epochs, validation_data=validation_ds)\n\n    # Extract the best hyperparameters after the search.\n    best_hp = tuner.get_best_hyperparameters()[0]\n    finetuning_model = tuner.get_best_models()[0]\n\n    print(\n        f\"The best hyperparameters found are:\\nLearning Rate: {best_hp['lr']}\"\n    )\n\n    if FLAGS.do_evaluation:\n        print(\"Evaluating on test set.\")\n        finetuning_model.evaluate(test_ds)\n\n    # TODO(mattdangerw): After incorporating keras_nlp tokenization, save an\n    # end-to-end model includeing preprocessing that operates on raw strings.\n    if FLAGS.saved_model_output:\n        print(f\"Saving to {FLAGS.saved_model_output}\")\n        finetuning_model.save(FLAGS.saved_model_output)\n\n\nif __name__ == \"__main__\":\n    flags.mark_flag_as_required(\"vocab_file\")\n    flags.mark_flag_as_required(\"bert_config_file\")\n    flags.mark_flag_as_required(\"saved_model_input\")\n    app.run(main)\n", "path": "examples/bert/run_glue_finetuning.py"}]}