handler.py

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


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

    def __call__(self, data: Dict[str, Any]) -> List[List[float]]:
        """
         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 not texts or not instruction:
            raise ValueError("Please provide texts and instruction")

        # make sure texts is a list
        if not isinstance(texts, list):
            texts = [texts]

        instructions = [[instruction, text] for text in texts]
        embeddings = self.model.encode(instructions)

        return embeddings.tolist()