Import only public API of TFDS

PiperOrigin-RevId: 756401612

Author: SeqIO Team

seqio/evaluation.py

@@ -33,7 +33,7 @@
 from seqio import metrics as metrics_lib
 from seqio import utils
 import tensorflow.compat.v2 as tf
-import tensorflow_datasets as tfds
+import tensorflow_datasets.public_api as tfds
 import typing_extensions
 
 Task = dataset_providers.Task

seqio/experimental_test.py

@@ -24,7 +24,7 @@
 from seqio import utils
 from seqio import vocabularies
 import tensorflow.compat.v2 as tf
-import tensorflow_datasets as tfds
+import tensorflow_datasets.public_api as tfds
 
 
 assert_dataset = test_utils.assert_dataset

seqio/loggers.py

@@ -27,7 +27,7 @@
 import numpy as np
 from seqio import metrics as metrics_lib
 import tensorflow as tf
-import tensorflow_datasets as tfds
+import tensorflow_datasets.public_api as tfds