misc(embeddings): clean up

handler.py

@@ -1,16 +1,18 @@
+import os
 import platform
 from typing import Union, Sequence, Sized
 
 import torch
-from loguru import logger
 from hfendpoints.openai import Context, run
 from hfendpoints.openai.embeddings import Embedding, EmbeddingEndpoint, EmbeddingRequest, EmbeddingResponse, Usage
-from sentence_transformers import SentenceTransformer
-
 from hfendpoints import EndpointConfig, Handler, __version__
+from loguru import logger
 from torch.backends.mkldnn import VERBOSE_ON_CREATION, VERBOSE_OFF
+from sentence_transformers import SentenceTransformer
 
-
+# Not used for now
+ENABLE_QUANTIZATION = bool(os.environ.get("HFENDPOINT_ENABLE_QUANTIZATION", "0"))
+SUPPORTED_AMP_DTYPES = {torch.float32, torch.bfloat16}
 
 def get_usage(tokens: Union[Sized, Sequence[Sized]], is_batched: bool) -> Usage:
     """
@@ -26,33 +28,39 @@ def get_usage(tokens: Union[Sized, Sequence[Sized]], is_batched: bool) -> Usage:
 
     return Usage(prompt_tokens=num_tokens, total_tokens=num_tokens)
 
-
 class SentenceTransformerHandler(Handler):
-    __slots__ = ("_config", "_model", "_model_config", "_model_name")
+    __slots__ = ("_config", "_dtype", "_model", "_model_name", "_use_amp")
 
     def __init__(self, config: EndpointConfig):
         self._config = config
-        self._model = SentenceTransformer(config.model_id, device="cpu", model_kwargs={"torch_dtype": "bfloat16"})
-        self._model_config = self._model._modules['0'].auto_model.config
+        self._dtype = torch.float32
         self._model_name = config.model_id
 
-        self._post_init()
+        self._allocate_model()
 
-    def _post_init(self):
-        self._model = self._model.to(memory_format=torch.channels_last)
+    def _allocate_model(self):
+        dtype = torch.bfloat16 if torch.cpu._is_avx512_bf16_supported() else torch.float32
+        model = SentenceTransformer(self._config.model_id, device="cpu", model_kwargs={"torch_dtype": dtype})
 
-        if "Intel" in platform.processor():
+        if platform.machine() == "x86_64":
             import intel_extension_for_pytorch as ipex
+            logger.info(f"x64 platform detected: {platform.processor()}")
+
             with torch.inference_mode():
-                self._model = self._model.eval()
-                self._model = ipex.optimize(self._model, dtype=torch.float32, weights_prepack=False)
-                self._model = torch.compile(self._model, backend="ipex")
+                model = model.eval()
+                model = model.to(memory_format=torch.channels_last)
+                model = ipex.optimize(model, dtype=dtype, weights_prepack=False, graph_mode=True, concat_linear=True)
+                model = torch.compile(model, dynamic=True, backend="ipex")
         else:
-            self._model = torch.compile(self._model)
+            model = torch.compile(model)
+
+        self._model = model
+        self._dtype = dtype
+        self._use_amp = dtype in SUPPORTED_AMP_DTYPES
 
     async def __call__(self, request: EmbeddingRequest, ctx: Context) -> EmbeddingResponse:
         with torch.backends.mkldnn.verbose(VERBOSE_ON_CREATION if self._config.is_debug else VERBOSE_OFF):
-            with torch.inference_mode(), torch.amp.autocast("cpu", dtype=torch.float32):
+            with torch.inference_mode(), torch.amp.autocast("cpu", dtype=self._dtype, enabled=self._use_amp):
                 tokens = self._model.tokenize(request.input)
                 vectors = self._model.encode(request.input)
 

requirements.txt

@@ -1,3 +1,5 @@
+# auto-round>=0.5.0
 loguru>=0.7.3
 sentence-transformers
+# transformers>=4.51
 torch>=2.5.0