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 transcribe(audio):
    if audio is None:
        return "No audio input provided. Please record or upload an audio file."

    # audio is now a file path, not a tuple
    try:
        array, sample_rate = librosa.load(audio, sr=16000)
    except Exception as e:
        return f"Error loading audio file: {str(e)}"

    # The rest of the function remains the same
    array = librosa.to_mono(array)
    input_features = processor(array, sampling_rate=sample_rate, return_tensors="pt").input_features

    # generate token ids
    predicted_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
    # decode token ids to text
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
    print(transcription)
    return transcription[0]  # Return the first (and only) transcription


demo = gr.Interface(
    fn=transcribe,
    inputs=[gr.Audio(sources=["microphone"], type="filepath")],
    outputs="text"
)

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