Update handler.py

#11
by ababeal - opened
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -24,6 +24,6 @@ class EndpointHandler():
24
  outputs = self.model(**batch_dict)
25
 
26
  embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
27
- embeddings = F.normalize(embeddings, p=2, dim=1).tolist()
28
-
29
- return embeddings
 
24
  outputs = self.model(**batch_dict)
25
 
26
  embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
27
+ embeddings = F.normalize(embeddings, p=2, dim=1)
28
+ scores = (embeddings[:2] @ embeddings[2:].T) * 100
29
+ return scores.tolist()