import torch
import gradio as gr
from transformers import pipeline

MODEL_NAME = "BlueRaccoon/whisper-small-kab"
lang = "uz"  # Used uz instead of kab

device = 0 if torch.cuda.is_available() else "cpu"
pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")


def transcribe(microphone):
    if microphone is None:
        return "ERROR: You need to record or upload an audio file."

    text = pipe(microphone)["text"]
    return text


with gr.Blocks() as demo:
    with gr.Tab("Transcribe Kabyle Audio"):
        gr.Markdown(
            f"""
            # Kabyle Whisper Demo: Transcribe Audio
            Transcribe Kabyle audio recorded from the microphone or uploaded as a file. This demo uses the fine-tuned
            checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe Kabyle audio
            files of arbitrary length.
            """
        )
        # Input for microphone recording only
        microphone_input = gr.Audio(type="filepath", label="Record or Upload Kabyle Audio")
        gr.Interface(
            fn=transcribe,
            inputs=[microphone_input],
            outputs=gr.Textbox(label="Transcription"),
        )

demo.launch()