from typing import Dict, List, Any
from transformers import ClapModel, ClapProcessor
import gc

class EndpointHandler:
    def __init__(self, path=""):
        model_name = "laion/larger_clap_general"
        self.model = ClapModel.from_pretrained(model_name)
        self.processor = ClapProcessor.from_pretrained(model_name)

    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
        """
        print(type(data))
        query = data['inputs']
        text_inputs = self.processor(text=query, return_tensors="pt")
        text_embed = self.model.get_text_features(**text_inputs)[0]
        gc.collect()
        return text_embed