Add alias for DeepSparseSentenceTransformer

neuralmagic · Oct 26, 2023 · 8ef1f86 · 8ef1f86
1 parent 1904691
commit 8ef1f86
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git 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
@@ -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
Original file line number	Diff line number	Diff line change
Expand Up		@@ -33,4 +33,4 @@
		)


		from .sentence_transformer import SentenceTransformer
		from .sentence_transformer import DeepSparseSentenceTransformer, SentenceTransformer