diff --git a/src/deepsparse/sentence_transformers/__init__.py b/src/deepsparse/sentence_transformers/__init__.py index 338a9f5dba..44c53314c5 100644 --- a/src/deepsparse/sentence_transformers/__init__.py +++ b/src/deepsparse/sentence_transformers/__init__.py @@ -33,4 +33,4 @@ ) -from .sentence_transformer import SentenceTransformer +from .sentence_transformer import DeepSparseSentenceTransformer, SentenceTransformer diff --git a/src/deepsparse/sentence_transformers/sentence_transformer.py b/src/deepsparse/sentence_transformers/sentence_transformer.py index c3454148d5..5e2dbd0855 100644 --- a/src/deepsparse/sentence_transformers/sentence_transformer.py +++ b/src/deepsparse/sentence_transformers/sentence_transformer.py @@ -28,7 +28,7 @@ DEFAULT_MODEL_NAME = "zeroshot/bge-small-en-v1.5-quant" -class SentenceTransformer: +class DeepSparseSentenceTransformer: """ Loads or creates a SentenceTransformer-compatible model that can be used to map text to embeddings. @@ -289,3 +289,6 @@ def mean_pooling( return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( input_mask_expanded.sum(1), min=1e-9 ) + +# for backwards compatibility +SentenceTransformer = DeepSparseSentenceTransformer