Fix Gemma3/3n/4 preprocessor crashes with Python-only tokenizer (keras-team#2773)

* Fix Gemma3/3n/4 preprocessor crashes with Python-only tokenizer

Use tf.shape(token_ids)[0] instead of tf.shape(prompts)[0] in
Gemma3CausalLMPreprocessor, Gemma3nCausalLMPreprocessor, and
Gemma4CausalLMPreprocessor after tokenization. This prevents crashes when
the Python-only tokenizer workflow returns list[list[int]] with varying
sequence lengths.

The tensor_utils.py ragged-batch fallback is intentionally omitted here;
PR keras-team#2769 adds _call_python paths that bypass convert_to_ragged_batch for
the Python-only workflow, so that part is not needed in this PR.

Adds regression tests for generate_preprocess() and __call__() with real
Python tokenizers and prompts that tokenize to different lengths.

* Strengthen call() regression tests for Gemma3/3n/4 python-only tokenizer fix

test_call_with_python_only_tokenizer previously constructed each
preprocessor with image_converter=None, which makes call() return from
the text-only branch before reaching the tf.shape(token_ids)[0] line
this fix touches, so the test passed regardless of the fix. Set an
image_converter (mirroring test_generate_preprocess_with_python_only_tokenizer)
so the test actually exercises that code path.

---------

Co-authored-by: RAHUL KUMAR <pctablet505@users.noreply.github.com>

Author: pctablet505

keras_hub/src/models/gemma3/gemma3_causal_lm_preprocessor.py

@@ -605,7 +605,7 @@ def call(
 
         # === Vision processing ===
 
-        batch_size = tf.shape(prompts)[0]
+        batch_size = tf.shape(token_ids)[0]
         desired_height = self.image_converter.image_size[0]
         desired_width = self.image_converter.image_size[1]
         if images is None:
@@ -759,7 +759,7 @@ def generate_preprocess(
 
         # === Vision processing ===
 
-        batch_size = tf.shape(prompts)[0]
+        batch_size = tf.shape(token_ids)[0]
         desired_height = self.image_converter.image_size[0]
         desired_width = self.image_converter.image_size[1]
         if images is None:

keras_hub/src/models/gemma3/gemma3_causal_lm_preprocessor_test.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 import pytest
 
@@ -7,6 +9,7 @@
 from keras_hub.src.models.gemma3.gemma3_image_converter import (
     Gemma3ImageConverter,
 )
+from keras_hub.src.models.gemma3.gemma3_tokenizer import Gemma3Tokenizer
 from keras_hub.src.tests.mocks.mock_gemma3_tokenizer import MockGemma3Tokenizer
 from keras_hub.src.tests.test_case import TestCase
 
@@ -181,6 +184,48 @@ def test_invalid_shape(self):
             }
             self.text_preprocessor(input_data)
 
+    def test_generate_preprocess_with_python_only_tokenizer(self):
+        proto_path = os.path.join(
+            self.get_test_data_dir(), "gemma3_test_vocab.spm"
+        )
+        tokenizer = Gemma3Tokenizer(proto=proto_path)
+        image_converter = Gemma3ImageConverter(image_size=(4, 4))
+        preprocessor = Gemma3CausalLMPreprocessor(
+            tokenizer=tokenizer,
+            image_converter=image_converter,
+            sequence_length=20,
+            max_images_per_prompt=2,
+            num_vision_tokens_per_image=5,
+        )
+        # Prompts tokenize to different lengths -> non-rectangular list.
+        prompts = ["the quick brown fox", "the quick"]
+        output = preprocessor.generate_preprocess(prompts, sequence_length=20)
+        self.assertEqual(output["token_ids"].shape[0], 2)
+        self.assertEqual(output["padding_mask"].shape[0], 2)
+
+    def test_call_with_python_only_tokenizer(self):
+        proto_path = os.path.join(
+            self.get_test_data_dir(), "gemma3_test_vocab.spm"
+        )
+        tokenizer = Gemma3Tokenizer(proto=proto_path)
+        # image_converter must be set: the text-only branch of `call()`
+        # returns before `token_ids` is used to compute `batch_size`.
+        image_converter = Gemma3ImageConverter(image_size=(4, 4))
+        preprocessor = Gemma3CausalLMPreprocessor(
+            tokenizer=tokenizer,
+            image_converter=image_converter,
+            sequence_length=20,
+            max_images_per_prompt=2,
+            num_vision_tokens_per_image=5,
+        )
+        # Prompts tokenize to different lengths -> non-rectangular list.
+        input_data = {
+            "prompts": ["the quick brown fox", "the quick"],
+            "responses": ["round", "round"],
+        }
+        output = preprocessor(input_data)
+        self.assertEqual(output[0]["token_ids"].shape[0], 2)
+
     @pytest.mark.kaggle_key_required
     @pytest.mark.extra_large
     def test_all_presets(self):

keras_hub/src/models/gemma3n/gemma3n_causal_lm_preprocessor.py

@@ -656,7 +656,7 @@ def call(
             return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
         # === Multimodal processing ===
-        batch_size = tf.shape(prompts)[0]
+        batch_size = tf.shape(token_ids)[0]
         desired_height = (
             self.image_converter.image_size[0] if self.image_converter else 0
         )
@@ -846,7 +846,7 @@ def generate_preprocess(
             }
 
         # === Multimodal processing ===
-        batch_size = tf.shape(prompts)[0]
+        batch_size = tf.shape(token_ids)[0]
         desired_height = (
             self.image_converter.image_size[0] if self.image_converter else 0
         )

keras_hub/src/models/gemma3n/gemma3n_causal_lm_preprocessor_test.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 
 from keras_hub.src.models.gemma3n.gemma3n_audio_converter import (
@@ -9,6 +11,7 @@
 from keras_hub.src.models.gemma3n.gemma3n_image_converter import (
     Gemma3nImageConverter,
 )
+from keras_hub.src.models.gemma3n.gemma3n_tokenizer import Gemma3nTokenizer
 from keras_hub.src.tests.mocks.mock_gemma3n_tokenizer import (
     MockGemma3nTokenizer,
 )
@@ -325,3 +328,51 @@ def test_text_only_with_audios_raises_error(self):
                 "audios": [np.ones((16000,))],
             }
             self.text_preprocessor(input_data)
+
+    def test_generate_preprocess_with_python_only_tokenizer(self):
+        proto_path = os.path.join(
+            self.get_test_data_dir(), "gemma3n_test_vocab.spm"
+        )
+        tokenizer = Gemma3nTokenizer(proto=proto_path)
+        image_converter = Gemma3nImageConverter(image_size=(4, 4))
+        preprocessor = Gemma3nCausalLMPreprocessor(
+            tokenizer=tokenizer,
+            image_converter=image_converter,
+            audio_converter=None,
+            sequence_length=20,
+            max_images_per_prompt=2,
+            num_vision_tokens_per_image=5,
+            max_audios_per_prompt=0,
+            num_audio_tokens_per_audio=0,
+        )
+        # Prompts tokenize to different lengths -> non-rectangular list.
+        prompts = ["the quick brown fox", "the quick"]
+        output = preprocessor.generate_preprocess(prompts, sequence_length=20)
+        self.assertEqual(output["token_ids"].shape[0], 2)
+        self.assertEqual(output["padding_mask"].shape[0], 2)
+
+    def test_call_with_python_only_tokenizer(self):
+        proto_path = os.path.join(
+            self.get_test_data_dir(), "gemma3n_test_vocab.spm"
+        )
+        tokenizer = Gemma3nTokenizer(proto=proto_path)
+        # image_converter must be set: the text-only branch of `call()`
+        # returns before `token_ids` is used to compute `batch_size`.
+        image_converter = Gemma3nImageConverter(image_size=(4, 4))
+        preprocessor = Gemma3nCausalLMPreprocessor(
+            tokenizer=tokenizer,
+            image_converter=image_converter,
+            audio_converter=None,
+            sequence_length=20,
+            max_images_per_prompt=2,
+            num_vision_tokens_per_image=5,
+            max_audios_per_prompt=0,
+            num_audio_tokens_per_audio=0,
+        )
+        # Prompts tokenize to different lengths -> non-rectangular list.
+        input_data = {
+            "prompts": ["the quick brown fox", "the quick"],
+            "responses": ["round", "round"],
+        }
+        output = preprocessor(input_data)
+        self.assertEqual(output[0]["token_ids"].shape[0], 2)

keras_hub/src/models/gemma4/gemma4_causal_lm_preprocessor.py

@@ -966,7 +966,7 @@ def call(
 
             return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
-        batch_size = tf.shape(prompts)[0]
+        batch_size = tf.shape(token_ids)[0]
 
         # === Vision & Audio processing ===
         if images is not None and self.image_converter is not None:
@@ -1149,7 +1149,7 @@ def generate_preprocess(
                 ),
             }
 
-        batch_size = tf.shape(prompts)[0]
+        batch_size = tf.shape(token_ids)[0]
 
         # === Vision & Audio processing ===
         if images is not None and self.image_converter is not None:

keras_hub/src/models/gemma4/gemma4_causal_lm_preprocessor_test.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 import pytest
 import tensorflow as tf
@@ -8,6 +10,7 @@
 from keras_hub.src.models.gemma4.gemma4_image_converter import (
     Gemma4ImageConverter,
 )
+from keras_hub.src.models.gemma4.gemma4_tokenizer import Gemma4Tokenizer
 from keras_hub.src.models.gemma4.gemma4_video_converter import (
     Gemma4VideoConverter,
 )
@@ -313,6 +316,48 @@ def test_video_metadata_timestamps(self):
             expanded_cleared.numpy()[0], expanded_default.numpy()[0]
         )
 
+    def test_generate_preprocess_with_python_only_tokenizer(self):
+        proto_path = os.path.join(
+            self.get_test_data_dir(), "gemma4_test_vocab.spm"
+        )
+        tokenizer = Gemma4Tokenizer(proto=proto_path)
+        image_converter = Gemma4ImageConverter(image_size=(4, 4), patch_size=4)
+        preprocessor = Gemma4CausalLMPreprocessor(
+            tokenizer=tokenizer,
+            image_converter=image_converter,
+            sequence_length=20,
+            max_images_per_prompt=2,
+            num_vision_tokens_per_image=5,
+        )
+        # Prompts tokenize to different lengths -> non-rectangular list.
+        prompts = ["the quick brown fox", "the quick"]
+        output = preprocessor.generate_preprocess(prompts, sequence_length=20)
+        self.assertEqual(output["token_ids"].shape[0], 2)
+        self.assertEqual(output["padding_mask"].shape[0], 2)
+
+    def test_call_with_python_only_tokenizer(self):
+        proto_path = os.path.join(
+            self.get_test_data_dir(), "gemma4_test_vocab.spm"
+        )
+        tokenizer = Gemma4Tokenizer(proto=proto_path)
+        # image_converter must be set: the text-only branch of `call()`
+        # returns before `token_ids` is used to compute `batch_size`.
+        image_converter = Gemma4ImageConverter(image_size=(4, 4), patch_size=4)
+        preprocessor = Gemma4CausalLMPreprocessor(
+            tokenizer=tokenizer,
+            image_converter=image_converter,
+            sequence_length=20,
+            max_images_per_prompt=2,
+            num_vision_tokens_per_image=5,
+        )
+        # Prompts tokenize to different lengths -> non-rectangular list.
+        input_data = {
+            "prompts": ["the quick brown fox", "the quick"],
+            "responses": ["round", "round"],
+        }
+        output = preprocessor(input_data)
+        self.assertEqual(output[0]["token_ids"].shape[0], 2)
+
     @pytest.mark.kaggle_key_required
     @pytest.mark.extra_large
     def test_all_presets(self):