Instructions to use mixedbread-ai/mxbai-edge-colbert-v0-32m with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- sentence-transformers
How to use mixedbread-ai/mxbai-edge-colbert-v0-32m with sentence-transformers:
from pylate import models queries = [ "Which planet is known as the Red Planet?", "What is the largest planet in our solar system?", ] documents = [ ["Mars is the Red Planet.", "Venus is Earth's twin."], ["Jupiter is the largest planet.", "Saturn has rings."], ] model = models.ColBERT(model_name_or_path="mixedbread-ai/mxbai-edge-colbert-v0-32m") queries_emb = model.encode(queries, is_query=True) docs_emb = model.encode(documents, is_query=False) - Transformers
How to use mixedbread-ai/mxbai-edge-colbert-v0-32m with Transformers:
# Load model directly from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("mixedbread-ai/mxbai-edge-colbert-v0-32m") model = AutoModel.from_pretrained("mixedbread-ai/mxbai-edge-colbert-v0-32m", device_map="auto") - Inference
- Notebooks
- Google Colab
- Kaggle
fix: Include all Dense projection layers in ONNX export (output dim 64)
Browse files## Problem
The current ONNX exports (model.onnx and model_int8.onnx) are missing the final Dense projection layers from PyLate's modules_list pipeline. This causes the ONNX models to output the intermediate hidden dimension instead of the correct final embedding dimension.
**mxbai-edge-colbert-v0-32m:**
- Current ONNX output: **768** dimensions (intermediate)
- Expected output: **64** dimensions (after all Dense projections)
## Root Cause
The ONNX export was done using the standard HuggingFace Transformers export pipeline, which only picks up the transformer backbone and possibly the first linear head. It does not include the additional Dense modules stored in the separate `1_Dense/`, `2_Dense/`, directories that are part of PyLate's `modules_list` architecture.
## Fix
Re-exported by wrapping the full pipeline (Transformer + all Dense layers) into a single module before ONNX export:
- **model.onnx**: fp32, opset 17, output dim = 64
- **model_int8.onnx**: int8 dynamic quantization with projection layers kept in fp32 (they are small and precision-sensitive)
Both files have been verified for:
- Correct output dimensions
- ONNX model validity (onnx.checker)
- Numerical consistency with PyTorch (fp32)
- model.onnx +2 -2
- model_int8.onnx +2 -2
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:01df15011c2a155eee7524613fb9125b1ffada9393c3eab02209b49ced84b368
|
| 3 |
+
size 131529322
|
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1f0fdd343a447c32bd6853f4162424b36bc8fe83eb5a82db18a6e555538a852f
|
| 3 |
+
size 36052322
|