app.py

"""TEA-ASR demo (Hugging Face Spaces, ZeroGPU). Taiwan-Mandarin ASR: Traditional script + Taiwanese lexicon +
Mandarin-English code-switch, adapted from Qwen3-ASR. No runtime post-processing (the Traditional decode is baked into the
model's own tokenizer).

ZeroGPU rules followed: `import spaces` before torch; models placed on cuda at MODULE level (CUDA-emulated at
startup); GPU-dependent fn decorated with @spaces.GPU. torch pinned to a ZeroGPU-supported version (see
requirements.txt) and Python pinned to 3.12 (see README).

Format tags (TEA-ASR-1.1 only): the second-generation flagship was trained with output-convention tags. When a tag
is selected the demo teacher-forces a decoder prefix `language {lang} format {tags}<asr_text>` — the same interface
as the vendor language hint, generalized. `keep-en` maximizes verbatim English on dense code-switch; the numeral
style forces Arabic (`digits`) or Chinese (`zh-num`) numerals. Untagged models ignore the controls."""
import spaces
import torch
import gradio as gr
import numpy as np
import librosa

# Models are public — no HF_TOKEN needed. TEA-ASR-1.1 is the second-gen flagship (default).
REPOS = {
    "TEA-ASR-1.1": "JacobLinCool/TEA-ASR-1.1",
    "TEA-ASR-1.1-mini": "JacobLinCool/TEA-ASR-1.1-mini",
    "TEA-ASR-1": "JacobLinCool/TEA-ASR-1",
    "TEA-ASR-1-mini": "JacobLinCool/TEA-ASR-1-mini",
}
# Models trained with output-convention format tags (the prefix control is meaningful only for these).
TAG_CAPABLE = {"TEA-ASR-1.1", "TEA-ASR-1.1-mini"}
LANGS = ["auto", "Chinese", "English"]
NUMERALS = ["auto", "123 (Arabic)", "一二三 (Chinese)"]

# load all models on cuda at module level (recommended ZeroGPU pattern)
MODELS, LOAD_ERR = {}, None
try:
    from qwen_asr import Qwen3ASRModel
    for _name, _repo in REPOS.items():
        MODELS[_name] = Qwen3ASRModel.from_pretrained(_repo, dtype=torch.bfloat16, device_map="cuda:0")
except Exception as e:
    LOAD_ERR = repr(e)
    print("model load failed:", LOAD_ERR)


def _build_tags(numerals, keep_en):
    """Selected UI controls -> canonical format-tag list (numeral tag first, then keep-en)."""
    tags = []
    if numerals == "123 (Arabic)":
        tags.append("digits")
    elif numerals == "一二三 (Chinese)":
        tags.append("zh-num")
    if keep_en:
        tags.append("keep-en")
    return tags


def _transcribe_prefixed(wrapper, wav, prefix, context):
    """Forced decoder-prefix transcription (mirrors qwen_asr's native path with a free-form prefix).

    The public `transcribe(language=...)` validates the language against a fixed list, so a format-tag
    prefix cannot pass through it. This replicates the native forced-language path — audio fold, chat
    prompt + `{prefix}`, generate, native output parse — with the prefix generalized.
    """
    from qwen_asr.inference.utils import float_range_normalize
    from qwen_asr import parse_asr_output

    hf_model, processor = wrapper.model, wrapper.processor
    wav = float_range_normalize(np.asarray(wav, dtype="float32"))
    msgs = [
        {"role": "system", "content": context or ""},
        {"role": "user", "content": [{"type": "audio", "audio": ""}]},
    ]
    base = processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
    if isinstance(base, list):
        base = base[0]
    inputs = processor(text=[base + prefix], audio=[wav], return_tensors="pt", padding=True)
    inputs = inputs.to(hf_model.device).to(hf_model.dtype)
    with torch.no_grad():
        generated = hf_model.generate(**inputs, max_new_tokens=wrapper.max_new_tokens)
    sequences = getattr(generated, "sequences", generated)
    continuation = processor.batch_decode(
        sequences[:, inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0]
    return parse_asr_output(continuation, user_language=prefix or None)[1]


@spaces.GPU(duration=120)
def transcribe(audio_path, model_choice, language, numerals, keep_en, context):
    if LOAD_ERR:
        return f"Model load failed.\n\n{LOAD_ERR}"
    if not audio_path:
        return "請提供音訊 / Please provide audio."
    wav, _ = librosa.load(audio_path, sr=16000, mono=True)
    tags = _build_tags(numerals, keep_en)
    if tags and model_choice in TAG_CAPABLE:
        # format-tag prefix path (tags are Chinese-space conventions; default the hint to Chinese)
        plang = language if language in ("Chinese", "English") else "Chinese"
        prefix = f"language {plang} format {' '.join(tags)}<asr_text>"
        return _transcribe_prefixed(MODELS[model_choice], wav, prefix, context)
    lang = None if language == "auto" else language
    out = MODELS[model_choice].transcribe(
        audio=[(np.asarray(wav, dtype="float32"), 16000)], context=(context or ""), language=lang)[0]
    return out.text


EXAMPLES = [
    # audio, model, language, numerals, keep_en, context
    ["examples/lecture_zh-TW.wav", "TEA-ASR-1.1", "Chinese", "auto", False, ""],
    ["examples/codeswitch_zh-en.wav", "TEA-ASR-1.1", "Chinese", "auto", True, ""],
    ["examples/mandarin_zh.wav", "TEA-ASR-1.1", "Chinese", "auto", False, ""],
]

DESC = (
    "**TEA-ASR (Taiwan Everyday Audio)** — Traditional-script + Taiwanese-lexicon ASR with robust "
    "Mandarin–English code-switch, adapted from Qwen3-ASR with a tokenizer-first procedure. "
    "Output is Traditional Chinese. Set the language hint to *Chinese* for Taiwan speech (best results).\n\n"
    "**Format tags** (TEA-ASR-1.1 and TEA-ASR-1.1-mini): tick **Keep English** to transcribe code-switch English "
    "verbatim instead of translating it, and pick a **numeral style** to force Arabic (123) or Chinese (一二三) "
    "numbers. First-gen models ignore these."
)

demo = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath", label="Audio (upload or record)"),
        gr.Dropdown(list(REPOS), value="TEA-ASR-1.1", label="Model"),
        gr.Dropdown(LANGS, value="Chinese", label="Language hint"),
        gr.Radio(NUMERALS, value="auto", label="Numeral style (format tag · TEA-ASR-1.1 models)"),
        gr.Checkbox(value=False, label="Keep English verbatim — don't translate code-switch (format tag · TEA-ASR-1.1 models)"),
        gr.Textbox(label="Context / hotwords (optional)", placeholder="例如：台積電 員工 名稱…"),
    ],
    outputs=gr.Textbox(label="Transcription (繁體中文)"),
    examples=EXAMPLES,
    cache_examples=False,
    title="TEA-ASR — Taiwan Mandarin ASR 🍵",
    description=DESC,
)

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