neuralmagic / deepsparse

Sparsity-aware deep learning inference runtime for CPUs
https://neuralmagic.com/deepsparse/
Other
3.01k stars 176 forks source link

Support for SentenceTransformers with `deepsparse.sentence_transformers.SentenceTransformer` #1301

Closed mgoin closed 1 year ago

mgoin commented 1 year ago

Direct replacement for SentenceTransformer implemented using DeepSparse and optimum-deepsparse

Performance optimizations with bucketing/batching will come in future work

Smoke Test:

pip install -e ".[sentence_transformers]"
# from sentence_transformers import SentenceTransformer
from deepsparse.sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", export=True)

# Sentences we'd like to encode
sentences = ['This framework generates embeddings for each input sentence',
    'Sentences are passed as a list of string.',
    'The quick brown fox jumps over the lazy dog.']

# Sentences are encoded by calling model.encode()
embeddings = model.encode(sentences)

# Print the embeddings
for sentence, embedding in zip(sentences, embeddings):
    print("Sentence:", sentence)
    print("Embedding:", embedding.shape)
    print("")

### Sentence: This framework generates embeddings for each input sentence
### Embedding: (384,)
### 
### Sentence: Sentences are passed as a list of string.
### Embedding: (384,)
### 
### Sentence: The quick brown fox jumps over the lazy dog.
### Embedding: (384,)

Full test through evaluation with MTEB

pip install mteb deepsparse-nightly[sentence_transformers] sentence-transformers
from mteb import MTEB

model_name = "TaylorAI/bge-micro-v2"

from deepsparse.sentence_transformers import SentenceTransformer
model = SentenceTransformer(model_name, export=True)
evaluation = MTEB(tasks=["Banking77Classification"])
results_ds = evaluation.run(model, output_folder=f"results/ds-{model_name}")
print(results_ds)

import sentence_transformers
model = sentence_transformers.SentenceTransformer(model_name)
evaluation = MTEB(tasks=["Banking77Classification"])
results_st = evaluation.run(model, output_folder=f"results/st-{model_name}")
print(results_st)

Output:

{'Banking77Classification': {'mteb_version': '1.1.1', 'dataset_revision': '0fd18e25b25c072e09e0d92ab615fda904d66300', 'mteb_dataset_name': 'Banking77Classification', 'test': {'accuracy': 0.8117207792207791, 'f1': 0.8109893836310513, 'accuracy_stderr': 0.007164150669501205, 'f1_stderr': 0.007346045502756079, 'main_score': 0.8117207792207791, 'evaluation_time': 26.05}}}

{'Banking77Classification': {'mteb_version': '1.1.1', 'dataset_revision': '0fd18e25b25c072e09e0d92ab615fda904d66300', 'mteb_dataset_name': 'Banking77Classification', 'test': {'accuracy': 0.8117207792207791, 'f1': 0.8109893836310513, 'accuracy_stderr': 0.007164150669501205, 'f1_stderr': 0.007346045502756079, 'main_score': 0.8117207792207791, 'evaluation_time': 12.21}}}