"""Kalenjin ASR — Gradio demo for Tonykip/whisper-kalenjin-v3-turbo.

Uses the CHUNKED long-form path (chunk_length_s=30, stride 5) + anti-loop decoding
(num_beams=5, no_repeat_ngram_size=3, repetition_penalty=1.15). The chunked path does NOT
invoke Whisper's sequential generate_with_fallback, so it sidesteps a transformers 5.x bug
in `_need_fallback` (IndexError on segment scores) while still preventing the long-form
"…ab usnatet…" looping. ZeroGPU-aware with a CPU fallback.

(The slightly-higher-accuracy "B3 robust-sequential" config — CER ~0.193 — is reserved for the
production endpoint, which needs word timestamps and uses return_timestamps='word', the code path
that avoids the same bug.)
"""
import os
import gradio as gr
import torch
from transformers import pipeline

if os.environ.get("SPACES_ZERO_GPU") is not None:
    import spaces
    gpu = spaces.GPU
else:
    def gpu(func=None, **_kw):
        return func if func is not None else (lambda f: f)

MODEL = "Tonykip/whisper-kalenjin-v3-turbo"
GEN = {"language": "sw", "task": "transcribe", "num_beams": 5,
       "no_repeat_ngram_size": 3, "repetition_penalty": 1.15}

_PIPE = None


def _pipe():
    global _PIPE
    if _PIPE is None:
        cuda = torch.cuda.is_available()
        dt = torch.float16 if cuda else torch.float32
        _PIPE = pipeline("automatic-speech-recognition", model=MODEL,
                         device=0 if cuda else -1, torch_dtype=dt,
                         chunk_length_s=30, stride_length_s=5)   # chunked → no sequential fallback
    return _PIPE


@gpu
def transcribe(path):
    if not path:
        return "🎙️ Record or upload Kalenjin audio, then press Transcribe."
    try:
        txt = _pipe()(path, generate_kwargs=GEN)["text"].strip()
    except Exception:
        txt = _pipe()(path)["text"].strip()
    return txt or "(no speech detected — try a clearer or longer clip)"


DESCRIPTION = """
# 🎙️ Kalenjin Speech Recognition

Transcribe spoken **Kalenjin** — covering **Kipsigis** and **Nandi**. Record a phrase or upload a
clip (any length) and press **Transcribe**.

*Powered by [`Tonykip/whisper-kalenjin-v3-turbo`](https://huggingface.co/Tonykip/whisper-kalenjin-v3-turbo),
a fine-tune of OpenAI Whisper-large-v3-turbo, with chunked long-form decoding + anti-repetition
controls so it stays stable on longer clips (no looping).*
"""

ABOUT = """
**Model:** [`Tonykip/whisper-kalenjin-v3-turbo`](https://huggingface.co/Tonykip/whisper-kalenjin-v3-turbo)
— fine-tune of Whisper-large-v3-turbo. The current best open Kalenjin ASR (~CER 0.21 on held-out
KaleBench-ASR).

**Decoding:** chunked long-form (30s windows) + beam-5 + `no_repeat_ngram_size=3` +
`repetition_penalty` — removes the long-form looping older configs showed.

**Tips & limits:** works best on clear Kipsigis/Nandi speech; English/Swahili code-switches are
usually kept; very noisy or overlapping audio is harder. On free CPU hardware transcription takes a
few seconds — switch the Space to a GPU/ZeroGPU runtime for near-instant results.

A separate experimental Parakeet fine-tune exists but is **not recommended for use**.
Built by **Tony Kipkemboi**, a native Kalenjin speaker.
"""

theme = gr.themes.Soft(primary_hue="green", secondary_hue="yellow")

with gr.Blocks(theme=theme, title="Kalenjin ASR") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(sources=["microphone", "upload"], type="filepath",
                                label="Kalenjin audio")
            go = gr.Button("Transcribe", variant="primary", size="lg")
        with gr.Column():
            out = gr.Textbox(label="Transcription", lines=6,
                             placeholder="The transcription will appear here…")
    go.click(transcribe, inputs=audio_in, outputs=out)
    audio_in.stop_recording(transcribe, inputs=audio_in, outputs=out)
    with gr.Accordion("About this model", open=False):
        gr.Markdown(ABOUT)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()