handler.py

from typing import Dict, List, Any
from InstructorEmbedding import INSTRUCTOR


INSTRUCTION_SEPARATOR = "|||"


class EndpointHandler:
    def __init__(self, path=""):
        # load model
        self.model = INSTRUCTOR(path, device="cuda")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str`)
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """
        # get inputs
        inputs: dict = data.pop("inputs", data)
        texts = inputs.pop("texts", None)
        instruction = inputs.pop("instruction", None)

        # if isinstance(inputs, str):
        #     inputs = [inputs]

        # run normal prediction
        # scores = self.model.predict_proba(inputs)[0]

        # return [{"label": self.id2label[i], "score": score.item()} for i, score in enumerate(scores)]
        instructions = [[instruction, text] for text in texts]
        embeddings = self.model.encode(instructions)
        return embeddings.tolist()