tomaarsen
commented
6 months ago Hello!
ONNX might be worth a shot:
CrossEncoder via transformers
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import time
model = AutoModelForSequenceClassification.from_pretrained("cross-encoder/ms-marco-MiniLM-L-6-v2")
tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-MiniLM-L-6-v2")
# 80 texts:
texts = [
["How many people live in Berlin?", "How many people live in Berlin?"],
[
"Berlin has a population of 3,520,031 registered inhabitants in an area of 891.82 square kilometers.",
"New York City is famous for the Metropolitan Museum of Art.",
],
] * 40
warmup_features = tokenizer(texts[:10], padding=True, return_tensors="pt")
features = tokenizer(texts, padding=True, return_tensors="pt")
with torch.no_grad():
scores = model(**warmup_features).logits
start_t = time.time()
scores = model(**features).logits
print(
f"Time: {len(texts) / (time.time() - start_t):.4f} sentences per second"
)Time: 601.4126 sentences per secondONNX:
Export via:
optimum-cli export onnx -m cross-encoder/ms-marco-MiniLM-L-6-v2 ce-ms-marco-MiniLM-L-6-v2
# Export with:
# optimum-cli export onnx -m cross-encoder/ms-marco-MiniLM-L-6-v2 ce-ms-marco-MiniLM-L-6-v2
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSequenceClassification
import torch
import time
model = ORTModelForSequenceClassification.from_pretrained("ce-ms-marco-MiniLM-L-6-v2")
tokenizer = AutoTokenizer.from_pretrained("ce-ms-marco-MiniLM-L-6-v2")
# 80 texts:
texts = [
["How many people live in Berlin?", "How many people live in Berlin?"],
[
"Berlin has a population of 3,520,031 registered inhabitants in an area of 891.82 square kilometers.",
"New York City is famous for the Metropolitan Museum of Art.",
],
] * 40
warmup_features = tokenizer(texts[:10], padding=True, return_tensors="pt")
features = tokenizer(texts, padding=True, return_tensors="pt")
with torch.no_grad():
scores = model(**warmup_features).logits
start_t = time.time()
scores = model(**features).logits
print(
f"Time: {len(texts) / (time.time() - start_t):.4f} sentences per second"
)Time: 905.0249 sentences per secondor
optimum-cli export onnx -m cross-encoder/ms-marco-MiniLM-L-6-v2 ce-ms-marco-MiniLM-L-6-v2 --optimize O3with the same script:
Time: 981.4480 sentences per secondBeyond that, multi-CPU reranking definitely should also work, but you'd have to implement it yourself manually. However, I think it might not fit your use case well, as it's primarily just for very large datasets (it's slow to make the pools).
Increasing the batch size might also work, indeed; but I'm not sure if that works well with CPU (it certainly does with GPU).
Another consideration is using the Intel Neural Compressor, e.g. with optimum-intel, but I am not very aware of this & I'm not sure if it would work on the CPU that you mentioned.
- Tom Aarsen
krumeto
snayan06
Hey team,
I am looking for ways to speed a cross-encoder (
cross-encoder/ms-marco-MiniLM-L-6-v2) on a CPU. For a variety of reasons, GPUs are off-the-table for now.We have the following setup:
sentence-transformers==2.2.2The model gets a short query and a list of 60 - 80 texts, typically above the 512 max_tokens (getting truncated).
Do you have any recommendations for speed-up?
Example ideas that we had:
batch_sizefrom the default of 32 to our expected texts (for example, 80)onnx- for example, based on https://www.philschmid.de/optimize-sentence-transformers or this model - https://huggingface.co/metarank/ce-msmarco-MiniLM-L6-v2Any ideas would be more than welcome. Thank you in advance!