|
from typing import Dict, List, Any |
|
from FlagEmbedding import BGEM3FlagModel |
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
self.model = BGEM3FlagModel(path, use_fp16=True) |
|
|
|
|
|
|
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
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 vector |
|
""" |
|
inputs = data.pop("inputs", data) |
|
parameters = data.pop("parameters", None) |
|
|
|
|
|
if parameters is not None: |
|
embeddings = self.model.encode(inputs, **parameters) |
|
else: |
|
embeddings = self.model.encode(inputs) |
|
|
|
return embeddings |
|
|
|
|
|
|