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)

        print(inputs)

        # add instruction
        query  = [['Retrieve documents that can help answer the question:',
                   inputs]]

        # encode
        embedding = self.model.encode(query)

        return embedding