|
import torch |
|
from typing import Dict |
|
from sentence_transformers import SentenceTransformer |
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
self.model = SentenceTransformer("all-MiniLM-L6-v2") |
|
|
|
|
|
def __call__(self, data: Dict[str, List[str]]) -> Dict[str, List[List[float]]]: |
|
""" |
|
Args: |
|
data (:obj:): |
|
includes the deserialized sentences as a list of strings |
|
Return: |
|
A :obj:`dict`: list of embeddings for each input sentence |
|
""" |
|
|
|
inputs = data.pop("inputs", data) |
|
|
|
|
|
embeddings = self.model.encode(inputs) |
|
|
|
|
|
return {"embeddings": embeddings.tolist()} |