We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Update our glue evaluation script to use tensorflow datasets
When it was originally written, glue for tfds was totally broken, so we used huggingface datasets.

This appears to be fixed, so let's used tfds instead.

We should update the line here as well to flip the dependency.
https://github.com/keras-team/keras-nlp/blob/master/setup.py#L55

--- 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.3.0",
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             "rouge-score",
52             "sentencepiece",
53         ],
54         "examples": [
55             "datasets",  # For GLUE in BERT example.
56             "nltk",
57             "wikiextractor",
58             "keras-tuner",
59         ],
60     },
61     classifiers=[
62         "Programming Language :: Python",
63         "Programming Language :: Python :: 3.7",
64         "Operating System :: Unix",
65         "Operating System :: Microsoft :: Windows",
66         "Operating System :: MacOS",
67         "Intended Audience :: Science/Research",
68         "Topic :: Scientific/Engineering",
69         "Topic :: Software Development",
70     ],
71     packages=find_packages(exclude=("*_test.py",)),
72 )
73
```

Path: `examples/bert/bert_finetune_glue.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 tempfile
17 
18 import datasets
19 import keras_tuner
20 import tensorflow as tf
21 from absl import app
22 from absl import flags
23 from tensorflow import keras
24 
25 import keras_nlp
26 from examples.bert.bert_config import FINETUNING_CONFIG
27 from examples.bert.bert_config import PREPROCESSING_CONFIG
28 
29 FLAGS = flags.FLAGS
30 
31 flags.DEFINE_string(
32     "vocab_file",
33     None,
34     "The vocabulary file for tokenization.",
35 )
36 
37 flags.DEFINE_string(
38     "saved_model_input",
39     None,
40     "The directory to load the pretrained model.",
41 )
42 
43 flags.DEFINE_string(
44     "saved_model_output",
45     None,
46     "The directory to save the finetuned model.",
47 )
48 
49 flags.DEFINE_string(
50     "task_name",
51     "mrpc",
52     "The name of the GLUE task to finetune on.",
53 )
54 
55 flags.DEFINE_bool(
56     "do_lower_case",
57     True,
58     "Whether to lower case the input text.",
59 )
60 
61 flags.DEFINE_bool(
62     "do_evaluation",
63     True,
64     "Whether to run evaluation on test data.",
65 )
66 
67 
68 def load_data(task_name):
69     if task_name in ("cola", "sst2"):
70         feature_names = ("sentence",)
71     elif task_name in ("mrpc", "stsb", "rte", "wnli"):
72         feature_names = ("sentence1", "sentence2")
73     elif task_name in ("mnli", "mnli_matched", "mnli_mismatched"):
74         feature_names = ("premise", "hypothesis")
75     elif task_name in "qnli":
76         feature_names = ("question", "sentence")
77     elif task_name in "qqp":
78         feature_names = ("question1", "question2")
79     else:
80         raise ValueError(f"Unkown task_name {task_name}.")
81 
82     test_suffix = ""
83     if task_name in ("mnli", "mnli_matched"):
84         # For "mnli", just run default to "mnli_matched".
85         task_name = "mnli"
86         test_suffix = "_matched"
87     elif task_name in ("mnli_mismatched",):
88         task_name = "mnli"
89         test_suffix = "_mismatched"
90 
91     def to_tf_dataset(split):
92         # Format each sample as a tuple of string features and an int label.
93         features = tuple([split[f] for f in feature_names])
94         label = tf.cast(split["label"], tf.int32)
95         return tf.data.Dataset.from_tensor_slices((features, label))
96 
97     data = datasets.load_dataset("glue", task_name)
98     data.set_format(type="tensorflow")
99     train_ds = to_tf_dataset(data["train"])
100     test_ds = to_tf_dataset(data["test" + test_suffix])
101     validation_ds = to_tf_dataset(data["validation" + test_suffix])
102     return train_ds, test_ds, validation_ds
103 
104 
105 class BertHyperModel(keras_tuner.HyperModel):
106     """Creates a hypermodel to help with the search space for finetuning."""
107 
108     def build(self, hp):
109         model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
110         finetuning_model = keras_nlp.models.BertClassifier(
111             base_model=model,
112             num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
113         )
114         finetuning_model.compile(
115             optimizer=keras.optimizers.Adam(
116                 learning_rate=hp.Choice(
117                     "lr", FINETUNING_CONFIG["learning_rates"]
118                 ),
119             ),
120             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
121             metrics=[keras.metrics.SparseCategoricalAccuracy()],
122         )
123         return finetuning_model
124 
125 
126 def main(_):
127     print(f"Reading input model from {FLAGS.saved_model_input}")
128 
129     tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
130         vocabulary=FLAGS.vocab_file,
131         lowercase=FLAGS.do_lower_case,
132     )
133     packer = keras_nlp.layers.MultiSegmentPacker(
134         sequence_length=PREPROCESSING_CONFIG["max_seq_length"],
135         start_value=tokenizer.token_to_id("[CLS]"),
136         end_value=tokenizer.token_to_id("[SEP]"),
137     )
138 
139     def preprocess_data(inputs, labels):
140         inputs = [tokenizer(x) for x in inputs]
141         token_ids, segment_ids = packer(inputs)
142         return {
143             "token_ids": token_ids,
144             "segment_ids": segment_ids,
145             "padding_mask": tf.cast(token_ids != 0, "int32"),
146         }, labels
147 
148     # Read and preprocess GLUE task data.
149     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
150 
151     batch_size = FINETUNING_CONFIG["batch_size"]
152     train_ds = train_ds.batch(batch_size).map(
153         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
154     )
155     validation_ds = validation_ds.batch(batch_size).map(
156         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
157     )
158     test_ds = test_ds.batch(batch_size).map(
159         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
160     )
161 
162     # Create a hypermodel object for a RandomSearch.
163     hypermodel = BertHyperModel()
164 
165     # Initialize the random search over the 4 learning rate parameters, for 4
166     # trials and 3 epochs for each trial.
167     tuner = keras_tuner.RandomSearch(
168         hypermodel=hypermodel,
169         objective=keras_tuner.Objective("val_loss", direction="min"),
170         max_trials=4,
171         overwrite=True,
172         project_name="hyperparameter_tuner_results",
173         directory=tempfile.mkdtemp(),
174     )
175 
176     tuner.search(
177         train_ds,
178         epochs=FINETUNING_CONFIG["epochs"],
179         validation_data=validation_ds,
180     )
181 
182     # Extract the best hyperparameters after the search.
183     best_hp = tuner.get_best_hyperparameters()[0]
184     finetuning_model = tuner.get_best_models()[0]
185 
186     print(
187         f"The best hyperparameters found are:\nLearning Rate: {best_hp['lr']}"
188     )
189 
190     if FLAGS.do_evaluation:
191         print("Evaluating on test set.")
192         finetuning_model.evaluate(test_ds)
193 
194     if FLAGS.saved_model_output:
195         print(f"Saving to {FLAGS.saved_model_output}")
196         finetuning_model.save(FLAGS.saved_model_output)
197 
198 
199 if __name__ == "__main__":
200     flags.mark_flag_as_required("vocab_file")
201     flags.mark_flag_as_required("saved_model_input")
202     app.run(main)
203
```
--- 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/bert_finetune_glue.py b/examples/bert/bert_finetune_glue.py
--- a/examples/bert/bert_finetune_glue.py
+++ b/examples/bert/bert_finetune_glue.py
@@ -15,9 +15,9 @@
 
 import tempfile
 
-import datasets
 import keras_tuner
 import tensorflow as tf
+import tensorflow_datasets as tfds
 from absl import app
 from absl import flags
 from tensorflow import keras
@@ -77,7 +77,7 @@
     elif task_name in "qqp":
         feature_names = ("question1", "question2")
     else:
-        raise ValueError(f"Unkown task_name {task_name}.")
+        raise ValueError(f"Unknown task_name {task_name}.")
 
     test_suffix = ""
     if task_name in ("mnli", "mnli_matched"):
@@ -88,17 +88,18 @@
         task_name = "mnli"
         test_suffix = "_mismatched"
 
-    def to_tf_dataset(split):
-        # Format each sample as a tuple of string features and an int label.
-        features = tuple([split[f] for f in feature_names])
-        label = tf.cast(split["label"], tf.int32)
-        return tf.data.Dataset.from_tensor_slices((features, label))
-
-    data = datasets.load_dataset("glue", task_name)
-    data.set_format(type="tensorflow")
-    train_ds = to_tf_dataset(data["train"])
-    test_ds = to_tf_dataset(data["test" + test_suffix])
-    validation_ds = to_tf_dataset(data["validation" + test_suffix])
+    def split_features(x):
+        return {feature_name: x[feature_name] for feature_name in feature_names}
+
+    train_ds, test_ds, validation_ds = tfds.load(
+        f"glue/{task_name}",
+        split=["train", "test" + test_suffix, "validation" + test_suffix],
+    )
+    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    test_ds = test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    validation_ds = validation_ds.map(
+        split_features, num_parallel_calls=tf.data.AUTOTUNE
+    )
     return train_ds, test_ds, validation_ds
 
 
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -52,7 +52,7 @@
             "sentencepiece",
         ],
         "examples": [
-            "datasets",  # For GLUE in BERT example.
+            "tensorflow_datasets",  # For GLUE in BERT example.
             "nltk",
             "wikiextractor",
             "keras-tuner",

{"golden_diff": "diff --git a/examples/bert/bert_finetune_glue.py b/examples/bert/bert_finetune_glue.py\n--- a/examples/bert/bert_finetune_glue.py\n+++ b/examples/bert/bert_finetune_glue.py\n@@ -15,9 +15,9 @@\n \n import tempfile\n \n-import datasets\n import keras_tuner\n import tensorflow as tf\n+import tensorflow_datasets as tfds\n from absl import app\n from absl import flags\n from tensorflow import keras\n@@ -77,7 +77,7 @@\n     elif task_name in \"qqp\":\n         feature_names = (\"question1\", \"question2\")\n     else:\n-        raise ValueError(f\"Unkown task_name {task_name}.\")\n+        raise ValueError(f\"Unknown task_name {task_name}.\")\n \n     test_suffix = \"\"\n     if task_name in (\"mnli\", \"mnli_matched\"):\n@@ -88,17 +88,18 @@\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+    def split_features(x):\n+        return {feature_name: x[feature_name] for feature_name in feature_names}\n+\n+    train_ds, test_ds, validation_ds = tfds.load(\n+        f\"glue/{task_name}\",\n+        split=[\"train\", \"test\" + test_suffix, \"validation\" + test_suffix],\n+    )\n+    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)\n+    test_ds = test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)\n+    validation_ds = validation_ds.map(\n+        split_features, num_parallel_calls=tf.data.AUTOTUNE\n+    )\n     return train_ds, test_ds, validation_ds\n \n \ndiff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -52,7 +52,7 @@\n             \"sentencepiece\",\n         ],\n         \"examples\": [\n-            \"datasets\",  # For GLUE in BERT example.\n+            \"tensorflow_datasets\",  # For GLUE in BERT example.\n             \"nltk\",\n             \"wikiextractor\",\n             \"keras-tuner\",\n", "issue": "Update our glue evaluation script to use tensorflow datasets\nWhen it was originally written, glue for tfds was totally broken, so we used huggingface datasets.\r\n\r\nThis appears to be fixed, so let's used tfds instead.\r\n\r\nWe should update the line here as well to flip the dependency.\r\nhttps://github.com/keras-team/keras-nlp/blob/master/setup.py#L55\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.3.0\",\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            \"rouge-score\",\n            \"sentencepiece\",\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 tempfile\n\nimport datasets\nimport keras_tuner\nimport tensorflow as tf\nfrom absl import app\nfrom absl import flags\nfrom tensorflow import keras\n\nimport keras_nlp\nfrom examples.bert.bert_config import FINETUNING_CONFIG\nfrom examples.bert.bert_config import PREPROCESSING_CONFIG\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    \"vocab_file\",\n    None,\n    \"The vocabulary file for tokenization.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_input\",\n    None,\n    \"The directory to load the pretrained model.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_output\",\n    None,\n    \"The directory to save the finetuned model.\",\n)\n\nflags.DEFINE_string(\n    \"task_name\",\n    \"mrpc\",\n    \"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.\",\n)\n\nflags.DEFINE_bool(\n    \"do_evaluation\",\n    True,\n    \"Whether to run evaluation on test data.\",\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 BertHyperModel(keras_tuner.HyperModel):\n    \"\"\"Creates a hypermodel to help with the search space for finetuning.\"\"\"\n\n    def build(self, hp):\n        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)\n        finetuning_model = keras_nlp.models.BertClassifier(\n            base_model=model,\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(\n                    \"lr\", FINETUNING_CONFIG[\"learning_rates\"]\n                ),\n            ),\n            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n            metrics=[keras.metrics.SparseCategoricalAccuracy()],\n        )\n        return finetuning_model\n\n\ndef main(_):\n    print(f\"Reading input model from {FLAGS.saved_model_input}\")\n\n    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(\n        vocabulary=FLAGS.vocab_file,\n        lowercase=FLAGS.do_lower_case,\n    )\n    packer = keras_nlp.layers.MultiSegmentPacker(\n        sequence_length=PREPROCESSING_CONFIG[\"max_seq_length\"],\n        start_value=tokenizer.token_to_id(\"[CLS]\"),\n        end_value=tokenizer.token_to_id(\"[SEP]\"),\n    )\n\n    def preprocess_data(inputs, labels):\n        inputs = [tokenizer(x) for x in inputs]\n        token_ids, segment_ids = packer(inputs)\n        return {\n            \"token_ids\": token_ids,\n            \"segment_ids\": segment_ids,\n            \"padding_mask\": tf.cast(token_ids != 0, \"int32\"),\n        }, labels\n\n    # Read and preprocess GLUE task data.\n    train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)\n\n    batch_size = FINETUNING_CONFIG[\"batch_size\"]\n    train_ds = train_ds.batch(batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    validation_ds = validation_ds.batch(batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    test_ds = test_ds.batch(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()\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        directory=tempfile.mkdtemp(),\n    )\n\n    tuner.search(\n        train_ds,\n        epochs=FINETUNING_CONFIG[\"epochs\"],\n        validation_data=validation_ds,\n    )\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    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(\"saved_model_input\")\n    app.run(main)\n", "path": "examples/bert/bert_finetune_glue.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.3.0\",\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            \"rouge-score\",\n            \"sentencepiece\",\n        ],\n        \"examples\": [\n            \"tensorflow_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 tempfile\n\nimport keras_tuner\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\nfrom absl import app\nfrom absl import flags\nfrom tensorflow import keras\n\nimport keras_nlp\nfrom examples.bert.bert_config import FINETUNING_CONFIG\nfrom examples.bert.bert_config import PREPROCESSING_CONFIG\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n    \"vocab_file\",\n    None,\n    \"The vocabulary file for tokenization.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_input\",\n    None,\n    \"The directory to load the pretrained model.\",\n)\n\nflags.DEFINE_string(\n    \"saved_model_output\",\n    None,\n    \"The directory to save the finetuned model.\",\n)\n\nflags.DEFINE_string(\n    \"task_name\",\n    \"mrpc\",\n    \"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.\",\n)\n\nflags.DEFINE_bool(\n    \"do_evaluation\",\n    True,\n    \"Whether to run evaluation on test data.\",\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\"Unknown 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 split_features(x):\n        return {feature_name: x[feature_name] for feature_name in feature_names}\n\n    train_ds, test_ds, validation_ds = tfds.load(\n        f\"glue/{task_name}\",\n        split=[\"train\", \"test\" + test_suffix, \"validation\" + test_suffix],\n    )\n    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)\n    test_ds = test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)\n    validation_ds = validation_ds.map(\n        split_features, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    return train_ds, test_ds, validation_ds\n\n\nclass BertHyperModel(keras_tuner.HyperModel):\n    \"\"\"Creates a hypermodel to help with the search space for finetuning.\"\"\"\n\n    def build(self, hp):\n        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)\n        finetuning_model = keras_nlp.models.BertClassifier(\n            base_model=model,\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(\n                    \"lr\", FINETUNING_CONFIG[\"learning_rates\"]\n                ),\n            ),\n            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n            metrics=[keras.metrics.SparseCategoricalAccuracy()],\n        )\n        return finetuning_model\n\n\ndef main(_):\n    print(f\"Reading input model from {FLAGS.saved_model_input}\")\n\n    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(\n        vocabulary=FLAGS.vocab_file,\n        lowercase=FLAGS.do_lower_case,\n    )\n    packer = keras_nlp.layers.MultiSegmentPacker(\n        sequence_length=PREPROCESSING_CONFIG[\"max_seq_length\"],\n        start_value=tokenizer.token_to_id(\"[CLS]\"),\n        end_value=tokenizer.token_to_id(\"[SEP]\"),\n    )\n\n    def preprocess_data(inputs, labels):\n        inputs = [tokenizer(x) for x in inputs]\n        token_ids, segment_ids = packer(inputs)\n        return {\n            \"token_ids\": token_ids,\n            \"segment_ids\": segment_ids,\n            \"padding_mask\": tf.cast(token_ids != 0, \"int32\"),\n        }, labels\n\n    # Read and preprocess GLUE task data.\n    train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)\n\n    batch_size = FINETUNING_CONFIG[\"batch_size\"]\n    train_ds = train_ds.batch(batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    validation_ds = validation_ds.batch(batch_size).map(\n        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE\n    )\n    test_ds = test_ds.batch(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()\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        directory=tempfile.mkdtemp(),\n    )\n\n    tuner.search(\n        train_ds,\n        epochs=FINETUNING_CONFIG[\"epochs\"],\n        validation_data=validation_ds,\n    )\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    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(\"saved_model_input\")\n    app.run(main)\n", "path": "examples/bert/bert_finetune_glue.py"}]}