fix: CI testing failures

Author: dataCenter430

keras/src/layers/preprocessing/text_vectorization.py

@@ -6,9 +6,6 @@
 from keras.src.layers.preprocessing.index_lookup import listify_tensors
 from keras.src.layers.preprocessing.string_lookup import StringLookup
 from keras.src.saving import serialization_lib
-from keras.src.trainers.data_adapters.grain_dataset_adapter import (
-    GrainDatasetAdapter,
-)
 from keras.src.utils import argument_validation
 from keras.src.utils import backend_utils
 from keras.src.utils import tf_utils
@@ -436,12 +433,12 @@ def adapt(self, data, batch_size=None, steps=None):
         elif grain.available and isinstance(
             data, (grain.MapDataset, grain.IterDataset, grain.DataLoader)
         ):
-            dataset_adapter = GrainDatasetAdapter(data)
-            tf_dataset = dataset_adapter.get_tf_dataset()
-            if steps is not None:
-                tf_dataset = tf_dataset.take(steps)
-            for batch in tf_dataset:
+            step = 0
+            for batch in data:
+                if steps is not None and step >= steps:
+                    break
                 self.update_state(_extract_adapt_batch(batch))
+                step += 1
         else:
             data = tf_utils.ensure_tensor(data, dtype="string")
             if data.shape.rank == 1:

keras/src/layers/preprocessing/text_vectorization_test.py

@@ -493,6 +493,10 @@ def test_adapt_with_steps(self):
         self.assertIn("bar", vocab)
         self.assertNotIn("unique_word", vocab)
 
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="TextVectorization and Grain adapt path require TensorFlow",
+    )
     def test_adapt_with_grain_dataset(self):
         pytest.importorskip("grain")
         import grain as grain_module
@@ -521,6 +525,10 @@ def __getitem__(self, index):
         self.assertIn("bar", vocab)
         self.assertIn("baz", vocab)
 
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="TextVectorization and Grain adapt path require TensorFlow",
+    )
     def test_adapt_with_grain_dataset_and_steps(self):
         pytest.importorskip("grain")
         import grain as grain_module