Repository: keras-team/keras-nlp

Issue:
Improve our continuous testing for model presets
Opening an issue to track the changes proposed on https://github.com/keras-team/keras-nlp/pull/357, as it has gotten slightly larger in scope.

I would like to propose the following changes to our "network_tests" for presets:

 - We collocate the preset testing within the model directory, and use test annotations to control how they are run.
 - We run the smallest available preset (per model) continuously on GCP, so we get some automated coverage for our preset code.
 - We actually test the output of our smallest available preset (with a relaxed float tolerance), so we can catch code updates that would break our checkpoints.


Discussion:

Resolving patch:
diff --git a/keras_nlp/conftest.py b/keras_nlp/conftest.py
--- a/keras_nlp/conftest.py
+++ b/keras_nlp/conftest.py
@@ -18,25 +18,48 @@
 
 def pytest_addoption(parser):
     parser.addoption(
-        "--runslow", action="store_true", default=False, help="run slow tests"
+        "--run_large",
+        action="store_true",
+        default=False,
+        help="run large tests",
+    )
+    parser.addoption(
+        "--run_extra_large",
+        action="store_true",
+        default=False,
+        help="run extra_large tests",
     )
 
 
 def pytest_configure(config):
-    config.addinivalue_line("markers", "slow: mark test as slow to run")
+    config.addinivalue_line(
+        "markers", "large: mark test as being slow or requiring a network"
+    )
+    config.addinivalue_line(
+        "markers",
+        "extra_large: mark test as being too large to run continuously",
+    )
 
 
 def pytest_collection_modifyitems(config, items):
-    if config.getoption("--runslow"):
-        # --runslow given in cli: do not skip slow tests
-        return
-    skip_slow = pytest.mark.skip(reason="need --runslow option to run")
+    run_extra_large_tests = config.getoption("--run_extra_large")
+    # Run large tests for --run_extra_large or --run_large.
+    run_large_tests = config.getoption("--run_large") or run_extra_large_tests
+
+    # Messages to annotate skipped tests with.
     skip_xla = pytest.mark.skipif(
         sys.platform == "darwin", reason="XLA unsupported on MacOS."
     )
-
+    skip_large = pytest.mark.skipif(
+        not run_large_tests, reason="need --run_large option to run"
+    )
+    skip_extra_large = pytest.mark.skipif(
+        not run_extra_large_tests, reason="need --run_extra_large option to run"
+    )
     for item in items:
-        if "slow" in item.keywords:
-            item.add_marker(skip_slow)
         if "jit_compile_true" in item.name:
             item.add_marker(skip_xla)
+        if "large" in item.keywords:
+            item.add_marker(skip_large)
+        if "extra_large" in item.keywords:
+            item.add_marker(skip_extra_large)


Test patch:
diff --git a/.cloudbuild/unit_test_jobs.jsonnet b/.cloudbuild/unit_test_jobs.jsonnet
--- a/.cloudbuild/unit_test_jobs.jsonnet
+++ b/.cloudbuild/unit_test_jobs.jsonnet
@@ -30,6 +30,7 @@ local unittest = base.BaseTest {
   command: [
     'pytest',
     'keras_nlp',
+    '--run_large',
   ],
 };
 
diff --git a/keras_nlp/tests/network_tests/test_load_presets.py b/keras_nlp/models/bert/bert_presets_test.py
similarity index 56%
rename from keras_nlp/tests/network_tests/test_load_presets.py
rename to keras_nlp/models/bert/bert_presets_test.py
--- a/keras_nlp/tests/network_tests/test_load_presets.py
+++ b/keras_nlp/models/bert/bert_presets_test.py
@@ -21,8 +21,64 @@
 from keras_nlp.models.bert.bert_tasks import BertClassifier
 
 
-@pytest.mark.slow
+@pytest.mark.large
+class BertPresetSmokeTest(tf.test.TestCase):
+    """
+    A smoke test for BERT presets we run continuously.
+
+    This only tests the smallest weights we have available. Run with:
+    `pytest keras_nlp/models/bert/bert_presets_test.py --run_large`
+    """
+
+    def test_preprocessor_output(self):
+        tokenizer = BertPreprocessor.from_preset(
+            "bert_tiny_uncased_en",
+            sequence_length=4,
+        )
+        outputs = tokenizer("The quick brown fox.")["token_ids"]
+        expected_outputs = [101, 1996, 4248, 102]
+        self.assertAllEqual(outputs, expected_outputs)
+
+    def test_backbone_output(self):
+        input_data = {
+            "token_ids": tf.constant([[101, 1996, 4248, 102]]),
+            "segment_ids": tf.constant([[0, 0, 0, 0]]),
+            "padding_mask": tf.constant([[1, 1, 1, 1]]),
+        }
+        model = Bert.from_preset(
+            "bert_tiny_uncased_en",
+        )
+        outputs = model(input_data)["sequence_output"][0, 0, :5]
+        # Outputs from our preset checkpoints should be stable!
+        # We should only update these numbers if we are updating a weights file,
+        # or have found a bug where output did not match the upstream source.
+        expected_outputs = [-1.38173, 0.16598, -2.92788, -2.66958, -0.61556]
+        # Keep a high tolerance, so we are robust to different hardware.
+        self.assertAllClose(outputs, expected_outputs, atol=0.01, rtol=0.01)
+
+    def test_classifier_output(self):
+        input_data = {
+            "token_ids": tf.constant([[101, 1996, 4248, 102]]),
+            "segment_ids": tf.constant([[0, 0, 0, 0]]),
+            "padding_mask": tf.constant([[1, 1, 1, 1]]),
+        }
+        model = BertClassifier(
+            "bert_tiny_uncased_en",
+        )
+        # We don't assert output values, as the head weights are random.
+        model(input_data)
+
+
+@pytest.mark.extra_large
 class BertPresetTest(tf.test.TestCase):
+    """
+    Test the full enumeration of our preset.
+
+    This only tests all enumeration of our presets and is only run manually.
+    Run with:
+    `pytest keras_nlp/models/bert_presets_test.py --run_extra_large`
+    """
+
     def test_load_bert(self):
         for preset in Bert.presets:
             model = Bert.from_preset(preset, load_weights=True)
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
--- a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
@@ -31,7 +31,7 @@
 )
 
 
-@pytest.mark.slow
+@pytest.mark.large
 class BytePairTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()