jpohhhh commited on
Commit
a154b3b
1 Parent(s): 2b69a61

Return embeddings as array instead of nested in JSON

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -17,7 +17,7 @@ class EndpointHandler():
17
  self.model = ORTModelForFeatureExtraction.from_pretrained(path, file_name="model-quantized.onnx")
18
  self.tokenizer = AutoTokenizer.from_pretrained(path)
19
 
20
- def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
21
  """
22
  Args:
23
  data (:obj:):
@@ -36,4 +36,4 @@ class EndpointHandler():
36
  # Normalize embeddings
37
  sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
38
  # postprocess the prediction
39
- return {"embeddings": sentence_embeddings.tolist()}
 
17
  self.model = ORTModelForFeatureExtraction.from_pretrained(path, file_name="model-quantized.onnx")
18
  self.tokenizer = AutoTokenizer.from_pretrained(path)
19
 
20
+ def __call__(self, data: Any) -> List[List[float]]:
21
  """
22
  Args:
23
  data (:obj:):
 
36
  # Normalize embeddings
37
  sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
38
  # postprocess the prediction
39
+ return sentence_embeddings.tolist()