from typing import Dict, List, Any
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np


class EndpointHandler():
    def __init__(self, path=""):
        self.model = hub.load(path)

    def __call__(self, data: Dict[str, Any]):
        inputs = data.pop("inputs")

        tf_inputs = tf.constant([inputs])
        
        embeddings = self.model(tf_inputs).numpy()

        return {"embeddings" : embeddings}