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:
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
