Add tests for the examples

PiperOrigin-RevId: 752272527

Author: Conchylicultor

gemma/gm/testing/__init__.py

@@ -16,5 +16,13 @@
 
 # pylint: disable=g-importing-member
 
-from gemma.gm.testing._dummy_model import DummyGemma
-from gemma.gm.testing._dummy_tokenizer import DummyTokenizer
+
+from etils import epy as _epy
+
+# pylint: disable=g-import-not-at-top,g-importing-member
+
+with _epy.lazy_api_imports(globals()):
+
+  from gemma.gm.testing._dummy_model import DummyGemma
+  from gemma.gm.testing._dummy_tokenizer import DummyTokenizer
+  from gemma.gm.testing._fixtures import use_hermetic_tokenizer

gemma/gm/testing/_fixtures.py

@@ -0,0 +1,43 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests fixtures."""
+
+import contextlib  # pylint: disable=unused-import
+from unittest import mock
+
+from etils import epath
+from gemma.gm.text import _tokenizer
+import pytest  # pytype: disable=import-error
+
+
+@pytest.fixture(autouse=True, scope='module')
+def use_hermetic_tokenizer():
+  """Use the local tokenizer, to avoid TFHub calls."""
+
+  new_path = epath.resource_path('gemma') / 'testdata/tokenizer_gemma3.model'
+
+  # We cannot mock `Gemma3Tokenizer.path` directly as dataclasses also
+  # set the value in the `__init__` default value.
+
+  old_init = _tokenizer.Gemma3Tokenizer.__init__
+
+  def mew_init(self, path=None, **kwargs):
+    del path
+    old_init(self, new_path, **kwargs)
+
+  with (
+      contextlib.nullcontext()
+  ):
+    yield

gemma/gm/tests/examples_test.py

@@ -0,0 +1,37 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gemma import gm
+from gemma.examples import seq2seq
+from kauldron import konfig
+import tensorflow_datasets as tfds
+
+# Activate the fixture
+use_hermetic_tokenizer = gm.testing.use_hermetic_tokenizer
+
+
+def test_examples():
+  cfg = seq2seq.get_config()
+  with konfig.mock_modules():
+    cfg.model = gm.testing.DummyGemma(
+        tokens='batch.input',
+    )
+  cfg.train_ds.num_workers = 0  # Disable multi-processing in tests.
+  cfg.workdir = '/tmp/gemma_test'
+
+  trainer = konfig.resolve(cfg)
+
+  # Resolve the training step, including the metrics, losses,...
+  with tfds.testing.mock_data(num_examples=10):
+    _ = trainer.context_specs

gemma/gm/text/_tokenizer_test.py

@@ -14,14 +14,14 @@
 
 import pickle
 
-from etils import epath
 from gemma import gm
 
+# Activate the fixture
+use_hermetic_tokenizer = gm.testing.use_hermetic_tokenizer
 
-def test_pickle():
 
-  tokenizer = gm.text.Gemma3Tokenizer(
-  )
+def test_pickle():
+  tokenizer = gm.text.Gemma3Tokenizer()
   tokenizer.encode('Hello world!')  # Trigger the lazy-loading of the tokenizer.
 
   pickle.dumps(tokenizer)