| from sentence_transformers import SentenceTransformer | |
| from typing import Dict, List, Any, Union | |
| class EndpointHandler: | |
| def __init__(self, model_path="bge-large-en/"): | |
| # Preload all the elements you are going to need at inference | |
| self.model = SentenceTransformer(model_path) | |
| def __call__(self, data: Dict[str, Any]) -> Union[List[List[float]], List[float]]: | |
| """ | |
| data args: | |
| inputs (:obj: `str` | `PIL.Image` | `np.array`) | |
| kwargs | |
| Return: | |
| A :obj:`list` | `dict`: will be serialized and returned | |
| """ | |
| # Extracting the inputs and kwargs | |
| inputs = data["inputs"] | |
| kwargs = data.get("kwargs", {}) | |
| normalize_embeddings = kwargs.get('normalize_embeddings', True) | |
| # Determine if the input is a query or a passage | |
| is_query = kwargs.get("is_query", False) | |
| if is_query: | |
| instruction = kwargs.get("query_instruction", "") | |
| if isinstance(inputs, list): | |
| inputs = [instruction + q for q in inputs] | |
| else: | |
| inputs = instruction + inputs | |
| # Encoding the inputs using the model | |
| embeddings = self.model.encode(inputs, normalize_embeddings=normalize_embeddings) | |
| # Return the serialized embeddings | |
| return embeddings.tolist() if isinstance(embeddings, list) else embeddings | |