from transformers import pipeline
import gradio as gr

# Load the fine-tuned Whisper model
pipe = pipeline(model="fawzanaramam/Whisper-Small-Finetuned-on-Surah-Fatiha")


def transcribe(audio):
    """Transcribes the given audio using Whisper"""
    return pipe(audio)["text"]


def save_transcription(text):
    """Creates a downloadable text file for the transcription"""
    file_path = "transcription.txt"
    with open(file_path, "w") as f:
        f.write(text)
    return file_path


# UI Layout
with gr.Blocks(
    css="""
    h1 { 
        font-size: 32px; 
        font-weight: bold;
        text-align: center;
    }
    h2, h3 {
        font-size: 26px; 
        font-weight: bold;
    }
    .gr-button {
        font-size: 18px;
        font-weight: bold;
    }
    .gr-textbox {
        font-size: 18px;
    }
"""
) as demo:

    gr.Markdown("# 📖 Whisper Small - Surah Fatiha Transcription")
    gr.Markdown(
        "Fine-tuned Whisper model for recognizing and transcribing Surah Fatiha."
    )

    with gr.Tab("🎤 Record Audio"):
        mic_input = gr.Audio(
            sources=["microphone"], type="filepath", label="Record Audio"
        )
        transcribe_btn_1 = gr.Button("Transcribe")
        output_text_1 = gr.Textbox(label="Transcription", interactive=False)
        download_btn_1 = gr.File(label="Download Transcription", visible=False)

    with gr.Tab("📥 Upload Audio"):
        file_input = gr.Audio(
            sources=["upload"], type="filepath", label="Upload Audio File"
        )
        transcribe_btn_2 = gr.Button("Transcribe")
        output_text_2 = gr.Textbox(label="Transcription", interactive=False)
        download_btn_2 = gr.File(label="Download Transcription", visible=False)

    # Function connections
    transcribe_btn_1.click(fn=transcribe, inputs=mic_input, outputs=output_text_1)
    transcribe_btn_2.click(fn=transcribe, inputs=file_input, outputs=output_text_2)

    transcribe_btn_1.click(
        fn=save_transcription, inputs=output_text_1, outputs=download_btn_1
    )
    transcribe_btn_2.click(
        fn=save_transcription, inputs=output_text_2, outputs=download_btn_2
    )

demo.launch()