import logging
from speechbrain.pretrained import EncoderClassifier
from typing import Dict, List, Any


class EndpointHandler:
    def __init__(self, path=""):
        self.model = EncoderClassifier.from_hparams("speechbrain/lang-id-voxlingua107-ecapa")
        print('model loaded')
        logging.info('model loaded')

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

        print('audio_url', inputs)
        logging.info(f'audio_url {inputs}')

        # run normal prediction
        output = self.model.classify_file(inputs)
        return {
            "prediction": float(output[1].exp()[0]),
            "language": output[3][0],
        }