import gradio as gr
from transformers import pipeline

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

# Define the transcription function for audio input
def transcribe_audio(audio):
    
    prediction = pipe(audio, batch_size=8, return_timestamps=True)["chunks"]
    
    transcription = "\n".join([f"[{chunk['timestamp'][0]:.2f}s - {chunk['timestamp'][1]:.2f}s] {chunk['text']}" for chunk in prediction])
    
    return transcription

# Create a Gradio interface
interface = gr.Interface(
    fn=transcribe_audio,  
    inputs=gr.Audio(type="filepath"), 
    outputs="text",  
    title="Whisper Small ASR",  
    description="Upload or record audio for transcription using Whisper Small."  
)

# Launch the Gradio app
interface.launch(share=True)