Nx vs. Python performance for sentence-transformer encoding

So to summarize this topic: Elixir + Nx perform quite well for generating sentence embeddings. It is very important to make sure that the model is compiled for the right sequence lengths though, as the input is padded.
For varying input lengths, the newest Bumblebee code on GitHub now supports specifying multiple sequence lengths, as seen in post 36.

Another option is using Ortex and an ONNX model. Here is an example livebook that uses this approach: Running the all-mpnet-base-v2 sentence transformer in Elixir using Ortex · GitHub. In that case, using a serving does not improve the performance that much, at least when running on CPU. Therefore one can also just call Ortex.run directly.

Thanks for all the replies and insights. I’m happy to see that this lead to some improvements in Bumblebee and Nx!

I’m also happy to report that we’re currently working on moving our production setup from Python to Elixir+Nx at my job :smiley: