import gradio as gr
import torch
from transformers import pipeline

device = "cuda:0" if torch.cuda.is_available() else "cpu"

pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-small",
    device=device,
)

def transcribe(audio):
    transcription = pipe(audio)["text"]
    return transcription

# Gradio interface
interface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(type="filepath"),
    outputs="text",
    title="Automatic Speech Recognition with Whisper Small",
    description="Upload an audio file to get the transcription."
)

if __name__ == "__main__":
    interface.launch(share=True)