handler.py

from typing import  Dict, List, Any
# from optimum.onnxruntime import ORTModelForSequenceClassification
# from transformers import pipeline, AutoTokenizer
from FlagEmbedding import BGEM3FlagModel


class EndpointHandler():
    def __init__(self, path="./"):
        # load the optimized model
        # モデルの準備
        self.model = BGEM3FlagModel("./", use_fp16=False)


    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
        """
        Args:
            data (:obj:):
                includes the input data and the parameters for the inference.
        Return:
            A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
                - "label": A string representing what the label/class is. There can be multiple labels.
                - "score": A score between 0 and 1 describing how confident the model is for this label/class.
        """
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None)
        
        sparse_embs = []
        
        # result = self.model.encode(inputs, return_dense=False, return_sparse=True)
        # sparse_vectors = result["lexical_weights"]
        # for sparse_vector in sparse_vectors:
        #     sparse_values = [value for value in sparse_vector.values()]
        #     sparse_dimensions = [int(key) for key in sparse_vector.keys()]
        #     sparse_embs.append((sparse_values, sparse_dimensions))

        # pass inputs with all kwargs in data
        # if parameters is not None:
        #     prediction = self.pipeline(inputs, **parameters)
        # else:
        #     prediction = self.pipeline(inputs)
        # postprocess the prediction
        return [
            [
                {"label": 0, "score": 0.5}
            ]
        ]