Fix save/load, cross-backend, and multi-image bugs in Mistral3

- Add missing build() to the vision encoder/projector layers, fixing
  weight reinitialization on save/load and a JAX param-count mismatch.
- Make Mistral3ImageConverter work inside tf.data on all three backends
  by using native tf ops for the graph path instead of keras.ops.
- Rework placeholder_indices from a single flattened batch index into
  per-row local positions, so prompts in the same batch can now have
  different numbers of images.
- Fill in preset kaggle handles/params, now that the checkpoints are
  uploaded.
- Round out test coverage across the model (backbone, causal LM,
  preprocessor, tokenizer, image converter, vision encoder).

Author: ahmed-damda

keras_hub/src/models/mistral3/mistral3_backbone.py

@@ -182,11 +182,9 @@ def __init__(
         image_sizes_input = keras.Input(
             shape=(2,), dtype="int32", name="image_sizes"
         )
-        # Flat positions of image placeholder tokens in the flattened
-        # `(batch * seq_length,)` sequence. Computed on the host (e.g. via
-        # `compute_image_placeholder_indices`) rather than derived in-graph,
-        # since a `nonzero`-style lookup has a data-dependent output shape
-        # that is incompatible with `jax.jit` tracing.
+        # Each example's own local image placeholder token positions,
+        # `-1`-padded to the batch's max count; see
+        # `compute_image_placeholder_indices`.
         placeholder_indices_input = keras.Input(
             shape=(None,),
             dtype="int32",

keras_hub/src/models/mistral3/mistral3_backbone_test.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 from keras import ops
 
 from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone
@@ -70,9 +71,7 @@ def setUp(self):
                 np.random.rand(2, 3, 8, 8).astype("float32")
             ),
             "image_sizes": ops.array([[8, 8], [8, 8]], dtype="int32"),
-            "placeholder_indices": ops.reshape(
-                ops.convert_to_tensor(placeholder_indices), (2, 4)
-            ),
+            "placeholder_indices": ops.convert_to_tensor(placeholder_indices),
         }
 
     def test_backbone_basics(self):
@@ -96,3 +95,62 @@ def test_backbone_basics(self):
             # its own vision-encoder-bearing backbone.
             run_quantization_check=False,
         )
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=Mistral3Backbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
+    def test_variable_images_per_prompt(self):
+        # One prompt with one image, one with two.
+        token_ids = ops.array(
+            [
+                [self.image_token_index] * 4 + [3, 0, 0, 0, 0],
+                [self.image_token_index] * 8 + [4],
+            ],
+            dtype="int32",
+        )
+        placeholder_indices = compute_image_placeholder_indices(
+            token_ids, image_token_index=self.image_token_index
+        )
+        input_data = {
+            "token_ids": token_ids,
+            "padding_mask": ops.ones((2, 9), dtype="int32"),
+            "pixel_values": ops.convert_to_tensor(
+                np.random.rand(3, 3, 8, 8).astype("float32")
+            ),
+            "image_sizes": ops.array([[8, 8], [8, 8], [8, 8]], dtype="int32"),
+            "placeholder_indices": ops.convert_to_tensor(placeholder_indices),
+        }
+        model = Mistral3Backbone(**self.init_kwargs)
+        output = model(input_data)
+        self.assertEqual(
+            ops.shape(output),
+            (2, 9, self.text_init_kwargs["hidden_dim"]),
+        )
+
+    def test_num_parameters(self):
+        model = Mistral3Backbone(**self.init_kwargs)
+        self.assertEqual(model.count_params(), 4016)
+        self.assertEqual(len(model.layers), 11)
+
+    @pytest.mark.kaggle_key_required
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        token_ids = ops.array([[1, 1824, 349, 524, 11234, 28804]])
+        input_data = {
+            "token_ids": token_ids,
+            "padding_mask": ops.ones_like(token_ids),
+            "pixel_values": ops.zeros((0, 3, 14, 14), dtype="float32"),
+            "image_sizes": ops.zeros((0, 2), dtype="int32"),
+            "placeholder_indices": ops.zeros((1, 0), dtype="int32"),
+        }
+        for preset in Mistral3Backbone.presets:
+            self.run_preset_test(
+                cls=Mistral3Backbone,
+                preset=preset,
+                input_data=input_data,
+            )

keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor.py

@@ -282,7 +282,7 @@ def call(
         placeholder_indices = compute_image_placeholder_indices(
             keras.ops.convert_to_numpy(model_token_ids),
             self.tokenizer.image_placeholder_token_id,
-        )[None, :]
+        )
 
         out_x = {
             "token_ids": model_token_ids,
@@ -351,7 +351,7 @@ def generate_preprocess(
             placeholder_indices = compute_image_placeholder_indices(
                 keras.ops.convert_to_numpy(token_ids),
                 self.tokenizer.image_placeholder_token_id,
-            )[None, :]
+            )
             out_x["pixel_values"] = pixel_values
             out_x["image_sizes"] = image_sizes
             out_x["placeholder_indices"] = placeholder_indices

keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor_test.py

@@ -1,4 +1,6 @@
 import numpy as np
+import pytest
+from keras import ops
 
 from keras_hub.src.models.mistral3.mistral3_causal_lm_preprocessor import (
     Mistral3CausalLMPreprocessor,
@@ -7,11 +9,61 @@
     Mistral3ImageConverter,
 )
 from keras_hub.src.models.mistral3.mistral3_tokenizer import Mistral3Tokenizer
-from keras_hub.src.models.mistral3.mistral3_tokenizer_test import (
-    _tekken_vision_init_kwargs,
-)
 from keras_hub.src.tests.test_case import TestCase
 
+# A tiktoken-style split pattern, matching the Tekken format.
+_TEKKEN_SPLIT_PATTERN = (
+    r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*"
+    r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|\p{N}| ?[^\s\p{L}\p{N}]+"
+    r"[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+"
+)
+
+
+def _bytes_to_unicode():
+    bs = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    return {b: chr(c) for b, c in zip(bs, cs)}
+
+
+def _tekken_vision_init_kwargs():
+    """Build a tiny Tekken (byte-level BPE) vocabulary with image tokens."""
+    byte_encoder = _bytes_to_unicode()
+    special_tokens = [
+        "<unk>",
+        "<s>",
+        "</s>",
+        "<pad>",
+        "[INST]",
+        "[IMG]",
+        "[IMG_BREAK]",
+        "[IMG_END]",
+    ]
+    vocabulary = {token: i for i, token in enumerate(special_tokens)}
+    offset = len(special_tokens)
+    for i in range(256):
+        vocabulary[byte_encoder[i]] = offset + i
+    merges = []
+    next_id = offset + 256
+    for a, b in [("t", "h"), ("th", "e"), ("i", "n")]:
+        vocabulary[a + b] = next_id
+        merges.append(f"{a} {b}")
+        next_id += 1
+    return {
+        "vocabulary": vocabulary,
+        "merges": merges,
+        "split_pattern": _TEKKEN_SPLIT_PATTERN,
+    }
+
 
 class Mistral3CausalLMPreprocessorTest(TestCase):
     def setUp(self):
@@ -26,6 +78,14 @@ def setUp(self):
             "spatial_merge_size": 1,
         }
 
+    def test_preprocessor_basics(self):
+        input_data = {"prompts": ["the tin", "in the"]}
+        self.run_preprocessor_test(
+            cls=Mistral3CausalLMPreprocessor,
+            init_kwargs=self.init_kwargs,
+            input_data=input_data,
+        )
+
     def test_generate_preprocess_with_images(self):
         preprocessor = Mistral3CausalLMPreprocessor(**self.init_kwargs)
         image = np.zeros((8, 8, 3), dtype="float32")
@@ -40,12 +100,40 @@ def test_generate_preprocess_with_images(self):
             "placeholder_indices",
         ):
             self.assertIn(key, x)
-        # An 8x8 image with `patch_size=4`, `spatial_merge_size=1` expands
-        # to a 2x2 grid of placeholder tokens.
         token_ids = np.array(x["token_ids"])
         num_placeholders = int(
             np.sum(
                 token_ids == preprocessor.tokenizer.image_placeholder_token_id
             )
         )
         self.assertEqual(num_placeholders, 4)
+
+    def test_generate_preprocess_text_only(self):
+        preprocessor = Mistral3CausalLMPreprocessor(**self.init_kwargs)
+        x = preprocessor.generate_preprocess("the tin")
+        self.assertEqual(set(x.keys()), {"token_ids", "padding_mask"})
+
+    def test_generate_postprocess(self):
+        preprocessor = Mistral3CausalLMPreprocessor(**self.init_kwargs)
+        input_data = {
+            "token_ids": ops.array([1, 265, 40, 124, 266, 0, 0, 0]),
+            "padding_mask": ops.array(
+                [True, True, True, True, True, False, False, False]
+            ),
+        }
+        x = preprocessor.generate_postprocess(input_data)
+        self.assertEqual(x, "the tin")
+
+    @pytest.mark.kaggle_key_required
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        input_data = {
+            "prompts": ["Describe the image. [IMG]"],
+            "images": [[self.load_test_image()]],
+        }
+        for preset in Mistral3CausalLMPreprocessor.presets:
+            self.run_preset_test(
+                cls=Mistral3CausalLMPreprocessor,
+                preset=preset,
+                input_data=input_data,
+            )