support mps backend.

Author: ninehills

.gitignore

@@ -162,3 +162,6 @@ cython_debug/
 .vscode
 .infinity_cache
 libs/infinity_emb/data/*
+
+# macOS
+.DS_Store

README.md

@@ -24,7 +24,7 @@ Infinity is a high-throughput, low-latency REST API for serving vector embedding
 ## Why Infinity:
 Infinity provides the following features:
 - **Deploy virtually any SentenceTransformer** - deploy the model you know from [SentenceTransformers](https://github.com/UKPLab/sentence-transformers/)
-- **Fast inference backends**: The inference server is built on top of [torch](https://github.com/pytorch/pytorch), [fastembed(onnx-cpu)](https://github.com/qdrant/fastembed) and [CTranslate2](https://github.com/OpenNMT/CTranslate2), getting most out of your **CUDA** or **CPU** hardware.
+- **Fast inference backends**: The inference server is built on top of [torch](https://github.com/pytorch/pytorch), [fastembed(onnx-cpu)](https://github.com/qdrant/fastembed) and [CTranslate2](https://github.com/OpenNMT/CTranslate2), getting most out of your **CUDA**, **CPU** or **MPS** hardware.
 - **Dynamic batching**: New embedding requests are queued while GPU is busy with the previous ones. New requests are squeezed intro your GPU/CPU as soon as ready. 
 - **Correct and tested implementation**: Unit and end-to-end tested. Embeddings via infinity are identical to [SentenceTransformers](https://github.com/UKPLab/sentence-transformers/) (up to numerical precision). Lets API users create embeddings till infinity and beyond.
 - **Easy to use**: The API is built on top of [FastAPI](https://fastapi.tiangolo.com/), [Swagger](https://swagger.io/) makes it fully documented. API are aligned to [OpenAI's Embedding specs](https://platform.openai.com/docs/guides/embeddings/what-are-embeddings). See below on how to get started.

libs/infinity_emb/infinity_emb/primitives.py

@@ -13,6 +13,7 @@
 class Device(enum.Enum):
     cpu = "cpu"
     cuda = "cuda"
+    mps = "mps"
     auto = None
 
 

libs/infinity_emb/infinity_emb/transformer/classifier/torch.py

@@ -31,7 +31,12 @@ def __init__(self, model_name_or_path, device: Optional[str] = None) -> None:
             top_k=None,
             torch_dtype=torch.float32 if used_device == "cpu" else torch.float16,
         )
-        self._pipe.model = to_bettertransformer(self._pipe.model, logger=logger)
+        if used_device == "mps":
+            logger.info(
+                "Disable Optimizations via Huggingface optimum for MPS Backend. "
+            )
+        else:
+            self._pipe.model = to_bettertransformer(self._pipe.model, logger)
 
         self._infinity_tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
 

libs/infinity_emb/infinity_emb/transformer/crossencoder/torch.py

@@ -47,7 +47,12 @@ def __init__(self, model_name_or_path, **kwargs):
         self._infinity_tokenizer = copy.deepcopy(self.tokenizer)
         self.model.eval()
 
-        self.model = to_bettertransformer(self.model, logger)
+        if self._target_device.type == "mps":
+            logger.info(
+                "Disable Optimizations via Huggingface optimum for MPS Backend. "
+            )
+        else:
+            self.model = to_bettertransformer(self.model, logger)
 
         if self._target_device.type == "cuda" and not os.environ.get(
             "INFINITY_DISABLE_HALF", ""

libs/infinity_emb/infinity_emb/transformer/embedder/sentence_transformer.py

@@ -54,7 +54,12 @@ def __init__(self, model_name_or_path, **kwargs):
         self._infinity_tokenizer = copy.deepcopy(fm.tokenizer)
         self.eval()
 
-        fm.auto_model = to_bettertransformer(fm.auto_model, logger)
+        if self._target_device.type == "mps":
+            logger.info(
+                "Disable Optimizations via Huggingface optimum for MPS Backend. "
+            )
+        else:
+            fm.auto_model = to_bettertransformer(fm.auto_model, logger)
 
         if self._target_device.type == "cuda" and not os.environ.get(
             "INFINITY_DISABLE_HALF", ""