"""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}` — 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)}" 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()