Skip to content

Add Mxbai(mixbread-ai) Embedding Model to Hub#2850

Open
laxmareddyp wants to merge 2 commits into
keras-team:masterfrom
laxmareddyp:add_mxbai_model
Open

Add Mxbai(mixbread-ai) Embedding Model to Hub#2850
laxmareddyp wants to merge 2 commits into
keras-team:masterfrom
laxmareddyp:add_mxbai_model

Conversation

@laxmareddyp

@laxmareddyp laxmareddyp commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Description of the change

Add 4 mixedbread-ai (mxbai) embedding model presets to KerasHub.

License: Apache-2.0

Presets

Preset Architecture Pooling Normalization Params
mxbai_embed_large_v1_en BERT (24L) CLS None 335M
mxbai_embed_2d_large_v1_en BERT (24L) CLS None 335M
mxbai_embed_xsmall_v1_en BERT (6L) Mean None 24M
deepset_mxbai_embed_de_large_v1 XLM-RoBERTa (24L) Mean L2 560M

Changes

  • tools/checkpoint_conversion/convert_mxbai_embed_checkpoints.py : Added conversion script for all 4 presets to use from_preset("hf://...") to leverage existing BERT/XLM-RoBERTa converters.
  • keras_hub/src/utils/transformers/convert_bert.py: Bug fix — added max_sequence_length to convert_backbone_config(). Without this, BERT models with non-512 max_position_embeddings (e.g., xsmall-v1 with 4096) fail during weight conversion due to position embedding shape mismatch.

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have followed the Keras Hub Model contribution guidelines in making these changes.
  • I have followed the Keras Hub API design guidelines in making these changes.
  • I have signed the Contributor License Agreement.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a checkpoint conversion script for HuggingFace mixedbread-ai mxbai-embed checkpoints to KerasHub format, supporting both BERT and XLM-RoBERTa architectures, and updates the BERT config converter to map the maximum sequence length. The review feedback recommends enhancing the conversion script's robustness by raising errors if numerical verification or semantic search ranking fails, rather than just printing diagnostics, to prevent saving corrupted presets.

Comment thread tools/checkpoint_conversion/convert_mxbai_embed_checkpoints.py
@laxmareddyp laxmareddyp added the kokoro:force-run Runs Tests on GPU label Jul 22, 2026
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Jul 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants