"""Japanese ASR demo for sbintuitions/kana-whisper on Hugging Face ZeroGPU.

The model is Whisper large-v3-turbo fine-tuned to transcribe Japanese speech
into *katakana* (pronunciation), e.g. キョーワイイテンキデスネ. It was built as
the ASR component for evaluating the reading accuracy of Japanese TTS systems.
"""

import numpy as np
import spaces
import torch
import gradio as gr
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

MODEL_ID = "sbintuitions/kana-whisper"

# ZeroGPU exposes CUDA at startup, so we can load straight onto the GPU. The
# @spaces.GPU decorator on transcribe() is what actually reserves the slice at
# call time; on non-ZeroGPU hardware the decorator is a no-op. We detect the
# device so the Space still builds/runs on plain CPU hardware (just slower).
if torch.cuda.is_available():
    device = "cuda"
    torch_dtype = torch.float16
else:
    device = "cpu"
    torch_dtype = torch.float32

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    MODEL_ID,
    dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
).to(device)

processor = AutoProcessor.from_pretrained(MODEL_ID)

asr = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    dtype=torch_dtype,
    device=device,
)

GENERATE_KWARGS = {"language": "ja", "task": "transcribe"}


def _to_float_mono(sampling_rate: int, data: np.ndarray) -> dict:
    """Normalize Gradio's (sr, ndarray) audio into what the ASR pipeline wants:
    a float32 mono waveform in [-1, 1] with its sampling rate (the feature
    extractor resamples to 16 kHz internally)."""
    data = np.asarray(data)
    if data.ndim > 1:  # stereo -> mono
        data = data.mean(axis=1)
    if np.issubdtype(data.dtype, np.integer):
        data = data.astype(np.float32) / np.iinfo(data.dtype).max
    else:
        data = data.astype(np.float32)
    return {"array": data, "sampling_rate": int(sampling_rate)}


@spaces.GPU
def transcribe(audio: tuple | None) -> str:
    if audio is None:
        return ""
    sampling_rate, data = audio
    inputs = _to_float_mono(sampling_rate, data)
    result = asr(inputs, generate_kwargs=GENERATE_KWARGS)
    return result["text"]


DESCRIPTION = """
# 🎙️ Kana Whisper — Japanese ASR → Katakana

Record or upload a short Japanese audio clip and get back its **katakana**
transcription (pronunciation, not kanji) using
[`sbintuitions/kana-whisper`](https://huggingface.co/sbintuitions/kana-whisper)
— a Whisper large-v3-turbo model fine-tuned by SB Intuitions.

> Example output: `キョーワイイテンキデスネ`
"""

with gr.Blocks(title="Kana Whisper — Japanese ASR") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        audio_in = gr.Audio(
            sources=["microphone", "upload"],
            type="numpy",
            label="Japanese audio (short clip)",
        )
        text_out = gr.Textbox(
            label="Katakana transcription",
            lines=4,
        )
    transcribe_btn = gr.Button("Transcribe", variant="primary")
    transcribe_btn.click(fn=transcribe, inputs=audio_in, outputs=text_out)

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