feat(embeddings): expose some more to Python and return corresponding embedding (with copy for now)

handler.py

@@ -1,4 +1,5 @@
 import platform
+from typing import Union, Sequence, Sized
 
 import torch
 from loguru import logger
@@ -7,6 +8,23 @@ from hfendpoints.openai.embeddings import Embedding, EmbeddingEndpoint, Embeddin
 from sentence_transformers import SentenceTransformer
 
 from hfendpoints import EndpointConfig, Handler, __version__
+from torch.backends.mkldnn import VERBOSE_ON_CREATION, VERBOSE_OFF
+
+
+
+def get_usage(tokens: Union[Sized, Sequence[Sized]], is_batched: bool) -> Usage:
+    """
+    Compute the number of processed tokens and return as Usage object matching OpenAI
+    :param tokens: List or nested List of tokens
+    :param is_batched: Flag indicating if the original request contained batched inputs
+    :return: Usage object matching OpenAI specifications
+    """
+    if is_batched:
+        num_tokens = sum(len(document) for document in tokens)
+    else:
+        num_tokens = len(tokens)
+
+    return Usage(prompt_tokens=num_tokens, total_tokens=num_tokens)
 
 
 class SentenceTransformerHandler(Handler):
@@ -32,21 +50,22 @@ class SentenceTransformerHandler(Handler):
         else:
             self._model = torch.compile(self._model)
 
-    @torch.compile
-    def forward(self, documents: str):
-        # TODO: Ask Tom how to do this better without tokenizing twice?
-        tokens = self._model.tokenize(documents)
-        vectors = self._model.encode(documents, output_value="sentence_embedding", normalize_embeddings=True)
-
-        return tokens, vectors
-
     async def __call__(self, request: EmbeddingRequest, ctx: Context) -> EmbeddingResponse:
-        with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON_CREATION):
+        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):
+                tokens = self._model.tokenize(request.input)
                 vectors = self._model.encode(request.input)
-                embedding = Embedding(index=0, embedding=vectors.tolist())
-                usage = Usage(prompt_tokens=len(request.input), total_tokens=len(request.input))
-                return EmbeddingResponse(model=self._model_name, embeddings=[embedding], usage=usage)
+
+                embeddings = [[None] * len(request)]
+                if not request.is_batched:
+                    embeddings[0] = Embedding(index=0, embedding=vectors.tolist())
+                else:
+                    for (index, embedding) in enumerate(vectors.tolist()):
+                        embedding = Embedding(index=index, embedding=embedding)
+                        embeddings[index] = embedding
+
+            usage = get_usage(tokens, request.is_batched)
+            return EmbeddingResponse(model=self._model_name, embeddings=embeddings, usage=usage)
 
 
 def entrypoint():