Hans Elias J commited on
Commit
e66ef1e
·
1 Parent(s): ee4ad34

return embeddings list directly

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -16,7 +16,7 @@ class EndpointHandler():
16
  self.tokenizer = AutoTokenizer.from_pretrained(path)
17
  self.model = AutoModel.from_pretrained(path)
18
 
19
- def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
20
  """
21
  data args:
22
  inputs (:obj: `List[str]`)
@@ -32,4 +32,4 @@ class EndpointHandler():
32
 
33
  embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask']).tolist()
34
 
35
- return {"embeddings": embeddings}
 
16
  self.tokenizer = AutoTokenizer.from_pretrained(path)
17
  self.model = AutoModel.from_pretrained(path)
18
 
19
+ def __call__(self, data: Dict[str, Any]) -> List[List[int]]:
20
  """
21
  data args:
22
  inputs (:obj: `List[str]`)
 
32
 
33
  embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask']).tolist()
34
 
35
+ return embeddings