from transformers import AutoProcessor, SeamlessM4Tv2ForSpeechToText
import torchaudio
from torch import Tensor
from base64 import b64encode


class EndpointHandler():
    def __init__(self, path=""):
        self.path = path
        self.processor = AutoProcessor.from_pretrained(
            "facebook/seamless-m4t-v2-large")
        self.model = SeamlessM4Tv2ForSpeechToText.from_pretrained(
            "facebook/seamless-m4t-v2-large")
        self.AUDIO_DEFAULT_SAMPLE_RATE = 16000

    def __call__(self, data: dict):
        input_audio_tensor = data["inputs"]["audio"]
        input_lang_code = data["inputs"]["input_lang_code"]
        output_lang_code = data["inputs"]["output_lang_code"]
        audio_inputs = self.processor(
            audios=input_audio_tensor, src_lang=input_lang_code, return_tensors="pt")
        id_array = self.model.generate(
            **audio_inputs, tgt_lang=output_lang_code)
        text_array_from_audio = self.processor.batch_decode(
            id_array, skip_special_tokens=True)
        return [{"text_array": text_array_from_audio}]