import spaces
import gradio as gr
from transformers import pipeline
from pydub import AudioSegment
import tempfile
import os

pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-large-v3",
    device_map="auto"
)

CHUNK_LENGTH_MS = 30 * 1000  # 30 seconds


@spaces.GPU
def transcribe(audio_path):
    audio = AudioSegment.from_file(audio_path)

    transcript = []

    for i in range(0, len(audio), CHUNK_LENGTH_MS):
        chunk = audio[i:i + CHUNK_LENGTH_MS]

        temp_file = tempfile.NamedTemporaryFile(
            suffix=".wav",
            delete=False
        )

        chunk.export(temp_file.name, format="wav")

        result = pipe(temp_file.name)

        transcript.append(result["text"])

        os.remove(temp_file.name)

    return "\n".join(transcript)


demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(type="filepath"),
    outputs=gr.Textbox(lines=20),
    title="Whisper Large V3 - Long Audio Support",
    description="Upload long audio files for transcription."
)

demo.launch()