Use SentenceTransformers instead

Per https://huggingface.co./sentence-transformers/multi-qa-MiniLM-L6-cos-v1/blob/main/README.md

handler.py

@@ -1,17 +1,8 @@
-from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModel
-import torch
-
-#Mean Pooling - Take attention mask into account for correct averaging
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+from sentence_transformers import SentenceTransformer, util
     
 class EndpointHandler():
     def __init__(self, path=""):
-        self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
-        self.model = AutoModel.from_pretrained('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
+        self.model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -22,11 +13,5 @@ class EndpointHandler():
             A :obj:`list` | `dict`: will be serialized and returned
         """
         sentences = data.pop("inputs",data)
-        encoded_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
-        # Compute token embeddings
-        with torch.no_grad():
-            model_output = self.model(**encoded_input)
-
-        # Perform pooling. In this case, max pooling.
-        sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
-        return sentence_embeddings.tolist()
+        embeddings = model.encode(sentences)
+        return embeddings.tolist()