fix: wrong usage computation

handler.py

@@ -1,7 +1,6 @@
 import platform
 from typing import Tuple, List
 
-import numpy as np
 import torch
 from hfendpoints.http import Context, run
 from hfendpoints.tasks import Usage
@@ -48,13 +47,13 @@ def get_cores_pinning_strategy() -> "CPUPool":
     return ipex.cpu.runtime.CPUPool(pinned_cpu_cores_id)
 
 
-def get_usage(mask: torch.IntTensor) -> Usage:
+def get_usage(mask: List[torch.IntTensor]) -> Usage:
     """
     Compute the number of processed tokens and return as Usage object matching OpenAI
     :param mask: Attention mask tensor, as returned by the model
     :return: Usage object matching OpenAI specifications
     """
-    num_tokens = np.sum(mask)
+    num_tokens = sum(x.sum().detach().item() for x in mask)
     return Usage(prompt_tokens=num_tokens, total_tokens=num_tokens)
 
 
@@ -68,7 +67,7 @@ class SentenceTransformerWithUsage(Module):
     def forward(self, sentences: list[str]) -> Tuple[List[List[int]], List[List[int]]]:
         vectors = self._model.encode(sentences, output_value=None)
         return (
-            [vector['attention_mask'].tolist() for vector in vectors],
+            [vector['attention_mask'] for vector in vectors],
             [vector['sentence_embedding'].tolist() for vector in vectors]
         )
 
@@ -124,6 +123,7 @@ class SentenceTransformerHandler(Handler):
 
             # TODO: Change the way we return usage
             usage = get_usage(mask)
+            vectors = vectors if request.is_batched else vectors[0]
             return EmbeddingResponse(embeddings=vectors, num_tokens=usage.total_tokens)
 
 
@@ -137,7 +137,7 @@ def entrypoint():
     handler = SentenceTransformerHandler(config)
 
     # Allocate endpoint
-    from hfendpoints.hfinference.embedding import EmbeddingEndpoint
+    from hfendpoints.openai.embedding import EmbeddingEndpoint
     endpoint = EmbeddingEndpoint(handler)
     run(endpoint, config.interface, config.port)