from typing import Dict, List, Any from InstructorEmbedding import INSTRUCTOR import torch class EndpointHandler(): def __init__(self, path=""): model = INSTRUCTOR(path) self.model = model if torch.cuda.is_available(): self.device = torch.device("cuda") self.model.to(self.device) else: self.device = torch.device("cpu") def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: """ data args: inputs (:obj: `str`) date (:obj: `str`) Return: A :obj:`list` | `dict`: will be serialized and returned """ # get inputs instruction = data.pop("instruction",data) text = data.pop("text", data) inputs = [[s, instruction] for s in text] if (self.device): inputs = torch.tensor(inputs).to(self.device) # Move inputs to the GPU embeddings = self.model.encode(inputs) return embeddings.tolist()