dabitbol/bge-m3-sparse-elastic

1from typing import Dict, List, Any2from FlagEmbedding import BGEM3FlagModel3 4class EndpointHandler():5    def __init__(self, path=""):6        self.model = BGEM3FlagModel('BAAI/bge-m3',  use_fp16=True)7 8    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:9        """10       data args:11            inputs (:`list`: `str`)12            kwargs13      Return:14            A :obj:`list` | `dict`: will be serialized and returned15        """16 17        results = []18        inputs = data.pop("inputs",data)19        for i in inputs:20            output = self.model.encode(i, return_dense=False, return_sparse=True, return_colbert_vecs=False)21            results.append(self._toJsonSerialisableFormat(self.model.convert_id_to_token(output['lexical_weights'])))22        23        return results24 25    def _toJsonSerialisableFormat(self, model_output):26        # convert the numpy float16 to a Python float27        # so it can be serialised as JSON28        return {key:float(value) for (key,value) in model_output.items()}29