import gradio as gr
from transformers import WhisperProcessor, FlaxWhisperForConditionalGeneration

# Load Whisper model
model_name = "openai/whisper-large"
processor = WhisperProcessor.from_pretrained(model_name)
model = FlaxWhisperForConditionalGeneration.from_pretrained(model_name)

def transcribe(audio):
    # Preprocess audio
    inputs = processor(audio, return_tensors="np", sampling_rate=16000)
    input_features = inputs.input_features

    # Generate transcription
    predicted_ids = model.generate(input_features)
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    
    return transcription

# Gradio Interface
interface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(source="upload", type="numpy"),
    outputs="text",
    title="Whisper JAX Transcription",
    description="Upload an audio file to transcribe using Whisper JAX."
)

interface.launch()