import torch
from transformers import pipeline
import gradio as gr

# Pick device (GPU if available, else CPU)
device = 0 if torch.cuda.is_available() else -1

# Initialize Whisper ASR pipeline
pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-small",
    chunk_length_s=30,
    device=device,
)

# Function to transcribe audio
def transcribe_audio(audio):
    if audio is None:
        return "Please record or upload an audio file."
    prediction = pipe(audio)["text"]
    return prediction

# Build Gradio interface
interface = gr.Interface(
    fn=transcribe_audio,
    inputs=gr.Audio(type="filepath", label="Record or Upload Audio"),
    outputs=gr.Textbox(label="Transcription"),
    title="Whisper Small - Automatic Speech Recognition",
    description="Record audio or upload a file and get a transcription using the OpenAI Whisper Small model."
)

# Run app
if __name__ == "__main__":
    interface.launch()