Nx vs. Python performance for sentence-transformer encoding

Questions & Help Questions
,
32

I experimented with running a model through ortex today, but I don’t see how this improves the situation with different sequence lengths. As soon as I send inputs with different sizes to my serving, I get an error that I cannot merge batches due to incompatible templates, which makes sense. Sending the full tokenizer output to the Ortex model does not offer a performance gain compared to running EXLA with the full sequence length.

Running different servings for each sequence length I see similar performance running the all-MiniLM-L6-v2 model on EXLA (see Nx vs. Python performance for sentence-transformer encoding - #18 by steffend) and Ortex:

image
image1200×676 36.4 KB

For small inputs, Ortex seems to perform a little better, for larger inputs a little worse.

Here is the Livebook I used:

Notice the part that uses the smallest input sequence length (by looking at the attention mask) in the client_preprocessing:

  def serving(model, tokenizer) do
    Nx.Serving.new(Ortex.Serving, model)
    |> Nx.Serving.client_preprocessing(fn inputs ->
      {:ok, encodings} = Tokenizers.Tokenizer.encode_batch(tokenizer, inputs)

      # get the maximum sequence length from the input by looking at the attention mask
      max_length =
        encodings
        |> Enum.map(&Tokenizers.Encoding.get_attention_mask/1)
        |> Enum.map(fn tensor -> Enum.sum(tensor) end)
        |> Enum.max(fn -> nil end)

      encodings =
        if max_length do
          for e <- encodings, do: Tokenizers.Encoding.truncate(e, max_length)
        else
          encodings
        end

      input_ids = for i <- encodings, do: Tokenizers.Encoding.get_ids(i)
      input_mask = for i <- encodings, do: Tokenizers.Encoding.get_attention_mask(i)
      token_type_ids = for i <- encodings, do: Tokenizers.Encoding.get_type_ids(i)

      inputs =
        Enum.zip_with([input_ids, input_mask, token_type_ids], fn [a, b, c] ->
          {Nx.tensor(a), Nx.tensor(b), Nx.tensor(c)}
        end)
        |> Nx.Batch.stack()

      {inputs, %{attention_mask: Nx.tensor(input_mask)}}
    end)
    |> Nx.Serving.client_postprocessing(fn {{output}, _meta}, client_info ->
      mean_pooling(output, client_info.attention_mask)
    end)
  end

I have to say that besides that Ortex works flawlessly and I’m able to run models that Bumblebee does not support yet (e.g. all-mpnet-base-v2): Running the all-mpnet-base-v2 sentence transformer in Elixir using Ortex · GitHub