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)

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

        instruction = 'Retrieve documents that can help answer the question:'

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

        return embeddings.tolist()