app.py

import os
import tempfile
import torch
import soundfile as sf
from transformers import pipeline
import gradio as gr
from pydub import AudioSegment

# =========================================================
# Whisper Tiny Hausa ASR
# =========================================================

MODEL_ID = "EYEDOL/whisper-tiny-hausa3"

DEVICE = 0 if torch.cuda.is_available() else -1

# Cache pipeline
ASR_PIPELINE = None

def get_asr_pipeline():
    global ASR_PIPELINE

    if ASR_PIPELINE is None:
        ASR_PIPELINE = pipeline(
            "automatic-speech-recognition",
            model=MODEL_ID,
            device=DEVICE
        )

    return ASR_PIPELINE

# =========================================================
# Utilities
# =========================================================

def save_numpy_to_wav(np_tuple):
    samplerate, data = np_tuple

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

    sf.write(tmp.name, data, samplerate)

    return tmp.name


def get_duration_seconds(path):
    try:
        info = sf.info(path)
        return info.duration
    except Exception:
        seg = AudioSegment.from_file(path)
        return len(seg) / 1000.0


def split_audio_file(path, chunk_length_ms=25000, overlap_ms=500):
    audio = AudioSegment.from_file(path)

    duration_ms = len(audio)

    chunks = []

    start = 0

    while start < duration_ms:

        end = min(start + chunk_length_ms, duration_ms)

        chunk = audio[start:end]

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

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

        chunks.append((tmp.name, start, end))

        start += max(1, chunk_length_ms - overlap_ms)

    return chunks


def transcribe_file(asr_pipeline, path, return_timestamps=False):

    if return_timestamps:
        return asr_pipeline(path, return_timestamps=True)

    return asr_pipeline(path)

# =========================================================
# Main transcription function
# =========================================================

def transcribe(
    audio_input,
    allow_longform_with_timestamps=False,
    chunk_length_seconds=25,
    overlap_seconds=0.5,
):

    if audio_input is None:
        return {"error": "No audio provided."}

    # Convert mic numpy input -> wav
    created_tmp_input = False

    if isinstance(audio_input, tuple):
        audio_path = save_numpy_to_wav(audio_input)
        created_tmp_input = True
    else:
        audio_path = audio_input

    duration_s = get_duration_seconds(audio_path)

    asr = get_asr_pipeline()

    # =====================================================
    # SHORT AUDIO
    # =====================================================

    if duration_s <= 30:

        out = transcribe_file(
            asr,
            audio_path,
            return_timestamps=False
        )

        text = out.get("text", out) if isinstance(out, dict) else str(out)

        segments = [{
            "start_s": 0.0,
            "end_s": duration_s,
            "text": text
        }]

        if created_tmp_input:
            try:
                os.unlink(audio_path)
            except:
                pass

        return {
            "full_text": text,
            "segments": segments
        }

    # =====================================================
    # LONG AUDIO WITH WHISPER TIMESTAMPS
    # =====================================================

    if allow_longform_with_timestamps:

        try:

            out = transcribe_file(
                asr,
                audio_path,
                return_timestamps=True
            )

            full_text = out.get("text", "")

            segments = []

            if "chunks" in out:

                for c in out["chunks"]:

                    ts = c.get("timestamp", [None, None])

                    segments.append({
                        "start_s": ts[0],
                        "end_s": ts[1],
                        "text": c.get("text", "")
                    })

            else:

                segments = [{
                    "start_s": 0.0,
                    "end_s": duration_s,
                    "text": full_text
                }]

            if created_tmp_input:
                try:
                    os.unlink(audio_path)
                except:
                    pass

            return {
                "full_text": full_text,
                "segments": segments
            }

        except Exception as e:
            print("Long-form failed. Falling back to chunking:", e)

    # =====================================================
    # CHUNKING FALLBACK
    # =====================================================

    chunk_length_ms = int(chunk_length_seconds * 1000)
    overlap_ms = int(overlap_seconds * 1000)

    chunks = split_audio_file(
        audio_path,
        chunk_length_ms=chunk_length_ms,
        overlap_ms=overlap_ms
    )

    stitched = []
    segments = []

    for chunk_path, start_ms, end_ms in chunks:

        try:

            out = transcribe_file(
                asr,
                chunk_path,
                return_timestamps=False
            )

            text = out.get("text", out) if isinstance(out, dict) else str(out)

        except Exception as e:

            text = f"[ERROR: {e}]"

        segments.append({
            "start_s": start_ms / 1000.0,
            "end_s": end_ms / 1000.0,
            "text": text
        })

        stitched.append(text)

        try:
            os.unlink(chunk_path)
        except:
            pass

    if created_tmp_input:
        try:
            os.unlink(audio_path)
        except:
            pass

    full_text = " ".join([x for x in stitched if x])

    return {
        "full_text": full_text,
        "segments": segments
    }

# =========================================================
# Gradio UI
# =========================================================

with gr.Blocks(title="Whisper Tiny Hausa ASR") as demo:

    gr.Markdown(
        """
        # Whisper Tiny Hausa ASR
        
        Upload audio or record with microphone.
        Supports long audio transcription.
        """
    )

    with gr.Row():

        with gr.Column(scale=2):

            mic_input = gr.Audio(
                label="Record Audio",
                type="numpy"
            )

            file_input = gr.Audio(
                label="Upload Audio File",
                type="filepath"
            )

            source = gr.Radio(
                ["Use microphone input", "Use uploaded file"],
                value="Use microphone input",
                label="Input source"
            )

            longform = gr.Checkbox(
                label="Use Whisper timestamps",
                value=True
            )

            chunk_len = gr.Slider(
                minimum=10,
                maximum=120,
                value=25,
                step=5,
                label="Chunk length (seconds)"
            )

            overlap = gr.Slider(
                minimum=0.0,
                maximum=5.0,
                value=0.5,
                step=0.5,
                label="Chunk overlap (seconds)"
            )

            transcribe_btn = gr.Button("Transcribe")

        with gr.Column(scale=3):

            full_text_out = gr.Textbox(
                label="Full transcription",
                lines=8
            )

            segments_out = gr.JSON(
                label="Segments"
            )

    def handle_transcription(
        mic_input,
        file_input,
        source_choice,
        use_longform,
        chunk_len_s,
        overlap_s
    ):

        audio_src = (
            mic_input
            if source_choice == "Use microphone input"
            else file_input
        )

        result = transcribe(
            audio_src,
            allow_longform_with_timestamps=use_longform,
            chunk_length_seconds=chunk_len_s,
            overlap_seconds=overlap_s
        )

        if "error" in result:
            return result["error"], []

        return result["full_text"], result["segments"]

    transcribe_btn.click(
        fn=handle_transcription,
        inputs=[
            mic_input,
            file_input,
            source,
            longform,
            chunk_len,
            overlap
        ],
        outputs=[
            full_text_out,
            segments_out
        ],
    )

# =========================================================
# Launch
# =========================================================

if __name__ == "__main__":
    demo.launch()