import torch
from transformers import pipeline
import gradio as gr
import os

MODEL_NAME = "HarshitJoshi/whisper-small-Hindi"
device = 0 if torch.cuda.is_available() else "cpu"
pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    device=device,
)

def transcribe_speech(filepath):
    output = pipe(
        filepath,
        max_new_tokens=256,
        generate_kwargs={
            "task": "transcribe",
            "language": "hindi",
        },
        chunk_length_s=10,
        batch_size=4,
    )
    return output["text"]

example_folder = "./examples"

demo = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Audio(label="Audio Input", type="filepath"),
    outputs=gr.Textbox(label="Transcription"),
    title="Hindi Speech Transcription",
    description=(
        "Upload an audio file or record using your microphone to transcribe Hindi speech."
    ),
    examples=example_folder,
    cache_examples=True,
    allow_flagging="never",
)

demo.launch(debug=True)