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")
        logging.info('model loaded')

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

        logging.info(f'audio_url {audio_url}')

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