[Bug]: Shape Mismatch During Expand Operation in ColBERT ONNX Model

[FaisalAliShah]

What happened?

I encountered a shape mismatch error when using ColBERT’s ONNX model to generate embeddings for a batch of text chunks. While most batches work fine, some cause the following error:
[E:onnxruntime:, sequential_executor.cc:516 ExecuteKernel] Non-zero status code returned while running Expand node. Name:'/bert/Expand' Status Message: /bert/Expand: left operand cannot broadcast on dim 1 LeftShape: {1,512}, RightShape: {18,513}

In this case:

• 18 is the batch size passed to the model.
• 512 and 513 refer to the token sequence lengths of the tensors being processed.
• The issue appears to be related to the tokenization and batching process inside the model, as the input to the ONNX model is plain text, and tokenization happens internally.

What is the expected behaviour?

model should generate embeddings for the provided text.

A minimal reproducible example

No response

What Python version are you on? e.g. python --version

Python 3.10.15

FastEmbed version

FastEmbed 0.3.6

What os are you seeing the problem on?

Linux

Relevant stack traces and/or logs

2024-11-21 10:10:48.454 | INFO     | app:embed_documents:61 - Received request to embed documents with 18 texts using model 'colbert'
2024-11-21 10:10:48.455 | INFO     | app:get_or_initialize_model:38 - Using cached model: colbert
2024-11-21 10:10:48.481 | ERROR    | app:embed_documents:71 - Error embedding documents: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Expand node. Name:'/bert/Expand' Status Message: /bert/Expand: left operand cannot broadcast on dim 1 LeftShape: {1,512}, RightShape: {18,513}
2024-11-21 11:05:09.729 | INFO     | app:embed_documents:61 - Received request to embed documents with 27 texts using model 'BM25'
2024-11-21 11:05:09.729 | INFO     | app:get_or_initialize_model:38 - Using cached model: BM25
2024-11-21 11:05:09.759 | INFO     | app:embed_documents:61 - Received request to embed documents with 27 texts using model 'colbert'
2024-11-21 11:05:09.759 | INFO     | app:get_or_initialize_model:38 - Using cached model: colbert
2024-11-21 11:05:10.213 | INFO     | app:embed_documents:66 - 128
2024-11-21 11:06:12.688 | INFO     | app:embed_documents:61 - Received request to embed documents with 27 texts using model 'BM25'
2024-11-21 11:06:12.689 | INFO     | app:get_or_initialize_model:38 - Using cached model: BM25

[abdelkareemkobo]

same error here also when trying to use AnswerdotAI with vepsa.onnx
InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Expand node. Name:'/bert/Expand' Status Message: invalid expand shape
but with Jinaai/Colbertv2 it worked fine

[AlterHoodie]

Hey,
Facing an issue while running the code snippet from the notebook qdrant/workshop-ultimate-hybrid-search, for batches having elements with large documents I get this error:
InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Expand node. Name:'/bert/Expand' Status Message: invalid expand shape

[atanudasgupta]

same here , InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running Expand node. Name:'/bert/Expand' Status Message: invalid expand shape.

looks like it is expecting chunk sizes to be smaller??

However - good news is It works with jinaai/jina-colbert-v2 embedding model without any issues.