from sentence_transformers import SentenceTransformer, util | |
class EndpointHandler(): | |
def __init__(self, path=""): | |
self.model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1') | |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
""" | |
data args: | |
inputs (:obj: `str` | `PIL.Image` | `np.array`) | |
kwargs | |
Return: | |
A :obj:`list` | `dict`: will be serialized and returned | |
""" | |
sentences = data.pop("inputs",data) | |
embeddings = model.encode(sentences) | |
return embeddings.tolist() |