No public description

PiperOrigin-RevId: 638306209
Change-Id: I29149616d8d244df0e1b28241221518aa8833dab

nanodo/data.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Data pipeline."""
 
+from collections.abc import Mapping, Sequence
 import dataclasses
 import enum
 import functools
@@ -112,11 +113,11 @@ def get_tokenizer(self) -> spm.SentencePieceProcessor:
 
 
 def _py_tokenize(
-    features: dict[str, str],
+    features: Mapping[str, str],
     spt: _SPTokenizer,
     pad_len: int | None = None,
     pad_id: int = PAD_ID,
-) -> list[int]:
+) -> Sequence[int]:
   """Tokenizes text into ids, optionally pads or truncates to pad_len."""
   text = features['text']
   tokenizer = spt.get_tokenizer()

pyproject.toml

@@ -60,7 +60,8 @@ dependencies = [
     "optax>=0.2.2",
     "orbax>=0.1.7",
     "sentencepiece>=0.2.0",
-    "tensorflow-datasets>=4.9.4",
+    # TODO: Temporary fix while waiting for TFDS to be released.
+    "tfds-nightly",
     "tensorflow>=2.16.1",
 ]
 

tests/train_test.py

@@ -166,7 +166,6 @@ def test_train_and_evaluate(self, preprocessing):
 
     c = _get_config(self)
     c.checkpoint = True
-    c.pygrain_worker_count = 0
 
     cfg = model.DoConfig(**c.model, V=c.V)
     m = model.TransformerDo(cfg)