Context
I am experimenting with text embedding with the hope of implementing semantic similarity search inside a Phoenix application.
My target use case involves a user writing a short sentence (typically 5 to 30 words). In less than a few seconds, I want to present the user with similar sentences out of a collection of equally short sentences previously written by other users.
The test that puzzles me
As a first quick test of feasibility, I am playing with the example posted by @jonatanklosko at Add text embedding serving · Issue #206 · elixir-nx/bumblebee · GitHub
{:ok, model_info} = Bumblebee.load_model({:hf, "bert-base-uncased"}, architecture: :base)
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "bert-base-uncased"})
text = "Hello, world!"
inputs = Bumblebee.apply_tokenizer(tokenizer, text)
Axon.predict(model_info.model, model_info.params, inputs).hidden_state[0]
The code executes without error but when I run it locally on my machine (MacBookAir <4yo), the last line Axon.predict(model_info.model, model_info.params, inputs).hidden_state[0] takes more than 1 minute to complete.
In contrast, the Python equivalent presented at the top of the same GH thread (Add text embedding serving · Issue #206 · elixir-nx/bumblebee · GitHub) completes almost instantaneously (fractions of a second) on the same machine:
from transformers import AutoTokenizer, AutoModel
import torch
# Load pre-trained model tokenizer and model weights
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")
# Tokenize input text
text = "Hello, world!"
tokens = tokenizer.encode(text, add_special_tokens=True, return_tensors="pt")
# Generate model embeddings
with torch.no_grad():
embeddings = model(tokens)[0].squeeze(0) # Remove batch dimension
# Print the embeddings for the first token
print(embeddings[0])
I am guessing this is not normal and I am doing something wrong. Any idea what that could be? Is there a better way to retrieve text embedding vectors than the script from Add text embedding serving · Issue #206 · elixir-nx/bumblebee · GitHub I am playing with?
Notes
- I am using:
{:bumblebee, "~> 0.5.3"},
{:nx, "~> 0.7.0"},
- The 1 minute runtime I am reporting above is for the last
Axon.predict(model_info.model, model_info.params, inputs).hidden_state[0]step alone (does not include the other mode/tokenizer loading steps). - I did read through Nx vs. Python performance for sentence-transformer encoding. I am guessing my issue is different than what’s discussed there since that post is “only” discussing 2x slower running time compared to equivalent Python code (much lower than the delta I am experiencing => for my application, I’d be more than happy with 2x slower than the equivalent Python runtime I am experiencing).
- Given my use case, since Python is fast enough, I realize I could let Python handle the embedding part and pick things up inside Phoenix after Python completes the embedding. But I’d prefer keeping it all in Elixir if possible.
