support for dimensions field like in OpenAI text-embedding-3, thanks

[ericg108]

Feature request

when initializing with SentenceTransformers, we can use the truncate_dim argument, like below:
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1", truncate_dim=dimensions)

and in calling OpenAI text-embedding-3, we can also pass a `` argument to get variant-length embeddings

dimensions integer Optional The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.
see also: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-dimensions

Motivation

more and more embedding models are supporting Matryoshka embeddings, namely allowing users to get dimensions of varying length, like mxbai-embed-large-v1, jina-embeddings-v3 etc.
this is very useful in scenarios with limited resources. hope it could be supported. Thanks.

Your contribution

I guess it's not a big modification. I may be able to add this feature when I'm told where to modify. Thanks.

[michaelfeil]

Hey, you are welcome to work on this. I have seen you never contributed, so here is a

Add a matryoshka_dim param:

This needs modifications in batch handler for embed,audio,video.

• https://github.com/michaelfeil/infinity/blob/607151b88d0ab8d35d65922f18b0723abd786599/libs/infinity_emb/infinity_emb/inference/batch_handler.py#L265C9-L265C33
• https://github.com/michaelfeil/infinity/blob/607151b88d0ab8d35d65922f18b0723abd786599/libs/infinity_emb/infinity_emb/engine.py#L236C9-L236C33
• https://github.com/michaelfeil/infinity/blob/607151b88d0ab8d35d65922f18b0723abd786599/libs/infinity_emb/infinity_emb/inference/batch_handler.py#L297C7-L297C33

There is also engine + syncengine, which needs three times the integration for above.
https://github.com/michaelfeil/infinity/blob/main/libs/infinity_emb/infinity_emb/engine.py
-> And a test for each of them in test_engine.py

-> Add as test in tests/end-to-end/test_dummyengine.py (will only be based on numpy model)
-> add test with openai client compatability. tests/end-to-end/test_ openai_compat.py

You also need to make the parameters / signature available to matryoska dim.
https://github.com/michaelfeil/infinity/blob/main/libs/infinity_emb/infinity_emb/sync_engine.py

Support in the pydantic model:

• infinity/libs/infinity_emb/infinity_emb/fastapi_schemas/pymodels.py

  Line 56 in 607151b

  user: Optional[str] = None

Should be a small code change (~50 LOC in 10 files), but its the first time adding a request time parameter and therefore needs extensive testing.
I think its easy, but I would enjoy if more people contribute! Let me know if this is helpful to get started.

[ericg108]

okay. Let me work on this.
But I'm not familiar with git operations which may take me some time to file a merge request.

[michaelfeil]

It will make a lot of sense to first know the basics of git & the feature branch pattern, as well as being familiar with inheritance & unit testing in Python.

[wirthual]

@ericg108 with the recent merge to master, you should be able now to truncate the returned embeddings by passing in a dimension params just like the openAI api.

Let me know if that works for you.

[ericg108]

@wirthual thank you for your effort. It took me a lot time to learn the codebase and still I didn't find the place to modify the SentenceTranformer initialization.
From your code, I see just slicing the dimension to get the matryoshka embeddings which I didn't expect. Does that mean the underlying implementation is not based on the sentencetransformer so we don't need to modify the SentenceTranformer initialization by passing the argument truncate_dim?
Thanks!

[michaelfeil]

@ericg108 yes, that's the feature you requested.
It's identical to what "big" inference providers, text-embedding-3 at OpenAI and the original matryoshka paper do. The reason why it works, is that the model was trained to sort the output by importance.

[ericg108]

@michaelfeil got it, thanks!

[wirthual]

Yes, its the same as passing truncate_dim to sentence_transformers . You can see the truncate_embeddings function which is called here is simply doing:

return embeddings[..., :truncate_dim]

However if you managed to pass truncate_dims to Sentence Transformers, this might be of interest for issue #482

Cheers