diff --git a/keras_hub/api/layers/__init__.py b/keras_hub/api/layers/__init__.py index 0d699da49d..736445a92f 100644 --- a/keras_hub/api/layers/__init__.py +++ b/keras_hub/api/layers/__init__.py @@ -129,6 +129,9 @@ from keras_hub.src.models.metaclip_2.metaclip_2_image_converter import ( MetaCLIP2ImageConverter as MetaCLIP2ImageConverter, ) +from keras_hub.src.models.mistral3.mistral3_image_converter import ( + Mistral3ImageConverter as Mistral3ImageConverter, +) from keras_hub.src.models.mit.mit_image_converter import ( MiTImageConverter as MiTImageConverter, ) diff --git a/keras_hub/api/models/__init__.py b/keras_hub/api/models/__init__.py index 67950bfb54..14568edbcb 100644 --- a/keras_hub/api/models/__init__.py +++ b/keras_hub/api/models/__init__.py @@ -517,6 +517,21 @@ from keras_hub.src.models.mistral.mistral_tokenizer import ( MistralTokenizer as MistralTokenizer, ) +from keras_hub.src.models.mistral3.mistral3_backbone import ( + Mistral3Backbone as Mistral3Backbone, +) +from keras_hub.src.models.mistral3.mistral3_causal_lm import ( + Mistral3CausalLM as Mistral3CausalLM, +) +from keras_hub.src.models.mistral3.mistral3_causal_lm_preprocessor import ( + Mistral3CausalLMPreprocessor as Mistral3CausalLMPreprocessor, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import ( + Mistral3Tokenizer as Mistral3Tokenizer, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoder as Mistral3VisionEncoder, +) from keras_hub.src.models.mit.mit_backbone import MiTBackbone as MiTBackbone from keras_hub.src.models.mit.mit_image_classifier import ( MiTImageClassifier as MiTImageClassifier, diff --git a/keras_hub/api/tokenizers/__init__.py b/keras_hub/api/tokenizers/__init__.py index 062a9973f1..3264ef55b9 100644 --- a/keras_hub/api/tokenizers/__init__.py +++ b/keras_hub/api/tokenizers/__init__.py @@ -77,6 +77,9 @@ from keras_hub.src.models.mistral.mistral_tokenizer import ( MistralTokenizer as MistralTokenizer, ) +from keras_hub.src.models.mistral3.mistral3_tokenizer import ( + Mistral3Tokenizer as Mistral3Tokenizer, +) from keras_hub.src.models.mixtral.mixtral_tokenizer import ( MixtralTokenizer as MixtralTokenizer, ) diff --git a/keras_hub/src/models/mistral/mistral_backbone.py b/keras_hub/src/models/mistral/mistral_backbone.py index f52935ad62..9985ae67b4 100644 --- a/keras_hub/src/models/mistral/mistral_backbone.py +++ b/keras_hub/src/models/mistral/mistral_backbone.py @@ -1,12 +1,10 @@ import keras from keras import ops from keras.layers import ReversibleEmbedding +from keras.layers import RMSNormalization from keras_hub.src.api_export import keras_hub_export from keras_hub.src.models.backbone import Backbone -from keras_hub.src.models.mistral.mistral_layer_norm import ( - MistralLayerNormalization, -) from keras_hub.src.models.mistral.mistral_transformer_decoder import ( MistralTransformerDecoder, ) @@ -136,7 +134,7 @@ def __init__( name=f"transformer_layer_{i}", ) self.transformer_layers.append(layer) - self.layer_norm = MistralLayerNormalization( + self.layer_norm = RMSNormalization( epsilon=layer_norm_epsilon, dtype=dtype, name="sequence_output_layernorm", diff --git a/keras_hub/src/models/mistral/mistral_layer_norm.py b/keras_hub/src/models/mistral/mistral_layer_norm.py deleted file mode 100644 index affca9c45f..0000000000 --- a/keras_hub/src/models/mistral/mistral_layer_norm.py +++ /dev/null @@ -1,35 +0,0 @@ -import keras -from keras import ops - - -# TODO: Deprecate this in favor of -# `keras.layers.LayerNormalization(rms_scaling=True)` once Keras 2 support is -# removed. -class MistralLayerNormalization(keras.layers.Layer): - """A normalization layer for Mistral that implements RMS normalization.""" - - def __init__(self, epsilon=1e-6, **kwargs): - super().__init__(**kwargs) - self.epsilon = epsilon - - def build(self, input_shape): - dim = input_shape[-1] - self.scale = self.add_weight( - name="scale", - trainable=True, - shape=(dim,), - initializer="ones", - dtype=self.variable_dtype, - ) - self.built = True - - def call(self, x): - x = ops.cast(x, "float32") - var = ops.mean(ops.power(x, 2), axis=-1, keepdims=True) - x = x * ops.rsqrt(var + self.epsilon) - return ops.cast(x * self.scale, self.compute_dtype) - - def get_config(self): - config = super().get_config() - config.update({"epsilon": self.epsilon}) - return config diff --git a/keras_hub/src/models/mistral/mistral_transformer_decoder.py b/keras_hub/src/models/mistral/mistral_transformer_decoder.py index 996c3459f4..71c5caf6be 100644 --- a/keras_hub/src/models/mistral/mistral_transformer_decoder.py +++ b/keras_hub/src/models/mistral/mistral_transformer_decoder.py @@ -1,5 +1,6 @@ import keras from keras import ops +from keras.layers import RMSNormalization from keras_hub.src.layers.modeling.transformer_layer_utils import ( compute_causal_mask, @@ -10,9 +11,6 @@ from keras_hub.src.models.mistral.mistral_attention import ( CachedMistralAttention, ) -from keras_hub.src.models.mistral.mistral_layer_norm import ( - MistralLayerNormalization, -) from keras_hub.src.utils.keras_utils import clone_initializer @@ -71,7 +69,7 @@ def build(self, decoder_sequence_shape): ) self._self_attention_layer.build(decoder_sequence_shape) - self._self_attention_layernorm = MistralLayerNormalization( + self._self_attention_layernorm = RMSNormalization( epsilon=self.layer_norm_epsilon, dtype=self.dtype_policy, name="self_attention_layernorm", @@ -116,7 +114,7 @@ def build(self, decoder_sequence_shape): ) ) - self._feedforward_layernorm = MistralLayerNormalization( + self._feedforward_layernorm = RMSNormalization( epsilon=self.layer_norm_epsilon, dtype=self.dtype_policy, name="feedforward_layernorm", diff --git a/keras_hub/src/models/mistral3/__init__.py b/keras_hub/src/models/mistral3/__init__.py new file mode 100644 index 0000000000..81190f2249 --- /dev/null +++ b/keras_hub/src/models/mistral3/__init__.py @@ -0,0 +1,5 @@ +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_presets import backbone_presets +from keras_hub.src.utils.preset_utils import register_presets + +register_presets(backbone_presets, Mistral3Backbone) diff --git a/keras_hub/src/models/mistral3/mistral3_backbone.py b/keras_hub/src/models/mistral3/mistral3_backbone.py new file mode 100644 index 0000000000..80dc3d6fd5 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_backbone.py @@ -0,0 +1,265 @@ +import keras +from keras import ops +from keras.layers import ReversibleEmbedding +from keras.layers import RMSNormalization + +from keras_hub.src.api_export import keras_hub_export +from keras_hub.src.models.backbone import Backbone +from keras_hub.src.models.mistral.mistral_transformer_decoder import ( + MistralTransformerDecoder, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3ImageFeatureExtractor, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3ImageTextEmbeddingMerger, +) + + +def _mistral_kernel_initializer(stddev=0.02): + return keras.initializers.RandomNormal(stddev=stddev) + + +@keras_hub_export("keras_hub.models.Mistral3Backbone") +class Mistral3Backbone(Backbone): + """ + The Mistral3 (Pixtral vision + Mistral text) core architecture. + + This network implements a multimodal Transformer-based decoder network, + Mistral3, as used by models such as Mistral Small 3.1/3.2. It includes + the token embedding lookups, a Pixtral-style vision encoder, and + transformer decoder layers. + + The default constructor gives a fully customizable, randomly initialized + Mistral3 model with any number of layers, heads, and embedding + dimensions. To load preset architectures and weights, use the + `from_preset` constructor. + + Args: + vocabulary_size: int. The size of the token vocabulary. + num_layers: int. The number of transformer layers. + num_query_heads: int. The number of query attention heads for + each transformer. + hidden_dim: int. The size of the transformer encoding and pooling + layers. + intermediate_dim: int. The output dimension of the first Dense layer + in a three-layer feedforward network for each transformer. + num_key_value_heads: int. The number of key and value attention heads + for each transformer. + vision_encoder: A `keras_hub.models.Mistral3VisionEncoder` instance. + multimodal_projector: A `Mistral3MultiModalProjector` instance. + rope_max_wavelength: int, optional. The maximum angular wavelength of + the sine/cosine curves, for rotary embeddings. Defaults to `10000`. + rope_scaling_factor: float, optional. The scaling factor for + calculation of rotary embedding. Defaults to `1.0`. + layer_norm_epsilon: float, optional. Epsilon for the layer + normalization layers in the transformer decoder. Defaults to `1e-6`. + sliding_window: int, optional. The sliding window for the mistral + attention layers. This controls the maximum cache size for the + attention layers in each transformer decoder. Only `sliding_window` + number of tokens are saved in the cache and used to generate the + next token. Defaults to `512`. Pass `None` to disable sliding + window attention entirely (e.g. Magistral). + head_dim: int, optional. The size of each attention head. When + `None` (the default), falls back to `hidden_dim // num_query_heads`. + Set explicitly when the model's head size is not equal to + `hidden_dim // num_query_heads` — e.g. Magistral uses + `head_dim=128` with `hidden_dim=5120` and `num_query_heads=32`. + image_token_index: int, optional. The token ID in `token_ids` that + marks image placeholder positions. Defaults to `10`. + dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use + for model computations and weights. Note that some computations, + such as softmax and layer normalization, will always be done at + float32 precision regardless of dtype. + + Examples: + + ```python + input_data = { + "token_ids": np.ones(shape=(1, 12), dtype="int32"), + "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]), + "pixel_values": np.ones(shape=(1, 3, 64, 64), dtype="float32"), + "image_sizes": np.array([[64, 64]], dtype="int32"), + "placeholder_indices": np.zeros(shape=(1, 16), dtype="int32"), + } + + # Pretrained Mistral3 decoder. + model = keras_hub.models.Mistral3Backbone.from_preset( + "mistral_small_3.1_24b_instruct_2503_en" + ) + model(input_data) + ``` + """ + + def __init__( + self, + vocabulary_size, + num_layers, + num_query_heads, + hidden_dim, + intermediate_dim, + num_key_value_heads, + vision_encoder, + multimodal_projector, + rope_max_wavelength=10000, + rope_scaling_factor=1.0, + layer_norm_epsilon=1e-6, + sliding_window=512, + head_dim=None, + dropout=0, + image_token_index=10, + dtype=None, + **kwargs, + ): + # === Layers === + self.token_embedding = ReversibleEmbedding( + input_dim=vocabulary_size, + output_dim=hidden_dim, + tie_weights=False, + embeddings_initializer=_mistral_kernel_initializer(stddev=0.01), + dtype=dtype, + name="token_embedding", + ) + self.transformer_layers = [] + for i in range(num_layers): + layer = MistralTransformerDecoder( + intermediate_dim=intermediate_dim, + num_query_heads=num_query_heads, + num_key_value_heads=num_key_value_heads, + rope_max_wavelength=rope_max_wavelength, + rope_scaling_factor=rope_scaling_factor, + layer_norm_epsilon=layer_norm_epsilon, + activation=ops.silu, + kernel_initializer=_mistral_kernel_initializer(stddev=0.02), + sliding_window=sliding_window, + head_dim=head_dim, + dropout=dropout, + dtype=dtype, + name=f"transformer_layer_{i}", + ) + self.transformer_layers.append(layer) + self.layer_norm = RMSNormalization( + epsilon=layer_norm_epsilon, + dtype=dtype, + name="sequence_output_layernorm", + ) + self.vision_encoder = vision_encoder + self.multimodal_projector = multimodal_projector + self.image_text_embedding_merger = Mistral3ImageTextEmbeddingMerger( + dtype=dtype, + name="image_text_embedding_merger", + ) + self.image_feature_extractor = Mistral3ImageFeatureExtractor( + vision_encoder, + multimodal_projector, + dtype=dtype, + name="image_feature_extractor", + ) + + # === Functional Model === + token_id_input = keras.Input( + shape=(None,), dtype="int32", name="token_ids" + ) + padding_mask_input = keras.Input( + shape=(None,), dtype="int32", name="padding_mask" + ) + # `None` spatial dims: HF's `PixtralImageProcessor` pads each batch + # to its own largest image, not to a fixed canvas, so the input + # canvas size varies per call. `image_sizes` carries each image's + # true (unpadded) `(height, width)` for cropping. + pixel_values_input = keras.Input( + shape=(vision_encoder.num_channels, None, None), + name="pixel_values", + ) + image_sizes_input = keras.Input( + shape=(2,), dtype="int32", name="image_sizes" + ) + # 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", + name="placeholder_indices", + ) + + x = self.token_embedding(token_id_input) + image_features = self.image_feature_extractor( + pixel_values_input, + image_sizes_input, + ) + x = self.image_text_embedding_merger( + x, image_features, placeholder_indices_input + ) + + for transformer_layer in self.transformer_layers: + x = transformer_layer(x, decoder_padding_mask=padding_mask_input) + sequence_output = self.layer_norm(x) + + super().__init__( + inputs={ + "token_ids": token_id_input, + "padding_mask": padding_mask_input, + "pixel_values": pixel_values_input, + "image_sizes": image_sizes_input, + "placeholder_indices": placeholder_indices_input, + }, + outputs=sequence_output, + dtype=dtype, + **kwargs, + ) + + # === Config === + self.vocabulary_size = vocabulary_size + self.num_layers = num_layers + self.num_query_heads = num_query_heads + self.hidden_dim = hidden_dim + self.intermediate_dim = intermediate_dim + self.rope_max_wavelength = rope_max_wavelength + self.num_key_value_heads = num_key_value_heads + self.rope_scaling_factor = rope_scaling_factor + self.sliding_window = sliding_window + self.head_dim = head_dim + self.layer_norm_epsilon = layer_norm_epsilon + self.dropout = dropout + self.image_token_index = image_token_index + + def get_config(self): + config = super().get_config() + config.update( + { + "vocabulary_size": self.vocabulary_size, + "num_layers": self.num_layers, + "num_query_heads": self.num_query_heads, + "hidden_dim": self.hidden_dim, + "intermediate_dim": self.intermediate_dim, + "rope_max_wavelength": self.rope_max_wavelength, + "rope_scaling_factor": self.rope_scaling_factor, + "num_key_value_heads": self.num_key_value_heads, + "sliding_window": self.sliding_window, + "head_dim": self.head_dim, + "layer_norm_epsilon": self.layer_norm_epsilon, + "dropout": self.dropout, + "image_token_index": self.image_token_index, + "vision_encoder": keras.layers.serialize(self.vision_encoder), + "multimodal_projector": keras.layers.serialize( + self.multimodal_projector + ), + } + ) + return config + + @classmethod + def from_config(cls, config): + config = dict(config) + config.update( + { + "vision_encoder": keras.layers.deserialize( + config["vision_encoder"] + ), + "multimodal_projector": keras.layers.deserialize( + config["multimodal_projector"] + ), + } + ) + return super().from_config(config) diff --git a/keras_hub/src/models/mistral3/mistral3_backbone_test.py b/keras_hub/src/models/mistral3/mistral3_backbone_test.py new file mode 100644 index 0000000000..3a91254964 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_backbone_test.py @@ -0,0 +1,156 @@ +import numpy as np +import pytest +from keras import ops + +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3MultiModalProjector, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoder, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + compute_image_placeholder_indices, +) +from keras_hub.src.tests.test_case import TestCase + + +class Mistral3BackboneTest(TestCase): + def setUp(self): + self.text_init_kwargs = { + "vocabulary_size": 10, + "num_layers": 2, + "num_query_heads": 8, + "num_key_value_heads": 4, + "hidden_dim": 16, + "intermediate_dim": 8, + "sliding_window": 2, + } + + vision_encoder = Mistral3VisionEncoder( + image_size=8, + patch_size=4, + hidden_dim=8, + num_layers=1, + num_heads=2, + head_dim=4, + intermediate_dim=8, + ) + multimodal_projector = Mistral3MultiModalProjector( + vision_hidden_dim=8, + text_hidden_dim=self.text_init_kwargs["hidden_dim"], + spatial_merge_size=1, + patch_size=4, + image_size=8, + ) + # Must stay inside `vocabulary_size`: the embedding lookup happens + # before the image-text merger overwrites these positions. + self.image_token_index = 9 + self.init_kwargs = { + **self.text_init_kwargs, + "vision_encoder": vision_encoder, + "multimodal_projector": multimodal_projector, + "image_token_index": self.image_token_index, + } + # Two 8x8 images, each a 2x2 patch grid; `spatial_merge_size=1` + # makes every patch its own merge window (4 rows per image). + token_ids = ops.array( + [ + [self.image_token_index] * 4 + [3], + [self.image_token_index] * 4 + [4], + ], + dtype="int32", + ) + placeholder_indices = compute_image_placeholder_indices( + token_ids, image_token_index=self.image_token_index + ) + self.input_data = { + "token_ids": token_ids, + "padding_mask": ops.ones((2, 5), dtype="int32"), + "pixel_values": ops.convert_to_tensor( + np.random.rand(2, 3, 8, 8).astype("float32") + ), + "image_sizes": ops.array([[8, 8], [8, 8]], dtype="int32"), + "placeholder_indices": ops.convert_to_tensor(placeholder_indices), + } + + def test_backbone_basics(self): + self.run_backbone_test( + cls=Mistral3Backbone, + init_kwargs=self.init_kwargs, + input_data=self.input_data, + expected_output_shape=( + 2, + 5, + self.text_init_kwargs["hidden_dim"], + ), + # Image inputs have no sequence axis to slice, so skip the + # default variable-length sweep. + variable_length_data=[self.input_data], + # `run_quantization_test` rebuilds `vision_encoder`/ + # `multimodal_projector` as standalone objects to apply a + # path-keyed `DTypePolicyMap`, but their sublayer paths change + # once they're no longer nested under the backbone -- the same + # structural mismatch `gemma3_backbone_test.py` works around for + # 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, + ) diff --git a/keras_hub/src/models/mistral3/mistral3_causal_lm.py b/keras_hub/src/models/mistral3/mistral3_causal_lm.py new file mode 100644 index 0000000000..24796fbe51 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_causal_lm.py @@ -0,0 +1,216 @@ +from keras import ops + +from keras_hub.src.api_export import keras_hub_export +from keras_hub.src.models.mistral.mistral_causal_lm import MistralCausalLM +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_causal_lm_preprocessor import ( + Mistral3CausalLMPreprocessor, +) +from keras_hub.src.utils.tensor_utils import any_equal + + +@keras_hub_export("keras_hub.models.Mistral3CausalLM") +class Mistral3CausalLM(MistralCausalLM): + """An end-to-end Mistral3 model for causal language modeling. + + A causal language model (LM) predicts the next token based on previous + tokens. This task setup can be used to train the model unsupervised on + plain text input, or to autoregressively generate text (optionally + grounded in one or more input images) similar to the data used for + training. This task can be used for pre-training or fine-tuning a + Mistral3 model, simply by calling `fit()`. + + This model has a `generate()` method, which generates text based on a + prompt. The generation strategy used is controlled by an additional + `sampler` argument on `compile()`. You can recompile the model with + different `keras_hub.samplers` objects to control the generation. By + default, `"top_k"` sampling will be used. + + Args: + backbone: A `keras_hub.models.Mistral3Backbone` instance. + preprocessor: A `keras_hub.models.Mistral3CausalLMPreprocessor` or + `None`. If `None`, this model will not apply preprocessing, and + inputs should be preprocessed before calling the model. + """ + + backbone_cls = Mistral3Backbone + preprocessor_cls = Mistral3CausalLMPreprocessor + + def call_with_cache( + self, + token_ids, + cache, + cache_update_index, + img_embeddings=None, + placeholder_indices=None, + ): + """Forward pass of `Mistral3CausalLM` with cache. + + `call_with_cache` adds an additional forward pass for the model for + autoregressive inference. Unlike calling the model directly, this method + allows caching previous key/value Tensors in multi-head attention layer, + and avoids recomputing the outputs of seen tokens. + + Args: + token_ids: a dense int Tensor with shape `(batch_size, max_length)`. + cache: a dense float Tensor, the cache of key and value. + cache_update_index: int, or int Tensor. The index of current inputs + in the whole sequence. + img_embeddings: a dense float Tensor of projected image features, + or `None` for a text-only forward pass. Scattered into + `token_ids`' embeddings at `placeholder_indices` before the + decoder layers run. + placeholder_indices: flat positions of image placeholder tokens + in `token_ids`. Required when `img_embeddings` is not + `None`. + + Returns: + A (logits, hidden_states, cache) tuple. Where `logits` is the + language model logits for the input token_ids, `hidden_states` is + the final hidden representation of the input tokens, and `cache` is + the decoding cache. + """ + x = self.backbone.token_embedding(token_ids) + if img_embeddings is not None: + x = self.backbone.image_text_embedding_merger( + x, img_embeddings, placeholder_indices + ) + # Each decoder layer has a cache; we update them separately. + updated_cache = [] + for i in range(self.backbone.num_layers): + current_cache = cache[:, i, ...] + x, next_cache = self.backbone.transformer_layers[i]( + x, + self_attention_cache=current_cache, + self_attention_cache_update_index=cache_update_index, + ) + updated_cache.append(next_cache) + cache = ops.stack(updated_cache, axis=1) + hidden_states = x = self.backbone.layer_norm(x) + logits = self.backbone.token_embedding(x, reverse=True) + return logits, hidden_states, cache + + def _build_cache( + self, token_ids, img_embeddings=None, placeholder_indices=None + ): + """Build an empty cache for use with `call_with_cache()`.""" + batch_size = ops.shape(token_ids)[0] + max_length = ops.shape(token_ids)[1] + num_layers = self.backbone.num_layers + num_key_value_heads = self.backbone.num_key_value_heads + head_dim = self.backbone.head_dim or ( + self.backbone.hidden_dim // self.backbone.num_query_heads + ) + shape = [ + batch_size, + num_layers, + 2, + max_length, + num_key_value_heads, + head_dim, + ] + cache = ops.zeros(shape, dtype=self.compute_dtype) + # Seed the cache. + _, hidden_states, cache = self.call_with_cache( + token_ids, + cache, + 0, + img_embeddings=img_embeddings, + placeholder_indices=placeholder_indices, + ) + return hidden_states, cache + + def generate_step( + self, + inputs, + stop_token_ids=None, + ): + """A compilable generation function for a single batch of inputs. + + This function represents the inner, XLA-compilable, generation function + for a single batch of inputs. + + Args: + inputs: A dictionary with keys `"token_ids"` and + `"padding_mask"`, and batched tensor values. When images are + present, also includes `"pixel_values"`, `"image_sizes"`, + and `"placeholder_indices"`. + stop_token_ids: List of id's of end token's to stop on. If all + sequences have produced a new stop token, generation + will stop. + """ + token_ids, padding_mask = inputs["token_ids"], inputs["padding_mask"] + pixel_values = inputs.get("pixel_values", None) + image_sizes = inputs.get("image_sizes", None) + placeholder_indices = inputs.get("placeholder_indices", None) + + # Compute image features once, at prefill, from a static (Python + # int, not tensor) shape check on the number of images. An unknown + # static shape (`None`) is treated as "no images". + img_embeddings = None + if pixel_values is not None and pixel_values.shape[0]: + img_embeddings = self.backbone.image_feature_extractor( + pixel_values, image_sizes + ) + else: + placeholder_indices = None + + # Create and seed cache with a single forward pass. + hidden_states, cache = self._build_cache( + token_ids, + img_embeddings=img_embeddings, + placeholder_indices=placeholder_indices, + ) + # Compute the lengths of all user inputted tokens ids. + row_lengths = ops.sum(ops.cast(padding_mask, "int32"), axis=-1) + # Start at the first index that has no user inputted id. + index = ops.min(row_lengths) + + def next(prompt, cache, index): + # The cache index is the index of our previous token. + cache_update_index = index - 1 + batch_size = ops.shape(prompt)[0] + prompt = ops.slice(prompt, [0, cache_update_index], [batch_size, 1]) + logits, hidden_states, cache = self.call_with_cache( + prompt, + cache, + cache_update_index, + ) + return ( + ops.squeeze(logits, axis=1), + ops.squeeze(hidden_states, axis=1), + cache, + ) + + token_ids = self.sampler( + next=next, + prompt=token_ids, + cache=cache, + index=index, + mask=padding_mask, + stop_token_ids=stop_token_ids, + hidden_states=hidden_states, + model=self, + ) + + # Compute an output padding mask with the token ids we updated. + if stop_token_ids is not None: + # Build a mask of stop_tokens locations not in the original + # prompt (not in locations where `padding_mask` is True). + end_locations = any_equal( + token_ids, stop_token_ids, ops.logical_not(padding_mask) + ) + + end_locations = ops.cast(end_locations, "int32") + # Use cumsum to get ones in all locations after end_locations. + cumsum = ops.cast(ops.cumsum(end_locations, axis=-1), "int32") + overflow = cumsum - end_locations + # Our padding mask is the inverse of these overflow locations. + padding_mask = ops.logical_not(ops.cast(overflow, "bool")) + else: + # Without early stopping, all locations will have been updated. + padding_mask = ops.ones_like(token_ids, dtype="bool") + return { + "token_ids": token_ids, + "padding_mask": padding_mask, + } diff --git a/keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor.py b/keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor.py new file mode 100644 index 0000000000..b639847554 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor.py @@ -0,0 +1,539 @@ +import keras + +from keras_hub.src.api_export import keras_hub_export +from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_image_converter import ( + Mistral3ImageConverter, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import Mistral3Tokenizer +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + compute_image_placeholder_indices, +) +from keras_hub.src.utils.tensor_utils import convert_to_numpy +from keras_hub.src.utils.tensor_utils import in_tf_function +from keras_hub.src.utils.tensor_utils import preprocessing_function + +try: + import tensorflow as tf +except ImportError: + tf = None + + +@keras_hub_export("keras_hub.models.Mistral3CausalLMPreprocessor") +class Mistral3CausalLMPreprocessor(CausalLMPreprocessor): + """Mistral3 Causal LM preprocessor. + + This preprocessing layer is meant for use with + `keras_hub.models.Mistral3CausalLM`. It takes in batches of prompts and + (optionally, per-prompt) images and returns outputs in a + `(x, y, sample_weight)` format, where the `y` label is the next token id + in the `x` sequence. + + `x` for `call()`/`generate_preprocess()` should be a dict with + `"prompts"` (and optionally `"responses"`) and, for multimodal inputs, an + `"images"` key. Images are matched to prompts by their `"[IMG]"` + placeholder occurrences, consumed in order — `"images"` can be any + reasonable nesting (a single image, a batched array, flat or + per-prompt-grouped lists), as long as the total image count matches the + total placeholder count. Omitting `"images"` (or passing `x` as a plain + string/list of strings) preprocesses as plain text, matching HF's + `Mistral3ForConditionalGeneration`, which also supports text-only calls. + + For use with generation, the layer also exposes two methods + `generate_preprocess()` and `generate_postprocess()`. When this preprocessor + is attached to a `keras_hub.models.Mistral3CausalLM` instance, these methods + will be called implicitly in `generate()`. They can also be called + standalone (e.g. to precompute preprocessing inputs for generation in a + separate process). + + Args: + tokenizer: A `keras_hub.models.Mistral3Tokenizer` instance. + image_converter: A `keras_hub.layers.Mistral3ImageConverter` + instance. + sequence_length: The length of the packed inputs. + add_start_token: If `True`, the preprocessor will prepend the tokenizer + start token to each input sequence. Default is `True`. + add_end_token: If `True`, the preprocessor will append the tokenizer + end token to each input sequence. Default is `True`. + spatial_merge_size: int. The multimodal projector's spatial merge + size, used to compute how many image placeholder tokens each + image expands to. Defaults to `2`. + + Call arguments: + x: A dict with `"prompts"` and, optionally, `"images"` keys. + y: Label data. Should always be `None` as the layer generates labels. + sample_weight: Label weights. Should always be `None` as the layer + generates label weights. + sequence_length: Pass to override the configured `sequence_length` of + the layer. + """ + + backbone_cls = Mistral3Backbone + tokenizer_cls = Mistral3Tokenizer + image_converter_cls = Mistral3ImageConverter + + def __init__( + self, + tokenizer, + image_converter, + sequence_length=1024, + add_start_token=True, + add_end_token=True, + spatial_merge_size=MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE, + **kwargs, + ): + super().__init__( + tokenizer=tokenizer, + sequence_length=sequence_length, + add_start_token=add_start_token, + add_end_token=add_end_token, + **kwargs, + ) + self.image_converter = image_converter + self.spatial_merge_size = spatial_merge_size + + def _compute_image_block_ids(self, height, width): + """Builds the token-ID block a single image expands to. + + Mirrors HF's Pixtral/Mistral3 processor: an image contributes one + `image_placeholder_token_id` per merged vision-patch row/column, + each row terminated by `image_break_token_id`, with the last row's + trailing break token swapped for `image_end_token_id`. + + Args: + height: int. The image's resized height, in pixels. + width: int. The image's resized width, in pixels. + + Returns: + list of int. The token IDs this image expands to. + """ + merge = self.image_converter.patch_size * self.spatial_merge_size + num_width_tokens = width // merge + num_height_tokens = height // merge + row = [self.tokenizer.image_placeholder_token_id] * num_width_tokens + row.append(self.tokenizer.image_break_token_id) + block = row * num_height_tokens + block[-1] = self.tokenizer.image_end_token_id + return block + + def _tokenize_base(self, prompt): + """Tokenizes `prompt` whole (not split around placeholders), since + SentencePiece's leading-space handling differs per call. + + Args: + prompt: str. The raw prompt text. + + Returns: + list of int. `prompt`'s token IDs, placeholders not expanded. + """ + base_ids = self.tokenizer(prompt) + return convert_to_numpy(base_ids).tolist() + + def _expand_image_blocks(self, base_ids, image_sizes): + """Splices each image's block token ids into `base_ids`. + + Args: + base_ids: list of int, from `_tokenize_base`. + image_sizes: list of `(height, width)` tuples, one per + placeholder occurrence in `base_ids`, in order. + + Returns: + list of int. The complete token ID sequence. + """ + placeholder_id = self.tokenizer.image_placeholder_token_id + token_ids = [] + image_idx = 0 + for token_id in base_ids: + if token_id == placeholder_id: + height, width = image_sizes[image_idx] + token_ids.extend(self._compute_image_block_ids(height, width)) + image_idx += 1 + else: + token_ids.append(token_id) + return token_ids + + def _tokenize_multimodal_prompts(self, prompts, image_sizes): + """Tokenizes `prompts` and splices in each image's block token ids. + + Args: + prompts: list of str. + image_sizes: list of `(height, width)` tuples, one per + placeholder occurrence across `prompts`, in order. + + Returns: + list of list of int. Token ids per prompt, placeholders + expanded. + """ + placeholder_id = self.tokenizer.image_placeholder_token_id + base_ids_per_prompt = [self._tokenize_base(p) for p in prompts] + occurrence_counts = [ + base_ids.count(placeholder_id) for base_ids in base_ids_per_prompt + ] + total_occurrences = sum(occurrence_counts) + if total_occurrences != len(image_sizes): + raise ValueError( + "The total number of image placeholder token occurrences " + "across `prompts` must match the number of images " + f"provided. Received: {total_occurrences} occurrence(s) " + f"across {len(prompts)} prompt(s), but {len(image_sizes)} " + "image(s)." + ) + + tokenized = [] + offset = 0 + for base_ids, num_occurrences in zip( + base_ids_per_prompt, occurrence_counts + ): + sizes_slice = image_sizes[offset : offset + num_occurrences] + offset += num_occurrences + tokenized.append(self._expand_image_blocks(base_ids, sizes_slice)) + return tokenized + + def _convert_images(self, flat_images): + """Runs `self.image_converter`, or signals an empty image batch. + + Args: + flat_images: list of raw images, or a `tf.Tensor` stacking + them on its leading axis (see `_flatten_images`). + + Returns: + `(pixel_values, image_sizes)`, or `(None, None)` if + `flat_images` is empty. + """ + # `len()` fails on a `tf.Tensor` with an unknown leading dim; use + # the static shape, treating unknown as non-empty. + if isinstance(flat_images, list): + num_images = len(flat_images) + else: + num_images = flat_images.shape[0] + if num_images == 0: + return None, None + return self.image_converter(flat_images) + + def _build_multimodal_inputs(self, prompts, flat_images): + """Tokenizes prompts and produces vision model inputs. + + Images are matched to prompts by consuming `flat_images` + left-to-right as placeholder tokens are encountered, not by any + caller-supplied grouping. Tokenization runs inside `tf.py_function` + since it needs concrete Python values, which `prompts` may not be + (e.g. inside `tf.data.Dataset.map`). + + Args: + prompts: list of str, or a `tf.Tensor` of str. The raw prompts. + flat_images: list of raw images, or a `tf.Tensor` stacking + them on its leading axis, in placeholder-occurrence order + across `prompts`. + + Returns: + `(tokenized, pixel_values, image_sizes)`. For an image-free + batch, `tokenized` is `prompts` unchanged and + `pixel_values`/`image_sizes` are `None`. Otherwise `tokenized` + is a ragged int32 tensor of token ids. + """ + pixel_values, image_sizes = self._convert_images(flat_images) + if pixel_values is None: + return prompts, None, None + + def _encode(prompts_tensor, image_sizes_tensor): + prompts_list = [p.decode("utf-8") for p in prompts_tensor.numpy()] + image_sizes_list = [ + tuple(size) for size in image_sizes_tensor.numpy().tolist() + ] + tokenized = self._tokenize_multimodal_prompts( + prompts_list, image_sizes_list + ) + return tf.ragged.constant(tokenized, dtype="int32") + + prompts_tensor = ( + prompts + if isinstance(prompts, tf.Tensor) + else tf.constant(prompts, dtype=tf.string) + ) + tokenized = tf.py_function( + _encode, + [prompts_tensor, image_sizes], + Tout=tf.RaggedTensorSpec( + shape=[None, None], dtype="int32", ragged_rank=1 + ), + ) + return tokenized, pixel_values, image_sizes + + def _flatten_images(self, images): + """Flattens `images` so all images sit on one leading axis. + + A `tf.Tensor` is folded via `tf.reshape` (graph safe); arbitrary + Python nesting is flattened by iteration (eager only). + + Returns: + Either a list of individual images or a single tensor + stacking every image on its leading axis. Both support + `len()` and are accepted by `self.image_converter`. + """ + if images is None: + return [] + if tf is not None and isinstance(images, tf.Tensor): + if images.shape.rank == 3: + return tf.expand_dims(images, axis=0) + image_shape = images.shape[-3:].as_list() + return tf.reshape(images, [-1] + image_shape) + if hasattr(images, "shape") and len(images.shape) == 3: + return [images] + if hasattr(images, "shape") and len(images.shape) == 4: + return list(images) + flat_images = [] + for item in images: + flat_images.extend(self._flatten_images(item)) + return flat_images + + def _build_multimodal_outputs(self, prompts, image_sizes, sequence_length): + """Builds `_call_multimodal_python`'s per-example outputs. + + Tokenization, packing, and placeholder-index computation each + depend on the previous step's concrete output, so they run + together in one `tf.py_function`. + + Args: + prompts: list of str, or a `tf.Tensor` of str. + image_sizes: int tensor `(num_images, 2)`, from + `self.image_converter`. + sequence_length: int. + + Returns: + `(model_token_ids, model_padding_mask, y, sample_weight, + placeholder_indices)`. + """ + + def _build(prompts_tensor, image_sizes_tensor): + prompts_list = [p.decode("utf-8") for p in prompts_tensor.numpy()] + image_sizes_list = [ + tuple(size) for size in image_sizes_tensor.numpy().tolist() + ] + tokenized = self._tokenize_multimodal_prompts( + prompts_list, image_sizes_list + ) + tokenized = tf.ragged.constant(tokenized, dtype="int32") + # Pad with one extra token to account for the truncation below. + token_ids, padding_mask = self.packer( + tokenized, + sequence_length=sequence_length + 1, + add_start_value=self.add_start_token, + add_end_value=self.add_end_token, + ) + model_token_ids = token_ids[..., :-1] + model_padding_mask = padding_mask[..., :-1] + y = token_ids[..., 1:] + sample_weight = padding_mask[..., 1:] + placeholder_indices = compute_image_placeholder_indices( + convert_to_numpy(model_token_ids), + self.tokenizer.image_placeholder_token_id, + ) + return ( + model_token_ids, + model_padding_mask, + y, + sample_weight, + placeholder_indices, + ) + + prompts_tensor = ( + prompts + if isinstance(prompts, tf.Tensor) + else tf.constant(prompts, dtype=tf.string) + ) + ( + model_token_ids, + model_padding_mask, + y, + sample_weight, + placeholder_indices, + ) = tf.py_function( + _build, + [prompts_tensor, image_sizes], + Tout=[tf.int32, tf.bool, tf.int32, tf.bool, tf.int32], + ) + # `tf.py_function` outputs have unknown rank unless set explicitly, + # which breaks the model's shape inference. `placeholder_indices`' + # last dim is data-dependent, so only its rank is fixed. + model_token_ids.set_shape([None, sequence_length]) + model_padding_mask.set_shape([None, sequence_length]) + y.set_shape([None, sequence_length]) + sample_weight.set_shape([None, sequence_length]) + placeholder_indices.set_shape([None, None]) + return ( + model_token_ids, + model_padding_mask, + y, + sample_weight, + placeholder_indices, + ) + + def _extract_multimodal_inputs(self, x): + """Normalizes `x` into `(prompts, flat_images, batched)`.""" + if isinstance(x, dict): + prompts = x["prompts"] + images = x.get("images", None) + else: + prompts = x + images = None + + batched = True + if isinstance(prompts, str): + batched = False + prompts = [prompts] + elif tf is not None and isinstance(prompts, tf.Tensor): + if prompts.shape.rank == 0: + batched = False + prompts = tf.expand_dims(prompts, 0) + else: + prompts = list(prompts) + + return prompts, self._flatten_images(images), batched + + def _call_multimodal_python( + self, x, y=None, sample_weight=None, sequence_length=None + ): + sequence_length = sequence_length or self.sequence_length + prompts, flat_images, batched = self._extract_multimodal_inputs(x) + pixel_values, image_sizes = self._convert_images(flat_images) + if pixel_values is None: + raise ValueError( + 'Mistral3\'s preprocessor was passed an `"images"` key but ' + "found zero images across the batch." + ) + + ( + model_token_ids, + model_padding_mask, + y, + sample_weight, + placeholder_indices, + ) = self._build_multimodal_outputs( + prompts, image_sizes, sequence_length + ) + + out_x = { + "token_ids": model_token_ids, + "padding_mask": model_padding_mask, + "pixel_values": pixel_values, + "image_sizes": image_sizes, + "placeholder_indices": placeholder_indices, + } + + if not batched: + out_x["token_ids"] = keras.ops.squeeze(out_x["token_ids"], axis=0) + out_x["padding_mask"] = keras.ops.squeeze( + out_x["padding_mask"], axis=0 + ) + y = keras.ops.squeeze(y, axis=0) + sample_weight = keras.ops.squeeze(sample_weight, axis=0) + + return keras.utils.pack_x_y_sample_weight(out_x, y, sample_weight) + + @preprocessing_function + def _call_multimodal_tf( + self, x, y=None, sample_weight=None, sequence_length=None + ): + return self._call_multimodal_python( + x, + y=y, + sample_weight=sample_weight, + sequence_length=sequence_length, + ) + + def call( + self, + x, + y=None, + sample_weight=None, + sequence_length=None, + ): + images = x.get("images") if isinstance(x, dict) else None + if images is None: + # Mistral3 (like the HF model it wraps) supports plain text-only + # calls: no image inputs are added to the output in that case. + prompts = x["prompts"] if isinstance(x, dict) else x + return super().call( + prompts, + y=y, + sample_weight=sample_weight, + sequence_length=sequence_length, + ) + + if not self._allow_python_workflow or in_tf_function(): + return self._call_multimodal_tf( + x, + y=y, + sample_weight=sample_weight, + sequence_length=sequence_length, + ) + return self._call_multimodal_python( + x, + y=y, + sample_weight=sample_weight, + sequence_length=sequence_length, + ) + + @preprocessing_function + def generate_preprocess( + self, + x, + sequence_length=None, + ): + """Convert prompts (and optional images) to model inputs for generation. + + `x` may be a string, list of strings, or a dict with a `"prompts"` + key and an `"images"` key. Returns a dict with `token_ids` and + `padding_mask`, plus `pixel_values`, `image_sizes`, and + `placeholder_indices` when images are present. + """ + images = x.get("images") if isinstance(x, dict) else None + if images is None: + # Mistral3 (like the HF model it wraps) supports plain text-only + # generation: no image inputs are added to the output in that + # case. + prompts = x["prompts"] if isinstance(x, dict) else x + return super().generate_preprocess( + prompts, sequence_length=sequence_length + ) + + if not self.built: + self.build(None) + + prompts, flat_images, batched = self._extract_multimodal_inputs(x) + tokenized, pixel_values, image_sizes = self._build_multimodal_inputs( + prompts, flat_images + ) + if pixel_values is None: + tokenized = self.tokenizer(tokenized) + token_ids, padding_mask = self.packer( + tokenized, sequence_length=sequence_length, add_end_value=False + ) + + out_x = { + "token_ids": token_ids, + "padding_mask": padding_mask, + } + if pixel_values is not None: + placeholder_indices = compute_image_placeholder_indices( + keras.ops.convert_to_numpy(token_ids), + self.tokenizer.image_placeholder_token_id, + ) + out_x["pixel_values"] = pixel_values + out_x["image_sizes"] = image_sizes + out_x["placeholder_indices"] = placeholder_indices + if not batched: + out_x["token_ids"] = keras.ops.squeeze(out_x["token_ids"], axis=0) + out_x["padding_mask"] = keras.ops.squeeze( + out_x["padding_mask"], axis=0 + ) + return out_x + + def get_config(self): + config = super().get_config() + config.update({"spatial_merge_size": self.spatial_merge_size}) + return config diff --git a/keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor_test.py b/keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor_test.py new file mode 100644 index 0000000000..26d441b5ed --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_causal_lm_preprocessor_test.py @@ -0,0 +1,135 @@ +import numpy as np +import pytest +from keras import ops + +from keras_hub.src.models.mistral3.mistral3_causal_lm_preprocessor import ( + Mistral3CausalLMPreprocessor, +) +from keras_hub.src.models.mistral3.mistral3_image_converter import ( + Mistral3ImageConverter, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import ( + MISTRAL3_TEKKEN_SPLIT_PATTERN as _TEKKEN_SPLIT_PATTERN, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import Mistral3Tokenizer +from keras_hub.src.tests.test_case import TestCase + + +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 = [ + "", + "", + "", + "", + "[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): + self.tokenizer = Mistral3Tokenizer(**_tekken_vision_init_kwargs()) + self.image_converter = Mistral3ImageConverter( + longest_edge=16, patch_size=4, spatial_merge_size=1 + ) + self.init_kwargs = { + "tokenizer": self.tokenizer, + "image_converter": self.image_converter, + "sequence_length": 32, + "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") + x = preprocessor.generate_preprocess( + {"prompts": "the [IMG] quick", "images": [image]} + ) + for key in ( + "token_ids", + "padding_mask", + "pixel_values", + "image_sizes", + "placeholder_indices", + ): + self.assertIn(key, x) + 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, + ) diff --git a/keras_hub/src/models/mistral3/mistral3_causal_lm_test.py b/keras_hub/src/models/mistral3/mistral3_causal_lm_test.py new file mode 100644 index 0000000000..ead81cd51b --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_causal_lm_test.py @@ -0,0 +1,280 @@ +from unittest.mock import patch + +import numpy as np +import pytest +from keras import ops +from keras import tree + +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_causal_lm import Mistral3CausalLM +from keras_hub.src.models.mistral3.mistral3_causal_lm_preprocessor import ( + Mistral3CausalLMPreprocessor, +) +from keras_hub.src.models.mistral3.mistral3_image_converter import ( + Mistral3ImageConverter, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import ( + MISTRAL3_TEKKEN_SPLIT_PATTERN as _TEKKEN_SPLIT_PATTERN, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import Mistral3Tokenizer +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3MultiModalProjector, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoder, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + compute_image_placeholder_indices, +) +from keras_hub.src.tests.test_case import TestCase + + +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 = [ + "", + "", + "", + "", + "[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 Mistral3CausalLMTest(TestCase): + def setUp(self): + self.tokenizer = Mistral3Tokenizer(**_tekken_vision_init_kwargs()) + self.image_converter = Mistral3ImageConverter( + longest_edge=8, patch_size=4, spatial_merge_size=1 + ) + self.preprocessor = Mistral3CausalLMPreprocessor( + tokenizer=self.tokenizer, + image_converter=self.image_converter, + sequence_length=16, + spatial_merge_size=1, + ) + vision_encoder = Mistral3VisionEncoder( + image_size=8, + patch_size=4, + hidden_dim=8, + num_layers=1, + num_heads=2, + head_dim=4, + intermediate_dim=8, + ) + multimodal_projector = Mistral3MultiModalProjector( + vision_hidden_dim=8, + text_hidden_dim=16, + spatial_merge_size=1, + patch_size=4, + image_size=8, + ) + self.backbone = Mistral3Backbone( + vocabulary_size=self.tokenizer.vocabulary_size(), + num_layers=2, + num_query_heads=8, + num_key_value_heads=4, + hidden_dim=16, + intermediate_dim=8, + sliding_window=2, + vision_encoder=vision_encoder, + multimodal_projector=multimodal_projector, + image_token_index=self.tokenizer.image_placeholder_token_id, + ) + self.init_kwargs = { + "backbone": self.backbone, + "preprocessor": self.preprocessor, + } + self.train_data = ( + { + "prompts": ["the [IMG] tin", "in [IMG] the"], + "images": [ + [np.zeros((8, 8, 3), dtype="float32")], + [np.ones((8, 8, 3), dtype="float32")], + ], + }, + ) + self.input_data = tree.map_structure( + ops.convert_to_tensor, self.preprocessor(*self.train_data)[0] + ) + + def test_causal_lm_basics(self): + self.run_task_test( + cls=Mistral3CausalLM, + init_kwargs=self.init_kwargs, + train_data=self.train_data, + expected_output_shape=(2, 16, self.tokenizer.vocabulary_size()), + ) + + def test_multimodal_generate(self): + vision_encoder = Mistral3VisionEncoder( + image_size=8, + patch_size=4, + hidden_dim=8, + num_layers=1, + num_heads=2, + head_dim=4, + intermediate_dim=8, + ) + multimodal_projector = Mistral3MultiModalProjector( + vision_hidden_dim=8, + text_hidden_dim=8, + spatial_merge_size=1, + patch_size=4, + image_size=8, + ) + # `image_token_index` must stay inside `vocabulary_size`, since the + # token embedding lookup happens before the image-text merger + # overwrites those positions. + image_token_index = 9 + backbone = Mistral3Backbone( + vocabulary_size=10, + num_layers=2, + num_query_heads=4, + num_key_value_heads=2, + hidden_dim=8, + intermediate_dim=16, + vision_encoder=vision_encoder, + multimodal_projector=multimodal_projector, + image_token_index=image_token_index, + ) + causal_lm = Mistral3CausalLM(backbone=backbone, preprocessor=None) + + # Two 8x8 images, each a 2x2 patch grid; `spatial_merge_size=1` + # makes every patch its own merge window (4 rows per image). + # Followed by one real token, then padding for incremental decoding. + token_ids = ops.array( + [ + [image_token_index] * 4 + [3, 0, 0], + [image_token_index] * 4 + [4, 0, 0], + ], + dtype="int32", + ) + padding_mask = ops.array( + [ + [1, 1, 1, 1, 1, 0, 0], + [1, 1, 1, 1, 1, 0, 0], + ], + ) + placeholder_indices = compute_image_placeholder_indices( + token_ids, image_token_index=image_token_index + ) + input_data = { + "token_ids": token_ids, + "padding_mask": padding_mask, + "pixel_values": ops.convert_to_tensor( + np.random.rand(2, 3, 8, 8).astype("float32") + ), + "image_sizes": ops.array([[8, 8], [8, 8]], dtype="int32"), + "placeholder_indices": ops.convert_to_tensor(placeholder_indices), + } + output = causal_lm.generate(input_data, stop_token_ids=None) + self.assertEqual(ops.shape(output["token_ids"]), (2, 7)) + self.assertEqual(ops.shape(output["padding_mask"]), (2, 7)) + + @pytest.mark.large + def test_saved_model(self): + self.run_model_saving_test( + cls=Mistral3CausalLM, + init_kwargs=self.init_kwargs, + input_data=self.input_data, + ) + + def test_generate(self): + causal_lm = Mistral3CausalLM(**self.init_kwargs) + prompt = "the tin" + output = causal_lm.generate(prompt) + self.assertTrue(prompt in output) + prompts = ["the tin", "in the"] + outputs = causal_lm.generate(prompts) + for prompt, output in zip(prompts, outputs): + self.assertTrue(prompt in output) + prompt_ids = self.preprocessor.generate_preprocess([prompt]) + causal_lm.preprocessor = None + outputs = causal_lm.generate(prompt_ids, stop_token_ids=None) + self.assertAllEqual( + outputs["token_ids"][:, :2], prompt_ids["token_ids"][:, :2] + ) + self.assertAllEqual( + outputs["padding_mask"][:, :2], prompt_ids["padding_mask"][:, :2] + ) + + def test_early_stopping(self): + causal_lm = Mistral3CausalLM(**self.init_kwargs) + call_with_cache = causal_lm.call_with_cache + + def wrapper(*args, **kwargs): + """Modify output logits to always favor end_token_id""" + logits, hidden_states, cache = call_with_cache(*args, **kwargs) + index = self.tokenizer.end_token_id + update = ops.ones_like(logits)[:, :, index] * 1.0e9 + update = ops.expand_dims(update, axis=-1) + logits = ops.slice_update(logits, (0, 0, index), update) + return logits, hidden_states, cache + + with patch.object(causal_lm, "call_with_cache", wraps=wrapper): + prompt = ["the tin", "in the"] + output = causal_lm.generate(prompt) + self.assertEqual(prompt, output) + + def test_generate_compilation(self): + causal_lm = Mistral3CausalLM(**self.init_kwargs) + causal_lm.generate("the tin") + first_fn = causal_lm.generate_function + causal_lm.generate("the tin") + second_fn = causal_lm.generate_function + self.assertEqual(first_fn, second_fn) + causal_lm.compile(sampler="greedy") + self.assertIsNone(causal_lm.generate_function) + + @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 Mistral3CausalLM.presets: + self.run_preset_test( + cls=Mistral3CausalLM, + preset=preset, + input_data=input_data, + ) diff --git a/keras_hub/src/models/mistral3/mistral3_image_converter.py b/keras_hub/src/models/mistral3/mistral3_image_converter.py new file mode 100644 index 0000000000..d67b90eb51 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_image_converter.py @@ -0,0 +1,191 @@ +import numpy as np +from keras import ops + +from keras_hub.src.api_export import keras_hub_export +from keras_hub.src.layers.preprocessing.image_converter import ImageConverter +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + compute_resize_size, +) +from keras_hub.src.utils.tensor_utils import preprocessing_function + +try: + import tensorflow as tf +except ImportError: + tf = None + +# CLIP normalization stats, in [0, 255] pixel-value units. +_CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +_CLIP_STD = [0.26862954, 0.26130258, 0.27577711] + + +@keras_hub_export("keras_hub.layers.Mistral3ImageConverter") +class Mistral3ImageConverter(ImageConverter): + """Converts raw images into `Mistral3Backbone`'s multimodal inputs. + + Each image is resized so its longest edge is at most `longest_edge` + pixels (aspect ratio preserved), then rounded up to a `patch_size` + multiple, matching HF's `PixtralImageProcessor`. Since every image in a + call can resize to a different shape, resizing runs in a Python loop + over `inputs` (a list of variable-size images) rather than through the + base class's single-`Resizing`-layer `call()`. + + Args: + longest_edge: int. The maximum size of an image's longer side after + resizing. Defaults to `1540`. + patch_size: int. The vision encoder's patch size. Defaults to `14`. + spatial_merge_size: int. The number of patches merged together per + side by the multimodal projector's patch merger. Resized image + dimensions are rounded up to a multiple of + `patch_size * spatial_merge_size`. Defaults to `2`. + scale: float, tuple of floats, or `None`. Per-channel scale applied + after resizing. Defaults to the CLIP normalization scale. + offset: float, tuple of floats, or `None`. Per-channel offset + applied after resizing. Defaults to the CLIP normalization + offset. + """ + + backbone_cls = Mistral3Backbone + + def __init__( + self, + longest_edge=1540, + patch_size=14, + spatial_merge_size=MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE, + scale=None, + offset=None, + **kwargs, + ): + if scale is None: + scale = [1.0 / 255.0 / s for s in _CLIP_STD] + if offset is None: + offset = [-m / s for m, s in zip(_CLIP_MEAN, _CLIP_STD)] + # `image_size=None` skips the base class's `Resizing` sublayer, + # since Pixtral's resize target is dynamic per image. `dtype` is + # always float32, independent of the model's compute dtype. Both + # are hardcoded below, so drop any incoming values (e.g. from a + # deserialized config) rather than conflict with them. + kwargs.pop("dtype", None) + kwargs.pop("image_size", None) + super().__init__( + image_size=None, + scale=scale, + offset=offset, + dtype="float32", + **kwargs, + ) + self.longest_edge = longest_edge + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + + def _call_python(self, inputs): + unbatched = ( + not isinstance(inputs, (list, tuple)) + and len(ops.shape(inputs)) == 3 + ) + if unbatched: + inputs = [inputs] + + # HF's `PixtralProcessor` rounds resized dimensions to a multiple of + # `patch_size * spatial_merge_size`, not `patch_size` alone, so that + # the patch grid divides evenly into the merged patch-merger grid. + merge_patch_size = self.patch_size * self.spatial_merge_size + resized_images = [] + image_sizes = [] + for image in inputs: + image = ops.convert_to_numpy(image).astype("float32") + height, width = image.shape[0], image.shape[1] + resized_height, resized_width = compute_resize_size( + height, width, self.longest_edge, merge_patch_size + ) + # HF's default resample for Mistral3 is bicubic, applied to the + # `uint8` pixel tensor (via torchvision) before it's ever cast to + # float, and torchvision rounds/clips its resize output back to + # the `uint8` grid. `antialias=True` matches its resize kernel + # exactly, but the output must also be rounded and clipped to + # `[0, 255]` here to reproduce that same `uint8` quantization, + # or the two pipelines diverge at every interpolated pixel. + image = ops.image.resize( + image, + size=(resized_height, resized_width), + interpolation="bicubic", + antialias=True, + ) + image = ops.convert_to_numpy(image) + image = np.clip(np.round(image), 0, 255).astype("float32") + scale = np.array(self.scale, dtype="float32") + offset = np.array(self.offset, dtype="float32") + image = image * scale + offset + # Channels-last `(H, W, 3)` -> channels-first `(3, H, W)`, to + # match `Mistral3Backbone`'s `pixel_values` input layout. + image = np.transpose(image, (2, 0, 1)) + resized_images.append(image) + image_sizes.append((resized_height, resized_width)) + + max_height = max(size[0] for size in image_sizes) + max_width = max(size[1] for size in image_sizes) + + padded_images = [] + for image, (resized_height, resized_width) in zip( + resized_images, image_sizes + ): + pad_height = max_height - resized_height + pad_width = max_width - resized_width + padded_images.append( + np.pad( + image, + ((0, 0), (0, pad_height), (0, pad_width)), + ) + ) + + pixel_values = np.stack(padded_images, axis=0).astype("float32") + image_sizes = np.array(image_sizes, dtype="int32") + if unbatched: + return pixel_values[0], image_sizes[0] + return pixel_values, image_sizes + + @preprocessing_function + def _call_tf(self, inputs): + images = tf.cast(inputs, "float32") + unbatched = len(images.shape) == 3 + if unbatched: + images = tf.expand_dims(images, axis=0) + + merge_patch_size = self.patch_size * self.spatial_merge_size + height, width = images.shape[1], images.shape[2] + resized_height, resized_width = compute_resize_size( + height, width, self.longest_edge, merge_patch_size + ) + images = tf.image.resize( + images, + size=(resized_height, resized_width), + method="bicubic", + antialias=True, + ) + images = tf.clip_by_value(tf.round(images), 0, 255) + scale = tf.constant(self.scale, dtype="float32") + offset = tf.constant(self.offset, dtype="float32") + images = images * scale + offset + images = tf.transpose(images, (0, 3, 1, 2)) + num_images = tf.shape(images)[0] + image_sizes = tf.tile( + tf.constant([[resized_height, resized_width]], dtype="int32"), + (num_images, 1), + ) + if unbatched: + return images[0], image_sizes[0] + return images, image_sizes + + def get_config(self): + config = super().get_config() + config.update( + { + "longest_edge": self.longest_edge, + "patch_size": self.patch_size, + "spatial_merge_size": self.spatial_merge_size, + } + ) + return config diff --git a/keras_hub/src/models/mistral3/mistral3_image_converter_test.py b/keras_hub/src/models/mistral3/mistral3_image_converter_test.py new file mode 100644 index 0000000000..bd041706b5 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_image_converter_test.py @@ -0,0 +1,78 @@ +import numpy as np + +from keras_hub.src.models.mistral3.mistral3_image_converter import ( + Mistral3ImageConverter, +) +from keras_hub.src.tests.test_case import TestCase + + +class Mistral3ImageConverterTest(TestCase): + def setUp(self): + # `spatial_merge_size=1` so the rounding multiple is `patch_size` + # alone, isolating resize-rounding tests from patch-merger granularity. + self.init_kwargs = { + "longest_edge": 16, + "patch_size": 4, + "spatial_merge_size": 1, + } + self.converter = Mistral3ImageConverter(**self.init_kwargs) + + def test_image_converter_basics(self): + image_a = np.full((8, 8, 3), 255.0, dtype="float32") + image_b = np.zeros((8, 8, 3), dtype="float32") + input_data = np.stack([image_a, image_b], axis=0) + self.run_preprocessing_layer_test( + cls=Mistral3ImageConverter, + init_kwargs=self.init_kwargs, + input_data=input_data, + ) + + def test_single_image_already_patch_multiple(self): + image = np.zeros((8, 8, 3), dtype="float32") + pixel_values, image_sizes = self.converter([image]) + self.assertEqual(pixel_values.shape, (1, 3, 8, 8)) + self.assertAllEqual(image_sizes, np.array([[8, 8]], dtype="int32")) + + def test_odd_sized_image_rounds_up_to_patch_multiple(self): + converter = Mistral3ImageConverter( + **{**self.init_kwargs, "longest_edge": 32} + ) + image = np.zeros((17, 17, 3), dtype="float32") + pixel_values, image_sizes = converter([image]) + # ratio = 17 / 32 < 1, so no downscale; each dim rounds up to the + # next multiple of `patch_size=4`: (17 - 1) // 4 + 1 = 5 -> 20. + self.assertAllEqual(image_sizes, np.array([[20, 20]], dtype="int32")) + self.assertEqual(pixel_values.shape, (1, 3, 20, 20)) + + def test_batch_with_different_sizes_reports_true_sizes_and_pads(self): + image_a = np.full((8, 8, 3), 255.0, dtype="float32") + image_b = np.zeros((12, 8, 3), dtype="float32") + pixel_values, image_sizes = self.converter([image_a, image_b]) + + self.assertAllEqual( + image_sizes, np.array([[8, 8], [12, 8]], dtype="int32") + ) + # Batch-local padding to this call's max (12, 8), not a fixed + # canvas. + self.assertEqual(pixel_values.shape, (2, 3, 12, 8)) + + pixel_values = np.array(pixel_values) + # Rows [8, 12) are zero padding added to reach the batch max height. + self.assertAllClose( + pixel_values[0, :, 8:, :], np.zeros((3, 4, 8), dtype="float32") + ) + # CLIP normalization: x * scale + offset, scale = 1/255/std, + # offset = -mean/std. + mean = np.array([0.48145466, 0.4578275, 0.40821073], dtype="float32") + std = np.array([0.26862954, 0.26130258, 0.27577711], dtype="float32") + expected_pixel = (255.0 / 255.0 - mean) / std + self.assertAllClose(pixel_values[0, :, 0, 0], expected_pixel, atol=1e-4) + + def test_spatial_merge_size_widens_rounding_multiple(self): + converter = Mistral3ImageConverter( + **{**self.init_kwargs, "longest_edge": 32, "spatial_merge_size": 2} + ) + image = np.zeros((9, 9, 3), dtype="float32") + pixel_values, image_sizes = converter([image]) + self.assertAllEqual(image_sizes, np.array([[16, 16]], dtype="int32")) + self.assertEqual(pixel_values.shape, (1, 3, 16, 16)) diff --git a/keras_hub/src/models/mistral3/mistral3_presets.py b/keras_hub/src/models/mistral3/mistral3_presets.py new file mode 100644 index 0000000000..41f610d85d --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_presets.py @@ -0,0 +1,42 @@ +"""Mistral3 model preset configurations.""" + +# Metadata for loading pretrained model weights. +backbone_presets = { + "mistral_small_3.1_24b_base_2503_en": { + "metadata": { + "description": ( + "24 billion parameter, 40-layer, pretrained Mistral3 model " + "with a Pixtral vision encoder for image input." + ), + "params": 24011361280, + "path": "mistral3", + }, + "kaggle_handle": "kaggle://keras/mistral3/keras/mistral_small_3.1_24b_base_2503_en/1", + }, + "mistral_small_3.1_24b_instruct_2503_en": { + "metadata": { + "description": ( + "24 billion parameter, 40-layer, instruction-tuned " + "Mistral3 model with a Pixtral vision encoder for image " + "input." + ), + "params": 24011361280, + "path": "mistral3", + }, + "kaggle_handle": "kaggle://keras/mistral3/keras/mistral_small_3.1_24b_instruct_2503_en/1", + }, + "mistral_small_3.2_24b_instruct_2506_en": { + "metadata": { + "description": ( + "24 billion parameter, 40-layer, instruction-tuned " + "Mistral3 model with a Pixtral vision encoder for image " + "input. An updated version of " + "mistral_small_3.1_24b_instruct_2503_en with improved " + "instruction-following and reduced repetition." + ), + "params": 24011361280, + "path": "mistral3", + }, + "kaggle_handle": "kaggle://keras/mistral3/keras/mistral_small_3.2_24b_instruct_2506_en/1", + }, +} diff --git a/keras_hub/src/models/mistral3/mistral3_tokenizer.py b/keras_hub/src/models/mistral3/mistral3_tokenizer.py new file mode 100644 index 0000000000..1d0bc83ac3 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_tokenizer.py @@ -0,0 +1,222 @@ +try: + import tensorflow as tf +except ImportError: + tf = None + +from keras_hub.src.api_export import keras_hub_export +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer +from keras_hub.src.utils.tensor_utils import preprocessing_function + +try: + import tokenizers as hf_tokenizers + from tokenizers import decoders + from tokenizers import models as hf_models + from tokenizers import pre_tokenizers +except ImportError: + hf_tokenizers = None + +# Tekken's pre-tokenization regex, shared by known Mistral3 checkpoints +# (mirrors `mistral_common`'s `Tekkenizer._pat_str`). Presets override this +# with the pattern read from the checkpoint's `tekken.json`. +MISTRAL3_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+" +) + + +@keras_hub_export( + [ + "keras_hub.tokenizers.Mistral3Tokenizer", + "keras_hub.models.Mistral3Tokenizer", + ] +) +class Mistral3Tokenizer(BytePairTokenizer): + """Mistral3 Tekken tokenizer layer based on byte-level BPE. + + Handles Mistral3's Tekken (tiktoken-style byte-level BPE) vocabulary. It + is based on `keras_hub.tokenizers.BytePairTokenizer`, uses the Tekken + pre-tokenization regex instead of the GPT-2/Llama3 pattern hardcoded in + the base class, and registers the special image tokens (`"[IMG]"`, + `"[IMG_BREAK]"`, `"[IMG_END]"`) used by multimodal Mistral3/Pixtral + models to mark image placeholder positions in a prompt. + + Not a subclass of `keras_hub.models.MistralTekkenTokenizer`: that class + builds its `unsplittable_tokens` list inline in `__init__` before calling + the `BytePairTokenizer` constructor, which doesn't leave a clean + extension point for adding the vision tokens via inheritance. + + The `vocabulary` and `merges` are usually produced from a `tekken.json` + file by the Hugging Face conversion path; see + `keras_hub.src.utils.transformers.convert_mistral3`. + + If input is a batch of strings (rank > 0), the layer will output a + `tf.RaggedTensor` where the last dimension of the output is ragged. + + If input is a scalar string (rank == 0), the layer will output a dense + `tf.Tensor` with static shape `[None]`. + + Args: + vocabulary: A dict mapping token strings to integer ids, or a path to a + vocabulary JSON file. + merges: A list of BPE merge rules, or a path to a merges file. + split_pattern: str, optional. The Tekken pre-tokenization regex. + Defaults to the pattern shared by known Mistral3 checkpoints. + control_tokens: list of str, optional. Extra reserved control tokens + (e.g. `"[INST]"`) to register as unsplittable, in addition to the + start/end/vision tokens. Defaults to `None`. + + Examples: + ```python + tokenizer = keras_hub.models.Mistral3Tokenizer.from_preset( + "hf://mistralai/Mistral-Small-3.2-24B-Instruct-2506", + ) + tokenizer("The quick brown fox jumped.") + tokenizer.detokenize(tokenizer("The quick brown fox jumped.")) + ``` + """ + + backbone_cls = Mistral3Backbone + + def __init__( + self, + vocabulary=None, + merges=None, + split_pattern=None, + control_tokens=None, + **kwargs, + ): + self.split_pattern = split_pattern or MISTRAL3_TEKKEN_SPLIT_PATTERN + self.control_tokens = list(control_tokens) if control_tokens else [] + self._add_special_token("", "start_token") + self._add_special_token("", "end_token") + self.pad_token_id = 0 + + # Tekken's control tokens (e.g. `"[INST]"`) occupy a reserved id + # block outside the BPE merges; register them as unsplittable, or + # literal occurrences in a prompt get shredded into byte-level + # tokens instead of mapping to their single reserved id. + unsplittable_tokens = [self.start_token, self.end_token] + for token in self.control_tokens: + if token not in unsplittable_tokens: + unsplittable_tokens.append(token) + + self._add_special_token("[IMG]", "image_placeholder_token") + self._add_special_token("[IMG_BREAK]", "image_break_token") + self._add_special_token("[IMG_END]", "image_end_token") + unsplittable_tokens += [ + self.image_placeholder_token, + self.image_break_token, + self.image_end_token, + ] + + super().__init__( + vocabulary=vocabulary, + merges=merges, + unsplittable_tokens=unsplittable_tokens, + **kwargs, + ) + + def _set_vocabulary_and_merges_tokenizers(self, vocabulary, merges): + self.vocabulary = vocabulary.copy() + self.merges = list(merges) + _merges = [] + for merge in self.merges: + if "#version:" in merge.lstrip(): + continue + a, b = str(merge).split(" ") + _merges.append((a, b)) + self._tokenizer = hf_tokenizers.Tokenizer( + hf_models.BPE(vocab=vocabulary, merges=_merges, fuse_unk=False) + ) + if self.unsplittable_tokens: + self._tokenizer.add_special_tokens(self.unsplittable_tokens) + self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split( + hf_tokenizers.Regex(self.split_pattern), + behavior="isolated", + ), + pre_tokenizers.ByteLevel( + add_prefix_space=self.add_prefix_space, use_regex=False + ), + ] + ) + self._tokenizer.decoder = decoders.ByteLevel() + + # Dummy attrs for serialization compatibility with the base class. + if not hasattr(self, "cache"): + self.byte2unicode = None + self.unicode2byte = None + self.cache = None + self.id_to_token_map = None + self.token_to_id_map = None + self.merge_ranks_lookup_default = None + self.merge_ranks = None + + def _set_vocabulary_and_merges_tf(self, vocabulary, merges): + # The base class hardcodes the GPT-2 split regex in its `tf.data` + # path, which does not match Tekken. We instead bridge to the + # `tokenizers` backend from within the graph (see `_tokenize_tf`), so + # there is nothing to build here. + self.vocabulary = vocabulary.copy() + self.merges = list(merges) + + def _maybe_initialized_tokenizers(self): + if getattr(self, "_tokenizer", None) is None: + self._set_vocabulary_and_merges_tokenizers( + self.vocabulary, self.merges + ) + + @preprocessing_function + def _tokenize_tf(self, inputs): + self._maybe_initialized_tokenizers() + + def _encode(string_tensor): + values = string_tensor.numpy() + strings = [v.decode("utf-8") for v in values.tolist()] + encodings = self._tokenizer.encode_batch( + strings, add_special_tokens=False + ) + return tf.ragged.constant( + [e.ids for e in encodings], dtype=self.compute_dtype + ) + + inputs = tf.convert_to_tensor(inputs) + unbatched = inputs.shape.rank == 0 + if unbatched: + inputs = tf.expand_dims(inputs, 0) + tokens = tf.py_function( + _encode, + [inputs], + Tout=tf.RaggedTensorSpec( + shape=[None, None], + dtype=self.compute_dtype, + ragged_rank=1, + ), + ) + + if self.sequence_length: + output_shape = tokens.shape.as_list() + output_shape[-1] = self.sequence_length + tokens = tokens.to_tensor( + shape=output_shape, + default_value=getattr(self, "pad_token_id", 0), + ) + if unbatched: + tokens = tokens[0] + return tokens + + def get_config(self): + config = super().get_config() + config.update( + { + "split_pattern": self.split_pattern, + "control_tokens": self.control_tokens, + } + ) + # `unsplittable_tokens` is derived from the special tokens in the + # constructor, so it is not a separate config argument. + del config["unsplittable_tokens"] + return config diff --git a/keras_hub/src/models/mistral3/mistral3_tokenizer_test.py b/keras_hub/src/models/mistral3/mistral3_tokenizer_test.py new file mode 100644 index 0000000000..a43c40c4f5 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_tokenizer_test.py @@ -0,0 +1,76 @@ +import pytest + +from keras_hub.src.models.mistral3.mistral3_tokenizer import ( + MISTRAL3_TEKKEN_SPLIT_PATTERN as _TEKKEN_SPLIT_PATTERN, +) +from keras_hub.src.models.mistral3.mistral3_tokenizer import Mistral3Tokenizer +from keras_hub.src.tests.test_case import TestCase + + +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 = [ + "", + "", + "", + "", + "[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 Mistral3TokenizerTest(TestCase): + def setUp(self): + self.init_kwargs = _tekken_vision_init_kwargs() + + def test_tokenizer_basics(self): + self.run_preprocessing_layer_test( + cls=Mistral3Tokenizer, + init_kwargs=self.init_kwargs, + input_data=["the tin", "in the"], + expected_output=[[265, 40, 124, 266], [266, 40, 265]], + ) + + @pytest.mark.kaggle_key_required + @pytest.mark.extra_large + def test_all_presets(self): + for preset in Mistral3Tokenizer.presets: + self.run_preset_test( + cls=Mistral3Tokenizer, + preset=preset, + input_data=["The quick brown fox jumped."], + ) diff --git a/keras_hub/src/models/mistral3/mistral3_vision_encoder.py b/keras_hub/src/models/mistral3/mistral3_vision_encoder.py new file mode 100644 index 0000000000..ef0b8949d0 --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_vision_encoder.py @@ -0,0 +1,1390 @@ +import keras +import numpy as np +from keras import ops + +from keras_hub.src.api_export import keras_hub_export + +MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE = 2 + + +def _mistral_kernel_initializer(stddev=0.02): + return keras.initializers.RandomNormal(stddev=stddev) + + +class Mistral3ImageFeatureExtractor(keras.layers.Layer): + """Computes projected Mistral3 image features. + + Wrapped in a `Layer` so shape queries see real tensors, not + graph-construction-time placeholders. + + Args: + vision_encoder: A `Mistral3VisionEncoder` instance. + multimodal_projector: A `Mistral3MultiModalProjector` instance. + vision_feature_layer: int. Which vision encoder hidden state to + project; only `-1` (the final hidden state) is supported. + """ + + def __init__( + self, + vision_encoder, + multimodal_projector, + vision_feature_layer=-1, + **kwargs, + ): + super().__init__(**kwargs) + if vision_feature_layer not in (-1, "last"): + raise NotImplementedError( + "`vision_feature_layer` only supports `-1` (the final " + "vision hidden state). `Mistral3VisionEncoder` does not " + f"expose intermediate hidden states. Received: " + f"{vision_feature_layer}." + ) + self.vision_encoder = vision_encoder + self.multimodal_projector = multimodal_projector + self.vision_feature_layer = vision_feature_layer + + def build(self, input_shape): + # `vision_encoder` is a `keras.Model`, so it must build through its + # own `__call__` rather than an external `.build()`. + merge_size = self.multimodal_projector.spatial_merge_size + patch_size = self.multimodal_projector.patch_size + num_channels = self.vision_encoder.num_channels + side = patch_size * merge_size + dummy_pixel_values = ops.zeros((1, num_channels, side, side)) + dummy_image_sizes = ops.convert_to_tensor([[side, side]], dtype="int32") + self.call(dummy_pixel_values, dummy_image_sizes) + self.built = True + + def call(self, pixel_values, image_sizes, training=None): + image_features = self.vision_encoder( + pixel_values, + image_sizes=image_sizes, + training=training, + ) + # `vision_encoder`'s padded-canvas output already matches what + # `Mistral3MultiModalProjector` expects; no further padding needed. + image_features = ops.squeeze(image_features, axis=0) + + patch_size = self.multimodal_projector.patch_size + max_patch_height = ops.shape(pixel_values)[2] // patch_size + max_patch_width = ops.shape(pixel_values)[3] // patch_size + + return self.multimodal_projector( + image_features, + image_sizes=image_sizes, + max_patch_height=max_patch_height, + max_patch_width=max_patch_width, + ) + + def compute_output_spec(self, pixel_values, image_sizes, training=None): + return keras.KerasTensor( + shape=(None, self.multimodal_projector.text_hidden_dim), + dtype=self.compute_dtype, + ) + + +def compute_image_placeholder_indices(token_ids, image_token_index): + """Compute per-example image placeholder positions for `token_ids`. + + Runs eagerly with NumPy in preprocessing, before `token_ids` reaches + the model. + + Args: + token_ids: int array `(batch, seq_length)`. + image_token_index: int. The token ID marking image placeholder + positions. + + Returns: + int32 NumPy array `(batch, max_placeholders)`: row `i` holds + example `i`'s own placeholder positions (local to its own + sequence), left-to-right, padded with `-1` up to the batch's max + count. + """ + token_ids = np.asarray(token_ids) + if token_ids.ndim == 1: + token_ids = token_ids[None, :] + rows = [ + np.nonzero(row == image_token_index)[0].astype("int32") + for row in token_ids + ] + max_count = max((len(row) for row in rows), default=0) + max_count = max(max_count, 1) + padded = np.full((len(rows), max_count), -1, dtype="int32") + for i, row in enumerate(rows): + padded[i, : len(row)] = row + return padded + + +def compute_resize_size(height, width, longest_edge, patch_size): + """Computes the resize target for one image. + + Scales `(height, width)` down (preserving aspect ratio) so its longest + edge is at most `longest_edge`, then rounds each dimension up to the + nearest multiple of `patch_size`. + + Args: + height: int. The image's original height. + width: int. The image's original width. + longest_edge: int. The maximum allowed size of the longer side. + patch_size: int. The patch size each output dimension must be a + multiple of. + + Returns: + `(resized_height, resized_width)` as plain Python ints. + """ + ratio = max(height / longest_edge, width / longest_edge) + if ratio > 1: + height = int(height / ratio) + width = int(width / ratio) + + resized_height = ((height - 1) // patch_size + 1) * patch_size + resized_width = ((width - 1) // patch_size + 1) * patch_size + return resized_height, resized_width + + +class Mistral3ImageTextEmbeddingMerger(keras.layers.Layer): + """Scatters projected image features into image placeholder positions. + + Replaces the token embeddings at `placeholder_indices` with the + concatenated, projected image features, matching HF's + `masked_scatter` fusion in `Mistral3Model`. + + `placeholder_indices` (each example's own local placeholder positions, + padded with `-1`) must be precomputed outside this layer — see + `compute_image_placeholder_indices` — since deriving them here via a + `nonzero`-style op would make the layer incompatible with `jax.jit` + tracing. Converting local positions to flat scatter targets is done + here with static-shape ops only (`arange`/`cumsum`), not a + data-dependent lookup. + """ + + def call(self, token_embeddings, image_features, placeholder_indices): + batch_size = ops.shape(token_embeddings)[0] + seq_length = ops.shape(token_embeddings)[1] + hidden_dim = ops.shape(token_embeddings)[2] + max_placeholders = ops.shape(placeholder_indices)[1] + + flat_embeddings = ops.reshape( + token_embeddings, + (batch_size * seq_length, hidden_dim), + ) + # Scratch row absorbs `placeholder_indices == -1` padding entries; + # sliced off below. + scratch_row = ops.zeros((1, hidden_dim), dtype=flat_embeddings.dtype) + flat_embeddings = ops.concatenate( + [flat_embeddings, scratch_row], axis=0 + ) + + placeholder_indices = ops.cast(placeholder_indices, "int32") + is_valid = placeholder_indices >= 0 + + row_index = ops.reshape( + ops.arange(batch_size, dtype="int32"), (batch_size, 1) + ) + scratch_row_index = batch_size * seq_length + global_indices = ops.where( + is_valid, + placeholder_indices + row_index * seq_length, + scratch_row_index, + ) + + # Row `i`'s own images sit in `image_features` at an offset equal + # to the number of real placeholders in preceding rows. + counts = ops.sum(ops.cast(is_valid, "int32"), axis=1) + row_starts = ops.cumsum(counts) - counts + local_positions = ops.reshape( + ops.arange(max_placeholders, dtype="int32"), (1, max_placeholders) + ) + feature_index = ops.reshape(row_starts, (batch_size, 1)) + ( + local_positions + ) + feature_index = ops.where(is_valid, feature_index, 0) + + image_features = ops.cast(image_features, flat_embeddings.dtype) + gathered_features = ops.take( + image_features, ops.reshape(feature_index, (-1,)), axis=0 + ) + + merged_embeddings = ops.scatter_update( + inputs=flat_embeddings, + indices=ops.expand_dims( + ops.reshape(global_indices, (-1,)), axis=-1 + ), + updates=gathered_features, + ) + merged_embeddings = merged_embeddings[: batch_size * seq_length] + + return ops.reshape( + merged_embeddings, + (batch_size, seq_length, hidden_dim), + ) + + def compute_output_shape(self, input_shape): + return input_shape + + +class Mistral3VisionRotaryEmbedding(keras.layers.Layer): + """2D rotary positional embedding for the Mistral3 vision encoder. + + Frequencies are built from 2D patch coordinates: the first half of the + frequency dims encodes height, the second half width. + """ + + def __init__( + self, + image_size, + patch_size, + head_dim, + rope_theta=10000.0, + **kwargs, + ): + super().__init__(**kwargs) + + self.image_size = image_size + self.patch_size = patch_size + self.head_dim = head_dim + self.rope_theta = rope_theta + + self.max_patches_per_side = image_size // patch_size + + def _compute_inv_freq(self): + max_patches = self.max_patches_per_side + + freq_indices = ops.arange(0, self.head_dim, 2, dtype="float32") + freqs = ops.divide( + 1.0, + ops.power(self.rope_theta, freq_indices / self.head_dim), + ) + + height_indices = ops.arange(max_patches, dtype="float32") + width_indices = ops.arange(max_patches, dtype="float32") + + freqs_h = ops.einsum( + "i,j->ij", + height_indices, + freqs[::2], + ) + + freqs_w = ops.einsum( + "i,j->ij", + width_indices, + freqs[1::2], + ) + + half_dim = self.head_dim // 4 + + # [H, 1, D/4] -> [H, W, D/4] + freqs_h = ops.broadcast_to( + ops.expand_dims(freqs_h, axis=1), + (max_patches, max_patches, half_dim), + ) + + # [1, W, D/4] -> [H, W, D/4] + freqs_w = ops.broadcast_to( + ops.expand_dims(freqs_w, axis=0), + (max_patches, max_patches, half_dim), + ) + + inv_freq = ops.concatenate( + [freqs_h, freqs_w], + axis=-1, + ) + + inv_freq = ops.reshape( + inv_freq, + (-1, self.head_dim // 2), + ) + + return ops.concatenate( + [inv_freq, inv_freq], + axis=-1, + ) + + def get_config(self): + config = super().get_config() + config.update( + { + "image_size": self.image_size, + "patch_size": self.patch_size, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + } + ) + return config + + def call(self, position_ids, dtype=None): + freqs = ops.take( + self._compute_inv_freq(), + position_ids, + axis=0, + ) + + cos = ops.cos(freqs) + sin = ops.sin(freqs) + + if dtype is not None: + cos = ops.cast(cos, dtype) + sin = ops.cast(sin, dtype) + + return cos, sin + + +def _rotate_half(x): + half = ops.shape(x)[-1] // 2 + + x1 = x[..., :half] + x2 = x[..., half:] + + return ops.concatenate( + [-x2, x1], + axis=-1, + ) + + +def _apply_rotary_pos_emb(q, k, cos, sin): + # q/k: [B, H, S, D]; cos/sin: [B, S, D]. + # Expand at axis 1 to broadcast cos/sin across attention heads. + + cos = ops.expand_dims(cos, axis=1) + sin = ops.expand_dims(sin, axis=1) + + q = q * cos + _rotate_half(q) * sin + k = k * cos + _rotate_half(k) * sin + + return q, k + + +class Mistral3VisionAttention(keras.layers.Layer): + """Multi-head self-attention used by the Mistral3 vision encoder. + + Args: + hidden_dim: int. The size of the attention layer's input/output. + num_heads: int. The number of attention heads. + head_dim: int. The size of each attention head. Defaults to + `hidden_dim // num_heads`. + dropout: float. The dropout probability applied to attention scores. + Defaults to `0.0`. + """ + + def __init__( + self, + hidden_dim, + num_heads, + head_dim=None, + dropout=0.0, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_dim = hidden_dim + self.num_heads = num_heads + self.head_dim = head_dim or hidden_dim // num_heads + self.dropout = dropout + + self.scaling = self.head_dim**-0.5 + + self.q_proj = keras.layers.Dense( + hidden_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="q_proj", + ) + self.k_proj = keras.layers.Dense( + hidden_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="k_proj", + ) + self.v_proj = keras.layers.Dense( + hidden_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="v_proj", + ) + self.o_proj = keras.layers.Dense( + hidden_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="o_proj", + ) + + self.attention_dropout = keras.layers.Dropout( + dropout, dtype=self.dtype_policy + ) + + def build(self, input_shape): + self.q_proj.build(input_shape) + self.k_proj.build(input_shape) + self.v_proj.build(input_shape) + self.o_proj.build(input_shape) + self.built = True + + def _reshape_to_heads(self, x): + batch_size = ops.shape(x)[0] + sequence_length = ops.shape(x)[1] + + x = ops.reshape( + x, + ( + batch_size, + sequence_length, + self.num_heads, + self.head_dim, + ), + ) + + return ops.transpose(x, (0, 2, 1, 3)) + + def call( + self, + inputs, + attention_mask=None, + position_embeddings=None, + training=None, + ): + q = self._reshape_to_heads(self.q_proj(inputs)) + k = self._reshape_to_heads(self.k_proj(inputs)) + v = self._reshape_to_heads(self.v_proj(inputs)) + + cos, sin = position_embeddings + q, k = _apply_rotary_pos_emb(q, k, cos, sin) + + # [B, H, S, D] @ [B, H, D, S] + attention_scores = ops.matmul( + q, + ops.transpose(k, (0, 1, 3, 2)), + ) + + attention_scores = attention_scores * self.scaling + + if attention_mask is not None: + attention_scores = attention_scores + attention_mask + + # HF explicitly performs softmax in float32. + attention_scores = ops.cast( + attention_scores, + "float32", + ) + + attention_scores = ops.softmax( + attention_scores, + axis=-1, + ) + + attention_scores = ops.cast( + attention_scores, + q.dtype, + ) + + attention_scores = self.attention_dropout( + attention_scores, + training=training, + ) + + attention_output = ops.matmul( + attention_scores, + v, + ) + + attention_output = ops.transpose( + attention_output, + (0, 2, 1, 3), + ) + + attention_output = ops.reshape( + attention_output, + ( + ops.shape(attention_output)[0], + ops.shape(attention_output)[1], + self.hidden_dim, + ), + ) + + return self.o_proj(attention_output) + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_dim": self.hidden_dim, + "num_heads": self.num_heads, + "head_dim": self.head_dim, + "dropout": self.dropout, + } + ) + return config + + +class Mistral3VisionMLP(keras.layers.Layer): + """SwiGLU MLP used by the Mistral3 vision encoder. + + Args: + hidden_dim: int. The size of the MLP's input/output. + intermediate_dim: int. The size of the MLP's intermediate layer. + activation: str or callable. The activation applied to the gate + projection. Defaults to `"silu"`. + """ + + def __init__( + self, + hidden_dim, + intermediate_dim, + activation="silu", + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_dim = hidden_dim + self.intermediate_dim = intermediate_dim + self.activation = activation + + self.gate_proj = keras.layers.Dense( + intermediate_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="gate_proj", + ) + + self.up_proj = keras.layers.Dense( + intermediate_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="up_proj", + ) + + self.down_proj = keras.layers.Dense( + hidden_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="down_proj", + ) + + self.activation_fn = keras.activations.get(activation) + + def build(self, input_shape): + self.gate_proj.build(input_shape) + self.up_proj.build(input_shape) + self.down_proj.build(input_shape[:-1] + (self.intermediate_dim,)) + self.built = True + + def call(self, inputs): + gate = self.activation_fn(self.gate_proj(inputs)) + up = self.up_proj(inputs) + + return self.down_proj(gate * up) + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_dim": self.hidden_dim, + "intermediate_dim": self.intermediate_dim, + "activation": self.activation, + } + ) + return config + + +class Mistral3VisionEncoderLayer(keras.layers.Layer): + """One Mistral3 vision transformer encoder layer. + + Args: + hidden_dim: int. The size of the transformer hidden state. + intermediate_dim: int. The size of the MLP's intermediate layer. + num_heads: int. The number of attention heads. + head_dim: int. The size of each attention head. Defaults to + `hidden_dim // num_heads`. + layer_norm_epsilon: float. The epsilon for RMS normalization. + Defaults to `1e-5`. + activation: str or callable. The MLP activation. Defaults to + `"silu"`. + dropout: float. The attention dropout probability. Defaults to + `0.0`. + """ + + def __init__( + self, + hidden_dim, + intermediate_dim, + num_heads, + head_dim=None, + layer_norm_epsilon=1e-5, + activation="silu", + dropout=0.0, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_dim = hidden_dim + self.intermediate_dim = intermediate_dim + self.num_heads = num_heads + self.head_dim = head_dim or hidden_dim // num_heads + self.layer_norm_epsilon = layer_norm_epsilon + self.activation = activation + self.dropout = dropout + + self.attention_norm = keras.layers.RMSNormalization( + epsilon=layer_norm_epsilon, + dtype=self.dtype_policy, + name="attention_norm", + ) + + self.attention = Mistral3VisionAttention( + hidden_dim=hidden_dim, + num_heads=num_heads, + head_dim=head_dim, + dropout=dropout, + dtype=self.dtype_policy, + name="attention", + ) + + self.ffn_norm = keras.layers.RMSNormalization( + epsilon=layer_norm_epsilon, + dtype=self.dtype_policy, + name="ffn_norm", + ) + + self.feed_forward = Mistral3VisionMLP( + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + activation=activation, + dtype=self.dtype_policy, + name="feed_forward", + ) + + def build(self, input_shape): + self.attention_norm.build(input_shape) + self.attention.build(input_shape) + self.ffn_norm.build(input_shape) + self.feed_forward.build(input_shape) + self.built = True + + def call( + self, + hidden_states, + attention_mask=None, + position_embeddings=None, + training=None, + ): + residual = hidden_states + + hidden_states = self.attention_norm(hidden_states) + + hidden_states = self.attention( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + training=training, + ) + + hidden_states = residual + hidden_states + + residual = hidden_states + + hidden_states = self.ffn_norm(hidden_states) + hidden_states = self.feed_forward(hidden_states) + + hidden_states = residual + hidden_states + + return hidden_states + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_dim": self.hidden_dim, + "intermediate_dim": self.intermediate_dim, + "num_heads": self.num_heads, + "head_dim": self.head_dim, + "layer_norm_epsilon": self.layer_norm_epsilon, + "activation": self.activation, + "dropout": self.dropout, + } + ) + return config + + +@keras_hub_export("keras_hub.models.Mistral3VisionEncoder") +class Mistral3VisionEncoder(keras.Model): + """Vision encoder used by Mistral3. + + This is not exposed as a standalone model. It is the vision tower + consumed by the Mistral3 multimodal architecture. + + `pixel_values` follows the Hugging Face Pixtral layout: + `(num_images, num_channels, height, width)`. + + When `image_sizes` is provided, each image is cropped to its effective + image size after patchification, the resulting patch sequences are + concatenated into a single sequence, and attention is restricted to + patches belonging to the same image. + """ + + def __init__( + self, + image_size=1540, + patch_size=14, + num_channels=3, + hidden_dim=1024, + num_layers=24, + num_heads=16, + head_dim=64, + intermediate_dim=4096, + rope_theta=10000.0, + layer_norm_epsilon=1e-5, + activation="silu", + attention_dropout=0.0, + **kwargs, + ): + super().__init__(**kwargs) + + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.num_heads = num_heads + self.head_dim = head_dim + self.intermediate_dim = intermediate_dim + self.rope_theta = rope_theta + self.layer_norm_epsilon = layer_norm_epsilon + self.activation = activation + self.attention_dropout = attention_dropout + + self.patch_conv = keras.layers.Conv2D( + filters=hidden_dim, + kernel_size=patch_size, + strides=patch_size, + padding="valid", + use_bias=False, + data_format="channels_last", + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="patch_conv", + ) + + self.ln_pre = keras.layers.RMSNormalization( + epsilon=layer_norm_epsilon, + dtype=self.dtype_policy, + name="ln_pre", + ) + + self.patch_positional_embedding = Mistral3VisionRotaryEmbedding( + image_size=image_size, + patch_size=patch_size, + head_dim=head_dim, + rope_theta=rope_theta, + dtype=self.dtype_policy, + name="patch_positional_embedding", + ) + + self.transformer_layers = [] + + for i in range(num_layers): + self.transformer_layers.append( + Mistral3VisionEncoderLayer( + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + num_heads=num_heads, + head_dim=head_dim, + layer_norm_epsilon=layer_norm_epsilon, + activation=activation, + dropout=attention_dropout, + dtype=self.dtype_policy, + name=f"transformer_layer_{i}", + ) + ) + + def build(self, input_shape): + self.patch_conv.build((None, None, None, self.num_channels)) + self.ln_pre.build((None, None, self.hidden_dim)) + for layer in self.transformer_layers: + layer.build((None, None, self.hidden_dim)) + self.built = True + + def _create_position_ids_for_images( + self, + image_sizes, + max_patch_height, + max_patch_width, + ): + """Create mesh-grid position IDs matching the compacted layout. + + Mirrors `_extract_patch_sequences`'s compaction (see that method for + why a cumulative-sum rank is used instead of `ops.nonzero`) so + position IDs line up with the patches they belong to. + """ + image_sizes = ops.cast(image_sizes, "int32") + patch_heights = image_sizes[:, 0] // self.patch_size + patch_widths = image_sizes[:, 1] // self.patch_size + num_images = ops.shape(image_sizes)[0] + + height_indices = ops.arange(max_patch_height, dtype="int32") + width_indices = ops.arange(max_patch_width, dtype="int32") + height_grid = ops.expand_dims(height_indices, axis=1) + width_grid = ops.expand_dims(width_indices, axis=0) + position_grid = ( + height_grid * (self.image_size // self.patch_size) + width_grid + ) + position_grid = ops.broadcast_to( + position_grid, + (num_images, max_patch_height, max_patch_width), + ) + + height_valid = ops.expand_dims( + ops.expand_dims(height_indices, axis=0) + < ops.expand_dims(patch_heights, axis=1), + axis=2, + ) + width_valid = ops.expand_dims( + ops.expand_dims(width_indices, axis=0) + < ops.expand_dims(patch_widths, axis=1), + axis=1, + ) + valid_patches = ops.logical_and(height_valid, width_valid) + position_grid = ops.reshape(position_grid, (-1,)) + valid_patches = ops.reshape(valid_patches, (-1,)) + + capacity = num_images * max_patch_height * max_patch_width + valid_int = ops.cast(valid_patches, "int32") + rank = ops.cumsum(valid_int) - valid_int + scatter_target = ops.where(valid_patches, rank, capacity) + buffer = ops.zeros((capacity + 1,), dtype="int32") + buffer = ops.scatter_update( + buffer, + ops.expand_dims(scatter_target, axis=1), + position_grid, + ) + return buffer[:capacity] + + def _create_block_attention_mask(self, patch_counts, dtype, capacity): + """Create a block-diagonal mask over the padded, compacted layout. + + `patch_counts` gives each image's real patch count; slots at or + past `sum(patch_counts)` are the tail padding produced by + `_extract_patch_sequences` and are assigned a sentinel block ID + (`num_images`) so they never attend to, or are attended to by, a + real patch. Computes each slot's image ID via a cumulative-sum + comparison rather than `ops.repeat(..., patch_counts)` (see + `_extract_patch_sequences` for why). + """ + patch_counts = ops.cast(patch_counts, "int32") + num_images = ops.shape(patch_counts)[0] + total_valid = ops.sum(patch_counts) + cumulative_ends = ops.cumsum(patch_counts) + + token_indices = ops.arange(capacity, dtype="int32") + image_ids = ops.sum( + ops.cast( + ops.expand_dims(token_indices, axis=1) + >= ops.expand_dims(cumulative_ends, axis=0), + "int32", + ), + axis=1, + ) + image_ids = ops.where( + token_indices >= total_valid, + num_images, + image_ids, + ) + + same_block = ops.equal( + ops.expand_dims(image_ids, axis=0), + ops.expand_dims(image_ids, axis=1), + ) + + dtype = keras.backend.standardize_dtype(dtype) + if dtype == "float16": + mask_value = -65504.0 + elif dtype == "bfloat16": + mask_value = -3.38953139e38 + else: + mask_value = -3.4028234663852886e38 + + mask = ops.where( + same_block, + ops.zeros_like(ops.cast(same_block, dtype)), + ops.cast(mask_value, dtype), + ) + return ops.expand_dims(ops.expand_dims(mask, axis=0), axis=0) + + def _normalize_image_sizes(self, image_sizes, num_images, height, width): + """Normalize image sizes as an integer tensor.""" + if image_sizes is None: + image_sizes = ops.stack([height, width]) + image_sizes = ops.broadcast_to( + image_sizes, + (num_images, 2), + ) + else: + image_sizes = ops.convert_to_tensor(image_sizes) + + return ops.cast(image_sizes, "int32") + + def _extract_patch_sequences(self, patch_embeds, image_sizes): + """Compact per-image patches into cumsum-offset order, tail-padded. + + Output has the same row count as the input + (`num_images * max_patch_height * max_patch_width`): valid patches + move to the front, in per-image row-major order (the layout + `Mistral3PatchMerger` expects), remainder zeroed. Uses a + cumulative-sum rank instead of `ops.nonzero`, since `nonzero`'s + output shape depends on tensor values and breaks `jax.jit` tracing. + """ + image_sizes = ops.cast(image_sizes, "int32") + patch_heights = image_sizes[:, 0] // self.patch_size + patch_widths = image_sizes[:, 1] // self.patch_size + num_images = ops.shape(patch_embeds)[0] + max_patch_height = ops.shape(patch_embeds)[1] + max_patch_width = ops.shape(patch_embeds)[2] + hidden_dim = ops.shape(patch_embeds)[-1] + + height_indices = ops.arange(max_patch_height, dtype="int32") + width_indices = ops.arange(max_patch_width, dtype="int32") + height_valid = ops.expand_dims( + ops.expand_dims(height_indices, axis=0) + < ops.expand_dims(patch_heights, axis=1), + axis=2, + ) + width_valid = ops.expand_dims( + ops.expand_dims(width_indices, axis=0) + < ops.expand_dims(patch_widths, axis=1), + axis=1, + ) + valid_patches = ops.logical_and(height_valid, width_valid) + patch_embeds = ops.reshape(patch_embeds, (-1, hidden_dim)) + valid_patches = ops.reshape(valid_patches, (-1,)) + + capacity = num_images * max_patch_height * max_patch_width + valid_int = ops.cast(valid_patches, "int32") + rank = ops.cumsum(valid_int) - valid_int + scatter_target = ops.where(valid_patches, rank, capacity) + buffer = ops.zeros((capacity + 1, hidden_dim), patch_embeds.dtype) + buffer = ops.scatter_update( + buffer, + ops.expand_dims(scatter_target, axis=1), + patch_embeds, + ) + return buffer[:capacity] + + def call( + self, + pixel_values, + image_sizes=None, + training=None, + ): + # HF input layout: [num_images, channels, height, width]. + pixel_values = ops.transpose(pixel_values, (0, 2, 3, 1)) + pixel_values = ops.cast( + pixel_values, + self.patch_conv.variable_dtype, + ) + patch_embeds = self.patch_conv(pixel_values) + max_patch_height = ops.shape(patch_embeds)[1] + max_patch_width = ops.shape(patch_embeds)[2] + + image_sizes = self._normalize_image_sizes( + image_sizes, + num_images=ops.shape(pixel_values)[0], + height=ops.shape(pixel_values)[1], + width=ops.shape(pixel_values)[2], + ) + patch_embeds = self._extract_patch_sequences( + patch_embeds, + image_sizes, + ) + patch_embeds = ops.expand_dims(patch_embeds, axis=0) + patch_embeds = self.ln_pre(patch_embeds) + + position_ids = self._create_position_ids_for_images( + image_sizes, + max_patch_height, + max_patch_width, + ) + position_ids = ops.expand_dims(position_ids, axis=0) + cos, sin = self.patch_positional_embedding( + position_ids, + dtype=patch_embeds.dtype, + ) + + patch_counts = (image_sizes[:, 0] // self.patch_size) * ( + image_sizes[:, 1] // self.patch_size + ) + capacity = ( + ops.shape(pixel_values)[0] * max_patch_height * max_patch_width + ) + attention_mask = self._create_block_attention_mask( + patch_counts, + patch_embeds.dtype, + capacity, + ) + + hidden_states = patch_embeds + for layer in self.transformer_layers: + hidden_states = layer( + hidden_states, + attention_mask=attention_mask, + position_embeddings=(cos, sin), + training=training, + ) + return hidden_states + + def compute_output_spec( + self, pixel_values, image_sizes=None, training=None + ): + """Declare the output shape without tracing `call()`. + + `MistralBackbone` builds this encoder into a functional model with + a dynamic `pixel_values` shape; tracing `call()`'s data-dependent + shape ops on symbolic dimensions isn't supported by the JAX + backend's graph tracer. + + Returns: + A `KerasTensor` with shape `(1, None, hidden_dim)`. + """ + return keras.KerasTensor( + shape=(1, None, self.hidden_dim), + dtype=self.compute_dtype, + ) + + def get_config(self): + config = super().get_config() + config.update( + { + "image_size": self.image_size, + "patch_size": self.patch_size, + "num_channels": self.num_channels, + "hidden_dim": self.hidden_dim, + "num_layers": self.num_layers, + "num_heads": self.num_heads, + "head_dim": self.head_dim, + "intermediate_dim": self.intermediate_dim, + "rope_theta": self.rope_theta, + "layer_norm_epsilon": self.layer_norm_epsilon, + "activation": self.activation, + "attention_dropout": self.attention_dropout, + } + ) + return config + + +class Mistral3PatchMerger(keras.layers.Layer): + """Spatially merge vision patches for Mistral3. + + Every `spatial_merge_size x spatial_merge_size` group of vision patches + is concatenated along the feature dimension and projected back down to + `hidden_dim`. + """ + + def __init__( + self, + hidden_dim=1024, + spatial_merge_size=MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE, + patch_size=14, + image_size=1540, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_dim = hidden_dim + self.image_size = image_size + self.spatial_merge_size = spatial_merge_size + self.patch_size = patch_size + + self.merging_layer = keras.layers.Dense( + hidden_dim, + use_bias=False, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="merging_layer", + ) + + def build(self, input_shape): + merge_input_dim = self.hidden_dim * self.spatial_merge_size**2 + self.merging_layer.build((None, merge_input_dim)) + self.built = True + + def call( + self, + image_features, + image_sizes, + max_patch_height, + max_patch_width, + ): + """Merge valid spatial neighborhoods using static-shape indexing. + + `image_features` is a padded canvas, sized + `num_images * max_patch_height * max_patch_width` (a function of + shapes only, `jax.jit`-safe): real per-image patches first, zero + rows after. `max_patch_height`/`max_patch_width` are passed in + rather than derived via `ops.max` so the output buffer size stays a + trace-time constant. + + Returns: + A tuple `(merged_padded, valid_count)`: `merged_padded` has + real merged windows as a prefix, zero-padded after; `valid_count` + is the scalar number of real windows in that prefix. + """ + image_features = ops.reshape( + image_features, + (-1, ops.shape(image_features)[-1]), + ) + image_sizes = ops.cast( + ops.convert_to_tensor(image_sizes), + "int32", + ) + # Clamped since Keras' torch auto-build traces this with all + # inputs filled with `1`s, which would otherwise divide by 0 + # below (`local_indices // widths_per_token`). + patch_heights = ops.maximum(image_sizes[:, 0] // self.patch_size, 1) + patch_widths = ops.maximum(image_sizes[:, 1] // self.patch_size, 1) + patch_counts = patch_heights * patch_widths + num_images = ops.shape(image_sizes)[0] + num_tokens = ops.shape(image_features)[0] + merge_size = self.spatial_merge_size + + image_offsets = ops.cumsum(patch_counts) - patch_counts + cumulative_ends = ops.cumsum(patch_counts) + total_valid_tokens = ops.sum(patch_counts) + + token_indices = ops.arange(num_tokens, dtype="int32") + # Static-shape stand-in for `ops.repeat(arange(num_images), + # patch_counts)`. + image_ids = ops.sum( + ops.cast( + ops.expand_dims(token_indices, 1) + >= ops.expand_dims(cumulative_ends, 0), + "int32", + ), + axis=1, + ) + image_ids = ops.clip(image_ids, 0, num_images - 1) + + widths_per_token = ops.take(patch_widths, image_ids, axis=0) + heights_per_token = ops.take(patch_heights, image_ids, axis=0) + offsets_per_token = ops.take(image_offsets, image_ids, axis=0) + local_indices = token_indices - offsets_per_token + local_rows = local_indices // widths_per_token + local_columns = local_indices % widths_per_token + + is_real = token_indices < total_valid_tokens + valid_windows = ops.logical_and( + is_real, + ops.logical_and( + local_rows % merge_size == 0, + local_columns % merge_size == 0, + ), + ) + valid_windows = ops.logical_and( + valid_windows, + local_rows + merge_size <= heights_per_token, + ) + valid_windows = ops.logical_and( + valid_windows, + local_columns + merge_size <= widths_per_token, + ) + + # The token itself is the window's top-left patch (floor-div/mod + # recombine exactly). Gather the `merge_size x merge_size` window in + # row-major order, matching the channel-major ordering of PyTorch's + # F.unfold (used by HF). + max_index = num_tokens - 1 + patch_indices = ops.stack( + [ + ops.clip( + token_indices + row * widths_per_token + col, + 0, + max_index, + ) + for row in range(merge_size) + for col in range(merge_size) + ], + axis=1, + ) + patches = ops.take(image_features, patch_indices, axis=0) + patches = ops.transpose(patches, (0, 2, 1)) + patches = ops.reshape( + patches, + (-1, self.hidden_dim * merge_size * merge_size), + ) + merged_all = self.merging_layer(patches) + + valid_int = ops.cast(valid_windows, "int32") + rank = ops.cumsum(valid_int) - valid_int + capacity = ( + num_images + * (max_patch_height // merge_size) + * (max_patch_width // merge_size) + ) + scatter_target = ops.where(valid_windows, rank, capacity) + buffer = ops.zeros((capacity + 1, self.hidden_dim), merged_all.dtype) + buffer = ops.scatter_update( + buffer, + ops.expand_dims(scatter_target, axis=1), + merged_all, + ) + merged_padded = buffer[:capacity] + valid_count = ops.sum(valid_int) + return merged_padded, valid_count + + def get_config(self): + config = super().get_config() + config.update( + { + "hidden_dim": self.hidden_dim, + "spatial_merge_size": self.spatial_merge_size, + "patch_size": self.patch_size, + "image_size": self.image_size, + } + ) + return config + + +class Mistral3MultiModalProjector(keras.layers.Layer): + """Multimodal projector used by Mistral3. + + Vision features are normalized, spatially merged, and projected into + the Mistral text-model hidden dimension. + """ + + def __init__( + self, + vision_hidden_dim=1024, + text_hidden_dim=5120, + spatial_merge_size=MISTRAL3_DEFAULT_SPATIAL_MERGE_SIZE, + patch_size=14, + layer_norm_epsilon=1e-5, + projector_hidden_act="gelu", + multimodal_projector_bias=False, + image_size=1540, + **kwargs, + ): + super().__init__(**kwargs) + + self.vision_hidden_dim = vision_hidden_dim + self.image_size = image_size + self.text_hidden_dim = text_hidden_dim + self.spatial_merge_size = spatial_merge_size + self.patch_size = patch_size + self.layer_norm_epsilon = layer_norm_epsilon + self.projector_hidden_act = projector_hidden_act + self.multimodal_projector_bias = multimodal_projector_bias + + self.norm = keras.layers.RMSNormalization( + epsilon=layer_norm_epsilon, + dtype=self.dtype_policy, + name="norm", + ) + + self.patch_merger = Mistral3PatchMerger( + hidden_dim=vision_hidden_dim, + spatial_merge_size=spatial_merge_size, + patch_size=patch_size, + image_size=image_size, + dtype=self.dtype_policy, + name="patch_merger", + ) + + self.linear_1 = keras.layers.Dense( + text_hidden_dim, + use_bias=multimodal_projector_bias, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="linear_1", + ) + + self.act = keras.activations.get( + projector_hidden_act, + ) + + self.linear_2 = keras.layers.Dense( + text_hidden_dim, + use_bias=multimodal_projector_bias, + kernel_initializer=_mistral_kernel_initializer(), + dtype=self.dtype_policy, + name="linear_2", + ) + + def build(self, input_shape): + self.norm.build((None, self.vision_hidden_dim)) + self.patch_merger.build((None, self.vision_hidden_dim)) + self.linear_1.build((None, self.vision_hidden_dim)) + self.linear_2.build((None, self.text_hidden_dim)) + self.built = True + + def call( + self, + image_features, + image_sizes, + max_patch_height, + max_patch_width, + ): + image_features = self.norm( + image_features, + ) + + # `valid_count` is unused: trimming to it needs a data-dependent + # slice, which breaks `jax.jit` tracing. The padded output is + # trimmed downstream instead, in `Mistral3ImageTextEmbeddingMerger`. + image_features, _ = self.patch_merger( + image_features, + image_sizes=image_sizes, + max_patch_height=max_patch_height, + max_patch_width=max_patch_width, + ) + + hidden_states = self.linear_1( + image_features, + ) + + hidden_states = self.act( + hidden_states, + ) + + hidden_states = self.linear_2( + hidden_states, + ) + + return hidden_states + + def compute_output_spec( + self, + image_features, + image_sizes, + max_patch_height, + max_patch_width, + ): + """Declare the output shape without tracing `call()`. + + Same reason as `Mistral3VisionEncoder.compute_output_spec`. + + Returns: + A `KerasTensor` with shape `(None, text_hidden_dim)`. + """ + return keras.KerasTensor( + shape=(None, self.text_hidden_dim), + dtype=self.compute_dtype, + ) + + def get_config(self): + config = super().get_config() + config.update( + { + "vision_hidden_dim": self.vision_hidden_dim, + "text_hidden_dim": self.text_hidden_dim, + "spatial_merge_size": self.spatial_merge_size, + "patch_size": self.patch_size, + "image_size": self.image_size, + "layer_norm_epsilon": self.layer_norm_epsilon, + "projector_hidden_act": self.projector_hidden_act, + "multimodal_projector_bias": (self.multimodal_projector_bias), + } + ) + return config diff --git a/keras_hub/src/models/mistral3/mistral3_vision_encoder_test.py b/keras_hub/src/models/mistral3/mistral3_vision_encoder_test.py new file mode 100644 index 0000000000..6045cf48ff --- /dev/null +++ b/keras_hub/src/models/mistral3/mistral3_vision_encoder_test.py @@ -0,0 +1,358 @@ +import numpy as np + +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3ImageFeatureExtractor, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3ImageTextEmbeddingMerger, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3MultiModalProjector, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3PatchMerger, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionAttention, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoder, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoderLayer, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionMLP, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionRotaryEmbedding, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + compute_image_placeholder_indices, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + compute_resize_size, +) +from keras_hub.src.tests.test_case import TestCase + + +class Mistral3VisionEncoderTest(TestCase): + def setUp(self): + self.image_size = 16 + self.patch_size = 4 + self.hidden_dim = 8 + self.num_heads = 2 + self.head_dim = self.hidden_dim // self.num_heads + self.intermediate_dim = 16 + self.num_layers = 2 + # Deliberately != num_heads to guard against a broadcast bug where + # the rotary embedding was expanded on the wrong axis. + self.batch_size = 3 + self.num_patches_per_side = self.image_size // self.patch_size + self.sequence_length = self.num_patches_per_side**2 + self.text_hidden_dim = 12 + self.spatial_merge_size = 2 + self.rope = Mistral3VisionRotaryEmbedding( + image_size=self.image_size, + patch_size=self.patch_size, + head_dim=self.head_dim, + ) + self.attention = Mistral3VisionAttention( + hidden_dim=self.hidden_dim, + num_heads=self.num_heads, + ) + self.encoder_layer = Mistral3VisionEncoderLayer( + hidden_dim=self.hidden_dim, + intermediate_dim=self.intermediate_dim, + num_heads=self.num_heads, + ) + self.encoder_init_kwargs = { + "image_size": self.image_size, + "patch_size": self.patch_size, + "hidden_dim": self.hidden_dim, + "num_layers": self.num_layers, + "num_heads": self.num_heads, + "head_dim": self.head_dim, + "intermediate_dim": self.intermediate_dim, + } + self.vision_encoder = Mistral3VisionEncoder(**self.encoder_init_kwargs) + self.patch_merger = Mistral3PatchMerger( + hidden_dim=self.hidden_dim, + spatial_merge_size=self.spatial_merge_size, + patch_size=self.patch_size, + image_size=self.image_size, + ) + self.projector = Mistral3MultiModalProjector( + vision_hidden_dim=self.hidden_dim, + text_hidden_dim=self.text_hidden_dim, + spatial_merge_size=self.spatial_merge_size, + patch_size=self.patch_size, + image_size=self.image_size, + ) + self.embedding_merger = Mistral3ImageTextEmbeddingMerger() + + # === Rotary embedding === + + def test_rotary_embedding_output_shape(self): + position_ids = np.zeros( + (self.batch_size, self.sequence_length), dtype="int32" + ) + cos, sin = self.rope(position_ids) + expected_shape = (self.batch_size, self.sequence_length, self.head_dim) + self.assertEqual(cos.shape, expected_shape) + self.assertEqual(sin.shape, expected_shape) + + def test_rotary_embedding_serialization(self): + self.run_serialization_test(self.rope) + + # === Attention === + + def test_attention_output_shape(self): + inputs = np.random.rand( + self.batch_size, self.sequence_length, self.hidden_dim + ) + position_ids = np.zeros( + (self.batch_size, self.sequence_length), dtype="int32" + ) + cos, sin = self.rope(position_ids) + output = self.attention(inputs, position_embeddings=(cos, sin)) + self.assertEqual( + output.shape, + (self.batch_size, self.sequence_length, self.hidden_dim), + ) + + def test_attention_serialization(self): + self.run_serialization_test(self.attention) + + # === MLP === + + def test_mlp_output_shape(self): + mlp = Mistral3VisionMLP( + hidden_dim=self.hidden_dim, + intermediate_dim=self.intermediate_dim, + ) + inputs = np.random.rand( + self.batch_size, self.sequence_length, self.hidden_dim + ) + output = mlp(inputs) + self.assertEqual( + output.shape, + (self.batch_size, self.sequence_length, self.hidden_dim), + ) + + # === Encoder layer === + + def test_encoder_layer_output_shape(self): + inputs = np.random.rand( + self.batch_size, self.sequence_length, self.hidden_dim + ) + position_ids = np.zeros( + (self.batch_size, self.sequence_length), dtype="int32" + ) + cos, sin = self.rope(position_ids) + output = self.encoder_layer(inputs, position_embeddings=(cos, sin)) + self.assertEqual( + output.shape, + (self.batch_size, self.sequence_length, self.hidden_dim), + ) + + def test_encoder_layer_serialization(self): + self.run_serialization_test(self.encoder_layer) + + # === Vision encoder === + + def test_encoder_output_shape(self): + # [num_images, channels, height, width] input, HF-Mistral style. + # All images are concatenated into a single sequence (block-diagonal + # attention keeps them from attending to each other), so the output + # always has batch dim 1. + pixel_values = np.random.rand( + self.batch_size, 3, self.image_size, self.image_size + ) + output = self.vision_encoder(pixel_values) + self.assertEqual( + output.shape, + (1, self.batch_size * self.sequence_length, self.hidden_dim), + ) + + def test_encoder_output_shape_with_variable_image_sizes(self): + # Two images sharing a common padded canvas but with different real + # sizes (one cropped in width). + pixel_values = np.random.rand(2, 3, self.image_size, self.image_size) + image_sizes = np.array([[16, 16], [16, 8]], dtype="int32") + output = self.vision_encoder(pixel_values, image_sizes=image_sizes) + # Padded to the full canvas capacity (real patches first, zero-padded + # tail): 2 images * 4x4 patch grid = 32 total tokens. + self.assertEqual(output.shape, (1, 32, self.hidden_dim)) + + def test_encoder_serialization(self): + self.run_serialization_test(self.vision_encoder) + + def test_encoder_saved_model(self): + pixel_values = np.random.rand( + self.batch_size, 3, self.image_size, self.image_size + ).astype("float32") + self.run_model_saving_test( + cls=Mistral3VisionEncoder, + init_kwargs=self.encoder_init_kwargs, + input_data=pixel_values, + ) + + # === Patch merger === + + def test_patch_merger_output_shape(self): + # Image 0: 4x4 patch grid (16 tokens, 4 merge windows). + # Image 1: 4x8 patch grid (32 tokens, 8 merge windows). + image_sizes = np.array([[16, 16], [16, 32]], dtype="int32") + max_patch_height, max_patch_width = 4, 8 + image_features = np.random.rand(16 + 32, self.hidden_dim) + merged, valid_count = self.patch_merger( + image_features, + image_sizes=image_sizes, + max_patch_height=max_patch_height, + max_patch_width=max_patch_width, + ) + capacity = 2 * (max_patch_height // 2) * (max_patch_width // 2) + self.assertEqual(merged.shape, (capacity, self.hidden_dim)) + self.assertEqual(int(valid_count), 12) + + def test_patch_merger_serialization(self): + self.run_serialization_test(self.patch_merger) + + # === Multimodal projector === + + def test_multimodal_projector_output_shape(self): + # Two images: 4x4 and 4x8 patch grids (16 + 32 = 48 tokens). + image_sizes = np.array([[16, 16], [16, 32]], dtype="int32") + image_features = np.random.rand(48, self.hidden_dim) + output = self.projector( + image_features, + image_sizes=image_sizes, + max_patch_height=4, + max_patch_width=8, + ) + # Padded to the full merge-window capacity (not sliced to the valid + # count): 2 images * (4//2) * (8//2) = 16 windows. + self.assertEqual(output.shape, (16, self.text_hidden_dim)) + + def test_multimodal_projector_serialization(self): + self.run_serialization_test(self.projector) + + # === Image feature extractor (encoder + projector, end to end) === + + def test_image_feature_extractor_output_shape(self): + # Two images padded to a common (16, 16) canvas; image 1 is cropped + # to half width. + pixel_values = np.random.rand(2, 3, self.image_size, self.image_size) + image_sizes = np.array([[16, 16], [16, 8]], dtype="int32") + extractor = Mistral3ImageFeatureExtractor( + self.vision_encoder, self.projector + ) + output = extractor(pixel_values, image_sizes) + # Padded to the full merge-window capacity (real windows first, + # zero-padded tail), not sliced to the valid count: 2 images * + # (4//2) * (4//2) = 8 windows. + self.assertEqual(output.shape, (8, self.text_hidden_dim)) + + def test_image_feature_extractor_rejects_unsupported_layer(self): + with self.assertRaises(NotImplementedError): + Mistral3ImageFeatureExtractor( + self.vision_encoder, + self.projector, + vision_feature_layer=-2, + ) + + # === Image/text embedding merger === + + def test_image_text_embedding_merger_scatters_features(self): + batch_size, seq_length, hidden_dim = 1, 5, 3 + token_embeddings = np.zeros((batch_size, seq_length, hidden_dim)) + image_features = np.array([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]) + # `(batch, max_placeholders)`: this example's own local positions. + placeholder_indices = np.array([[1, 3]], dtype="int32") + output = self.embedding_merger( + token_embeddings, image_features, placeholder_indices + ) + expected = np.zeros((batch_size, seq_length, hidden_dim)) + expected[0, 1] = 1.0 + expected[0, 3] = 2.0 + self.assertAllClose(output, expected) + + def test_image_text_embedding_merger_batched_with_padding(self): + # Row 0 has two images (local positions 1, 3); row 1 has only one + # (local position 2, with a `-1`-padded second column) -- the case + # `placeholder_indices`' padding exists to support. + batch_size, seq_length, hidden_dim = 2, 5, 3 + token_embeddings = np.zeros((batch_size, seq_length, hidden_dim)) + image_features = np.array( + [[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [3.0, 3.0, 3.0]] + ) + placeholder_indices = np.array([[1, 3], [2, -1]], dtype="int32") + output = self.embedding_merger( + token_embeddings, image_features, placeholder_indices + ) + expected = np.zeros((batch_size, seq_length, hidden_dim)) + expected[0, 1] = 1.0 + expected[0, 3] = 2.0 + expected[1, 2] = 3.0 + self.assertAllClose(output, expected) + + def test_image_text_embedding_merger_serialization(self): + self.run_serialization_test(self.embedding_merger) + + # === compute_image_placeholder_indices === + + def test_compute_image_placeholder_indices(self): + token_ids = np.array([[1, 10, 3, 10], [10, 2, 3, 4]]) + indices = compute_image_placeholder_indices( + token_ids, image_token_index=10 + ) + # `(batch, max_placeholders)`: each row's own local positions, + # padded with `-1` up to the batch's max count. + self.assertAllEqual(indices, np.array([[1, 3], [0, -1]])) + + def test_compute_image_placeholder_indices_none_present(self): + token_ids = np.array([[1, 2, 3]]) + indices = compute_image_placeholder_indices( + token_ids, image_token_index=10 + ) + self.assertAllEqual(indices, np.array([[-1]])) + + # === compute_resize_size === + + def test_compute_resize_size_exact_multiple(self): + # Already a multiple of `patch_size`, well under `longest_edge`: + # no scaling, no rounding. + size = compute_resize_size( + height=16, width=16, longest_edge=32, patch_size=4 + ) + self.assertEqual(size, (16, 16)) + + def test_compute_resize_size_rounds_up(self): + # Under `longest_edge`, but not a `patch_size` multiple: each dim + # rounds up independently, (17 - 1) // 4 + 1 = 5 -> 20. + size = compute_resize_size( + height=17, width=17, longest_edge=32, patch_size=4 + ) + self.assertEqual(size, (20, 20)) + + def test_compute_resize_size_clamps_wide_image(self): + # Wide image: width is the longest edge, so it is scaled down to + # `longest_edge` and height is scaled by the same ratio, preserving + # aspect ratio, before rounding up to a `patch_size` multiple. + # ratio = 40 / 16 = 2.5 -> height = floor(20 / 2.5) = 8, + # width = floor(40 / 2.5) = 16 (both already patch multiples). + size = compute_resize_size( + height=20, width=40, longest_edge=16, patch_size=4 + ) + self.assertEqual(size, (8, 16)) + + def test_compute_resize_size_clamps_tall_image(self): + # Tall image: height is the longest edge. Same scale factor is + # applied to both dimensions. + # ratio = 40 / 16 = 2.5 -> height = floor(40 / 2.5) = 16, + # width = floor(20 / 2.5) = 8. + size = compute_resize_size( + height=40, width=20, longest_edge=16, patch_size=4 + ) + self.assertEqual(size, (16, 8)) diff --git a/keras_hub/src/utils/transformers/convert_mistral3.py b/keras_hub/src/utils/transformers/convert_mistral3.py new file mode 100644 index 0000000000..a68c3dba01 --- /dev/null +++ b/keras_hub/src/utils/transformers/convert_mistral3.py @@ -0,0 +1,413 @@ +import numpy as np + +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3MultiModalProjector, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoder, +) +from keras_hub.src.utils.preset_utils import check_file_exists +from keras_hub.src.utils.preset_utils import get_file +from keras_hub.src.utils.preset_utils import load_json +from keras_hub.src.utils.transformers.convert_mistral import ( + _convert_tekken_tokenizer as _convert_mistral_tekken_tokenizer, +) +from keras_hub.src.utils.transformers.convert_mistral import ( + convert_backbone_config as convert_text_backbone_config, +) + +backbone_cls = Mistral3Backbone + + +_PIXTRAL_DEFAULT_RESCALE_FACTOR = 1 / 255 + + +def _get_rope_theta(config, default=10000.0): + rope_theta = config.get("rope_parameters", {}).get("rope_theta") + if rope_theta is None: + rope_theta = config.get("rope_theta", default) + return rope_theta + + +def _convert_tekken_tokenizer(path): + """Like `convert_mistral._convert_tekken_tokenizer`, but also returns + `control_tokens`: the reserved-block special-token pieces, which + `Mistral3Tokenizer` registers as unsplittable alongside its vision + tokens. + """ + try: + from mistral_common.tokens.tokenizers.tekken import Tekkenizer + except ImportError: + raise ImportError( + "Converting a Tekken (`tekken.json`) tokenizer requires the " + "`mistral_common` package. Please install it with " + "`pip install mistral-common`." + ) + + vocabulary, merges, split_pattern = _convert_mistral_tekken_tokenizer(path) + + tokenizer = Tekkenizer.from_file(path) + control_tokens = [ + tokenizer.id_to_piece(rank) + for rank in range(tokenizer.num_special_tokens) + ] + + return vocabulary, merges, split_pattern, control_tokens + + +def _load_pixtral_defaults_from_mistral_common(): + # Some checkpoints (e.g. Mistral Small 3.2) ship no + # `preprocessor_config.json`; fall back to `mistral_common`'s fixed + # constants instead of duplicating the numbers here. + try: + from mistral_common.tokens.tokenizers.image import DATASET_MEAN + from mistral_common.tokens.tokenizers.image import DATASET_STD + except ImportError: + raise ImportError( + "Converting a Mistral3 checkpoint with no " + "`preprocessor_config.json` requires the `mistral_common` " + "package. Please install it with `pip install mistral-common`." + ) + return list(DATASET_MEAN), list(DATASET_STD) + + +def load_image_converter_config(preset, transformers_config): + vision_config = transformers_config["vision_config"] + if check_file_exists(preset, "preprocessor_config.json"): + preprocessor_config = load_json(preset, "preprocessor_config.json") + mean = preprocessor_config["image_mean"] + std = preprocessor_config["image_std"] + rescale_factor = preprocessor_config["rescale_factor"] + patch_size = preprocessor_config["patch_size"] + if isinstance(patch_size, dict): + patch_size = patch_size.get("height") or patch_size.get("width") + size = preprocessor_config["size"] + longest_edge = ( + size.get("longest_edge") if isinstance(size, dict) else None + ) + else: + mean, std = _load_pixtral_defaults_from_mistral_common() + rescale_factor = _PIXTRAL_DEFAULT_RESCALE_FACTOR + patch_size = vision_config["patch_size"] + longest_edge = vision_config["image_size"] + + config = {} + if mean is not None and std is not None: + config["offset"] = [-m / s for m, s in zip(mean, std)] + config["scale"] = [rescale_factor / s for s in std] + if patch_size is not None: + config["patch_size"] = patch_size + if longest_edge is not None: + config["longest_edge"] = longest_edge + config["spatial_merge_size"] = transformers_config["spatial_merge_size"] + return config + + +def convert_backbone_config(transformers_config): + text_config = transformers_config["text_config"] + vision_config = transformers_config["vision_config"] + backbone_config = convert_text_backbone_config(text_config) + + vision_hidden_dim = vision_config["hidden_size"] + vision_num_heads = vision_config["num_attention_heads"] + vision_head_dim = vision_config.get("head_dim") or ( + vision_hidden_dim // vision_num_heads + ) + vision_image_size = vision_config["image_size"] + vision_patch_size = vision_config["patch_size"] + vision_encoder = Mistral3VisionEncoder( + image_size=vision_image_size, + patch_size=vision_patch_size, + num_channels=vision_config["num_channels"], + hidden_dim=vision_hidden_dim, + num_layers=vision_config["num_hidden_layers"], + num_heads=vision_num_heads, + head_dim=vision_head_dim, + intermediate_dim=vision_config["intermediate_size"], + rope_theta=_get_rope_theta(vision_config), + layer_norm_epsilon=vision_config.get("rms_norm_eps", 1e-5), + activation=vision_config["hidden_act"], + attention_dropout=vision_config["attention_dropout"], + ) + + multimodal_projector = Mistral3MultiModalProjector( + vision_hidden_dim=vision_hidden_dim, + text_hidden_dim=text_config["hidden_size"], + spatial_merge_size=transformers_config["spatial_merge_size"], + patch_size=vision_patch_size, + layer_norm_epsilon=text_config["rms_norm_eps"], + projector_hidden_act=transformers_config["projector_hidden_act"], + multimodal_projector_bias=transformers_config[ + "multimodal_projector_bias" + ], + image_size=vision_image_size, + ) + + image_token_index = transformers_config["image_token_index"] + + backbone_config.update( + { + "vision_encoder": vision_encoder, + "multimodal_projector": multimodal_projector, + "image_token_index": image_token_index, + } + ) + return backbone_config + + +def _port_text_weights(backbone, loader, tie_word_embeddings): + # Embeddings + loader.port_weight( + keras_variable=backbone.token_embedding.embeddings, + hf_weight_key="language_model.model.embed_tokens.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + # When embeddings are tied, `lm_head.weight` is not saved as a separate + # tensor in the checkpoint; reuse the embedding weights instead. + lm_head_key = ( + "language_model.model.embed_tokens.weight" + if tie_word_embeddings + else "lm_head.weight" + ) + loader.port_weight( + keras_variable=backbone.token_embedding.reverse_embeddings, + hf_weight_key=lm_head_key, + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + + # Attention blocks + for index in range(backbone.num_layers): + decoder_layer = backbone.transformer_layers[index] + + # Norm layers + loader.port_weight( + keras_variable=decoder_layer._self_attention_layernorm.scale, + hf_weight_key=f"language_model.model.layers.{index}.input_layernorm.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + loader.port_weight( + keras_variable=decoder_layer._feedforward_layernorm.scale, + hf_weight_key=( + f"language_model.model.layers.{index}.post_attention_layernorm.weight" + ), + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + + # Attention layers + loader.port_weight( + keras_variable=decoder_layer._self_attention_layer._query_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.self_attn.q_proj.weight", + hook_fn=lambda hf_tensor, keras_shape: np.reshape( + np.transpose(hf_tensor.astype(np.float32)), keras_shape + ), + ) + loader.port_weight( + keras_variable=decoder_layer._self_attention_layer._key_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.self_attn.k_proj.weight", + hook_fn=lambda hf_tensor, keras_shape: np.reshape( + np.transpose(hf_tensor.astype(np.float32)), keras_shape + ), + ) + loader.port_weight( + keras_variable=decoder_layer._self_attention_layer._value_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.self_attn.v_proj.weight", + hook_fn=lambda hf_tensor, keras_shape: np.reshape( + np.transpose(hf_tensor.astype(np.float32)), keras_shape + ), + ) + loader.port_weight( + keras_variable=decoder_layer._self_attention_layer._output_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.self_attn.o_proj.weight", + hook_fn=lambda hf_tensor, keras_shape: np.reshape( + np.transpose(hf_tensor.astype(np.float32)), keras_shape + ), + ) + + # MLP layers + loader.port_weight( + keras_variable=decoder_layer._feedforward_gate_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.mlp.gate_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=decoder_layer._feedforward_intermediate_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.mlp.up_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=decoder_layer._feedforward_output_dense.kernel, + hf_weight_key=f"language_model.model.layers.{index}.mlp.down_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + + # Normalization + loader.port_weight( + keras_variable=backbone.layer_norm.scale, + hf_weight_key="language_model.model.norm.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + + +def _port_vision_weights(backbone, loader): + vision_encoder = backbone.vision_encoder + projector = backbone.multimodal_projector + + loader.port_weight( + keras_variable=vision_encoder.patch_conv.kernel, + hf_weight_key="vision_tower.patch_conv.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(2, 3, 1, 0) + ), + ) + loader.port_weight( + keras_variable=vision_encoder.ln_pre.scale, + hf_weight_key="vision_tower.ln_pre.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + + for index in range(vision_encoder.num_layers): + layer = vision_encoder.transformer_layers[index] + layer_prefix = f"vision_tower.transformer.layers.{index}" + + loader.port_weight( + keras_variable=layer.attention_norm.scale, + hf_weight_key=f"{layer_prefix}.attention_norm.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + loader.port_weight( + keras_variable=layer.ffn_norm.scale, + hf_weight_key=f"{layer_prefix}.ffn_norm.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + + loader.port_weight( + keras_variable=layer.attention.q_proj.kernel, + hf_weight_key=f"{layer_prefix}.attention.q_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=layer.attention.k_proj.kernel, + hf_weight_key=f"{layer_prefix}.attention.k_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=layer.attention.v_proj.kernel, + hf_weight_key=f"{layer_prefix}.attention.v_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=layer.attention.o_proj.kernel, + hf_weight_key=f"{layer_prefix}.attention.o_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + + loader.port_weight( + keras_variable=layer.feed_forward.gate_proj.kernel, + hf_weight_key=f"{layer_prefix}.feed_forward.gate_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=layer.feed_forward.up_proj.kernel, + hf_weight_key=f"{layer_prefix}.feed_forward.up_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=layer.feed_forward.down_proj.kernel, + hf_weight_key=f"{layer_prefix}.feed_forward.down_proj.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + + # Multimodal projector + loader.port_weight( + keras_variable=projector.norm.scale, + hf_weight_key="multi_modal_projector.norm.weight", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + loader.port_weight( + keras_variable=projector.patch_merger.merging_layer.kernel, + hf_weight_key="multi_modal_projector.patch_merger.merging_layer.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=projector.linear_1.kernel, + hf_weight_key="multi_modal_projector.linear_1.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + loader.port_weight( + keras_variable=projector.linear_2.kernel, + hf_weight_key="multi_modal_projector.linear_2.weight", + hook_fn=lambda hf_tensor, _: np.transpose( + hf_tensor.astype(np.float32), axes=(1, 0) + ), + ) + if projector.linear_1.use_bias: + loader.port_weight( + keras_variable=projector.linear_1.bias, + hf_weight_key="multi_modal_projector.linear_1.bias", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + loader.port_weight( + keras_variable=projector.linear_2.bias, + hf_weight_key="multi_modal_projector.linear_2.bias", + hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float32), + ) + + +def convert_weights(backbone, loader, transformers_config): + tie_word_embeddings = transformers_config.get("tie_word_embeddings", False) + _port_text_weights( + backbone, + loader, + tie_word_embeddings=tie_word_embeddings, + ) + _port_vision_weights(backbone, loader) + + +def convert_tokenizer(cls, preset, **kwargs): + # Mistral3 checkpoints always ship a Tekken (byte-level BPE) `tekken.json` + # tokenizer; there is currently no SentencePiece Mistral3 preset to + # support. + if not check_file_exists(preset, "tekken.json"): + raise ValueError( + f"Could not find a `tekken.json` file for preset '{preset}'. " + "Mistral3 checkpoint conversion currently only supports Tekken " + "(byte-level BPE) tokenizers." + ) + tekken_path = get_file(preset, "tekken.json") + vocabulary, merges, split_pattern, control_tokens = ( + _convert_tekken_tokenizer(tekken_path) + ) + return cls( + vocabulary=vocabulary, + merges=merges, + split_pattern=split_pattern, + control_tokens=control_tokens, + **kwargs, + ) diff --git a/keras_hub/src/utils/transformers/convert_mistral3_test.py b/keras_hub/src/utils/transformers/convert_mistral3_test.py new file mode 100644 index 0000000000..4585c08119 --- /dev/null +++ b/keras_hub/src/utils/transformers/convert_mistral3_test.py @@ -0,0 +1,175 @@ +import json +import os +import tempfile + +import numpy as np +import pytest +from keras import ops + +from keras_hub.src.models.mistral3.mistral3_backbone import Mistral3Backbone +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3MultiModalProjector, +) +from keras_hub.src.models.mistral3.mistral3_vision_encoder import ( + Mistral3VisionEncoder, +) +from keras_hub.src.tests.test_case import TestCase +from keras_hub.src.utils.transformers import convert_mistral3 + + +class TestTask(TestCase): + @pytest.mark.large + def test_convert_multimodal_preset_matches_hf(self): + # Build a tiny Mistral3 (Pixtral vision tower + Mistral text model) + # checkpoint and check that the converted `Mistral3Backbone` matches + # HF's reference forward pass end to end, including the vision + # tower, multimodal projector, and image/text embedding merge. + torch = pytest.importorskip("torch") + transformers = pytest.importorskip("transformers") + + text_config = transformers.MistralConfig( + vocab_size=100, + hidden_size=16, + intermediate_size=24, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + sliding_window=None, + rope_theta=1_000_000.0, + rms_norm_eps=1e-5, + ) + vision_config = transformers.PixtralVisionConfig( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_channels=3, + image_size=16, + patch_size=4, + rope_parameters={"rope_theta": 10000.0}, + ) + config = transformers.Mistral3Config( + text_config=text_config, + vision_config=vision_config, + image_token_index=10, + spatial_merge_size=2, + ) + torch.manual_seed(0) + hf_model = transformers.Mistral3ForConditionalGeneration(config).eval() + + with tempfile.TemporaryDirectory() as preset_dir: + hf_model.save_pretrained(preset_dir) + keras_backbone = Mistral3Backbone.from_preset(preset_dir) + + self.assertIsNotNone(keras_backbone.vision_encoder) + + # A single 16x16 image with a 4x4 patch size produces a 4x4 patch + # grid (16 patches); a spatial merge size of 2 merges these into a + # 2x2 grid of 4 tokens, so 4 placeholder tokens at id 10 are needed. + input_ids = np.array([[1, 10, 10, 10, 10, 3, 4]], dtype="int32") + padding_mask = np.ones_like(input_ids) + pixel_values = np.random.rand(1, 3, 16, 16).astype("float32") + image_sizes = np.array([[16, 16]], dtype="int32") + placeholder_indices = np.array([[1, 2, 3, 4]], dtype="int32") + + keras_out = ops.convert_to_numpy( + keras_backbone( + { + "token_ids": input_ids, + "padding_mask": padding_mask, + "pixel_values": pixel_values, + "image_sizes": image_sizes, + "placeholder_indices": placeholder_indices, + } + ) + ) + with torch.no_grad(): + hf_out = ( + hf_model.model( + input_ids=torch.tensor(input_ids), + attention_mask=torch.tensor(padding_mask), + pixel_values=torch.tensor(pixel_values), + image_sizes=torch.tensor(image_sizes), + ) + .last_hidden_state.detach() + .cpu() + .numpy() + ) + self.assertEqual(keras_out.shape, hf_out.shape) + # fp16 weight storage dominates the parity bound, as in the + # text-only converter test. + self.assertAllClose(keras_out, hf_out, atol=1e-2) + + def test_convert_backbone_config_detects_mistral3(self): + transformers_config = { + "text_config": { + "vocab_size": 100, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "hidden_size": 32, + "intermediate_size": 48, + "num_key_value_heads": 2, + "rope_theta": 1_000_000.0, + "rms_norm_eps": 1e-5, + "sliding_window": None, + }, + "vision_config": { + "hidden_size": 16, + "intermediate_size": 24, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_channels": 3, + "image_size": 32, + "patch_size": 8, + "hidden_act": "gelu", + "attention_dropout": 0.0, + "rope_parameters": {"rope_theta": 10000.0}, + }, + "image_token_index": 10, + "spatial_merge_size": 2, + "projector_hidden_act": "gelu", + "multimodal_projector_bias": False, + } + keras_config = convert_mistral3.convert_backbone_config( + transformers_config + ) + self.assertIsInstance( + keras_config["vision_encoder"], Mistral3VisionEncoder + ) + self.assertIsInstance( + keras_config["multimodal_projector"], + Mistral3MultiModalProjector, + ) + self.assertEqual(keras_config["image_token_index"], 10) + self.assertEqual(keras_config["rope_max_wavelength"], 1_000_000.0) + self.assertEqual( + keras_config["vision_encoder"].get_config()["rope_theta"], + 10000.0, + ) + + def test_load_image_converter_config_without_preprocessor_config(self): + # Some checkpoints (e.g. Mistral Small 3.2) ship no + # `preprocessor_config.json`; the image normalization mean/std + # should come from `mistral_common`, not a local hardcoded copy. + pytest.importorskip("mistral_common") + from mistral_common.tokens.tokenizers.image import DATASET_MEAN + from mistral_common.tokens.tokenizers.image import DATASET_STD + + transformers_config = { + "vision_config": {"patch_size": 14, "image_size": 1540}, + "spatial_merge_size": 2, + } + with tempfile.TemporaryDirectory() as dir_path: + with open(os.path.join(dir_path, "config.json"), "w") as f: + json.dump(transformers_config, f) + config = convert_mistral3.load_image_converter_config( + dir_path, transformers_config + ) + expected_offset = [-m / s for m, s in zip(DATASET_MEAN, DATASET_STD)] + expected_scale = [(1 / 255) / s for s in DATASET_STD] + self.assertAllClose(config["offset"], expected_offset) + self.assertAllClose(config["scale"], expected_scale) + self.assertEqual(config["patch_size"], 14) + self.assertEqual(config["longest_edge"], 1540) + self.assertEqual(config["spatial_merge_size"], 2) diff --git a/keras_hub/src/utils/transformers/preset_loader.py b/keras_hub/src/utils/transformers/preset_loader.py index eaf2b92aa2..300de3c3e8 100644 --- a/keras_hub/src/utils/transformers/preset_loader.py +++ b/keras_hub/src/utils/transformers/preset_loader.py @@ -24,6 +24,7 @@ from keras_hub.src.utils.transformers import convert_llama3 from keras_hub.src.utils.transformers import convert_metaclip_2 from keras_hub.src.utils.transformers import convert_mistral +from keras_hub.src.utils.transformers import convert_mistral3 from keras_hub.src.utils.transformers import convert_mixtral from keras_hub.src.utils.transformers import convert_pali_gemma from keras_hub.src.utils.transformers import convert_qwen @@ -85,6 +86,8 @@ def __init__(self, preset, config): self.converter = convert_metaclip_2 elif model_type == "mistral": self.converter = convert_mistral + elif model_type == "mistral3": + self.converter = convert_mistral3 elif model_type == "paligemma": self.converter = convert_pali_gemma elif model_type == "vit": diff --git a/tools/checkpoint_conversion/convert_mistral3_checkpoints.py b/tools/checkpoint_conversion/convert_mistral3_checkpoints.py new file mode 100644 index 0000000000..0afa561ad9 --- /dev/null +++ b/tools/checkpoint_conversion/convert_mistral3_checkpoints.py @@ -0,0 +1,442 @@ +"""Convert multimodal Mistral3 HuggingFace checkpoints to KerasHub presets. + +Usage: + python tools/checkpoint_conversion/convert_mistral3_checkpoints.py \ + --preset mistral_small_3.1_24b_instruct_2503_en +""" + +import gc +import os + +os.environ["KERAS_BACKEND"] = "torch" +os.environ["CUDA_VISIBLE_DEVICES"] = "-1" + +import numpy as np +import requests +import torch + +device = torch.device("cpu") +torch.set_default_device(device) + +from absl import app # noqa: E402 +from absl import flags # noqa: E402 +from huggingface_hub import hf_hub_download # noqa: E402 +from keras import ops # noqa: E402 +from PIL import Image # noqa: E402 +from transformers import AutoProcessor # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 +from transformers import Mistral3ForConditionalGeneration # noqa: E402 + +import keras_hub # noqa: E402 + +_IMAGE_URL = "http://images.cocodataset.org/val2017/000000039769.jpg" + +PRESET_MAP = { + "mistral_small_3.1_24b_base_2503_en": ( + "mistralai/Mistral-Small-3.1-24B-Base-2503" + ), + "mistral_small_3.1_24b_instruct_2503_en": ( + "mistralai/Mistral-Small-3.1-24B-Instruct-2503" + ), + "mistral_small_3.2_24b_instruct_2506_en": ( + "mistralai/Mistral-Small-3.2-24B-Instruct-2506" + ), +} + +MAX_NEW_TOKENS = 64 + +TEXT_PROMPT = "What is Keras?" + +IMAGE_PROMPT = "What is in this image?" + +FLAGS = flags.FLAGS +flags.DEFINE_string( + "preset", None, f"Must be one of {','.join(PRESET_MAP.keys())}" +) +flags.DEFINE_boolean( + "skip_generate", + False, + "Skip the generation comparison step. Useful for large models where " + "generation is slow or unnecessary (numerics verification is sufficient).", +) + + +def load_reference_image(): + return Image.open(requests.get(_IMAGE_URL, stream=True).raw).convert("RGB") + + +def build_text_inputs(hf_preset, text): + # Some checkpoints (e.g. Mistral Small 3.2) ship only `tekken.json`, + # with no `tokenizer_config.json` for `AutoTokenizer` to resolve a + # class from; fall back to `mistral_common` reading it directly. + try: + hf_tokenizer = AutoTokenizer.from_pretrained(hf_preset) + return hf_tokenizer([text], return_tensors="pt"), hf_tokenizer + except OSError: + pass + + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer + + tekken_path = hf_hub_download(hf_preset, "tekken.json") + raw_tokenizer = MistralTokenizer.from_file( + tekken_path + ).instruct_tokenizer.tokenizer + token_ids = raw_tokenizer.encode(text, bos=True, eos=False) + inputs = { + "input_ids": torch.tensor([token_ids]), + "attention_mask": torch.ones(1, len(token_ids), dtype=torch.long), + } + return inputs, raw_tokenizer + + +def run_hf_text_forward(hf_model, hf_preset, skip_generate=False): + hf_inputs, hf_tokenizer = build_text_inputs(hf_preset, TEXT_PROMPT) + with torch.no_grad(): + hf_outputs = hf_model(**hf_inputs) + results = { + "prompt": TEXT_PROMPT, + "token_ids": hf_inputs["input_ids"].detach().cpu().numpy(), + "logits": hf_outputs.logits.detach().cpu().numpy(), + } + if not skip_generate: + with torch.no_grad(): + generated_ids = hf_model.generate( + **hf_inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False + ) + prompt_length = hf_inputs["input_ids"].shape[1] + generated_token_ids = generated_ids[0, prompt_length:].tolist() + try: + results["generated_text"] = hf_tokenizer.decode( + generated_token_ids, skip_special_tokens=True + ) + except TypeError: + # `mistral_common`'s raw tokenizer (used in the fallback path + # above) doesn't accept `skip_special_tokens`. + results["generated_text"] = hf_tokenizer.decode(generated_token_ids) + return results + + +def build_multimodal_inputs(hf_preset, hf_config, image): + # Falls back to `mistral_common`'s `encode_chat_completion` when a + # checkpoint has no chat template (e.g. a base model) or no + # `preprocessor_config.json` for `AutoProcessor` to resolve a class + # from (e.g. Mistral Small 3.2). + image_token_index = hf_config.image_token_index + + try: + hf_processor = AutoProcessor.from_pretrained(hf_preset) + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": IMAGE_PROMPT}, + ], + } + ] + prompt = hf_processor.apply_chat_template( + messages, add_generation_prompt=True + ) + inputs = hf_processor(text=prompt, images=image, return_tensors="np") + token_ids = inputs["input_ids"].astype("int32") + padding_mask = inputs["attention_mask"].astype("int32") + pixel_values = inputs["pixel_values"].astype("float32") + image_sizes = inputs["image_sizes"].astype("int32") + tokenizer = hf_processor.tokenizer + except (OSError, ValueError): + from mistral_common.protocol.instruct.chunk import ImageChunk + from mistral_common.protocol.instruct.chunk import TextChunk + from mistral_common.protocol.instruct.messages import UserMessage + from mistral_common.protocol.instruct.request import ( + ChatCompletionRequest, + ) + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer + + tekken_path = hf_hub_download(hf_preset, "tekken.json") + request = ChatCompletionRequest( + messages=[ + UserMessage( + content=[ + ImageChunk(image=image), + TextChunk(text=IMAGE_PROMPT), + ] + ) + ] + ) + mistral_tokenizer = MistralTokenizer.from_file(tekken_path) + tokenized = mistral_tokenizer.encode_chat_completion(request) + token_ids = np.array([tokenized.tokens], dtype="int32") + padding_mask = np.ones_like(token_ids) + pixel_values = tokenized.images[0][None, ...].astype("float32") + image_sizes = np.array([pixel_values.shape[-2:]], dtype="int32") + prompt = f"[INST][IMG]{IMAGE_PROMPT}[/INST]" + tokenizer = mistral_tokenizer.instruct_tokenizer.tokenizer + + flat_ids = token_ids.reshape(-1) + placeholder_indices = np.where(flat_ids == image_token_index)[0].astype( + "int32" + )[None, :] + + return { + "prompt": prompt, + "image": np.asarray(image), + "token_ids": token_ids, + "padding_mask": padding_mask, + "pixel_values": pixel_values, + "image_sizes": image_sizes, + "placeholder_indices": placeholder_indices, + "tokenizer": tokenizer, + } + + +def precompute_hf_outputs(hf_preset, hf_config, skip_generate=False): + hf_model = Mistral3ForConditionalGeneration.from_pretrained( + hf_preset, device_map="cpu", torch_dtype=torch.float32 + ) + hf_model.eval() + + text_results = run_hf_text_forward( + hf_model, hf_preset, skip_generate=skip_generate + ) + + image = load_reference_image() + inputs = build_multimodal_inputs(hf_preset, hf_config, image) + with torch.no_grad(): + hf_image_outputs = hf_model( + input_ids=torch.tensor(inputs["token_ids"]), + attention_mask=torch.tensor(inputs["padding_mask"]), + pixel_values=torch.tensor(inputs["pixel_values"]), + image_sizes=torch.tensor(inputs["image_sizes"]), + ) + hf_results = { + "num_parameters": hf_model.num_parameters(), + "text": text_results, + "image": { + **inputs, + "logits": hf_image_outputs.logits.detach().cpu().numpy(), + }, + } + if not skip_generate: + with torch.no_grad(): + generated_ids = hf_model.generate( + input_ids=torch.tensor(inputs["token_ids"]), + attention_mask=torch.tensor(inputs["padding_mask"]), + pixel_values=torch.tensor(inputs["pixel_values"]), + image_sizes=torch.tensor(inputs["image_sizes"]), + max_new_tokens=MAX_NEW_TOKENS, + do_sample=False, + ) + prompt_length = inputs["token_ids"].shape[1] + generated_token_ids = generated_ids[0, prompt_length:].tolist() + try: + generated_text = inputs["tokenizer"].decode( + generated_token_ids, skip_special_tokens=True + ) + except TypeError: + # `mistral_common`'s raw tokenizer (used in the fallback path + # above) doesn't accept `skip_special_tokens`. + generated_text = inputs["tokenizer"].decode(generated_token_ids) + hf_results["image"]["generated_text"] = generated_text + + del hf_model + gc.collect() + return hf_results + + +def check_param_count(keras_model, hf_results): + keras_params = keras_model.backbone.count_params() + hf_params = hf_results["num_parameters"] + print(f"\nKerasHub params: {keras_params:,}") + print(f"HF params: {hf_params:,}") + np.testing.assert_equal(keras_params, hf_params) + print("✅ Parameter count matches.") + + +def test_numerics(label, keras_logits, hf_logits): + keras_logits = ops.convert_to_numpy(keras_logits).astype("float32") + abs_diff = np.abs(keras_logits - hf_logits) + max_diff = float(np.max(abs_diff)) + mean_diff = float(np.mean(abs_diff)) + print(f"KerasHub logits [{label}]:", keras_logits[0, 0, :5]) + print(f"HF logits [{label}]: ", hf_logits[0, 0, :5]) + try: + np.testing.assert_allclose( + keras_logits, hf_logits, atol=1e-3, rtol=1e-3 + ) + print( + f"✅ [{label}] Logits within 1e-3 tolerance " + f"(max={max_diff:.6f}, mean={mean_diff:.6f})." + ) + except AssertionError: + tol = 1e-3 + 1e-3 * np.abs(hf_logits) + mismatched = int(np.sum(abs_diff > tol)) + total = hf_logits.size + matched_pct = 100 * (1.0 - mismatched / total) + print( + f"⚠️ [{label}] Logits exceed 1e-3 tolerance — " + f"max={max_diff:.6f}, mean={mean_diff:.6f}, " + f"matching={matched_pct:.2f}% ({total - mismatched}/{total}).\n" + ) + + +def test_generate( + label, + keras_model, + prompt, + hf_generated_text, + prompt_token_count, + image=None, +): + x = {"prompts": [prompt]} + if image is not None: + x["images"] = [[image]] + max_length = prompt_token_count + MAX_NEW_TOKENS + kh_output = keras_model.generate(x, max_length=max_length) + kh_text = kh_output[0] if isinstance(kh_output, list) else kh_output + if isinstance(kh_text, str): + if kh_text.startswith(prompt): + kh_text = kh_text[len(prompt) :] + else: + # `[IMG]` placeholders expand into real image tokens during + # preprocessing and decode back to nothing, so the decoded text + # won't literally start with `prompt` for image inputs. Strip + # everything up to and including the last `[/INST]` instead. + idx = kh_text.rfind("[/INST]") + if idx != -1: + kh_text = kh_text[idx + len("[/INST]") :] + print(f"\n[{label}] HF generated: {hf_generated_text}") + print(f"[{label}] KH generated: {kh_text}") + + +def run_kh_forward(backbone, backbone_inputs): + with torch.no_grad(): + hidden_states = backbone(backbone_inputs) + return backbone.token_embedding(hidden_states, reverse=True) + + +def test_token_ids(label, preprocessor, prompt, hf_token_ids, image=None): + x = {"prompts": [prompt]} + if image is not None: + x["images"] = [[image]] + keras_inputs = preprocessor.generate_preprocess( + x, sequence_length=hf_token_ids.shape[1] + ) + keras_token_ids = ops.convert_to_numpy(keras_inputs["token_ids"]) + np.testing.assert_array_equal(keras_token_ids, hf_token_ids) + print(f"✅ [{label}] Token IDs match.") + + +def validate_output(keras_model, hf_results, skip_generate=False): + check_param_count(keras_model, hf_results) + backbone = keras_model.backbone + preprocessor = keras_model.preprocessor + text_results = hf_results["text"] + image_results = hf_results["image"] + + test_token_ids("text", preprocessor, TEXT_PROMPT, text_results["token_ids"]) + + # The backbone always declares `pixel_values`/`image_sizes`/ + # `placeholder_indices` as graph inputs, so a text-only forward pass + # feeds it empty-batched image tensors rather than omitting them — + # this is a no-op through the image-merge layer. + vision_encoder = backbone.vision_encoder + patch_size = vision_encoder.patch_size + token_ids = ops.convert_to_tensor(text_results["token_ids"].astype("int32")) + backbone_inputs = { + "token_ids": token_ids, + "padding_mask": ops.ones_like(token_ids), + "pixel_values": ops.zeros( + (0, vision_encoder.num_channels, patch_size, patch_size), + dtype="float32", + ), + "image_sizes": ops.zeros((0, 2), dtype="int32"), + "placeholder_indices": ops.zeros((1, 0), dtype="int32"), + } + keras_logits = run_kh_forward(backbone, backbone_inputs) + test_numerics("text", keras_logits, text_results["logits"]) + + test_token_ids( + "image", + preprocessor, + image_results["prompt"], + image_results["token_ids"], + image=image_results["image"], + ) + + # Feed HF's preprocessed `pixel_values` directly, rather than re-running + # the Keras preprocessor, to avoid PIL vs `ops.image.resize` divergence. + backbone_inputs = { + "token_ids": ops.convert_to_tensor( + image_results["token_ids"].astype("int32") + ), + "padding_mask": ops.convert_to_tensor( + image_results["padding_mask"].astype("int32") + ), + "pixel_values": ops.convert_to_tensor(image_results["pixel_values"]), + "image_sizes": ops.convert_to_tensor( + image_results["image_sizes"].astype("int32") + ), + "placeholder_indices": ops.convert_to_tensor( + image_results["placeholder_indices"].astype("int32") + ), + } + keras_logits = run_kh_forward(backbone, backbone_inputs) + test_numerics("image", keras_logits, image_results["logits"]) + + if not skip_generate: + keras_model.compile(sampler="greedy") + test_generate( + "text", + keras_model, + text_results["prompt"], + text_results.get("generated_text"), + text_results["token_ids"].shape[1], + ) + test_generate( + "image", + keras_model, + image_results["prompt"], + image_results.get("generated_text"), + image_results["token_ids"].shape[1], + image=image_results["image"], + ) + + +def main(_): + if FLAGS.preset not in PRESET_MAP: + raise ValueError( + f"Invalid preset {FLAGS.preset}. Must be one " + f"of {','.join(PRESET_MAP.keys())}" + ) + preset = FLAGS.preset + hf_preset = PRESET_MAP[preset] + + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(hf_preset) + + hf_results = precompute_hf_outputs( + hf_preset, hf_config, skip_generate=FLAGS.skip_generate + ) + print("\n-> Huggingface model loaded and reference outputs computed") + + keras_model = keras_hub.models.Mistral3CausalLM.from_preset( + f"hf://{hf_preset}", dtype="float32" + ) + print("\n-> KerasHub model loaded") + + validate_output(keras_model, hf_results, skip_generate=FLAGS.skip_generate) + print("\n✅ Tests passed!") + + del keras_model + gc.collect() + keras_model = keras_hub.models.Mistral3CausalLM.from_preset( + f"hf://{hf_preset}", dtype="bfloat16" + ) + keras_model.save_to_preset(f"./{preset}") + print("\n✅ Saved the model preset in bfloat16") + + +if __name__ == "__main__": + flags.mark_flag_as_required("preset") + app.run(main)