handler.py

from transformers import pipeline
import numpy as np 

class EndpointHandler():
    def __init__(self, path=""):
        self.pipeline = pipeline("sentence-embeddings",model=path)

    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
        """

        def normalize_vectors(vectors):
            """
            Normalize a single vector or a list of vectors.
        
            Parameters:
                vectors (Union[np.ndarray, List[np.ndarray]]): Input vector or list of vectors.
        
            Returns:
                Union[np.ndarray, List[np.ndarray]]: Normalized vector or list of normalized vectors.
            """
            if isinstance(vectors, np.ndarray):
                # If it's a single vector, normalize it
                return vectors / np.linalg.norm(vectors)
            elif isinstance(vectors, list):
                # If it's a list of vectors, normalize each vector in the list
                return [vector / np.linalg.norm(vector) for vector in vectors]
            else:
                raise ValueError("Input must be a numpy array or a list of numpy arrays.")
        
        inputs = data.pop("inputs",data)
        prediction = self.pipeline(inputs)
        normalized_prediction = normalize_vectors(prediction)
        
        return {"embeddings":  normalized_prediction}