Merge pull request #5001 from jburnim:jburnim_python3.13

PiperOrigin-RevId: 816318427

Author: Flax Authors

examples/gemma/input_pipeline_test.py

@@ -14,6 +14,7 @@
 
 import os
 import pathlib
+import sys
 import tempfile
 
 from absl.testing import absltest
@@ -32,6 +33,8 @@ class InputPipelineTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    if sys.version_info >= (3, 13):
+      self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')
     self.train_ds, self.eval_ds = self._get_datasets()
 
   def _get_datasets(self):

examples/gemma/tokenizer.py

@@ -16,14 +16,16 @@
 
 import dataclasses
 import os
+import sys
 import tempfile
 import time
 from typing import Any
 from collections.abc import Iterable
 
 import jax
 import tensorflow as tf
-import tensorflow_text as tftxt
+if sys.version_info < (3, 13):
+  import tensorflow_text as tftxt
 from absl import logging
 from sentencepiece import SentencePieceTrainer, SentencePieceProcessor
 

examples/lm1b/input_pipeline_test.py

@@ -14,6 +14,7 @@
 
 import os
 import pathlib
+import sys
 import tempfile
 
 from absl.testing import absltest
@@ -33,6 +34,8 @@ class InputPipelineTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    if sys.version_info >= (3, 13):
+      self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')
     self.train_ds, self.eval_ds, self.predict_ds = self._get_datasets()
 
   def _get_datasets(self):

examples/lm1b/tokenizer.py

@@ -15,6 +15,7 @@
 """Provides op for tokenizing a dataset."""
 
 import os
+import sys
 import tempfile
 import time
 from typing import Any, Dict, Tuple
@@ -25,7 +26,8 @@
 import jax
 from sentencepiece import SentencePieceTrainer
 import tensorflow as tf
-import tensorflow_text as tftxt
+if sys.version_info < (3, 13):
+  import tensorflow_text as tftxt
 
 Features = dict[str, tf.Tensor]
 

examples/lm1b/train_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pathlib
+import sys
 import tempfile
 
 from absl import logging
@@ -33,6 +34,8 @@ class TrainTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    if sys.version_info >= (3, 13):
+      self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')
     tf.config.experimental.set_visible_devices([], 'GPU')
 
   def test_train_and_evaluate(self):

examples/lm1b_nnx/input_pipeline_test.py

@@ -14,6 +14,7 @@
 
 import os
 import pathlib
+import sys
 import tempfile
 
 from absl.testing import absltest

examples/lm1b_nnx/tokenizer.py

@@ -16,14 +16,16 @@
 
 import dataclasses
 import os
+import sys
 import tempfile
 import time
 from typing import Any
 from collections.abc import Iterable
 
 import jax
 import tensorflow as tf
-import tensorflow_text as tftxt
+if sys.version_info < (3, 13):
+  import tensorflow_text as tftxt
 from absl import logging
 from sentencepiece import SentencePieceTrainer
 

examples/lm1b_nnx/train_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pathlib
+import sys
 import tempfile
 
 import jax
@@ -31,6 +32,8 @@ class TrainTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    if sys.version_info >= (3, 13):
+      self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')
     tf.config.experimental.set_visible_devices([], 'GPU')
 
   def test_train_and_evaluate(self):

examples/sst2/input_pipeline.py

@@ -14,13 +14,15 @@
 
 """SST-2 input pipeline."""
 
+import sys
 from typing import Any, Dict, Optional
 
 from absl import logging
 import numpy as np
 import tensorflow as tf
 import tensorflow_datasets as tfds
-import tensorflow_text as text
+if sys.version_info < (3, 13):
+  import tensorflow_text as text
 
 import vocabulary
 
@@ -177,7 +179,7 @@ def __init__(
       self,
       tfds_name: str = 'glue/sst2',
       vocab_path: str = 'vocab.txt',
-      tokenizer: text.Tokenizer = text.WhitespaceTokenizer(),
+      tokenizer=None,
       split='train',
   ):
     """Initializes the SST2 data source."""
@@ -194,6 +196,8 @@ def __init__(
     self.vocab = vocabulary.Vocabulary(vocab_path=vocab_path)
 
     # Convert the sentences to sequences of token IDs and compute length.
+    if tokenizer is None:
+       tokenizer = text.WhitespaceTokenizer()
     self.tokenizer = tokenizer
     self.tf_vocab = vocab_to_hashtable(self.vocab, unk_idx=self.vocab.unk_idx)
     self.examples = self.dataset.map(

examples/sst2/input_pipeline_test.py

@@ -14,6 +14,7 @@
 
 import os
 import pathlib
+import sys
 import tempfile
 
 from absl.testing import absltest
@@ -27,6 +28,8 @@ class InputPipelineTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    if sys.version_info >= (3, 13):
+      self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')
     self.vocab_path = self._get_vocab_path()
     self.dataset = self._get_dataset(self.vocab_path)