Add Support for BAAI/bge-m3

[Ya-shh]

Canonical vector values: https://colab.research.google.com/drive/1OTZcLxWWAmJqV1ZDS57W5x6zJPRzku4n?usp=sharing

Onnx model: https://huggingface.co/yashvardhan7/bge-m3-onnx/tree/main

Contributed ONNX weights to the BAAI/bge-m3 HF model repo :https://huggingface.co/BAAI/bge-m3/tree/main/onnx

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[Ya-shh]

@Anush008 ,

This one has successfully passed all the tests conducted locally !

Additionally I have updated onnx_embedding.py file as it was generating this error :
test_embedding():
is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")

    for model_desc in TextEmbedding.list_supported_models():
        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
            continue

        dim = model_desc["dim"]
        model = TextEmbedding(model_name=model_desc["model"])

        docs = ["hello world", "flag embedding"]

      embeddings = list(model.embed(docs))

fastembed/tests/test_text_onnx_embeddings.py:48:

---

fastembed/fastembed/text/text_embedding.py:89: in embed
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
fastembed/fastembed/text/onnx_embedding.py:223: in embed
yield from self._embed_documents(
fastembed/fastembed/common/onnx_model.py:97: in _embed_documents
yield from self._post_process_onnx_output(self.onnx_embed(batch))
fastembed/fastembed/common/onnx_model.py:70: in onnx_embed
model_output = self.model.run(None, onnx_input)

---

self = <onnxruntime.capi.onnxruntime_inference_collection.InferenceSession object at 0x10217eb90>
output_names = ['token_embeddings', 'sentence_embedding']
input_feed = {'attention_mask': array([[1, 1, 1, 1, 1, 0],
[1, 1, 1, 1, 1, 1]]), 'input_ids': array([[ 0, 33600, 31, ...[ 0, 49938, 6, 55720, 59725, 2]]), 'token_type_ids': array([[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0]])}
run_options = None

def run(self, output_names, input_feed, run_options=None):
    """
    Compute the predictions.

    :param output_names: name of the outputs
    :param input_feed: dictionary ``{ input_name: input_value }``
    :param run_options: See :class:`onnxruntime.RunOptions`.
    :return: list of results, every result is either a numpy array,
        a sparse tensor, a list or a dictionary.

    ::

        sess.run([output_name], {input_name: x})
    """
    self._validate_input(list(input_feed.keys()))
    if not output_names:
        output_names = [output.name for output in self._outputs_meta]
    try:

      return self._sess.run(output_names, input_feed, run_options)

E onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Invalid input name: token_type_ids

../.pyenv/versions/3.10.10/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:220: InvalidArgument

[Ya-shh]

Fixes #107

[Anush008]

CI passed. Awesome.

I see you've contributed the ONNX weights to the BAAI/bge-m3 HF model repo.

https://huggingface.co/BAAI/bge-m3

We can now use the original repo as the source?

[Ya-shh]

CI passed. Awesome.

I see you've contributed the ONNX weights to the BAAI/bge-m3 HF model repo.

https://huggingface.co/BAAI/bge-m3

We can now use the original repo as the source?

Done !

[Anush008]

🚀🚀

[NirantK]

As mentioned in the issue — we'd want to support BGE-M3 with not just dense vectors, but all 3: SPLADE and ColBERT.

It's not very useful as just a dense model and adds to the confusion.

Would appreciate it if you can add a PR with a new category of models e.g. multifunc — from multi-functionality which supports all three.

[Ya-shh]

@NirantK , I found a bge-m3-onnx model on HF which supports both dense and ColBERT vectors. I even tested the outputs on colab :https://colab.research.google.com/drive/1wLfKa3zdPkVo3ewnrTu6BHyyo1RrLINL?usp=sharing

[Ya-shh]

however , when I tested it with my bge-m3 onnx model repo for ColBERT vectors .This was the output:https://colab.research.google.com/drive/1SNhC089nCiLixhQVGT-VmYQeOu7WHdba?usp=sharing

[Ya-shh]

Could we test bge-m3-onnx with SPLADE vectors like this ? Notebook:https://colab.research.google.com/drive/1phjnxXzUzOMd_x3_015ZImjnBuio5IbW?usp=sharing

[Ya-shh]

@Anush008 could you please review this ?

[Anush008]

We'll have to wait for Nirant to take a look at those.

[Ya-shh]

@NirantK could you please review this ?

[NirantK]

First off, thanks a ton for the love and contributions @Ya-shh — really appreciate.

Thank you for the investigations and looking around for more ONNX exports e.g. aapot, ddmitov, official BGE-M3 — none of them do all 3 as far as I can tell

Given the time and effort you've put across PRs, I believe it'd be helpful to expand how I'm thinking about this right now:

BGE-M3

1. BGE-M3 should exist in a separate, new category e.g. multi like we do for sparse
2. It should support all 3 — not just dense and ColBERT. This means FastEmbed will have to store and do some of the normalisation, loadings heads (for ColBERT and SPLADE) and passing attention which FlagEmbedding is doing right now. This PR should include those changes.

Tests and Notebooks

The notebooks are a bit difficult to make sense of:

1. SPLADE notebook has PyTorch outputs, not ONNX. We'd need both to compare and use tests for both SPLADE, Dense and write new ones for ColBERT

2. ColBERT notebooks: I don't see a comparison or test in both of them? But the dimensions looks different and I assume that's because of different normalisation or export mismatch between the two models

Suggestions

1. Perhaps contribute a simple ColBERT notebook in a new category first? That will allow us to collaborate on ColBERT tests as well. I suspect this is the most popular model and a great starting point: https://github.com/stanford-futuredata/ColBERT

2. Please look into closing existing open PRs before opening new ones? It's not a great idea to have lot of open streams of work. If the ONNX export needs more investigation. I would love it if you took that up and contributed the solution here as an example notebook, I'll be happy and excited to review that!

[Ya-shh]

@NirantK Thank you immensely for your kind words and recognition. Your appreciation means a lot to me!

Your detailed feedback is incredibly helpful in guiding my efforts further. I'm deeply grateful for the opportunity to expand on my contributions and improve the ONNX exports. I wholeheartedly agree with your suggestions and will prioritize them accordingly.

⎯Regarding your thoughts on BGE-M3, creating a separate category multi , for that I will definitely implement the necessary changes to support all three aspects efficiently.

As for the notebooks, I understand the importance of clarity and consistency. Rest assured, I'll work diligently to ensure they provide comprehensive comparisons and tests for all three vectors.

Your suggestion to begin with a ColBERT notebook is well-taken, and I'll proceed with enthusiasm. Collaborating on tests for this popular model is a fantastic opportunity to enhance our collective efforts.

Lastly, I appreciate your reminder to close existing PRs before opening new ones. Streamlining our workflow will undoubtedly lead to greater efficiency and effectiveness. I'll make it a priority to address any outstanding tasks promptly.

Once again, thank you for your unwavering support and encouragement. I'm too excited to delve deeper into the ONNX export investigation and share my findings with you through an example notebook. Your review and feedback are eagerly anticipated!