feat(embedding): move to common I/O representation

handler.py

@@ -1,10 +1,12 @@
 import platform
-from operator import itemgetter
-from typing import Generator, Tuple
+from typing import Tuple, List
 
+import numpy as np
 import torch
 from hfendpoints.http import Context, run
-from hfendpoints.http.embeddings import Embedding, EmbeddingEndpoint, EmbeddingRequest, EmbeddingResponse, Usage
+from hfendpoints.openai.embedding import EmbeddingEndpoint
+from hfendpoints.tasks import Usage
+from hfendpoints.tasks.embedding import EmbeddingRequest, EmbeddingResponse
 from intel_extension_for_pytorch.cpu.runtime import pin
 from loguru import logger
 from sentence_transformers import SentenceTransformer
@@ -45,7 +47,6 @@ def get_cores_pinning_strategy() -> "CPUPool":
 
     logger.info(f"Pinning CPU cores to {pinned_cpu_cores_id}")
     return ipex.cpu.runtime.CPUPool(pinned_cpu_cores_id)
-    # return ipex.cpu.runtime.CPUPool(node_id=0)
 
 
 def get_usage(mask: torch.IntTensor) -> Usage:
@@ -54,7 +55,7 @@ def get_usage(mask: torch.IntTensor) -> Usage:
     :param mask: Attention mask tensor, as returned by the model
     :return: Usage object matching OpenAI specifications
     """
-    num_tokens = sum(m.sum().item() for m in mask)
+    num_tokens = np.sum(mask)
     return Usage(prompt_tokens=num_tokens, total_tokens=num_tokens)
 
 
@@ -65,9 +66,12 @@ class SentenceTransformerWithUsage(Module):
         super().__init__()
         self._model = model
 
-    def forward(self, sentences: list[str]) -> Tuple[Generator[torch.Tensor], Generator[torch.Tensor]]:
+    def forward(self, sentences: list[str]) -> Tuple[List[List[int]], List[List[int]]]:
         vectors = self._model.encode(sentences, output_value=None)
-        return map(itemgetter('attention_mask'), vectors), map(itemgetter('sentence_embedding'), vectors)
+        return (
+            [vector['attention_mask'].tolist() for vector in vectors],
+            [vector['sentence_embedding'].tolist() for vector in vectors]
+        )
 
 
 class SentenceTransformerHandler(Handler):
@@ -119,13 +123,9 @@ class SentenceTransformerHandler(Handler):
                 with pin(self._pinned_cores):
                     mask, vectors = self._model(request.input if request.is_batched else [request.input])
 
-                embeddings = [None] * len(request)
-                for (index, embedding) in enumerate(vectors):
-                    embedding = Embedding(index=index, embedding=embedding.tolist())
-                    embeddings[index] = embedding
-
+            # TODO: Change the way we return usage
             usage = get_usage(mask)
-            return EmbeddingResponse(model=self._model_name, embeddings=embeddings, usage=usage)
+            return EmbeddingResponse(embeddings=vectors, num_tokens=usage.total_tokens)
 
 
 def entrypoint():