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):
    # Transcribe the uploaded or recorded audio
    prediction = pipe(audio, batch_size=8, return_timestamps=True)["chunks"]
    
    # Format the output to show text with timestamps
    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,  # The function to be applied to the audio input
    inputs=gr.Audio(type="filepath"),  # Users can record or upload audio
    outputs="text",  # The output is the transcription (text with timestamps)
    title="Whisper Small ASR",  # Title of your app
    description="Upload or record audio for transcription using Whisper Small."  # Description of your app
)

# Launch the Gradio app
interface.launch()