Add Mistral3 model to Hub - #3020
Conversation
…d image converter and update preprocessor and causal lm
…tch HF numerics; update checkpoint conversion script to validate text and image paths separately.
…or configs Some Mistral3 checkpoints (e.g. Mistral Small 3.2) ship only tekken.json, without tokenizer_config.json or preprocessor_config.json. Fall back to mistral_common directly for tokenization and image preprocessing in convert_mistral.py and the checkpoint conversion script, and trim docstrings/comments down to the essential rationale.
- Unify MistralCausalLMPreprocessor.generate_preprocess() into a single method for text-only and multimodal presets: calls with no images delegate straight to the base CausalLMPreprocessor (matching Qwen3.5's pattern), while both text-only and multimodal presets now accept a {"prompts": ..., "images": ...} dict uniformly.
- Add @preprocessing_function to generate_preprocess() so it correctly handles the tf.Tensor decorator produces.
- Fix Mistral3ImageConverter: dtype and image_size were hardcoded in __init__ but also passed during deserialization, causing "got multiple values for keyword argument". Remove the hardcoded call.
- Add test_multimodal_serialization and other preprocessor test coverage for the now-removed _expand_image_placeholders in favor of testing the unified generate_preprocess().
- Simplify convert_mistral_checkpoints.py's test_token_ids to build its own {"prompts", "images"} pair, removing the duplicated dict construction at both call sites.
Merge the two precompute_hf_outputs functions and drop the redundant multimodal flag in favor of an image key check; trim stale comments.
…gation - Trim redundant/private-method tests across mistral test files, hoist shared setup into setUp, add missing serialization coverage. - Split multimodal backbone tests into MistralMultimodalBackboneTest using run_backbone_test, with run_quantization_check=False (same workaround gemma3_backbone_test.py uses for nested-submodel quantization path mismatches). - Fix real bug: ~20 sublayer constructions in mistral_vision_encoder.py never forwarded dtype=self.dtype_policy, so mixed-precision policies never reached the vision encoder's sublayers. - Rename compute_pixtral_resize_size to compute_resize_size.
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces multimodal (image + text) support for the Mistral3/Pixtral architecture in KerasHub, adding a vision encoder, multimodal projector, image converter, and updating the preprocessor, causal LM, tokenizers, and checkpoint conversion utilities. The review feedback highlights critical issues with batching in the multimodal preprocessing pipeline; specifically, the current implementation of compute_image_placeholder_indices flattens the batch dimension, which prevents correct mapping of placeholder indices for batch sizes greater than 1 and causes shape mismatches in the preprocessor and test suite.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for multimodal Mistral3 models (Pixtral) by introducing a vision encoder, multimodal projector, and updated preprocessing logic. The review identified critical issues regarding the use of backend-specific code, a logic bug in the placeholder index calculation for batch sizes greater than 1, and a minor docstring formatting violation.
| elif tf is not None and isinstance(prompts, tf.Tensor): | ||
| # `@preprocessing_function` converts raw inputs to `tf.Tensor`s | ||
| # eagerly; decode back to Python strings, since image-placeholder | ||
| # expansion below is per-example and variable-length. | ||
| if prompts.shape.rank == 0: | ||
| batched = False | ||
| prompts = [prompts] | ||
| if images is not None: | ||
| images = [images] | ||
| prompts = [p.numpy().decode("utf-8") for p in prompts] | ||
|
|
||
| if images is None: | ||
| images_per_prompt = [[] for _ in prompts] | ||
| else: | ||
| images_per_prompt = [ | ||
| list(per_prompt_images) for per_prompt_images in images | ||
| ] | ||
| return list(prompts), images_per_prompt, batched | ||
|
|
||
| def call( | ||
| self, | ||
| x, | ||
| y=None, | ||
| sample_weight=None, | ||
| sequence_length=None, | ||
| ): | ||
| if self.text_only_model: | ||
| return super().call( | ||
| x, | ||
| y=y, | ||
| sample_weight=sample_weight, | ||
| sequence_length=sequence_length, | ||
| ) | ||
|
|
||
| sequence_length = sequence_length or self.sequence_length | ||
| prompts, images_per_prompt, batched = self._extract_multimodal_inputs(x) | ||
| tokenized, pixel_values, image_sizes = self._build_multimodal_inputs( | ||
| prompts, images_per_prompt | ||
| ) | ||
| if pixel_values is None: | ||
| raise ValueError( | ||
| "Mistral3's multimodal preprocessor requires at least one " | ||
| "image per batch when `image_converter` is set; got a " | ||
| "batch with zero images." | ||
| ) | ||
|
|
||
| tokenized = tf.ragged.constant(tokenized, dtype="int32") |
There was a problem hiding this comment.
This part of the code has operations that are specific to the TensorFlow backend, which violates the repository's style guide for backend-agnostic code.
- On lines 275-284,
tf.Tensorand.numpy()are used. This should be replaced with backend-agnostic alternatives likekeras.backend.is_tensorandkeras.ops.convert_to_numpy. - On line 321,
tf.ragged.constantis used. This is also TensorFlow-specific. Theself.packerlayer can handle a list of lists directly, so this conversion is unnecessary.
A similar issue with tf.ragged.constant exists on line 392 in the generate_preprocess method.
References
- Avoid backend-specific operations (like TensorFlow-specific classes/methods) to maintain backend-agnostic code. (link)
There was a problem hiding this comment.
This is expected here, not a violation, preprocessors run inside tf.data via @preprocessing_function, so tf ops are the norm in this file. tf.ragged.constant, isinstance(x, tf.Tensor), and .numpy().decode() are the same pattern used by blip2, gemma3, gemma3n, gemma4, qwen3_5, and rwkv7's preprocessors.
| placeholder_indices = compute_image_placeholder_indices( | ||
| keras.ops.convert_to_numpy(model_token_ids), | ||
| self.tokenizer.image_placeholder_token_id, | ||
| )[None, :] |
There was a problem hiding this comment.
The calculation of placeholder_indices is incorrect for a batch size greater than 1. compute_image_placeholder_indices returns a flat array of indices for the entire batch. Adding a new leading dimension with [None, :] results in a shape of (1, num_total_placeholders), which will not work correctly when the batch size is greater than 1. The indices should be reshaped to (batch_size, num_placeholders_per_item).
To avoid creating a divergence between parallel execution paths (the standard preprocess path and the generate_preprocess path), please ensure this fix is applied consistently to both paths.
| placeholder_indices = compute_image_placeholder_indices( | |
| keras.ops.convert_to_numpy(model_token_ids), | |
| self.tokenizer.image_placeholder_token_id, | |
| )[None, :] | |
| placeholder_indices = compute_image_placeholder_indices( | |
| keras.ops.convert_to_numpy(model_token_ids), | |
| self.tokenizer.image_placeholder_token_id, | |
| ) | |
| batch_size = keras.ops.shape(model_token_ids)[0] | |
| if keras.ops.shape(placeholder_indices)[0] == 0: | |
| num_placeholders_per_item = 0 | |
| else: | |
| num_placeholders_per_item = ( | |
| keras.ops.shape(placeholder_indices)[0] // batch_size | |
| ) | |
| placeholder_indices = keras.ops.reshape( | |
| placeholder_indices, (batch_size, num_placeholders_per_item) | |
| ) |
References
- When aligning parallel execution paths (such as a Python path and a framework/TF path), avoid introducing path-specific fixes or improvements that would create a divergence between them. Apply such changes consistently to both paths, preferably in a separate pull request.
There was a problem hiding this comment.
Same as before, placeholder_indices are global flat indices into (batch*seq_length,), not per-example, so (1, N) is correct as-is. The suggested fix also assumes an even placeholder count per example (total // batch_size), which breaks once prompts have different numbers of images.
There was a problem hiding this comment.
Update: reworked this. placeholder_indices is no longer a flattened (1, N) global-index array, it's now (batch, max_placeholders), each row's own local positions, padded with -1. The old scheme couldn't actually handle prompts in the same batch with different numbers of images, so this was worth fixing. The suggested reshape still wouldn't have worked as-is (assumed an even split), but the underlying concern was valid.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces multimodal support for the Mistral model architecture, enabling the integration of a vision encoder and multimodal projector. Key changes include the addition of a Mistral3ImageConverter for image preprocessing, a Mistral3VisionEncoder for feature extraction, and updates to the MistralBackbone and MistralCausalLM to handle image-text embedding merging. The PR also updates the tokenizer to support special vision tokens and includes comprehensive testing and checkpoint conversion utilities for multimodal presets. The review comments correctly identify an API inconsistency regarding the add_end_token default value and suggest a more robust structural approach for defensive type conversion in the preprocessor, both of which are actionable and valuable for maintaining repository standards.
| image_converter=None, | ||
| sequence_length=1024, | ||
| add_start_token=True, | ||
| add_end_token=True, |
There was a problem hiding this comment.
The default value for add_end_token is changed to True here, which differs from the False default in the base CausalLMPreprocessor class.
This inconsistency in the API can be surprising for users and may lead to unexpected behavior. For API consistency, it's generally best to keep default values the same as the parent class unless there's a strong, model-specific reason for the change.
If this change is intentional, please document the reasoning in the docstring. Otherwise, I'd recommend reverting to the base class default: add_end_token=False.
There was a problem hiding this comment.
Checked this, it's actually not a change. The base CausalLMPreprocessor.__init__ has defaulted add_end_token=True and before this PR, Mistral had no __init__ override at all, so it already inherited True. The old docstring saying "Default is False" was just stale and never matched runtime behavior, this PR's explicit default actually corrects that inconsistency rather than introducing one.
| def call( | ||
| self, | ||
| x, | ||
| y=None, | ||
| sample_weight=None, | ||
| sequence_length=None, | ||
| ): | ||
| if self.text_only_model: | ||
| return super().call( | ||
| x, | ||
| y=y, | ||
| sample_weight=sample_weight, | ||
| sequence_length=sequence_length, | ||
| ) | ||
|
|
||
| sequence_length = sequence_length or self.sequence_length | ||
| prompts, images_per_prompt, batched = self._extract_multimodal_inputs(x) | ||
| tokenized, pixel_values, image_sizes = self._build_multimodal_inputs( | ||
| prompts, images_per_prompt | ||
| ) | ||
| if pixel_values is None: | ||
| raise ValueError( | ||
| "Mistral3's multimodal preprocessor requires at least one " | ||
| "image per batch when `image_converter` is set; got a " | ||
| "batch with zero images." | ||
| ) | ||
|
|
||
| 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] | ||
| placeholder_indices = compute_image_placeholder_indices( | ||
| keras.ops.convert_to_numpy(model_token_ids), | ||
| self.tokenizer.image_placeholder_token_id, | ||
| )[None, :] | ||
|
|
||
| out_x = { | ||
| "token_ids": model_token_ids, | ||
| "padding_mask": model_padding_mask, | ||
| "pixel_values": pixel_values, | ||
| "image_sizes": image_sizes, | ||
| "placeholder_indices": placeholder_indices, | ||
| } | ||
| # Target `y` will be the next token. | ||
| y = token_ids[..., 1:] | ||
| sample_weight = padding_mask[..., 1:] | ||
|
|
||
| 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) | ||
|
|
There was a problem hiding this comment.
When handling inputs that may be passed as tensors in a Python-only execution path, please perform defensive type conversion at the entry point of the path rather than deep inside helper functions (such as _extract_multimodal_inputs which decodes tensors to Python strings). This prevents potential issues when downstream helper functions perform operations on these values.
References
- When handling inputs that may be passed as tensors in a Python-only execution path, perform defensive type conversion (e.g., converting to int) at the entry point of the path rather than deep inside helper functions, especially if multiple downstream helper functions perform arithmetic on the value.
There was a problem hiding this comment.
_extract_multimodal_inputs is actually the very first thing called in both call() and generate_preprocess(), and the tensor-to-string decode happens right at the top of it, before anything else touches the value. So this is already happening at the entry point, it's just tucked into a small helper instead of inlined. Don't see a concrete case where this causes a problem, but happy to inline it if there's any specific scenario in mind.
|
hi @ahmed-damda thank you for the work you've begun on adding Mistral3. Note, however, that the Keras Team is currently working on adding Mistral 3 to support multi-modal functionality, so we recommend you to hold off on any further work on this PR and encourage you to try out our implementation when it's complete. |
|
Hi @rwitcher , @vijay-dabur (Myself) and @ahmed-damda have been added to Keras Team for implementation of models and we were assigned this model by @divyashreepathihalli and @laxmareddyp . We are mostly done with the model work required here, we will be sharing this PR for review today. Let me know if there is any confusion. |
Hey @vijay-dabur There was some confusion regarding the GitHub handle on this PR since this task was originally assigned to you, and we saw it opened under a different name without prior notice. We wanted to hold off on reviewing/merging to prevent duplicate effort between you and the other contributor. Additionally, as previously discussed, Mistral 3 needs to be in a separate model folder rather than inside the existing mistral folder. This structure issue is also a reason to @rwitcher commented on the PR. |
Multimodal Mistral3 code had been grafted directly onto the text-only Mistral classes. Mirrors the Gemma/Gemma3 and Mistral/Mixtral precedent of one self-contained directory per model. - keras_hub/src/models/mistral/: reverted to text-only. - keras_hub/src/models/mistral3/: new model directory. Mistral3Backbone is standalone but reuses CachedMistralAttention/MistralTransformerDecoder from mistral/ (identical text decoder); Mistral3CausalLM subclasses MistralCausalLM; Mistral3CausalLMPreprocessor and Mistral3Tokenizer (Tekken-only) are standalone. - Split HF conversion (convert_mistral.py/convert_mistral3.py) and the checkpoint validation script accordingly. - Updated public API exports.
ahmed-damda
changed the title
…te in checkpoint conversion
- Add missing build() to the vision encoder/projector layers, fixing weight reinitialization on save/load and a JAX param-count mismatch. - Make Mistral3ImageConverter work inside tf.data on all three backends by using native tf ops for the graph path instead of keras.ops. - Rework placeholder_indices from a single flattened batch index into per-row local positions, so prompts in the same batch can now have different numbers of images. - Fill in preset kaggle handles/params, now that the checkpoints are uploaded. - Round out test coverage across the model (backbone, causal LM, preprocessor, tokenizer, image converter, vision encoder).
a4a13ac to
961f8f7
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces the multimodal Mistral3 (Pixtral vision + Mistral text) architecture to KerasHub, implementing the backbone, causal LM task, preprocessor, tokenizer, image converter, and HuggingFace conversion utilities. It also refactors the existing Mistral model to use Keras's native RMSNormalization. The review feedback highlights several critical issues that should be addressed: an XLA tracing bug in generate_step when the batch dimension is dynamic, a potential crash in the tokenizer when split_pattern is None, a serialization failure in the preprocessor due to a missing image_converter in get_config(), incorrect handling of 4D image inputs in the preprocessor, and dead code in Mistral3PatchMerger. All review comments are valid and should be kept.
| 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 | ||
| ) |
There was a problem hiding this comment.
Under XLA compilation or tracing (e.g., in JAX or TensorFlow graph mode), the batch dimension of pixel_values is often None (dynamic batch size). Checking pixel_values.shape[0] directly in Python will evaluate to False when it is None, which completely disables the vision encoder during compiled execution. We should check if pixel_values.shape[0] is None or pixel_values.shape[0] > 0 to support dynamic batch sizes during tracing.
| 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 | |
| ) | |
| img_embeddings = None | |
| if pixel_values is not None and ( | |
| pixel_values.shape[0] is None or pixel_values.shape[0] > 0 | |
| ): | |
| img_embeddings = self.backbone.image_feature_extractor( | |
| pixel_values, image_sizes | |
| ) |
There was a problem hiding this comment.
This is intentional, same static-shape convention Gemma3 uses in its generate_step (None shape = no images).
| def __init__( | ||
| self, | ||
| vocabulary=None, | ||
| merges=None, | ||
| split_pattern=None, | ||
| control_tokens=None, | ||
| **kwargs, | ||
| ): | ||
| self.split_pattern = split_pattern |
There was a problem hiding this comment.
If split_pattern is None (which is the default), instantiating the tokenizer with just vocabulary and merges will crash during _set_vocabulary_and_merges_tokenizers because hf_tokenizers.Regex(self.split_pattern) on line 128 expects a string, not None. We should default split_pattern to the Tekken split pattern if it is None.
| def __init__( | |
| self, | |
| vocabulary=None, | |
| merges=None, | |
| split_pattern=None, | |
| control_tokens=None, | |
| **kwargs, | |
| ): | |
| self.split_pattern = split_pattern | |
| def __init__( | |
| self, | |
| vocabulary=None, | |
| merges=None, | |
| split_pattern=None, | |
| control_tokens=None, | |
| **kwargs, | |
| ): | |
| if split_pattern is None: | |
| 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+" | |
| ) | |
| self.split_pattern = split_pattern |
There was a problem hiding this comment.
Yeah this would crash if someone instantiated the tokenizer directly without a split_pattern. Added a MISTRAL3_TEKKEN_SPLIT_PATTERN default, still overridable.
| def get_config(self): | ||
| config = super().get_config() | ||
| config.update({"spatial_merge_size": self.spatial_merge_size}) | ||
| return config |
There was a problem hiding this comment.
The image_converter attribute is a required argument of __init__ but is missing from get_config(). This will cause serialization to fail when saving and loading the preprocessor, as from_config will not be able to deserialize the preprocessor without image_converter. We should serialize image_converter using keras.layers.serialize(self.image_converter) and implement from_config to deserialize it.
| def get_config(self): | |
| config = super().get_config() | |
| config.update({"spatial_merge_size": self.spatial_merge_size}) | |
| return config | |
| def get_config(self): | |
| config = super().get_config() | |
| config.update( | |
| { | |
| "spatial_merge_size": self.spatial_merge_size, | |
| "image_converter": keras.layers.serialize(self.image_converter), | |
| } | |
| ) | |
| return config | |
| @classmethod | |
| def from_config(cls, config): | |
| if "image_converter" in config and config["image_converter"] is not None: | |
| config["image_converter"] = keras.layers.deserialize( | |
| config["image_converter"] | |
| ) | |
| return super().from_config(config) |
There was a problem hiding this comment.
This is already handled, image_converter gets serialized/deserialized generically by the base Preprocessor.get_config/from_config, and it's covered by the existing round-trip test. No change needed here.
| else: | ||
| images_per_prompt = [ | ||
| list(per_prompt_images) for per_prompt_images in images | ||
| ] |
There was a problem hiding this comment.
If images is passed as a 4D NumPy array or tf.Tensor of shape (batch_size, height, width, channels) (representing one image per prompt), per_prompt_images will be a 3D array/tensor of shape (height, width, channels). Calling list(per_prompt_images) will slice the image along its height dimension, resulting in a list of 2D slices of shape (width, channels) instead of a list containing the single 3D image. We should handle 4D inputs explicitly by expanding their dimensions or checking the rank of the input.
| else: | |
| images_per_prompt = [ | |
| list(per_prompt_images) for per_prompt_images in images | |
| ] | |
| else: | |
| first_img = images[0] if len(images) > 0 else None | |
| if (hasattr(images, "shape") and len(images.shape) == 4) or ( | |
| isinstance(images, list) | |
| and first_img is not None | |
| and hasattr(first_img, "shape") | |
| and len(first_img.shape) == 3 | |
| ): | |
| images_per_prompt = [[img] for img in images] | |
| else: | |
| images_per_prompt = [ | |
| list(per_prompt_images) for per_prompt_images in images | |
| ] |
There was a problem hiding this comment.
Fair point, this was a real gap. Reworked the preprocessor so images are matched to prompts by placeholder-token count instead of requiring a specific nesting, now accepts a single image, a flat list, a batched array, or the original list-of-lists, all in one pass.
| def _merge_image( | ||
| self, | ||
| image_features, | ||
| height, | ||
| width, | ||
| ): |
There was a problem hiding this comment.
The helper method _merge_image is defined but never called anywhere in Mistral3PatchMerger or the rest of the codebase. It is dead code and should be removed to improve maintainability.
|
The zizmor failure seems unrelated to this PR |
ahmed-damda
changed the title
4 participants
Description of the change
This PR adds Mistral3, a multimodal (image + text) model using the Pixtral vision architecture that Mistral3 models (e.g. Mistral Small 3.1/3.2) are built on.
What's new
Mistral3VisionEncoderkeras_hub.models.Mistral3VisionEncoderMistral3MultiModalProjector/Mistral3PatchMergerMistral3ImageConverterkeras_hub.layers.Mistral3ImageConverterMistral3Backbonecall()takestoken_ids/padding_maskpluspixel_values,image_sizes, andplaceholder_indicesfor image inputsMistral3CausalLMcall_with_cache/generate_stepcompute image features once at prefill and scatter them into the token embeddings before caching, so incremental decode steps stay text-onlyMistral3CausalLMPreprocessorimage_converter;xcan be a dict of{"prompts", "images"}and prompts get their image placeholder tokens expanded automaticallyMistral3Tokenizer[IMG],[IMG_BREAK],[IMG_END]as special tokensconvert_mistral3.pymistral_commonfallback for checkpoints (like Mistral Small 3.2) that ship nopreprocessor_config.jsonconvert_mistral3_checkpoints.pyPresets
mistral_small_3.1_24b_base_2503_enmistral_small_3.1_24b_instruct_2503_enmistral_small_3.2_24b_instruct_2506_enmistral_small_3.1_24b_instruct_2503_en, with improved instruction-following and reduced repetition.Design notes
pixel_valueshas dynamic spatial dimensions rather than a fixed canvas, matching HF'sPixtralImageProcessor, which pads each batch to its own largest image.image_sizescarries each image's true(height, width)so the encoder can crop correctly after patchification.RMSNormalizationrather than a custom implementation.Testing
Checkpoint conversion was validated against reference HF checkpoints for numerical parity on both text and image generation paths, alongside unit tests for the vision encoder, image converter, and multimodal backbone/preprocessor paths. Using Keras's built-in
RMSNormalizationreduced the max logit mismatch against the HF reference from 28 to 8.Mistral Small 3.2 Instruct: numerics comparison:
Mistral Small 3.2 Instruct: generation
Checklist