import gradio as gr
from transformers import WhisperProcessor, WhisperForConditionalGeneration
import librosa

processor = WhisperProcessor.from_pretrained("Neurai/NeuraSpeech_WhisperBase")
model = WhisperForConditionalGeneration.from_pretrained("Neurai/NeuraSpeech_WhisperBase")
forced_decoder_ids = processor.get_decoder_prompt_ids(language="fa", task="transcribe")


def reverse_audio(audio):
    array, sample_rate = audio
    sr = 16000
    array = librosa.to_mono(array)
    array = librosa.resample(array, orig_sr=sample_rate, target_sr=16000)
    input_features = processor(array, sampling_rate=sr, return_tensors="pt").input_features

    # generate token ids
    predicted_ids = model.generate(input_features)
    # decode token ids to text
    transcription = processor.batch_decode(predicted_ids, )
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
    return transcription


input_audio = gr.Audio(
    sources=["microphone"],
    waveform_options=gr.WaveformOptions(
        waveform_color="#01C6FF",
        waveform_progress_color="#0066B4",
        skip_length=2,
        show_controls=True,
    ),
)
demo = gr.Interface(
    fn=reverse_audio,
    inputs=input_audio,
    outputs="text"
)

if __name__ == "__main__":
    demo.launch()