Squash history
Browse filesCo-authored-by: Fouad Oubari <foubari@users.noreply.huggingface.co>
- .gitattributes +6 -0
- README.md +251 -0
- assets/jais-asr-benchmarks.png +3 -0
- chat_template.jinja +102 -0
- config.json +161 -0
- configuration_jais_asr.py +270 -0
- generation_config.json +14 -0
- language.py +328 -0
- mlx-vlm-jais-asr.patch +583 -0
- model-00001-of-00002.safetensors +3 -0
- model-00002-of-00002.safetensors +3 -0
- model.safetensors.index.json +0 -0
- modeling_jais_asr.py +1555 -0
- pipeline_jais_asr.py +187 -0
- postprocess.py +404 -0
- preprocessor_config.json +14 -0
- processing_jais_asr.py +390 -0
- processor_config.json +9 -0
- special_tokens_map.json +23 -0
- tokenizer.json +3 -0
- tokenizer_config.json +18 -0
.gitattributes
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.gguf filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
assets/jais-asr-benchmarks.png filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,251 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
library_name: mlx
|
| 4 |
+
pipeline_tag: automatic-speech-recognition
|
| 5 |
+
base_model: IFM/Jais-ASR
|
| 6 |
+
base_model_relation: quantized
|
| 7 |
+
model_name: Jais-ASR
|
| 8 |
+
language:
|
| 9 |
+
- ar
|
| 10 |
+
- arb # Modern Standard Arabic
|
| 11 |
+
- arz # Egyptian
|
| 12 |
+
- apd # Sudanese
|
| 13 |
+
- ary # Moroccan
|
| 14 |
+
- arq # Algerian
|
| 15 |
+
- aeb # Tunisian
|
| 16 |
+
- ayl # Libyan
|
| 17 |
+
- mey # Mauritanian (Hassaniya)
|
| 18 |
+
- apc # Levantine: Syrian, Lebanese, Jordanian, Palestinian
|
| 19 |
+
- acm # Iraqi
|
| 20 |
+
- afb # Gulf: Kuwaiti, Qatari, Bahraini, Emirati
|
| 21 |
+
- ars # Saudi (Najdi)
|
| 22 |
+
- acw # Saudi (Hijazi)
|
| 23 |
+
- acx # Omani
|
| 24 |
+
- ayn # Yemeni (Sanaani)
|
| 25 |
+
- en
|
| 26 |
+
- fr
|
| 27 |
+
- es
|
| 28 |
+
- it
|
| 29 |
+
tags:
|
| 30 |
+
- jais
|
| 31 |
+
- ifm
|
| 32 |
+
- arabic
|
| 33 |
+
- arabic-dialects
|
| 34 |
+
- speech-to-text
|
| 35 |
+
- code-switching
|
| 36 |
+
- diacritization
|
| 37 |
+
- language-identification
|
| 38 |
+
- open-weights
|
| 39 |
+
- custom_code
|
| 40 |
+
- mlx
|
| 41 |
+
- mlx-vlm
|
| 42 |
+
---
|
| 43 |
+
|
| 44 |
+
# Jais-ASR-MLX-8bit
|
| 45 |
+
|
| 46 |
+
> [!NOTE]
|
| 47 |
+
> This repository contains the 8-bit MLX version of [IFM/Jais-ASR](https://huggingface.co/IFM/Jais-ASR) for Apple silicon, for use with `mlx-vlm`. The 4-bit version is [IFM/Jais-ASR-MLX-4bit](https://huggingface.co/IFM/Jais-ASR-MLX-4bit); see the [table below](#mlx-configurations). For NVIDIA or AMD GPUs, use [IFM/Jais-ASR-GGUF](https://huggingface.co/IFM/Jais-ASR-GGUF). The Quickstart further down applies to the original model.
|
| 48 |
+
>
|
| 49 |
+
> **Compatibility:** requires [mlx-vlm](https://github.com/Blaizzy/mlx-vlm) at commit `6ecadd7` with `mlx-vlm-jais-asr.patch` from this repository, and PyTorch (used by the audio processor).
|
| 50 |
+
|
| 51 |
+
```bash
|
| 52 |
+
# 1. Install mlx-vlm with Jais-ASR support (once)
|
| 53 |
+
git clone https://github.com/Blaizzy/mlx-vlm.git && cd mlx-vlm
|
| 54 |
+
git checkout 6ecadd767ca1c7c54763289c750dd180d7945037
|
| 55 |
+
# Recommended: use a dedicated environment (Python 3.10+).
|
| 56 |
+
# Skip the next line if you already activated one with conda, uv, pyenv, etc.
|
| 57 |
+
python3 -m venv .venv && source .venv/bin/activate
|
| 58 |
+
pip install -e . torch
|
| 59 |
+
hf download IFM/Jais-ASR-MLX-8bit mlx-vlm-jais-asr.patch --local-dir . && git apply mlx-vlm-jais-asr.patch && cd ..
|
| 60 |
+
|
| 61 |
+
# 2. Transcribe (the model downloads on first use)
|
| 62 |
+
mlx_vlm.generate --model IFM/Jais-ASR-MLX-8bit --trust-remote-code --audio speech.wav --prompt "" \
|
| 63 |
+
--temperature 0 --repetition-penalty 1.05 --max-tokens 256
|
| 64 |
+
|
| 65 |
+
# 3. Optional: set the language in advance (ar, msa, EG, MA, ...; see Language and Dialect Settings)
|
| 66 |
+
mlx_vlm.generate --model IFM/Jais-ASR-MLX-8bit --trust-remote-code --audio speech.wav --prompt "" \
|
| 67 |
+
--temperature 0 --repetition-penalty 1.05 --max-tokens 256 --language EG
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
## MLX Configurations
|
| 71 |
+
|
| 72 |
+
| Configuration | Repository | Text decoder | Audio encoder | WER % | CER % | Size (GB) |
|
| 73 |
+
|---|---|---|---|---:|---:|---:|
|
| 74 |
+
| Original model (BF16) | [IFM/Jais-ASR](https://huggingface.co/IFM/Jais-ASR) | BF16 | BF16 | 22.85 | 8.37 | 11.46 |
|
| 75 |
+
| 8-bit | [IFM/Jais-ASR-MLX-8bit](https://huggingface.co/IFM/Jais-ASR-MLX-8bit) | 8-bit (same values as `Jais-ASR-Q8_0.gguf`) | BF16 | 22.76 | 8.30 | 7.03 |
|
| 76 |
+
| 4-bit GPTQ | [IFM/Jais-ASR-MLX-4bit](https://huggingface.co/IFM/Jais-ASR-MLX-4bit) | 4-bit GPTQ (same values as `Jais-ASR-Q4_K_S.gguf`) | BF16 | 22.94 | 8.34 | 4.50 |
|
| 77 |
+
|
| 78 |
+
WER / CER (%) are averaged over the six test sets, evaluated with `mlx-vlm` (greedy decoding, repetition penalty 1.05, `code_switch=False`, at most 256 generated tokens); lower is better. The original model row reports its published scores. Size is the total download size of each repository.
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
Jais-ASR is IFM's state-of-the-art speech recognition model for Arabic as it is spoken: Modern Standard Arabic and dialects from 19 Arab countries, with natural code-switching into English, French, Spanish, and Italian. Based on Jais-3-4B, it adds optional diacritization, built-in language and dialect identification, and an explicit signal when there is no speech to transcribe. The weights are released under Apache 2.0 with a vLLM plugin, and quantized versions are available in the [Jais Speech collection](https://huggingface.co/collections/IFM/jais-speech).
|
| 82 |
+
|
| 83 |
+
<p align="center">
|
| 84 |
+
<img src="assets/jais-asr-benchmarks.png" alt="Average word and character error rates on the Open Universal Arabic ASR Leaderboard. Jais-ASR has the lowest of both, 22.85% WER and 8.37% CER, ahead of Audar-ASR-V1-Turbo, Speechmatics Melia 1, ElevenLabs Scribe v2, cohere-transcribe-arabic-07-2026, omniASR-LLM-7B, Qwen3-Omni-30B-A3B, Qwen3-ASR-1.7B and Voxtral-Small-24B." width="100%">
|
| 85 |
+
</p>
|
| 86 |
+
|
| 87 |
+
## Jais-ASR Highlights
|
| 88 |
+
|
| 89 |
+
- **Leading performance in Arabic ASR.** Jais-ASR has the lowest average error rates on the [Open Universal Arabic ASR Leaderboard](https://github.com/Natural-Language-Processing-Elm/open_universal_arabic_asr_leaderboard), 22.85% WER and 8.37% CER, ahead of every open-weight model compared and the commercial APIs from Speechmatics and ElevenLabs. See [Benchmark Results](#benchmark-results).
|
| 90 |
+
- **Built for Arabic as it is spoken in real life.** MSA and dialects across 19 countries, natural code-switching with English, French, Spanish, and Italian, and robust transcription in noisy, real-world recordings.
|
| 91 |
+
- **Accurate and faithful.** Optional diacritization, built-in language and dialect identification, and transcripts designed to avoid memorized boilerplate, repetition loops, and translating speech instead of transcribing it. When there is no speech, the model says so instead of inventing text.
|
| 92 |
+
- **Ready for production.** Run it with Transformers, or serve it with the included vLLM plugin behind an OpenAI-compatible `/v1/audio/transcriptions` endpoint. The bf16 weights are 11.4 GB.
|
| 93 |
+
|
| 94 |
+
## Benchmark Results
|
| 95 |
+
|
| 96 |
+
The [Open Universal Arabic ASR Leaderboard](https://github.com/Natural-Language-Processing-Elm/open_universal_arabic_asr_leaderboard) covers six public test sets spanning read speech, broadcast media, and conversational dialects, in clean and noisy conditions. Evaluated with the leaderboard's own protocol, Jais-ASR leads on both average metrics.
|
| 97 |
+
|
| 98 |
+
<!-- TABLE:START leaderboard -->
|
| 99 |
+
<div style="font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;margin:0 auto;padding:8px 0 16px;overflow-x:auto"><table style="display:table;width:100%;table-layout:fixed;border-collapse:collapse;font-size:12px;margin:0"><thead><tr><th style="width:18%;border-bottom:none"></th><th colspan="1" style="padding:6px 4px 2px;text-align:center;font-size:11px;font-weight:600;letter-spacing:0.04em;text-transform:uppercase;opacity:0.65;border-bottom:1px solid rgba(128,128,128,0.25)"></th><th colspan="6" style="padding:6px 4px 2px;text-align:center;font-size:11px;font-weight:600;letter-spacing:0.04em;text-transform:uppercase;opacity:0.65;border-bottom:1px solid rgba(128,128,128,0.25)">Open-weight models</th><th colspan="2" style="padding:6px 4px 2px;text-align:center;font-size:11px;font-weight:600;letter-spacing:0.04em;text-transform:uppercase;opacity:0.65;border-bottom:1px solid rgba(128,128,128,0.25)">Closed models</th></tr><tr><th style="padding:10px 6px;text-align:left;border-bottom:2px solid #0b6b4f"></th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;background:rgba(11,107,79,0.08);">Jais-ASR</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">Audar-ASR-V1-Turbo</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">cohere-transcribe-arabic-07-2026</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">omniASR-LLM-7B</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">Qwen3-Omni-30B-A3B</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">Qwen3-ASR-1.7B</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">Voxtral-Small-24B</th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">Melia 1<div style="margin-top:2px;font-size:10px;opacity:0.65;font-weight:400">Speechmatics</div></th><th style="padding:10px 3px;text-align:center;font-weight:600;border-bottom:2px solid #0b6b4f;color:#0b6b4f;font-size:12.5px;line-height:1.2;width:9.11%;overflow-wrap:anywhere;">Scribe v2<div style="margin-top:2px;font-size:10px;opacity:0.65;font-weight:400">ElevenLabs</div></th></tr></thead><tbody><tr><td colspan="10" style="padding:6px 10px;font-weight:600;font-size:12.5px;color:#0b6b4f;border-bottom:1px solid rgba(11,107,79,0.25);background:rgba(11,107,79,0.12)">Word error rate (WER, %) ↓</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">Average</div><div style="margin-top:2px;font-size:10px;opacity:0.65">Macro-average, six test sets</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>22.85</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">23.17</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">25.87</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">28.32</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">30.71</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">33.36</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">34.47</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">23.51</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">24.54</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">SADA</div><div style="margin-top:2px;font-size:10px;opacity:0.65">Saudi dialectal broadcast</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);">30.85</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;"><strong>28.92</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">37.47</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">41.61</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">44.82</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">45.53</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">50.82</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">34.74</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">33.96</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">Common Voice 18</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA read speech</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);">8.56</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">8.09</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.82</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">8.75</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">11.46</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.90</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">15.25</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;"><strong>5.19</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">8.71</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">MASC clean</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA lectures, clean</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>16.12</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.73</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">19.60</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">19.69</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">21.47</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">24.37</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">23.96</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.23</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.91</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">MASC noisy</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA lectures, noisy</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>25.44</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">27.19</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">27.07</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">29.29</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">30.85</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">34.29</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">34.43</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">25.89</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">27.15</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">MGB-2</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA broadcast news</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>9.59</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">11.08</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">15.54</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">14.13</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">13.09</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.57</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.03</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">12.61</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">10.30</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">Casablanca</div><div style="margin-top:2px;font-size:10px;opacity:0.65">8 dialects, conversational</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);">46.56</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">47.02</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">49.71</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">56.46</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">62.55</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">64.47</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">66.30</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;"><strong>46.41</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">50.20</td></tr><tr><td colspan="10" style="padding:6px 10px;font-weight:600;font-size:12.5px;color:#0b6b4f;border-bottom:1px solid rgba(11,107,79,0.25);background:rgba(11,107,79,0.12)">Character error rate (CER, %) ↓</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">Average</div><div style="margin-top:2px;font-size:10px;opacity:0.65">Macro-average, six test sets</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>8.37</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">9.23</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">11.80</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">12.52</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">13.67</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">12.33</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">15.29</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">10.62</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">9.96</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">SADA</div><div style="margin-top:2px;font-size:10px;opacity:0.65">Saudi dialectal broadcast</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);">14.33</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;"><strong>13.40</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">23.53</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">24.95</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">26.11</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">19.90</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">28.85</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">21.21</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">16.48</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">Common Voice 18</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA read speech</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);">2.70</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">2.46</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;"><strong>1.62</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">2.71</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">4.28</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.06</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.54</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;"><strong>1.62</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">3.00</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">MASC clean</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA lectures, clean</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>3.76</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.32</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">6.45</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.76</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.59</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.72</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">7.06</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.11</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">4.97</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">MASC noisy</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA lectures, noisy</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>7.51</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">9.78</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">10.13</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">10.66</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">11.28</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">10.84</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">12.22</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">9.47</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">9.36</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">MGB-2</div><div style="margin-top:2px;font-size:10px;opacity:0.65">MSA broadcast news</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>3.92</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.64</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">8.40</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">7.10</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">6.20</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">6.25</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">7.41</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">5.88</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">4.78</td></tr><tr><td style="padding:6px 4px 6px 10px;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle"><div style="font-size:12.5px;font-weight:600;line-height:1.2;color:inherit">Casablanca</div><div style="margin-top:2px;font-size:10px;opacity:0.65">8 dialects, conversational</div></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;background:rgba(11,107,79,0.08);"><strong>18.02</strong></td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">18.75</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">20.66</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">23.96</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">28.53</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">26.23</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">30.64</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">20.41</td><td style="padding:6px 2px;text-align:center;border-bottom:1px solid rgba(128,128,128,0.15);vertical-align:middle;font-size:12.5px;line-height:1.2;">21.18</td></tr></tbody></table></div>
|
| 100 |
+
<!-- TABLE:END leaderboard -->
|
| 101 |
+
|
| 102 |
+
<p style="font-size:1.25em;line-height:1.6">Error rates in %; lower is better. Average is the macro-average over the six test sets. Jais-ASR was measured with the leaderboard's evaluation script. Scores for the other open-weight models are taken from the leaderboard; Speechmatics Melia 1 and ElevenLabs Scribe v2 were measured by IFM.</p>
|
| 103 |
+
|
| 104 |
+
## Capabilities
|
| 105 |
+
|
| 106 |
+
Jais-ASR brings Arabic speech transcription, code-switching, diacritization, and language identification into one model, with controls to tailor the output to your application.
|
| 107 |
+
|
| 108 |
+
- **Transcribe multilingual speech as spoken.** Enable `code_switch=True` to keep English, French, Spanish, or Italian spans in their original script alongside Arabic, useful for interviews and everyday conversations.
|
| 109 |
+
- **Add diacritics for clearer Arabic pronunciation.** Enable `diacritics=True` to add vowel marks to Arabic transcripts, making them easier to read and supporting language learning. Diacritics and code-switching can be enabled together.
|
| 110 |
+
- **Detect the language and dialect as you transcribe.** Each transcript comes with its recognized language or dialect label, so you can organize recordings and route text to downstream tools without a separate identification model.
|
| 111 |
+
- **Handle non-speech audio explicitly.** When audio contains no speech (e.g., music or noise), the model returns `[IFM-NoTranscript]` with the flag `refused=True` instead of hallucinated content, giving applications a clear signal to handle the case separately.
|
| 112 |
+
- **Process multiple recordings at once.** Batch transcription supports files with different durations and sample rates.
|
| 113 |
+
|
| 114 |
+
### Language and Dialect Settings
|
| 115 |
+
|
| 116 |
+
Automatic language and dialect detection is enabled by default. If the language of the audio is known in advance, set it with one of the codes below: `ar` allows Modern Standard Arabic and all 19 supported dialects, `msa` selects Modern Standard Arabic, and a dialect code such as `EG` selects that dialect.
|
| 117 |
+
|
| 118 |
+
| Language | Setting |
|
| 119 |
+
|---|---|
|
| 120 |
+
| Arabic, including all supported dialects | `ar` |
|
| 121 |
+
| Modern Standard Arabic | `msa` |
|
| 122 |
+
| English | `en` |
|
| 123 |
+
| French | `fr` |
|
| 124 |
+
| Spanish | `es` |
|
| 125 |
+
| Italian | `it` |
|
| 126 |
+
|
| 127 |
+
The 19 Arabic dialect codes are:
|
| 128 |
+
|
| 129 |
+
Algerian (`DZ`) · Bahraini (`BH`) · Egyptian (`EG`) · Emirati (`AE`) · Iraqi (`IQ`) · Jordanian (`JO`) · Kuwaiti (`KW`) · Lebanese (`LB`) · Libyan (`LY`) · Mauritanian (`MR`) · Moroccan (`MA`) · Omani (`OM`) · Palestinian (`PS`) · Qatari (`QA`) · Saudi (`SA`) · Sudanese (`SD`) · Syrian (`SY`) · Tunisian (`TN`) · Yemeni (`YE`)
|
| 130 |
+
|
| 131 |
+
## Quickstart
|
| 132 |
+
|
| 133 |
+
Install PyTorch for your hardware, then install the inference dependencies:
|
| 134 |
+
|
| 135 |
+
```bash
|
| 136 |
+
pip install 'transformers==5.16.1' accelerate soundfile soxr
|
| 137 |
+
```
|
| 138 |
+
|
| 139 |
+
The examples below use a GPU with bfloat16 support. The bf16 weights are 11.4 GB; plan on roughly 12 GB of GPU memory.
|
| 140 |
+
|
| 141 |
+
### Transformers Pipeline
|
| 142 |
+
|
| 143 |
+
```python
|
| 144 |
+
import torch
|
| 145 |
+
from transformers import pipeline
|
| 146 |
+
|
| 147 |
+
asr = pipeline(
|
| 148 |
+
"automatic-speech-recognition",
|
| 149 |
+
model="IFM/Jais-ASR",
|
| 150 |
+
trust_remote_code=True,
|
| 151 |
+
device_map="auto",
|
| 152 |
+
dtype=torch.bfloat16,
|
| 153 |
+
)
|
| 154 |
+
result = asr("speech.wav", max_new_tokens=256)
|
| 155 |
+
print(result["text"])
|
| 156 |
+
```
|
| 157 |
+
|
| 158 |
+
The pipeline accepts an audio file path or a waveform dictionary such as `{"array": waveform, "sampling_rate": 16000}`. For batches and structured metadata, use the model and processor directly as shown below.
|
| 159 |
+
|
| 160 |
+
### Batch Transcription
|
| 161 |
+
|
| 162 |
+
This example transcribes two files together, including clips with different durations or sample rates.
|
| 163 |
+
|
| 164 |
+
```python
|
| 165 |
+
import soundfile as sf
|
| 166 |
+
import torch
|
| 167 |
+
from transformers import AutoModelForCausalLM, AutoProcessor
|
| 168 |
+
|
| 169 |
+
model_id = "IFM/Jais-ASR"
|
| 170 |
+
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
| 171 |
+
processor.tokenizer.padding_side = "left"
|
| 172 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 173 |
+
model_id,
|
| 174 |
+
trust_remote_code=True,
|
| 175 |
+
device_map="auto",
|
| 176 |
+
dtype=torch.bfloat16,
|
| 177 |
+
).eval()
|
| 178 |
+
|
| 179 |
+
paths = ["speech_1.wav", "speech_2.wav"]
|
| 180 |
+
audio = [sf.read(path, dtype="float32") for path in paths]
|
| 181 |
+
waveforms = [waveform for waveform, _ in audio]
|
| 182 |
+
sample_rates = [sample_rate for _, sample_rate in audio]
|
| 183 |
+
|
| 184 |
+
prompt = processor.tokenizer.apply_chat_template(
|
| 185 |
+
[{"role": "user", "content": [{"type": "audio"}]}],
|
| 186 |
+
tokenize=False,
|
| 187 |
+
add_generation_prompt=True,
|
| 188 |
+
# diacritics=True, # Add Arabic vowel marks.
|
| 189 |
+
# code_switch=True, # Keep other languages in their original script.
|
| 190 |
+
)
|
| 191 |
+
inputs = processor(
|
| 192 |
+
text=[prompt] * len(waveforms),
|
| 193 |
+
audio=waveforms,
|
| 194 |
+
sampling_rate=sample_rates,
|
| 195 |
+
padding=True,
|
| 196 |
+
return_tensors="pt",
|
| 197 |
+
).to(model.device)
|
| 198 |
+
|
| 199 |
+
with torch.inference_mode():
|
| 200 |
+
output = model.generate(
|
| 201 |
+
**inputs,
|
| 202 |
+
do_sample=False,
|
| 203 |
+
max_new_tokens=256,
|
| 204 |
+
# language="EG", # Set the target language to override automatic detection.
|
| 205 |
+
)
|
| 206 |
+
|
| 207 |
+
transcripts = processor.post_process_transcription(
|
| 208 |
+
output, input_lengths=inputs["input_ids"].shape[1],
|
| 209 |
+
)
|
| 210 |
+
for path, result in zip(paths, transcripts):
|
| 211 |
+
print(path, {
|
| 212 |
+
"text": result.text,
|
| 213 |
+
"language": result.language,
|
| 214 |
+
"refused": result.refused,
|
| 215 |
+
})
|
| 216 |
+
```
|
| 217 |
+
|
| 218 |
+
`language` contains a language or dialect label, such as `msa`, `EG`, or `en`. When the model emits `[IFM-NoTranscript]`, postprocessing sets `refused=True` and returns an empty transcript. Uncomment either prompt option above to request diacritics or code-switching; both can be enabled together.
|
| 219 |
+
|
| 220 |
+
> [!TIP]
|
| 221 |
+
> Pass each recording's original sample rate; the processor resamples it to 16 kHz automatically. Increase `max_new_tokens` for longer recordings: a result with `truncated=True` means generation ended before the transcript was complete.
|
| 222 |
+
|
| 223 |
+
### Serving with vLLM
|
| 224 |
+
|
| 225 |
+
Jais-ASR ships with a vLLM plugin in [`vllm_plugin/`](https://huggingface.co/IFM/Jais-ASR/tree/main/vllm_plugin). Install it in its own environment, then start the server:
|
| 226 |
+
|
| 227 |
+
```bash
|
| 228 |
+
pip install 'vllm==0.23.0' 'transformers==5.5.4' accelerate soundfile soxr
|
| 229 |
+
hf download IFM/Jais-ASR --include "vllm_plugin/*" --local-dir Jais-ASR
|
| 230 |
+
pip install ./Jais-ASR/vllm_plugin
|
| 231 |
+
vllm serve IFM/Jais-ASR --trust-remote-code --reasoning-parser jais_asr_tag --max-model-len 2048
|
| 232 |
+
```
|
| 233 |
+
|
| 234 |
+
Transcribe a file through the OpenAI-compatible endpoint:
|
| 235 |
+
|
| 236 |
+
```bash
|
| 237 |
+
curl http://localhost:8000/v1/audio/transcriptions \
|
| 238 |
+
-F file=@clip.wav \
|
| 239 |
+
-F model=IFM/Jais-ASR
|
| 240 |
+
```
|
| 241 |
+
|
| 242 |
+
## Citation
|
| 243 |
+
|
| 244 |
+
```bibtex
|
| 245 |
+
@misc{jaisasr2026,
|
| 246 |
+
title = {Introducing Jais-ASR: Arabic Speech Recognition for Every Dialect},
|
| 247 |
+
author = {{IFM Team}},
|
| 248 |
+
year = {2026},
|
| 249 |
+
url = {https://huggingface.co/IFM/Jais-ASR},
|
| 250 |
+
}
|
| 251 |
+
```
|
assets/jais-asr-benchmarks.png
ADDED
|
Git LFS Details
|
chat_template.jinja
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{%- if language is defined and language is not none -%}
|
| 2 |
+
{{- raise_exception("language is a generation-time pin. Pass it to model.generate(), the pipeline, or the API; use language_prefill here to insert a tag.") -}}
|
| 3 |
+
{%- endif -%}
|
| 4 |
+
{#- ===================================================================================
|
| 5 |
+
Jais-ASR — audio transcription chat template.
|
| 6 |
+
|
| 7 |
+
Optional flags (keyword arguments to apply_chat_template, or
|
| 8 |
+
chat_template_kwargs on a server):
|
| 9 |
+
diacritics=True -> fully diacritized transcript
|
| 10 |
+
code_switch=True -> keep foreign-language spans verbatim
|
| 11 |
+
language_prefill="EG" -> prefill the language tag (see the model card)
|
| 12 |
+
An explicit system message always wins over the mode flags.
|
| 13 |
+
=================================================================================== -#}
|
| 14 |
+
{{- bos_token }}
|
| 15 |
+
{#- ---- mode: the system message content is the switch -------------------------- -#}
|
| 16 |
+
{%- set _m = namespace(modes='') -%}
|
| 17 |
+
{%- if code_switch is defined and code_switch -%}
|
| 18 |
+
{%- set _m.modes = '[Code-switch]' -%}
|
| 19 |
+
{%- endif -%}
|
| 20 |
+
{%- if diacritics is defined and diacritics -%}
|
| 21 |
+
{%- set _m.modes = (_m.modes ~ ' [Diacritics]') if _m.modes else '[Diacritics]' -%}
|
| 22 |
+
{%- endif -%}
|
| 23 |
+
{%- set _has_system = messages and messages[0].role == 'system' -%}
|
| 24 |
+
{%- if _has_system -%}
|
| 25 |
+
{%- if messages[0].content is string -%}
|
| 26 |
+
{%- set system_content = messages[0].content -%}
|
| 27 |
+
{%- elif messages[0].content -%}
|
| 28 |
+
{%- set _sys = namespace(text='') -%}
|
| 29 |
+
{%- for part in messages[0].content -%}
|
| 30 |
+
{%- if part.type == 'text' -%}
|
| 31 |
+
{%- set _sys.text = _sys.text + part.text -%}
|
| 32 |
+
{%- endif -%}
|
| 33 |
+
{%- endfor -%}
|
| 34 |
+
{%- set system_content = _sys.text -%}
|
| 35 |
+
{%- else -%}
|
| 36 |
+
{%- set system_content = '' -%}
|
| 37 |
+
{%- endif -%}
|
| 38 |
+
{%- else -%}
|
| 39 |
+
{%- set system_content = _m.modes -%}
|
| 40 |
+
{%- endif -%}
|
| 41 |
+
{%- if system_content -%}
|
| 42 |
+
{{- '<|ifm|im_start|>system\n' + system_content + '<|ifm|im_end|>' }}
|
| 43 |
+
{%- endif -%}
|
| 44 |
+
{#- ---- turns ------------------------------------------------------------------ -#}
|
| 45 |
+
{%- for message in messages -%}
|
| 46 |
+
{%- if message.content is string -%}
|
| 47 |
+
{%- set content = message.content -%}
|
| 48 |
+
{%- elif not message.content -%}
|
| 49 |
+
{%- set content = '' -%}
|
| 50 |
+
{%- else -%}
|
| 51 |
+
{#- content parts: {"type": "audio"} renders the audio placeholder,
|
| 52 |
+
which the processor expands -#}
|
| 53 |
+
{%- set _ns = namespace(text='') -%}
|
| 54 |
+
{%- for part in message.content -%}
|
| 55 |
+
{%- if part.type == 'audio' -%}
|
| 56 |
+
{%- set _ns.text = _ns.text + '<|ifm|audio|>' -%}
|
| 57 |
+
{%- elif part.type == 'text' -%}
|
| 58 |
+
{%- set _ns.text = _ns.text + part.text -%}
|
| 59 |
+
{%- else -%}
|
| 60 |
+
{{- raise_exception("audio-only model: unsupported content part type '" ~ part.type ~ "' (use 'audio' or 'text')") -}}
|
| 61 |
+
{%- endif -%}
|
| 62 |
+
{%- endfor -%}
|
| 63 |
+
{%- set content = _ns.text -%}
|
| 64 |
+
{%- endif -%}
|
| 65 |
+
{%- if message.role == 'system' -%}
|
| 66 |
+
{%- if not loop.first -%}
|
| 67 |
+
{{- '<|ifm|im_start|>system\n' + content + '<|ifm|im_end|>' }}
|
| 68 |
+
{%- endif -%}
|
| 69 |
+
{%- elif message.role == 'user' -%}
|
| 70 |
+
{{- '<|ifm|im_start|>user\n' + content + '<|ifm|im_end|>' }}
|
| 71 |
+
{%- elif message.role == 'assistant' -%}
|
| 72 |
+
{%- set thinking = message.reasoning if message.reasoning is defined and message.reasoning is string else '' -%}
|
| 73 |
+
{{- '<|ifm|im_start|>assistant' }}
|
| 74 |
+
{%- generation -%}
|
| 75 |
+
{{- '\n<ifm|think>\n' + thinking + '</ifm|think>' + content }}
|
| 76 |
+
{{- '<|ifm|im_end|>' -}}
|
| 77 |
+
{%- endgeneration -%}
|
| 78 |
+
{%- else -%}
|
| 79 |
+
{{- raise_exception("unsupported role '" ~ message.role ~ "' (use system, user or assistant)") -}}
|
| 80 |
+
{%- endif -%}
|
| 81 |
+
{%- endfor -%}
|
| 82 |
+
{#- ---- generation prompt ------------------------------------------------------- -#}
|
| 83 |
+
{%- if add_generation_prompt -%}
|
| 84 |
+
{%- if enable_thinking is defined and enable_thinking -%}
|
| 85 |
+
{%- if language_prefill is defined and language_prefill -%}
|
| 86 |
+
{{- raise_exception("language_prefill= prefill is not supported with enable_thinking=True: the prefill cannot follow an open think block. Drop language_prefill= or disable thinking.") -}}
|
| 87 |
+
{%- endif -%}
|
| 88 |
+
{{- '<|ifm|im_start|>assistant\n<ifm|think>\n' }}
|
| 89 |
+
{%- else -%}
|
| 90 |
+
{{- '<|ifm|im_start|>assistant\n<ifm|think>\n</ifm|think>' }}
|
| 91 |
+
{#- language_prefill= prefills the decision tag into the prompt: a transcript is then
|
| 92 |
+
guaranteed and [IFM-NoTranscript] cannot be produced. Unset = the model
|
| 93 |
+
decides. To pin a language while keeping abstention, use language.LanguagePin. -#}
|
| 94 |
+
{%- if language_prefill is defined and language_prefill -%}
|
| 95 |
+
{%- set _tags = ['msa','en','fr','es','it','EG','SA','MA','IQ','JO','SY','QA','DZ','LB','OM','YE','TN','MR','KW','SD','LY','AE','PS','BH'] -%}
|
| 96 |
+
{%- if language_prefill not in _tags -%}
|
| 97 |
+
{{- raise_exception("language_prefill must be one of " ~ (_tags | join(', ')) ~ " (got '" ~ language_prefill ~ "'). An ISO code such as 'ar' covers several tags and cannot be prefilled; use the LanguagePin logits processor for that.") -}}
|
| 98 |
+
{%- endif -%}
|
| 99 |
+
{{- '[' ~ language_prefill ~ ']' }}
|
| 100 |
+
{%- endif -%}
|
| 101 |
+
{%- endif -%}
|
| 102 |
+
{%- endif -%}
|
config.json
ADDED
|
@@ -0,0 +1,161 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"JaisASRForConditionalGeneration"
|
| 4 |
+
],
|
| 5 |
+
"audio_config": {
|
| 6 |
+
"activation_dropout": 0,
|
| 7 |
+
"activation_function": "gelu",
|
| 8 |
+
"attention_dropout": 0,
|
| 9 |
+
"conv_chunksize": 500,
|
| 10 |
+
"d_model": 1280,
|
| 11 |
+
"downsample_hidden_size": 480,
|
| 12 |
+
"dropout": 0,
|
| 13 |
+
"dtype": "bfloat16",
|
| 14 |
+
"encoder_attention_heads": 20,
|
| 15 |
+
"encoder_ffn_dim": 5120,
|
| 16 |
+
"encoder_layers": 32,
|
| 17 |
+
"initializer_range": 0.02,
|
| 18 |
+
"max_source_positions": 1500,
|
| 19 |
+
"model_type": "qwen3_omni_moe_audio_encoder",
|
| 20 |
+
"n_window": 50,
|
| 21 |
+
"n_window_infer": 800,
|
| 22 |
+
"num_mel_bins": 128,
|
| 23 |
+
"output_dim": 2048,
|
| 24 |
+
"scale_embedding": false,
|
| 25 |
+
"tie_word_embeddings": true,
|
| 26 |
+
"torch_dtype": "bfloat16"
|
| 27 |
+
},
|
| 28 |
+
"audio_end_token_id": 250036,
|
| 29 |
+
"audio_start_token_id": 250035,
|
| 30 |
+
"audio_token_id": 250062,
|
| 31 |
+
"auto_map": {
|
| 32 |
+
"AutoConfig": "configuration_jais_asr.JaisASRConfig",
|
| 33 |
+
"AutoModel": "modeling_jais_asr.JaisASRModel",
|
| 34 |
+
"AutoModelForCausalLM": "modeling_jais_asr.JaisASRForConditionalGeneration",
|
| 35 |
+
"AutoProcessor": "processing_jais_asr.JaisASRProcessor"
|
| 36 |
+
},
|
| 37 |
+
"bos_token_id": 0,
|
| 38 |
+
"custom_pipelines": {
|
| 39 |
+
"automatic-speech-recognition": {
|
| 40 |
+
"impl": "pipeline_jais_asr.JaisASRTranscriptionPipeline",
|
| 41 |
+
"pt": [
|
| 42 |
+
"AutoModelForCausalLM"
|
| 43 |
+
],
|
| 44 |
+
"tf": []
|
| 45 |
+
}
|
| 46 |
+
},
|
| 47 |
+
"do_sample": false,
|
| 48 |
+
"dtype": "bfloat16",
|
| 49 |
+
"eos_token_id": [
|
| 50 |
+
1,
|
| 51 |
+
250019
|
| 52 |
+
],
|
| 53 |
+
"generation_config": {
|
| 54 |
+
"bos_token_id": 0,
|
| 55 |
+
"eos_token_id": [
|
| 56 |
+
1,
|
| 57 |
+
250019
|
| 58 |
+
],
|
| 59 |
+
"pad_token_id": 0,
|
| 60 |
+
"do_sample": false,
|
| 61 |
+
"temperature": 0.0,
|
| 62 |
+
"num_beams": 1,
|
| 63 |
+
"max_new_tokens": 256,
|
| 64 |
+
"repetition_penalty": 1.05,
|
| 65 |
+
"transformers_version": "5.5.4"
|
| 66 |
+
},
|
| 67 |
+
"model_type": "jais_asr",
|
| 68 |
+
"pad_token_id": 0,
|
| 69 |
+
"projector_hidden_act": "gelu",
|
| 70 |
+
"quantization": {
|
| 71 |
+
"group_size": 32,
|
| 72 |
+
"bits": 8,
|
| 73 |
+
"mode": "affine"
|
| 74 |
+
},
|
| 75 |
+
"quantization_config": {
|
| 76 |
+
"group_size": 32,
|
| 77 |
+
"bits": 8,
|
| 78 |
+
"mode": "affine"
|
| 79 |
+
},
|
| 80 |
+
"temperature": 0.0,
|
| 81 |
+
"text_config": {
|
| 82 |
+
"attention_bias": false,
|
| 83 |
+
"attention_dropout": 0.0,
|
| 84 |
+
"bos_token_id": 0,
|
| 85 |
+
"dtype": "bfloat16",
|
| 86 |
+
"eos_token_id": 1,
|
| 87 |
+
"head_dim": 128,
|
| 88 |
+
"hidden_act": "silu",
|
| 89 |
+
"hidden_size": 2560,
|
| 90 |
+
"initializer_range": 0.02,
|
| 91 |
+
"intermediate_size": 10240,
|
| 92 |
+
"layernorm_num_groups": 2,
|
| 93 |
+
"max_position_embeddings": 524288,
|
| 94 |
+
"model_type": "k2_horizon",
|
| 95 |
+
"num_attention_heads": 32,
|
| 96 |
+
"num_hidden_layers": 36,
|
| 97 |
+
"num_key_value_heads": 8,
|
| 98 |
+
"pad_token_id": 0,
|
| 99 |
+
"query_key_norm": false,
|
| 100 |
+
"rms_norm_eps": 1e-06,
|
| 101 |
+
"rope_head_dim": 128,
|
| 102 |
+
"rope_parameters": {
|
| 103 |
+
"rope_theta": 10000000.0,
|
| 104 |
+
"rope_type": "default"
|
| 105 |
+
},
|
| 106 |
+
"rope_theta": 10000000.0,
|
| 107 |
+
"sliding_window": null,
|
| 108 |
+
"tie_word_embeddings": false,
|
| 109 |
+
"torch_dtype": "bfloat16",
|
| 110 |
+
"use_cache": true,
|
| 111 |
+
"use_sliding_window": false,
|
| 112 |
+
"vocab_size": 250624,
|
| 113 |
+
"num_experts": 0,
|
| 114 |
+
"num_experts_per_tok": 0,
|
| 115 |
+
"num_shared_experts": 0,
|
| 116 |
+
"mova_num_experts": 0,
|
| 117 |
+
"mova_num_experts_per_tok": 0,
|
| 118 |
+
"moe_intermediate_size": 0,
|
| 119 |
+
"attention_gate_func": null,
|
| 120 |
+
"mlp_only_layers": [
|
| 121 |
+
0,
|
| 122 |
+
1,
|
| 123 |
+
2,
|
| 124 |
+
3,
|
| 125 |
+
4,
|
| 126 |
+
5,
|
| 127 |
+
6,
|
| 128 |
+
7,
|
| 129 |
+
8,
|
| 130 |
+
9,
|
| 131 |
+
10,
|
| 132 |
+
11,
|
| 133 |
+
12,
|
| 134 |
+
13,
|
| 135 |
+
14,
|
| 136 |
+
15,
|
| 137 |
+
16,
|
| 138 |
+
17,
|
| 139 |
+
18,
|
| 140 |
+
19,
|
| 141 |
+
20,
|
| 142 |
+
21,
|
| 143 |
+
22,
|
| 144 |
+
23,
|
| 145 |
+
24,
|
| 146 |
+
25,
|
| 147 |
+
26,
|
| 148 |
+
27,
|
| 149 |
+
28,
|
| 150 |
+
29,
|
| 151 |
+
30,
|
| 152 |
+
31,
|
| 153 |
+
32,
|
| 154 |
+
33,
|
| 155 |
+
34,
|
| 156 |
+
35
|
| 157 |
+
]
|
| 158 |
+
},
|
| 159 |
+
"tie_word_embeddings": false,
|
| 160 |
+
"transformers_version": "5.5.4"
|
| 161 |
+
}
|
configuration_jais_asr.py
ADDED
|
@@ -0,0 +1,270 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2024-2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Jais-ASR configuration, including its Qwen3-Omni audio and dense K2 components.
|
| 15 |
+
|
| 16 |
+
Component class names identify their upstream architecture; source credits appear
|
| 17 |
+
above each section. Config fields and checkpoint metadata retain their existing API.
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
import transformers as _tf
|
| 21 |
+
|
| 22 |
+
try:
|
| 23 |
+
from transformers.configuration_utils import PreTrainedConfig
|
| 24 |
+
except ImportError:
|
| 25 |
+
from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
|
| 26 |
+
|
| 27 |
+
_MIN = (4, 57)
|
| 28 |
+
_v = tuple(int(x) for x in _tf.__version__.split(".")[:2] if x.isdigit())
|
| 29 |
+
if _v and _v < _MIN:
|
| 30 |
+
raise ImportError(
|
| 31 |
+
f"Jais-ASR requires transformers >= 4.57 (found {_tf.__version__}). "
|
| 32 |
+
"Tested on 4.57.3, 5.3.0 and 5.16.1."
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
# Qwen3-Omni audio encoder: adapted from Hugging Face Transformers v5.16.1.
|
| 36 |
+
# https://github.com/huggingface/transformers/tree/v5.16.1/src/transformers/models/qwen3_omni_moe
|
| 37 |
+
# Retains Jais-ASR's packed chunking, generalized lengths and positional-buffer fixes.
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
class Qwen3OmniMoeAudioEncoderConfig(PreTrainedConfig):
|
| 41 |
+
"""Configuration of the MawJais-stt audio tower."""
|
| 42 |
+
|
| 43 |
+
model_type = "qwen3_omni_moe_audio_encoder"
|
| 44 |
+
|
| 45 |
+
def __init__(
|
| 46 |
+
self,
|
| 47 |
+
num_mel_bins=128,
|
| 48 |
+
encoder_layers=32,
|
| 49 |
+
encoder_attention_heads=20,
|
| 50 |
+
encoder_ffn_dim=5120,
|
| 51 |
+
d_model=1280,
|
| 52 |
+
dropout=0,
|
| 53 |
+
attention_dropout=0,
|
| 54 |
+
activation_function="gelu",
|
| 55 |
+
activation_dropout=0,
|
| 56 |
+
scale_embedding=False,
|
| 57 |
+
initializer_range=0.02,
|
| 58 |
+
max_source_positions=1500,
|
| 59 |
+
n_window=50,
|
| 60 |
+
output_dim=2048,
|
| 61 |
+
n_window_infer=800,
|
| 62 |
+
conv_chunksize=500,
|
| 63 |
+
downsample_hidden_size=480,
|
| 64 |
+
**kwargs,
|
| 65 |
+
):
|
| 66 |
+
super().__init__(**kwargs)
|
| 67 |
+
|
| 68 |
+
self.num_mel_bins = num_mel_bins
|
| 69 |
+
self.d_model = d_model
|
| 70 |
+
self.encoder_layers = encoder_layers
|
| 71 |
+
self.encoder_attention_heads = encoder_attention_heads
|
| 72 |
+
self.encoder_ffn_dim = encoder_ffn_dim
|
| 73 |
+
self.dropout = dropout
|
| 74 |
+
self.attention_dropout = attention_dropout
|
| 75 |
+
self.activation_function = activation_function
|
| 76 |
+
self.activation_dropout = activation_dropout
|
| 77 |
+
self.num_hidden_layers = encoder_layers
|
| 78 |
+
self.initializer_range = initializer_range
|
| 79 |
+
self.scale_embedding = (
|
| 80 |
+
scale_embedding # scale factor will be sqrt(d_model) if True
|
| 81 |
+
)
|
| 82 |
+
self.max_source_positions = max_source_positions
|
| 83 |
+
self.n_window = n_window
|
| 84 |
+
self.output_dim = output_dim
|
| 85 |
+
self.n_window_infer = n_window_infer
|
| 86 |
+
self.conv_chunksize = conv_chunksize
|
| 87 |
+
self.downsample_hidden_size = downsample_hidden_size
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
# Dense K2-Horizon decoder: adapted from IFM/K2-Horizon-3.7B.
|
| 91 |
+
# Source revision: 633f52ad28b17edeabd82afc61d2d13b4c59a561.
|
| 92 |
+
# https://huggingface.co/IFM/K2-Horizon-3.7B/tree/633f52ad28b17edeabd82afc61d2d13b4c59a561
|
| 93 |
+
# Retains grouped RMSNorm, partial RoPE and Transformers 4.57/5.x compatibility.
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
class K2HorizonConfig(PreTrainedConfig):
|
| 97 |
+
"""Dense K2-Horizon configuration used by Jais-ASR.
|
| 98 |
+
|
| 99 |
+
Adapted from IFM/K2-Horizon-3.7B revision 633f52ad28b17edeabd82afc61d2d13b4c59a561.
|
| 100 |
+
Uses an explicit constructor for Transformers 4.57/5.x compatibility.
|
| 101 |
+
|
| 102 |
+
Defaults here are NOT this model's values — always load from config.json.
|
| 103 |
+
The two fields that carry the architecture:
|
| 104 |
+
layernorm_num_groups : groups for the grouped RMSNorm (2 in this checkpoint)
|
| 105 |
+
rope_head_dim : RoPE width per head; None or == head_dim -> full RoPE
|
| 106 |
+
"""
|
| 107 |
+
|
| 108 |
+
model_type = "k2_horizon"
|
| 109 |
+
keys_to_ignore_at_inference = ["past_key_values"]
|
| 110 |
+
|
| 111 |
+
# tensor-parallel / pipeline plans (also read by the vLLM Transformers backend)
|
| 112 |
+
base_model_tp_plan = {
|
| 113 |
+
"layers.*.self_attn.q_proj": "colwise",
|
| 114 |
+
"layers.*.self_attn.k_proj": "colwise",
|
| 115 |
+
"layers.*.self_attn.v_proj": "colwise",
|
| 116 |
+
"layers.*.self_attn.o_proj": "rowwise",
|
| 117 |
+
"layers.*.mlp.gate_proj": "colwise",
|
| 118 |
+
"layers.*.mlp.up_proj": "colwise",
|
| 119 |
+
"layers.*.mlp.down_proj": "rowwise",
|
| 120 |
+
}
|
| 121 |
+
base_model_pp_plan = {
|
| 122 |
+
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
| 123 |
+
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
| 124 |
+
"norm": (["hidden_states"], ["hidden_states"]),
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
def __init__(
|
| 128 |
+
self,
|
| 129 |
+
vocab_size=151936,
|
| 130 |
+
hidden_size=2048,
|
| 131 |
+
intermediate_size=6144,
|
| 132 |
+
num_hidden_layers=24,
|
| 133 |
+
num_attention_heads=32,
|
| 134 |
+
num_key_value_heads=4,
|
| 135 |
+
hidden_act="silu",
|
| 136 |
+
max_position_embeddings=32768,
|
| 137 |
+
initializer_range=0.02,
|
| 138 |
+
rms_norm_eps=1e-6,
|
| 139 |
+
use_cache=True,
|
| 140 |
+
tie_word_embeddings=False,
|
| 141 |
+
rope_theta=10000.0,
|
| 142 |
+
rope_scaling=None,
|
| 143 |
+
attention_bias=False,
|
| 144 |
+
use_sliding_window=False,
|
| 145 |
+
sliding_window=4096,
|
| 146 |
+
attention_dropout=0.0,
|
| 147 |
+
query_key_norm=True,
|
| 148 |
+
layernorm_num_groups=1,
|
| 149 |
+
rope_head_dim=None,
|
| 150 |
+
head_dim=None,
|
| 151 |
+
pad_token_id=None,
|
| 152 |
+
**kwargs,
|
| 153 |
+
):
|
| 154 |
+
# Only the dense K2-Horizon architecture is included in this repository.
|
| 155 |
+
for field in (
|
| 156 |
+
"num_experts",
|
| 157 |
+
"num_experts_per_tok",
|
| 158 |
+
"num_shared_experts",
|
| 159 |
+
"mova_num_experts",
|
| 160 |
+
"mova_num_experts_per_tok",
|
| 161 |
+
"moe_intermediate_size",
|
| 162 |
+
):
|
| 163 |
+
value = kwargs.pop(field, 0)
|
| 164 |
+
if value != 0:
|
| 165 |
+
raise ValueError(f"Jais-ASR includes dense K2 only; {field} must be 0")
|
| 166 |
+
setattr(self, field, 0)
|
| 167 |
+
self.attention_gate_func = kwargs.pop("attention_gate_func", None)
|
| 168 |
+
if self.attention_gate_func is not None:
|
| 169 |
+
raise ValueError("Jais-ASR does not use attention gating")
|
| 170 |
+
self.mlp_only_layers = list(range(num_hidden_layers))
|
| 171 |
+
kwargs.pop("mlp_only_layers", None)
|
| 172 |
+
self.vocab_size = vocab_size
|
| 173 |
+
self.max_position_embeddings = max_position_embeddings
|
| 174 |
+
self.hidden_size = hidden_size
|
| 175 |
+
self.intermediate_size = intermediate_size
|
| 176 |
+
self.num_hidden_layers = num_hidden_layers
|
| 177 |
+
self.num_attention_heads = num_attention_heads
|
| 178 |
+
self.head_dim = (
|
| 179 |
+
head_dim if head_dim is not None else hidden_size // num_attention_heads
|
| 180 |
+
)
|
| 181 |
+
self.use_sliding_window = use_sliding_window
|
| 182 |
+
self.sliding_window = sliding_window if use_sliding_window else None
|
| 183 |
+
|
| 184 |
+
self.num_key_value_heads = num_key_value_heads
|
| 185 |
+
self.hidden_act = hidden_act
|
| 186 |
+
self.initializer_range = initializer_range
|
| 187 |
+
self.rms_norm_eps = rms_norm_eps
|
| 188 |
+
self.use_cache = use_cache
|
| 189 |
+
self.rope_theta = rope_theta
|
| 190 |
+
self.rope_scaling = rope_scaling
|
| 191 |
+
self.attention_bias = attention_bias
|
| 192 |
+
self.attention_dropout = attention_dropout
|
| 193 |
+
# Validate the correctness of rotary position embeddings parameters
|
| 194 |
+
# BC: if there is a 'type' field, move it to 'rope_type'.
|
| 195 |
+
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
| 196 |
+
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
| 197 |
+
if hasattr(PreTrainedConfig, "validate_rope"):
|
| 198 |
+
# transformers>=5: rope_parameters is the source of truth
|
| 199 |
+
self.standardize_rope_params()
|
| 200 |
+
self.validate_rope()
|
| 201 |
+
else:
|
| 202 |
+
from transformers.modeling_rope_utils import rope_config_validation
|
| 203 |
+
|
| 204 |
+
rope_config_validation(self)
|
| 205 |
+
|
| 206 |
+
self.query_key_norm = query_key_norm
|
| 207 |
+
self.layernorm_num_groups = layernorm_num_groups
|
| 208 |
+
self.rope_head_dim = rope_head_dim
|
| 209 |
+
|
| 210 |
+
super().__init__(
|
| 211 |
+
tie_word_embeddings=tie_word_embeddings,
|
| 212 |
+
pad_token_id=pad_token_id,
|
| 213 |
+
**kwargs,
|
| 214 |
+
)
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
# Jais-ASR composition: audio encoder + learned projector + dense text decoder.
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
class JaisASRConfig(PreTrainedConfig):
|
| 221 |
+
model_type = "jais_asr"
|
| 222 |
+
sub_configs = {
|
| 223 |
+
"text_config": K2HorizonConfig,
|
| 224 |
+
"audio_config": Qwen3OmniMoeAudioEncoderConfig,
|
| 225 |
+
}
|
| 226 |
+
|
| 227 |
+
def __init__(
|
| 228 |
+
self,
|
| 229 |
+
text_config=None,
|
| 230 |
+
audio_config=None,
|
| 231 |
+
audio_token_id=250062,
|
| 232 |
+
audio_start_token_id=250035,
|
| 233 |
+
audio_end_token_id=250036,
|
| 234 |
+
projector_hidden_act="gelu",
|
| 235 |
+
**kwargs,
|
| 236 |
+
):
|
| 237 |
+
if isinstance(text_config, dict) or text_config is None:
|
| 238 |
+
text_config = K2HorizonConfig(**(text_config or {}))
|
| 239 |
+
if isinstance(audio_config, dict) or audio_config is None:
|
| 240 |
+
audio_config = Qwen3OmniMoeAudioEncoderConfig(**(audio_config or {}))
|
| 241 |
+
self.text_config = text_config
|
| 242 |
+
self.audio_config = audio_config
|
| 243 |
+
# `<|ifm|audio|>`: the repeated per-frame placeholder the model masked_scatters onto.
|
| 244 |
+
# `<audio>` / `</audio>`: span delimiters only (never scattered); one variable-length run
|
| 245 |
+
# of audio_token_id sits between them — no fixed count, no padding, no maximum.
|
| 246 |
+
self.audio_token_id = audio_token_id
|
| 247 |
+
self.audio_start_token_id = audio_start_token_id
|
| 248 |
+
self.audio_end_token_id = audio_end_token_id
|
| 249 |
+
self.projector_hidden_act = projector_hidden_act
|
| 250 |
+
# the text decoder is NOT tied; mirror whatever the text backbone declares (False here).
|
| 251 |
+
kwargs.setdefault(
|
| 252 |
+
"tie_word_embeddings", getattr(text_config, "tie_word_embeddings", False)
|
| 253 |
+
)
|
| 254 |
+
super().__init__(**kwargs)
|
| 255 |
+
|
| 256 |
+
# convenience passthroughs read by some frameworks
|
| 257 |
+
@property
|
| 258 |
+
def hidden_size(self):
|
| 259 |
+
return self.text_config.hidden_size
|
| 260 |
+
|
| 261 |
+
@property
|
| 262 |
+
def vocab_size(self):
|
| 263 |
+
return self.text_config.vocab_size
|
| 264 |
+
|
| 265 |
+
@vocab_size.setter
|
| 266 |
+
def vocab_size(self, v):
|
| 267 |
+
self.text_config.vocab_size = v
|
| 268 |
+
|
| 269 |
+
|
| 270 |
+
__all__ = ["JaisASRConfig", "K2HorizonConfig", "Qwen3OmniMoeAudioEncoderConfig"]
|
generation_config.json
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token_id": 0,
|
| 3 |
+
"eos_token_id": [
|
| 4 |
+
1,
|
| 5 |
+
250019
|
| 6 |
+
],
|
| 7 |
+
"pad_token_id": 0,
|
| 8 |
+
"do_sample": false,
|
| 9 |
+
"temperature": 0.0,
|
| 10 |
+
"num_beams": 1,
|
| 11 |
+
"max_new_tokens": 256,
|
| 12 |
+
"repetition_penalty": 1.05,
|
| 13 |
+
"transformers_version": "5.5.4"
|
| 14 |
+
}
|
language.py
ADDED
|
@@ -0,0 +1,328 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Choosing the transcription language in advance.
|
| 2 |
+
|
| 3 |
+
By default Jais-ASR decides for itself, at one generated position, between abstaining
|
| 4 |
+
(`[IFM-NoTranscript]`) and a language/dialect tag. This module lets a caller pin that
|
| 5 |
+
language WITHOUT taking away the model's ability to abstain.
|
| 6 |
+
|
| 7 |
+
Why that needs care
|
| 8 |
+
-------------------
|
| 9 |
+
The obvious implementation -- restrict the decision to {pinned_tag, IF} -- silently
|
| 10 |
+
changes the abstention rate: when the pinned tag is not the model's own first choice,
|
| 11 |
+
the refusal token can outscore it even on clear speech, and the model abstains where
|
| 12 |
+
it would have transcribed. So the two decisions are read separately, from the same
|
| 13 |
+
untouched logits:
|
| 14 |
+
|
| 15 |
+
abstain iff logit(IF) > max over ALL tags <- compare against all language tags
|
| 16 |
+
else argmax over the ALLOWED tags <- the caller's pin
|
| 17 |
+
|
| 18 |
+
The abstain/transcribe decision is compared against the model's OWN best tag, never the
|
| 19 |
+
pinned one, so pinning does not change this comparison at a fixed decision prefix. Only which tag gets written changes.
|
| 20 |
+
|
| 21 |
+
Two ways to use it, with different guarantees:
|
| 22 |
+
|
| 23 |
+
constrained `LanguagePin` logits processor. Preserves the strict score comparison
|
| 24 |
+
above (ties transcribe), not unconstrained sampling probabilities. Needs the
|
| 25 |
+
processor to be passed to generate().
|
| 26 |
+
prefill write `[EG]` into the prompt (chat template `language_prefill=`). Works in any
|
| 27 |
+
engine with no extra code, but the tag is then in the PROMPT, so
|
| 28 |
+
`[IFM-NoTranscript]` can never be produced -- you always get text, even
|
| 29 |
+
for silence. Use it when you want a transcript no matter what.
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
from __future__ import annotations
|
| 33 |
+
|
| 34 |
+
import logging
|
| 35 |
+
import math
|
| 36 |
+
|
| 37 |
+
try:
|
| 38 |
+
from .postprocess import (
|
| 39 |
+
LANGUAGE_TAG_IDS,
|
| 40 |
+
BRACKET_OPEN_ID,
|
| 41 |
+
BRACKET_CLOSE_ID,
|
| 42 |
+
has_prefilled_tag,
|
| 43 |
+
THINK_OPEN_ID,
|
| 44 |
+
THINK_CLOSE_ID,
|
| 45 |
+
)
|
| 46 |
+
except ImportError:
|
| 47 |
+
if __package__:
|
| 48 |
+
raise
|
| 49 |
+
from postprocess import (
|
| 50 |
+
LANGUAGE_TAG_IDS,
|
| 51 |
+
BRACKET_OPEN_ID,
|
| 52 |
+
BRACKET_CLOSE_ID,
|
| 53 |
+
has_prefilled_tag,
|
| 54 |
+
THINK_OPEN_ID,
|
| 55 |
+
THINK_CLOSE_ID,
|
| 56 |
+
)
|
| 57 |
+
|
| 58 |
+
__all__ = [
|
| 59 |
+
"SUPPORTED_LANGUAGES",
|
| 60 |
+
"resolve_language",
|
| 61 |
+
"LanguagePin",
|
| 62 |
+
"prefill_for",
|
| 63 |
+
"resolve_language_controls",
|
| 64 |
+
]
|
| 65 |
+
|
| 66 |
+
REFUSAL_DECISION_ID = 5126 # 'IF', the continuation of [IFM-NoTranscript]
|
| 67 |
+
TAG_TO_ID = {t: i for i, t in LANGUAGE_TAG_IDS.items()}
|
| 68 |
+
|
| 69 |
+
#: the 19 country codes plus MSA -- everything the model treats as Arabic
|
| 70 |
+
ARABIC_TAGS = (
|
| 71 |
+
"msa",
|
| 72 |
+
"EG",
|
| 73 |
+
"SA",
|
| 74 |
+
"MA",
|
| 75 |
+
"IQ",
|
| 76 |
+
"JO",
|
| 77 |
+
"SY",
|
| 78 |
+
"QA",
|
| 79 |
+
"DZ",
|
| 80 |
+
"LB",
|
| 81 |
+
"OM",
|
| 82 |
+
"YE",
|
| 83 |
+
"TN",
|
| 84 |
+
"MR",
|
| 85 |
+
"KW",
|
| 86 |
+
"SD",
|
| 87 |
+
"LY",
|
| 88 |
+
"AE",
|
| 89 |
+
"PS",
|
| 90 |
+
"BH",
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
#: ISO 639-1 -> the tags that count as that language. `ar` deliberately maps to ALL the
|
| 94 |
+
#: Arabic tags: pinning "Arabic" should let the model keep choosing the dialect.
|
| 95 |
+
ISO_TO_TAGS = {
|
| 96 |
+
"ar": ARABIC_TAGS,
|
| 97 |
+
"en": ("en",),
|
| 98 |
+
"fr": ("fr",),
|
| 99 |
+
"es": ("es",),
|
| 100 |
+
"it": ("it",),
|
| 101 |
+
}
|
| 102 |
+
|
| 103 |
+
#: what vLLM's SupportsTranscription advertises on /v1/audio/transcriptions --
|
| 104 |
+
#: the ISO codes plus every concrete Jais-ASR tag, so dialect pinning is reachable
|
| 105 |
+
SUPPORTED_LANGUAGES = {
|
| 106 |
+
"ar": "Arabic",
|
| 107 |
+
"en": "English",
|
| 108 |
+
"fr": "French",
|
| 109 |
+
"es": "Spanish",
|
| 110 |
+
"it": "Italian",
|
| 111 |
+
"msa": "Arabic (Modern Standard)",
|
| 112 |
+
**{t: f"Arabic ({t})" for t in ARABIC_TAGS if t != "msa"},
|
| 113 |
+
}
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def resolve_language(language):
|
| 117 |
+
"""Normalise a caller's language spec into the set of allowed tags.
|
| 118 |
+
|
| 119 |
+
Accepts an ISO 639-1 code ("ar"), a model tag ("EG", "msa"), a list of either, or
|
| 120 |
+
None for "no restriction". Raises on anything unrecognised rather than silently
|
| 121 |
+
falling back to auto -- a typo that quietly disables pinning is worse than an error.
|
| 122 |
+
"""
|
| 123 |
+
if language is None:
|
| 124 |
+
return None
|
| 125 |
+
if not isinstance(language, (str, list, tuple)):
|
| 126 |
+
raise ValueError(
|
| 127 |
+
"language must be a tag, ISO code, non-empty list/tuple, or None"
|
| 128 |
+
)
|
| 129 |
+
specs = [language] if isinstance(language, str) else list(language)
|
| 130 |
+
if not specs:
|
| 131 |
+
raise ValueError("language list must not be empty; use None for no restriction")
|
| 132 |
+
tags = []
|
| 133 |
+
for s in specs:
|
| 134 |
+
if not isinstance(s, str):
|
| 135 |
+
raise ValueError("each language must be a string tag or ISO code")
|
| 136 |
+
s = s.strip()
|
| 137 |
+
if s in ISO_TO_TAGS:
|
| 138 |
+
tags.extend(ISO_TO_TAGS[s])
|
| 139 |
+
elif s in TAG_TO_ID:
|
| 140 |
+
tags.append(s)
|
| 141 |
+
elif s.lower() in ISO_TO_TAGS:
|
| 142 |
+
tags.extend(ISO_TO_TAGS[s.lower()])
|
| 143 |
+
else:
|
| 144 |
+
raise ValueError(
|
| 145 |
+
f"unknown language {s!r}. Use an ISO 639-1 code "
|
| 146 |
+
f"({', '.join(sorted(ISO_TO_TAGS))}) or a model tag "
|
| 147 |
+
f"({', '.join(sorted(TAG_TO_ID))})."
|
| 148 |
+
)
|
| 149 |
+
seen, out = set(), []
|
| 150 |
+
for t in tags:
|
| 151 |
+
if t not in seen:
|
| 152 |
+
seen.add(t)
|
| 153 |
+
out.append(t)
|
| 154 |
+
return out
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def prefill_for(language):
|
| 158 |
+
"""The literal prefix for the PREFILL path, e.g. 'EG' -> '[EG]'.
|
| 159 |
+
|
| 160 |
+
Only a single concrete tag can be prefilled -- you cannot write a set into a prompt.
|
| 161 |
+
Returns None when the spec resolves to more than one tag. The prefill ends at ']':
|
| 162 |
+
the space before the first transcript word is the model's own to produce.
|
| 163 |
+
"""
|
| 164 |
+
tags = resolve_language(language)
|
| 165 |
+
if not tags or len(tags) != 1:
|
| 166 |
+
return None
|
| 167 |
+
return f"[{tags[0]}]"
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class LanguagePin:
|
| 171 |
+
"""Pin the language while leaving the abstention decision untouched.
|
| 172 |
+
|
| 173 |
+
Engine-agnostic: `apply(generated_ids, logits)` is the whole rule. Thin adapters for
|
| 174 |
+
transformers (`__call__`) and for any engine that hands over (token_ids, logits).
|
| 175 |
+
|
| 176 |
+
pin = LanguagePin("EG") # or "ar", or ["EG", "MA"]
|
| 177 |
+
model.generate(**inputs, logits_processor=[pin.for_transformers(prompt_len)])
|
| 178 |
+
"""
|
| 179 |
+
|
| 180 |
+
def __init__(self, language, force_bracket=True):
|
| 181 |
+
self.tags = resolve_language(language)
|
| 182 |
+
self.allowed_ids = [TAG_TO_ID[t] for t in self.tags] if self.tags else None
|
| 183 |
+
self.all_tag_ids = list(LANGUAGE_TAG_IDS)
|
| 184 |
+
self.force_bracket = force_bracket
|
| 185 |
+
|
| 186 |
+
def apply(self, generated_ids, logits, *, thinking=False):
|
| 187 |
+
"""Constrain the decision zone in-place and return the score vector.
|
| 188 |
+
|
| 189 |
+
Prompt tokens must be excluded. If the prompt opens a thinking block,
|
| 190 |
+
pass thinking=True: scores are untouched until the generated closing
|
| 191 |
+
marker, and the decision zone starts immediately after that marker.
|
| 192 |
+
The adapters infer this flag from the prompt automatically.
|
| 193 |
+
"""
|
| 194 |
+
if self.allowed_ids is None:
|
| 195 |
+
return logits
|
| 196 |
+
if thinking:
|
| 197 |
+
for index, token_id in enumerate(generated_ids):
|
| 198 |
+
if int(token_id) == THINK_CLOSE_ID:
|
| 199 |
+
generated_ids = generated_ids[index + 1 :]
|
| 200 |
+
break
|
| 201 |
+
else:
|
| 202 |
+
return logits
|
| 203 |
+
n_generated = len(generated_ids)
|
| 204 |
+
if n_generated > 2:
|
| 205 |
+
return logits
|
| 206 |
+
if n_generated == 0:
|
| 207 |
+
return _keep_only(logits, BRACKET_OPEN_ID) if self.force_bracket else logits
|
| 208 |
+
# With force_bracket=False, non-tag output must remain untouched.
|
| 209 |
+
if int(generated_ids[0]) != BRACKET_OPEN_ID:
|
| 210 |
+
return logits
|
| 211 |
+
if n_generated == 2:
|
| 212 |
+
if int(generated_ids[1]) in self.allowed_ids:
|
| 213 |
+
return _keep_only(logits, BRACKET_CLOSE_ID)
|
| 214 |
+
return logits # Let [IFM-NoTranscript] continue naturally.
|
| 215 |
+
|
| 216 |
+
# Read both decisions before masking. Refusal must strictly beat every tag;
|
| 217 |
+
# ties transcribe. This is a deterministic score rule, not preservation of
|
| 218 |
+
# the probability distribution of unconstrained sampling.
|
| 219 |
+
best_tag = max(self.all_tag_ids, key=lambda i: float(logits[i]))
|
| 220 |
+
abstains = float(logits[REFUSAL_DECISION_ID]) > float(logits[best_tag])
|
| 221 |
+
keep_id = (
|
| 222 |
+
REFUSAL_DECISION_ID
|
| 223 |
+
if abstains
|
| 224 |
+
else max(self.allowed_ids, key=lambda i: float(logits[i]))
|
| 225 |
+
)
|
| 226 |
+
return _keep_only(logits, keep_id)
|
| 227 |
+
|
| 228 |
+
def for_transformers(self, prompt_len):
|
| 229 |
+
"""Bind to the padded prompt width; supports batch and expanded beam rows."""
|
| 230 |
+
if not isinstance(prompt_len, int) or prompt_len <= 0:
|
| 231 |
+
raise ValueError("prompt_len must be the positive padded prompt width")
|
| 232 |
+
outer = self
|
| 233 |
+
|
| 234 |
+
class _Proc:
|
| 235 |
+
def __call__(self, input_ids, scores):
|
| 236 |
+
if input_ids.shape[1] < prompt_len:
|
| 237 |
+
raise ValueError("input_ids is shorter than the bound prompt_len")
|
| 238 |
+
if outer.allowed_ids is None:
|
| 239 |
+
return scores
|
| 240 |
+
# Stateless normal-prompt fast path after the decision zone.
|
| 241 |
+
if input_ids.shape[1] - prompt_len > 2:
|
| 242 |
+
if bool((input_ids[:, prompt_len - 1] == THINK_CLOSE_ID).all()):
|
| 243 |
+
return scores
|
| 244 |
+
for b in range(scores.shape[0]):
|
| 245 |
+
prompt = input_ids[b, :prompt_len]
|
| 246 |
+
if has_prefilled_tag(prompt):
|
| 247 |
+
continue
|
| 248 |
+
outer.apply(
|
| 249 |
+
input_ids[b, prompt_len:],
|
| 250 |
+
scores[b],
|
| 251 |
+
thinking=has_open_thinking(prompt),
|
| 252 |
+
)
|
| 253 |
+
return scores
|
| 254 |
+
|
| 255 |
+
return _Proc()
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
def _keep_only(logits, token_id):
|
| 259 |
+
keep = logits[token_id].clone() if hasattr(logits, "clone") else logits[token_id]
|
| 260 |
+
if not math.isfinite(float(keep)):
|
| 261 |
+
raise ValueError(
|
| 262 |
+
"language pin selected a token without a finite score; check conflicting "
|
| 263 |
+
"token constraints or invalid logits"
|
| 264 |
+
)
|
| 265 |
+
logits[:] = float("-inf")
|
| 266 |
+
logits[token_id] = keep
|
| 267 |
+
return logits
|
| 268 |
+
|
| 269 |
+
|
| 270 |
+
def has_open_thinking(prompt_ids):
|
| 271 |
+
"""Whether the final thinking delimiter in the prompt is an opening one."""
|
| 272 |
+
# Tensor reversal copies the prompt and iteration unbinds every element.
|
| 273 |
+
# Index backward so a final closing delimiter costs constant work.
|
| 274 |
+
for index in range(len(prompt_ids) - 1, -1, -1):
|
| 275 |
+
token_id = int(prompt_ids[index])
|
| 276 |
+
if token_id == THINK_CLOSE_ID:
|
| 277 |
+
return False
|
| 278 |
+
if token_id == THINK_OPEN_ID:
|
| 279 |
+
return True
|
| 280 |
+
return False
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
#: The exact prompt contract, rendered here so engines that build their own prompt
|
| 284 |
+
#: (vLLM's transcription endpoint) do not have to run the jinja template. Kept in
|
| 285 |
+
#: exact agreement with chat_template.jinja: the system turn is rendered only when
|
| 286 |
+
#: there is system content.
|
| 287 |
+
GENERATION_PROMPT = (
|
| 288 |
+
"<|ifm|begin_of_text|>"
|
| 289 |
+
"{system_turn}"
|
| 290 |
+
"<|ifm|im_start|>user\n{audio}<|ifm|im_end|>"
|
| 291 |
+
"<|ifm|im_start|>assistant\n<ifm|think>\n</ifm|think>{prefill}"
|
| 292 |
+
)
|
| 293 |
+
SYSTEM_TURN = "<|ifm|im_start|>system\n{system}<|ifm|im_end|>"
|
| 294 |
+
AUDIO_TOKEN = "<|ifm|audio|>"
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
def resolve_language_controls(
|
| 298 |
+
language=None, language_prefill=None, *, enable_thinking=False
|
| 299 |
+
):
|
| 300 |
+
"""Validate mutually exclusive pin/prefill options and canonicalize prefill."""
|
| 301 |
+
if language is not None and language_prefill is not None:
|
| 302 |
+
raise ValueError("Use either language (pinning) or language_prefill, not both")
|
| 303 |
+
if language is not None:
|
| 304 |
+
resolve_language(language)
|
| 305 |
+
if language_prefill is not None:
|
| 306 |
+
if not isinstance(language_prefill, str):
|
| 307 |
+
raise ValueError("language_prefill must be one concrete language tag")
|
| 308 |
+
tags = resolve_language(language_prefill)
|
| 309 |
+
if len(tags) != 1:
|
| 310 |
+
raise ValueError(
|
| 311 |
+
"language_prefill must resolve to one tag; use language for a family"
|
| 312 |
+
)
|
| 313 |
+
if enable_thinking:
|
| 314 |
+
raise ValueError(
|
| 315 |
+
"language_prefill cannot be combined with enable_thinking=True"
|
| 316 |
+
)
|
| 317 |
+
language_prefill = tags[0]
|
| 318 |
+
return language, language_prefill
|
| 319 |
+
|
| 320 |
+
|
| 321 |
+
def build_prompt(system="", language_prefill=None, audio_token=AUDIO_TOKEN):
|
| 322 |
+
"""Render the ASR prompt; only language_prefill inserts a decision tag."""
|
| 323 |
+
_, tag = resolve_language_controls(language_prefill=language_prefill)
|
| 324 |
+
return GENERATION_PROMPT.format(
|
| 325 |
+
system_turn=SYSTEM_TURN.format(system=system) if system else "",
|
| 326 |
+
audio=audio_token,
|
| 327 |
+
prefill=f"[{tag}]" if tag else "",
|
| 328 |
+
)
|
mlx-vlm-jais-asr.patch
ADDED
|
@@ -0,0 +1,583 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
diff --git a/mlx_vlm/generate/ar.py b/mlx_vlm/generate/ar.py
|
| 2 |
+
index f2f9ae50..f91e06a3 100644
|
| 3 |
+
--- a/mlx_vlm/generate/ar.py
|
| 4 |
+
+++ b/mlx_vlm/generate/ar.py
|
| 5 |
+
@@ -321,6 +321,12 @@ def generate_step(
|
| 6 |
+
)
|
| 7 |
+
if logits_processors is not None:
|
| 8 |
+
processors.extend(logits_processors)
|
| 9 |
+
+ # Model-specific generate options, e.g. the Jais-ASR language pin.
|
| 10 |
+
+ generation_processors = getattr(model, "generation_processors", None)
|
| 11 |
+
+ if generation_processors is not None:
|
| 12 |
+
+ processors.extend(generation_processors(input_ids, kwargs))
|
| 13 |
+
+ elif kwargs.get("language") is not None:
|
| 14 |
+
+ raise ValueError(f"{type(model).__name__} does not support the language option")
|
| 15 |
+
|
| 16 |
+
y = input_ids
|
| 17 |
+
tokens = mx.array([], dtype=input_ids.dtype)
|
| 18 |
+
diff --git a/mlx_vlm/generate/dispatch.py b/mlx_vlm/generate/dispatch.py
|
| 19 |
+
index 585eadeb..f98eb64f 100644
|
| 20 |
+
--- a/mlx_vlm/generate/dispatch.py
|
| 21 |
+
+++ b/mlx_vlm/generate/dispatch.py
|
| 22 |
+
@@ -417,6 +417,13 @@ def parse_arguments():
|
| 23 |
+
default=None,
|
| 24 |
+
help="Penalty factor for previously generated tokens.",
|
| 25 |
+
)
|
| 26 |
+
+ parser.add_argument(
|
| 27 |
+
+ "--language",
|
| 28 |
+
+ type=str,
|
| 29 |
+
+ default=None,
|
| 30 |
+
+ help="Transcription language for models that support pinning it (Jais-ASR): "
|
| 31 |
+
+ "an ISO 639-1 code (ar, en, fr, es, it) or a model tag (EG, MA, msa, ...).",
|
| 32 |
+
+ )
|
| 33 |
+
parser.add_argument(
|
| 34 |
+
"--repetition-context-size",
|
| 35 |
+
type=int,
|
| 36 |
+
@@ -1550,6 +1557,7 @@ def main():
|
| 37 |
+
"min_p": args.min_p,
|
| 38 |
+
"repetition_penalty": args.repetition_penalty,
|
| 39 |
+
"repetition_context_size": args.repetition_context_size,
|
| 40 |
+
+ **({"language": args.language} if args.language else {}),
|
| 41 |
+
"presence_penalty": args.presence_penalty,
|
| 42 |
+
"presence_context_size": args.presence_context_size,
|
| 43 |
+
"frequency_penalty": args.frequency_penalty,
|
| 44 |
+
@@ -1599,6 +1607,7 @@ def main():
|
| 45 |
+
"max_tokens": args.max_tokens,
|
| 46 |
+
"repetition_penalty": args.repetition_penalty,
|
| 47 |
+
"repetition_context_size": args.repetition_context_size,
|
| 48 |
+
+ **({"language": args.language} if args.language else {}),
|
| 49 |
+
"presence_penalty": args.presence_penalty,
|
| 50 |
+
"presence_context_size": args.presence_context_size,
|
| 51 |
+
"frequency_penalty": args.frequency_penalty,
|
| 52 |
+
diff --git a/mlx_vlm/models/jais_asr/__init__.py b/mlx_vlm/models/jais_asr/__init__.py
|
| 53 |
+
new file mode 100644
|
| 54 |
+
index 00000000..baf44b0a
|
| 55 |
+
--- /dev/null
|
| 56 |
+
+++ b/mlx_vlm/models/jais_asr/__init__.py
|
| 57 |
+
@@ -0,0 +1,13 @@
|
| 58 |
+
+from ..k2_horizon.language import LanguageModel
|
| 59 |
+
+from ..qwen3_omni_moe.audio import AudioModel
|
| 60 |
+
+from .config import AudioConfig, ModelConfig, TextConfig
|
| 61 |
+
+from .jais_asr import Model
|
| 62 |
+
+
|
| 63 |
+
+__all__ = [
|
| 64 |
+
+ "AudioConfig",
|
| 65 |
+
+ "AudioModel",
|
| 66 |
+
+ "LanguageModel",
|
| 67 |
+
+ "Model",
|
| 68 |
+
+ "ModelConfig",
|
| 69 |
+
+ "TextConfig",
|
| 70 |
+
+]
|
| 71 |
+
diff --git a/mlx_vlm/models/jais_asr/config.py b/mlx_vlm/models/jais_asr/config.py
|
| 72 |
+
new file mode 100644
|
| 73 |
+
index 00000000..d2c304b7
|
| 74 |
+
--- /dev/null
|
| 75 |
+
+++ b/mlx_vlm/models/jais_asr/config.py
|
| 76 |
+
@@ -0,0 +1,27 @@
|
| 77 |
+
+from dataclasses import dataclass
|
| 78 |
+
+from typing import List, Optional, Union
|
| 79 |
+
+
|
| 80 |
+
+from ..base import BaseModelConfig
|
| 81 |
+
+from ..k2_horizon.config import ModelConfig as TextConfig
|
| 82 |
+
+from ..qwen3_omni_moe.config import AudioConfig
|
| 83 |
+
+
|
| 84 |
+
+
|
| 85 |
+
+@dataclass
|
| 86 |
+
+class ModelConfig(BaseModelConfig):
|
| 87 |
+
+ text_config: TextConfig = None
|
| 88 |
+
+ audio_config: AudioConfig = None
|
| 89 |
+
+ model_type: str = "jais_asr"
|
| 90 |
+
+ audio_token_id: int = 250062
|
| 91 |
+
+ audio_start_token_id: int = 250035
|
| 92 |
+
+ audio_end_token_id: int = 250036
|
| 93 |
+
+ projector_hidden_act: str = "gelu"
|
| 94 |
+
+ eos_token_id: Optional[Union[int, List[int]]] = None
|
| 95 |
+
+
|
| 96 |
+
+ def __post_init__(self):
|
| 97 |
+
+ if self.text_config is None or isinstance(self.text_config, dict):
|
| 98 |
+
+ self.text_config = TextConfig.from_dict(self.text_config or {})
|
| 99 |
+
+ if self.audio_config is None or isinstance(self.audio_config, dict):
|
| 100 |
+
+ self.audio_config = AudioConfig.from_dict(self.audio_config or {})
|
| 101 |
+
+
|
| 102 |
+
+
|
| 103 |
+
+__all__ = ["AudioConfig", "ModelConfig", "TextConfig"]
|
| 104 |
+
diff --git a/mlx_vlm/models/jais_asr/jais_asr.py b/mlx_vlm/models/jais_asr/jais_asr.py
|
| 105 |
+
new file mode 100644
|
| 106 |
+
index 00000000..e5a79f66
|
| 107 |
+
--- /dev/null
|
| 108 |
+
+++ b/mlx_vlm/models/jais_asr/jais_asr.py
|
| 109 |
+
@@ -0,0 +1,141 @@
|
| 110 |
+
+"""Jais-ASR: Qwen3-Omni audio encoder + projector + K2-Horizon decoder.
|
| 111 |
+
+
|
| 112 |
+
+Mirrors ``JaisASRForConditionalGeneration`` of MBZUAI-Paris/Jais-ASR: the audio
|
| 113 |
+
+encoder output is projected to the decoder width and written into the
|
| 114 |
+
+``<|ifm|audio|>`` placeholder positions of the prompt, in order.
|
| 115 |
+
+"""
|
| 116 |
+
+
|
| 117 |
+
+from typing import Optional
|
| 118 |
+
+
|
| 119 |
+
+import mlx.core as mx
|
| 120 |
+
+import mlx.nn as nn
|
| 121 |
+
+
|
| 122 |
+
+from ..base import InputEmbeddingsFeatures, LanguageModelOutput
|
| 123 |
+
+from ..k2_horizon.language import LanguageModel
|
| 124 |
+
+from ..qwen3_omni_moe.audio import AudioModel
|
| 125 |
+
+from .config import ModelConfig
|
| 126 |
+
+from .language import LanguagePin
|
| 127 |
+
+
|
| 128 |
+
+
|
| 129 |
+
+class Projector(nn.Module):
|
| 130 |
+
+ """LayerNorm -> Linear -> GELU -> Linear (audio width -> decoder width)."""
|
| 131 |
+
+
|
| 132 |
+
+ def __init__(self, config: ModelConfig):
|
| 133 |
+
+ super().__init__()
|
| 134 |
+
+ in_dim = config.audio_config.output_dim
|
| 135 |
+
+ out_dim = config.text_config.hidden_size
|
| 136 |
+
+ if config.projector_hidden_act != "gelu":
|
| 137 |
+
+ raise ValueError(f"Unsupported projector activation: {config.projector_hidden_act}")
|
| 138 |
+
+ self.norm = nn.LayerNorm(in_dim)
|
| 139 |
+
+ self.linear_1 = nn.Linear(in_dim, out_dim)
|
| 140 |
+
+ self.linear_2 = nn.Linear(out_dim, out_dim)
|
| 141 |
+
+
|
| 142 |
+
+ def __call__(self, x: mx.array) -> mx.array:
|
| 143 |
+
+ return self.linear_2(nn.gelu(self.linear_1(self.norm(x))))
|
| 144 |
+
+
|
| 145 |
+
+
|
| 146 |
+
+class Model(nn.Module):
|
| 147 |
+
+ def __init__(self, config: ModelConfig):
|
| 148 |
+
+ super().__init__()
|
| 149 |
+
+ self.config = config
|
| 150 |
+
+ self.model_type = config.model_type
|
| 151 |
+
+ self.audio_tower = AudioModel(config.audio_config)
|
| 152 |
+
+ # Named like other mlx-vlm projectors so quantization leaves it in full precision.
|
| 153 |
+
+ self.multi_modal_projector = Projector(config)
|
| 154 |
+
+ self.language_model = LanguageModel(config.text_config)
|
| 155 |
+
+
|
| 156 |
+
+ def get_audio_features(
|
| 157 |
+
+ self,
|
| 158 |
+
+ input_features: mx.array,
|
| 159 |
+
+ feature_attention_mask: Optional[mx.array] = None,
|
| 160 |
+
+ ) -> mx.array:
|
| 161 |
+
+ """Pack the clips of a batch, encode them, project to the decoder width."""
|
| 162 |
+
+ if input_features.ndim == 2:
|
| 163 |
+
+ input_features = input_features[None]
|
| 164 |
+
+ batch, _, frames = input_features.shape
|
| 165 |
+
+ if feature_attention_mask is None:
|
| 166 |
+
+ lengths = [frames] * batch
|
| 167 |
+
+ else:
|
| 168 |
+
+ lengths = [int(n) for n in feature_attention_mask.sum(axis=-1).tolist()]
|
| 169 |
+
+ packed = mx.concatenate(
|
| 170 |
+
+ [input_features[i, :, :n] for i, n in enumerate(lengths)], axis=-1
|
| 171 |
+
+ )
|
| 172 |
+
+ dtype = self.audio_tower.conv_out.weight.dtype
|
| 173 |
+
+ hidden = self.audio_tower(
|
| 174 |
+
+ packed.astype(dtype), feature_lens=mx.array(lengths, dtype=mx.int32)
|
| 175 |
+
+ )
|
| 176 |
+
+ return self.multi_modal_projector(hidden)
|
| 177 |
+
+
|
| 178 |
+
+ def get_input_embeddings(
|
| 179 |
+
+ self,
|
| 180 |
+
+ input_ids: Optional[mx.array] = None,
|
| 181 |
+
+ pixel_values: Optional[mx.array] = None,
|
| 182 |
+
+ **kwargs,
|
| 183 |
+
+ ) -> InputEmbeddingsFeatures:
|
| 184 |
+
+ embeds = self.language_model.model.embed_tokens(input_ids)
|
| 185 |
+
+ input_features = kwargs.get("input_features")
|
| 186 |
+
+ if input_features is None:
|
| 187 |
+
+ return InputEmbeddingsFeatures(inputs_embeds=embeds)
|
| 188 |
+
+ placeholders = input_ids == self.config.audio_token_id
|
| 189 |
+
+ n_placeholders = int(placeholders.sum())
|
| 190 |
+
+ if n_placeholders == 0:
|
| 191 |
+
+ return InputEmbeddingsFeatures(inputs_embeds=embeds)
|
| 192 |
+
+
|
| 193 |
+
+ audio = self.get_audio_features(
|
| 194 |
+
+ input_features, kwargs.get("feature_attention_mask")
|
| 195 |
+
+ ).astype(embeds.dtype)
|
| 196 |
+
+ if audio.shape[0] != n_placeholders:
|
| 197 |
+
+ raise ValueError(
|
| 198 |
+
+ f"Audio placeholder tokens ({n_placeholders}) != audio frames "
|
| 199 |
+
+ f"({audio.shape[0]}); build prompts with the Jais-ASR processor."
|
| 200 |
+
+ )
|
| 201 |
+
+ # Placeholders are filled in sequence order, as masked_scatter does.
|
| 202 |
+
+ flat = placeholders.reshape(-1)
|
| 203 |
+
+ order = mx.maximum(mx.cumsum(flat.astype(mx.int32)) - 1, 0)
|
| 204 |
+
+ hidden = embeds.shape[-1]
|
| 205 |
+
+ merged = mx.where(flat[:, None], audio[order], embeds.reshape(-1, hidden))
|
| 206 |
+
+ return InputEmbeddingsFeatures(inputs_embeds=merged.reshape(embeds.shape))
|
| 207 |
+
+
|
| 208 |
+
+ def __call__(
|
| 209 |
+
+ self,
|
| 210 |
+
+ input_ids: mx.array,
|
| 211 |
+
+ pixel_values: Optional[mx.array] = None,
|
| 212 |
+
+ mask: Optional[mx.array] = None,
|
| 213 |
+
+ cache=None,
|
| 214 |
+
+ **kwargs,
|
| 215 |
+
+ ) -> LanguageModelOutput:
|
| 216 |
+
+ inputs_embeds = self.get_input_embeddings(input_ids, **kwargs).inputs_embeds
|
| 217 |
+
+ return self.language_model(input_ids, cache=cache, inputs_embeds=inputs_embeds)
|
| 218 |
+
+
|
| 219 |
+
+ def sanitize(self, weights):
|
| 220 |
+
+ """Map Jais-ASR (Transformers) names to this module tree; idempotent."""
|
| 221 |
+
+ renamed = {}
|
| 222 |
+
+ for key, value in weights.items():
|
| 223 |
+
+ if key.startswith("model.audio_tower."):
|
| 224 |
+
+ key = key.replace("model.audio_tower.", "audio_tower.", 1)
|
| 225 |
+
+ elif key.startswith("model.embed_audio."):
|
| 226 |
+
+ key = key.replace("model.embed_audio.", "multi_modal_projector.", 1)
|
| 227 |
+
+ elif key.startswith("model.language_model."):
|
| 228 |
+
+ key = key.replace("model.language_model.", "language_model.model.", 1)
|
| 229 |
+
+ elif key == "lm_head.weight":
|
| 230 |
+
+ key = "language_model.lm_head.weight"
|
| 231 |
+
+ renamed[key] = value
|
| 232 |
+
+ return renamed
|
| 233 |
+
+
|
| 234 |
+
+ def generation_processors(self, input_ids, options):
|
| 235 |
+
+ """Logits processors for the generate options of this model.
|
| 236 |
+
+
|
| 237 |
+
+ `language` pins the transcription language (see language.py). It is removed
|
| 238 |
+
+ from `options` so it never reaches the forward pass.
|
| 239 |
+
+ """
|
| 240 |
+
+ language = options.pop("language", None)
|
| 241 |
+
+ if language is None:
|
| 242 |
+
+ return []
|
| 243 |
+
+ return [LanguagePin(language).for_mlx(input_ids.reshape(-1).tolist())]
|
| 244 |
+
+
|
| 245 |
+
+ def make_cache(self):
|
| 246 |
+
+ return self.language_model.make_cache()
|
| 247 |
+
+
|
| 248 |
+
+ @property
|
| 249 |
+
+ def layers(self):
|
| 250 |
+
+ return self.language_model.layers
|
| 251 |
+
diff --git a/mlx_vlm/models/jais_asr/language.py b/mlx_vlm/models/jais_asr/language.py
|
| 252 |
+
new file mode 100644
|
| 253 |
+
index 00000000..174a49d7
|
| 254 |
+
--- /dev/null
|
| 255 |
+
+++ b/mlx_vlm/models/jais_asr/language.py
|
| 256 |
+
@@ -0,0 +1,242 @@
|
| 257 |
+
+"""Choosing the Jais-ASR transcription language in advance (language pin).
|
| 258 |
+
+
|
| 259 |
+
+Same rule as `language.py` of the Jais-ASR model repository, for MLX arrays.
|
| 260 |
+
+
|
| 261 |
+
+By default Jais-ASR decides for itself, at one generated position, between abstaining
|
| 262 |
+
+(`[IFM-NoTranscript]`) and a language/dialect tag. Pinning restricts which tag is written
|
| 263 |
+
+without taking away the model's ability to abstain: the two decisions are read
|
| 264 |
+
+separately, from the same untouched scores,
|
| 265 |
+
+
|
| 266 |
+
+ abstain iff score(IF) > max over ALL tags <- the model's own best tag
|
| 267 |
+
+ else argmax over the ALLOWED tags <- the caller's pin
|
| 268 |
+
+
|
| 269 |
+
+so the abstain/transcribe decision does not change; only which tag gets written does.
|
| 270 |
+
+The other way to choose the language, prefill (`[EG]` written at the end of the prompt),
|
| 271 |
+
+needs no code but can never abstain.
|
| 272 |
+
+
|
| 273 |
+
+Usage: `generate(model, processor, prompt, audio=[...], language="EG")`, or
|
| 274 |
+
+`mlx_vlm.generate ... --language EG`. `language` accepts an ISO 639-1 code ("ar" keeps
|
| 275 |
+
+every Arabic tag, so the model still chooses the dialect), a model tag ("EG", "msa"),
|
| 276 |
+
+or a list of either.
|
| 277 |
+
+"""
|
| 278 |
+
+
|
| 279 |
+
+import math
|
| 280 |
+
+
|
| 281 |
+
+import mlx.core as mx
|
| 282 |
+
+
|
| 283 |
+
+# Token ids of the Jais-ASR tokenizer (model repository `postprocess.py`).
|
| 284 |
+
+BRACKET_OPEN_ID = 60 # '['
|
| 285 |
+
+BRACKET_CLOSE_ID = 62 # ']'
|
| 286 |
+
+THINK_OPEN_ID = 250029 # <ifm|think>
|
| 287 |
+
+THINK_CLOSE_ID = 250030 # '</ifm|think>', the last token of every rendered prompt
|
| 288 |
+
+REFUSAL_DECISION_ID = 5126 # 'IF', the continuation of [IFM-NoTranscript]
|
| 289 |
+
+
|
| 290 |
+
+#: decision-token id -> tag surface. 24 tags, all distinct.
|
| 291 |
+
+LANGUAGE_TAG_IDS = {
|
| 292 |
+
+ 226606: "msa", # Modern Standard Arabic
|
| 293 |
+
+ 272: "en",
|
| 294 |
+
+ 2075: "fr",
|
| 295 |
+
+ 273: "es",
|
| 296 |
+
+ 286: "it",
|
| 297 |
+
+ 7646: "EG",
|
| 298 |
+
+ 8048: "SA",
|
| 299 |
+
+ 6101: "MA",
|
| 300 |
+
+ 29354: "IQ",
|
| 301 |
+
+ 23129: "JO",
|
| 302 |
+
+ 20056: "SY",
|
| 303 |
+
+ 27526: "QA",
|
| 304 |
+
+ 64489: "DZ",
|
| 305 |
+
+ 23852: "LB",
|
| 306 |
+
+ 5594: "OM",
|
| 307 |
+
+ 28472: "YE",
|
| 308 |
+
+ 43852: "TN",
|
| 309 |
+
+ 17604: "MR",
|
| 310 |
+
+ 25893: "KW",
|
| 311 |
+
+ 9373: "SD",
|
| 312 |
+
+ 12933: "LY",
|
| 313 |
+
+ 5402: "AE",
|
| 314 |
+
+ 4295: "PS",
|
| 315 |
+
+ 41905: "BH",
|
| 316 |
+
+}
|
| 317 |
+
+TAG_TO_ID = {t: i for i, t in LANGUAGE_TAG_IDS.items()}
|
| 318 |
+
+
|
| 319 |
+
+#: the 19 country codes plus MSA -- everything the model treats as Arabic
|
| 320 |
+
+ARABIC_TAGS = (
|
| 321 |
+
+ "msa",
|
| 322 |
+
+ "EG",
|
| 323 |
+
+ "SA",
|
| 324 |
+
+ "MA",
|
| 325 |
+
+ "IQ",
|
| 326 |
+
+ "JO",
|
| 327 |
+
+ "SY",
|
| 328 |
+
+ "QA",
|
| 329 |
+
+ "DZ",
|
| 330 |
+
+ "LB",
|
| 331 |
+
+ "OM",
|
| 332 |
+
+ "YE",
|
| 333 |
+
+ "TN",
|
| 334 |
+
+ "MR",
|
| 335 |
+
+ "KW",
|
| 336 |
+
+ "SD",
|
| 337 |
+
+ "LY",
|
| 338 |
+
+ "AE",
|
| 339 |
+
+ "PS",
|
| 340 |
+
+ "BH",
|
| 341 |
+
+)
|
| 342 |
+
+
|
| 343 |
+
+#: ISO 639-1 -> the tags that count as that language. `ar` deliberately maps to ALL the
|
| 344 |
+
+#: Arabic tags: pinning "Arabic" should let the model keep choosing the dialect.
|
| 345 |
+
+ISO_TO_TAGS = {
|
| 346 |
+
+ "ar": ARABIC_TAGS,
|
| 347 |
+
+ "en": ("en",),
|
| 348 |
+
+ "fr": ("fr",),
|
| 349 |
+
+ "es": ("es",),
|
| 350 |
+
+ "it": ("it",),
|
| 351 |
+
+}
|
| 352 |
+
+
|
| 353 |
+
+
|
| 354 |
+
+def resolve_language(language):
|
| 355 |
+
+ """Normalise a caller's language spec into the list of allowed tags.
|
| 356 |
+
+
|
| 357 |
+
+ Accepts an ISO 639-1 code ("ar"), a model tag ("EG", "msa"), a list of either, or
|
| 358 |
+
+ None for "no restriction". Raises on anything unrecognised rather than silently
|
| 359 |
+
+ falling back to auto -- a typo that quietly disables pinning is worse than an error.
|
| 360 |
+
+ """
|
| 361 |
+
+ if language is None:
|
| 362 |
+
+ return None
|
| 363 |
+
+ if not isinstance(language, (str, list, tuple)):
|
| 364 |
+
+ raise ValueError(
|
| 365 |
+
+ "language must be a tag, ISO code, non-empty list/tuple, or None"
|
| 366 |
+
+ )
|
| 367 |
+
+ specs = [language] if isinstance(language, str) else list(language)
|
| 368 |
+
+ if not specs:
|
| 369 |
+
+ raise ValueError("language list must not be empty; use None for no restriction")
|
| 370 |
+
+ tags = []
|
| 371 |
+
+ for s in specs:
|
| 372 |
+
+ if not isinstance(s, str):
|
| 373 |
+
+ raise ValueError("each language must be a string tag or ISO code")
|
| 374 |
+
+ s = s.strip()
|
| 375 |
+
+ if s in ISO_TO_TAGS:
|
| 376 |
+
+ tags.extend(ISO_TO_TAGS[s])
|
| 377 |
+
+ elif s in TAG_TO_ID:
|
| 378 |
+
+ tags.append(s)
|
| 379 |
+
+ elif s.lower() in ISO_TO_TAGS:
|
| 380 |
+
+ tags.extend(ISO_TO_TAGS[s.lower()])
|
| 381 |
+
+ else:
|
| 382 |
+
+ raise ValueError(
|
| 383 |
+
+ f"unknown language {s!r}. Use an ISO 639-1 code "
|
| 384 |
+
+ f"({', '.join(sorted(ISO_TO_TAGS))}) or a model tag "
|
| 385 |
+
+ f"({', '.join(sorted(TAG_TO_ID))})."
|
| 386 |
+
+ )
|
| 387 |
+
+ seen, out = set(), []
|
| 388 |
+
+ for t in tags:
|
| 389 |
+
+ if t not in seen:
|
| 390 |
+
+ seen.add(t)
|
| 391 |
+
+ out.append(t)
|
| 392 |
+
+ return out
|
| 393 |
+
+
|
| 394 |
+
+
|
| 395 |
+
+def has_prefilled_tag(prompt_ids):
|
| 396 |
+
+ """Only a complete known [language] suffix counts as a prompt prefill."""
|
| 397 |
+
+ return (
|
| 398 |
+
+ len(prompt_ids) >= 3
|
| 399 |
+
+ and int(prompt_ids[-3]) == BRACKET_OPEN_ID
|
| 400 |
+
+ and int(prompt_ids[-2]) in LANGUAGE_TAG_IDS
|
| 401 |
+
+ and int(prompt_ids[-1]) == BRACKET_CLOSE_ID
|
| 402 |
+
+ )
|
| 403 |
+
+
|
| 404 |
+
+
|
| 405 |
+
+def has_open_thinking(prompt_ids):
|
| 406 |
+
+ """Whether the final thinking delimiter in the prompt is an opening one."""
|
| 407 |
+
+ for index in range(len(prompt_ids) - 1, -1, -1):
|
| 408 |
+
+ token_id = int(prompt_ids[index])
|
| 409 |
+
+ if token_id == THINK_CLOSE_ID:
|
| 410 |
+
+ return False
|
| 411 |
+
+ if token_id == THINK_OPEN_ID:
|
| 412 |
+
+ return True
|
| 413 |
+
+ return False
|
| 414 |
+
+
|
| 415 |
+
+
|
| 416 |
+
+def _keep_only(logits, token_id):
|
| 417 |
+
+ keep = logits[token_id]
|
| 418 |
+
+ if not math.isfinite(keep.item()):
|
| 419 |
+
+ raise ValueError(
|
| 420 |
+
+ "language pin selected a token without a finite score; check conflicting "
|
| 421 |
+
+ "token constraints or invalid logits"
|
| 422 |
+
+ )
|
| 423 |
+
+ return mx.where(mx.arange(logits.shape[-1]) == token_id, logits, -mx.inf)
|
| 424 |
+
+
|
| 425 |
+
+
|
| 426 |
+
+class LanguagePin:
|
| 427 |
+
+ """Pin the language while leaving the abstention decision untouched."""
|
| 428 |
+
+
|
| 429 |
+
+ def __init__(self, language, force_bracket=True):
|
| 430 |
+
+ self.tags = resolve_language(language)
|
| 431 |
+
+ self.allowed_ids = [TAG_TO_ID[t] for t in self.tags] if self.tags else None
|
| 432 |
+
+ self.all_tag_ids = list(LANGUAGE_TAG_IDS)
|
| 433 |
+
+ self.force_bracket = force_bracket
|
| 434 |
+
+
|
| 435 |
+
+ def apply(self, generated_ids, logits, *, thinking=False):
|
| 436 |
+
+ """Constrain the decision zone and return the score vector.
|
| 437 |
+
+
|
| 438 |
+
+ `generated_ids` excludes the prompt. If the prompt opens a thinking block, pass
|
| 439 |
+
+ thinking=True: scores are untouched until the generated closing marker, and
|
| 440 |
+
+ the decision zone starts immediately after that marker.
|
| 441 |
+
+ """
|
| 442 |
+
+ if self.allowed_ids is None:
|
| 443 |
+
+ return logits
|
| 444 |
+
+ if thinking:
|
| 445 |
+
+ for index, token_id in enumerate(generated_ids):
|
| 446 |
+
+ if int(token_id) == THINK_CLOSE_ID:
|
| 447 |
+
+ generated_ids = generated_ids[index + 1 :]
|
| 448 |
+
+ break
|
| 449 |
+
+ else:
|
| 450 |
+
+ return logits
|
| 451 |
+
+ n_generated = len(generated_ids)
|
| 452 |
+
+ if n_generated > 2:
|
| 453 |
+
+ return logits
|
| 454 |
+
+ if n_generated == 0:
|
| 455 |
+
+ return _keep_only(logits, BRACKET_OPEN_ID) if self.force_bracket else logits
|
| 456 |
+
+ # With force_bracket=False, non-tag output must remain untouched.
|
| 457 |
+
+ if int(generated_ids[0]) != BRACKET_OPEN_ID:
|
| 458 |
+
+ return logits
|
| 459 |
+
+ if n_generated == 2:
|
| 460 |
+
+ if int(generated_ids[1]) in self.allowed_ids:
|
| 461 |
+
+ return _keep_only(logits, BRACKET_CLOSE_ID)
|
| 462 |
+
+ return logits # Let [IFM-NoTranscript] continue naturally.
|
| 463 |
+
+
|
| 464 |
+
+ # Read both decisions before masking. Refusal must strictly beat every tag;
|
| 465 |
+
+ # ties transcribe. This is a deterministic score rule, not preservation of
|
| 466 |
+
+ # the probability distribution of unconstrained sampling.
|
| 467 |
+
+ ids = self.all_tag_ids + [REFUSAL_DECISION_ID]
|
| 468 |
+
+ scores = dict(zip(ids, logits[mx.array(ids)].astype(mx.float32).tolist()))
|
| 469 |
+
+ best_tag = max(self.all_tag_ids, key=lambda i: scores[i])
|
| 470 |
+
+ abstains = scores[REFUSAL_DECISION_ID] > scores[best_tag]
|
| 471 |
+
+ keep_id = (
|
| 472 |
+
+ REFUSAL_DECISION_ID
|
| 473 |
+
+ if abstains
|
| 474 |
+
+ else max(self.allowed_ids, key=lambda i: scores[i])
|
| 475 |
+
+ )
|
| 476 |
+
+ return _keep_only(logits, keep_id)
|
| 477 |
+
+
|
| 478 |
+
+ def for_mlx(self, prompt_ids):
|
| 479 |
+
+ """mlx-vlm logits processor: (tokens, logits [1, vocab]) -> logits.
|
| 480 |
+
+
|
| 481 |
+
+ mlx-vlm passes the end of the prompt (what its chunked prefill left) on the first
|
| 482 |
+
+ call and one more token per call after that, so the generated tokens are the ones
|
| 483 |
+
+ after the length seen on the first call.
|
| 484 |
+
+ """
|
| 485 |
+
+ if has_prefilled_tag(prompt_ids):
|
| 486 |
+
+ raise ValueError("language cannot be combined with a prefilled language tag")
|
| 487 |
+
+ thinking = has_open_thinking(prompt_ids)
|
| 488 |
+
+ prompt_len = []
|
| 489 |
+
+
|
| 490 |
+
+ def processor(tokens, logits):
|
| 491 |
+
+ if not prompt_len:
|
| 492 |
+
+ prompt_len.append(len(tokens))
|
| 493 |
+
+ generated = tokens[prompt_len[0]:]
|
| 494 |
+
+ if self.allowed_ids is None or len(generated) > 2 and not thinking:
|
| 495 |
+
+ return logits
|
| 496 |
+
+ return self.apply(generated.tolist(), logits[0], thinking=thinking)[None]
|
| 497 |
+
+
|
| 498 |
+
+ return processor
|
| 499 |
+
diff --git a/mlx_vlm/models/k2_horizon/config.py b/mlx_vlm/models/k2_horizon/config.py
|
| 500 |
+
index bf1dbc6b..550c5697 100644
|
| 501 |
+
--- a/mlx_vlm/models/k2_horizon/config.py
|
| 502 |
+
+++ b/mlx_vlm/models/k2_horizon/config.py
|
| 503 |
+
@@ -20,6 +20,7 @@ class ModelConfig(BaseModelConfig):
|
| 504 |
+
rope_scaling: Optional[Dict[str, Any]] = None
|
| 505 |
+
attention_bias: bool = False
|
| 506 |
+
tie_word_embeddings: bool = False
|
| 507 |
+
+ layernorm_num_groups: int = 1
|
| 508 |
+
|
| 509 |
+
@classmethod
|
| 510 |
+
def from_dict(cls, params: Dict[str, Any]):
|
| 511 |
+
diff --git a/mlx_vlm/models/k2_horizon/language.py b/mlx_vlm/models/k2_horizon/language.py
|
| 512 |
+
index f3e4d5ef..eb32e0c8 100644
|
| 513 |
+
--- a/mlx_vlm/models/k2_horizon/language.py
|
| 514 |
+
+++ b/mlx_vlm/models/k2_horizon/language.py
|
| 515 |
+
@@ -14,6 +14,30 @@ from ..rope_utils import initialize_rope
|
| 516 |
+
from .config import ModelConfig
|
| 517 |
+
|
| 518 |
+
|
| 519 |
+
+class GroupedRMSNorm(nn.Module):
|
| 520 |
+
+ """RMSNorm over ``groups`` equal slices of the last axis (K2HorizonRMSNorm).
|
| 521 |
+
+
|
| 522 |
+
+ Statistics and the weight product are computed in float32, then cast back,
|
| 523 |
+
+ as in the reference implementation.
|
| 524 |
+
+ """
|
| 525 |
+
+
|
| 526 |
+
+ def __init__(self, dims: int, groups: int = 1, eps: float = 1e-6):
|
| 527 |
+
+ super().__init__()
|
| 528 |
+
+ if dims % groups:
|
| 529 |
+
+ raise ValueError(f"hidden size {dims} is not divisible by {groups} groups")
|
| 530 |
+
+ self.weight = mx.ones((dims,))
|
| 531 |
+
+ self.groups = groups
|
| 532 |
+
+ self.eps = eps
|
| 533 |
+
+
|
| 534 |
+
+ def __call__(self, x: mx.array) -> mx.array:
|
| 535 |
+
+ if self.groups == 1:
|
| 536 |
+
+ return mx.fast.rms_norm(x, self.weight, self.eps)
|
| 537 |
+
+ shape = x.shape
|
| 538 |
+
+ y = x.astype(mx.float32).reshape(*shape[:-1], self.groups, shape[-1] // self.groups)
|
| 539 |
+
+ y = y * mx.rsqrt(mx.mean(mx.square(y), axis=-1, keepdims=True) + self.eps)
|
| 540 |
+
+ return (self.weight.astype(mx.float32) * y.reshape(shape)).astype(x.dtype)
|
| 541 |
+
+
|
| 542 |
+
+
|
| 543 |
+
class Attention(nn.Module):
|
| 544 |
+
def __init__(self, args: ModelConfig):
|
| 545 |
+
super().__init__()
|
| 546 |
+
@@ -74,9 +98,11 @@ class TransformerBlock(nn.Module):
|
| 547 |
+
super().__init__()
|
| 548 |
+
self.self_attn = Attention(args)
|
| 549 |
+
self.mlp = SwiGLUMLP(args.hidden_size, args.intermediate_size)
|
| 550 |
+
- self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
| 551 |
+
- self.post_attention_layernorm = nn.RMSNorm(
|
| 552 |
+
- args.hidden_size, eps=args.rms_norm_eps
|
| 553 |
+
+ self.input_layernorm = GroupedRMSNorm(
|
| 554 |
+
+ args.hidden_size, args.layernorm_num_groups, eps=args.rms_norm_eps
|
| 555 |
+
+ )
|
| 556 |
+
+ self.post_attention_layernorm = GroupedRMSNorm(
|
| 557 |
+
+ args.hidden_size, args.layernorm_num_groups, eps=args.rms_norm_eps
|
| 558 |
+
)
|
| 559 |
+
|
| 560 |
+
def __call__(
|
| 561 |
+
@@ -95,7 +121,9 @@ class K2HorizonModel(nn.Module):
|
| 562 |
+
self.args = args
|
| 563 |
+
self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
|
| 564 |
+
self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
|
| 565 |
+
- self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
|
| 566 |
+
+ self.norm = GroupedRMSNorm(
|
| 567 |
+
+ args.hidden_size, args.layernorm_num_groups, eps=args.rms_norm_eps
|
| 568 |
+
+ )
|
| 569 |
+
|
| 570 |
+
def __call__(
|
| 571 |
+
self,
|
| 572 |
+
diff --git a/mlx_vlm/prompt_utils.py b/mlx_vlm/prompt_utils.py
|
| 573 |
+
index 96afd7de..857d61a0 100644
|
| 574 |
+
--- a/mlx_vlm/prompt_utils.py
|
| 575 |
+
+++ b/mlx_vlm/prompt_utils.py
|
| 576 |
+
@@ -54,6 +54,7 @@ MODEL_CONFIG = {
|
| 577 |
+
"qwen3_5_moe": MessageFormat.LIST_WITH_IMAGE_FIRST,
|
| 578 |
+
"qwen4_exp": MessageFormat.LIST_WITH_IMAGE_FIRST,
|
| 579 |
+
"qwen3_omni_moe": MessageFormat.LIST_WITH_IMAGE_FIRST,
|
| 580 |
+
+ "jais_asr": MessageFormat.LIST_WITH_IMAGE_FIRST,
|
| 581 |
+
"minicpmo": MessageFormat.IMAGE_TOKEN,
|
| 582 |
+
"minicpmv4_6": MessageFormat.IMAGE_TOKEN_WRAPPED,
|
| 583 |
+
"mistral3": MessageFormat.LIST_WITH_IMAGE_FIRST,
|
model-00001-of-00002.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0d614b3a08298dcf3a4f8e5126fb5a8269c5cd26ea32bf36909d826ea6eefe84
|
| 3 |
+
size 5362405082
|
model-00002-of-00002.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:62996500bd1feb1c780e47178f029b9fa0beaab7bb9364a6192f6efe00c7d8fe
|
| 3 |
+
size 1647930813
|
model.safetensors.index.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
modeling_jais_asr.py
ADDED
|
@@ -0,0 +1,1555 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2024-2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
|
| 2 |
+
#
|
| 3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
| 4 |
+
# you may not use this file except in compliance with the License.
|
| 5 |
+
# You may obtain a copy of the License at
|
| 6 |
+
#
|
| 7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
| 8 |
+
#
|
| 9 |
+
# Unless required by applicable law or agreed to in writing, software
|
| 10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
| 11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 12 |
+
# See the License for the specific language governing permissions and
|
| 13 |
+
# limitations under the License.
|
| 14 |
+
"""Jais-ASR model with explicitly attributed Qwen3-Omni and K2-Horizon components.
|
| 15 |
+
|
| 16 |
+
Sections: shared attention helper, Qwen audio encoder, K2 dense text decoder,
|
| 17 |
+
and Jais-ASR composition. Parameter names and generation APIs are unchanged.
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
import copy
|
| 21 |
+
import inspect as _inspect
|
| 22 |
+
import math
|
| 23 |
+
from typing import Callable, Optional
|
| 24 |
+
|
| 25 |
+
import numpy as np
|
| 26 |
+
import torch
|
| 27 |
+
from torch import nn
|
| 28 |
+
from torch.nn import functional as F
|
| 29 |
+
from transformers.activations import ACT2FN
|
| 30 |
+
from transformers.cache_utils import Cache, DynamicCache
|
| 31 |
+
from transformers.generation import GenerationMixin
|
| 32 |
+
from transformers.masking_utils import (
|
| 33 |
+
create_causal_mask,
|
| 34 |
+
create_sliding_window_causal_mask,
|
| 35 |
+
)
|
| 36 |
+
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
|
| 37 |
+
from transformers.modeling_layers import GradientCheckpointingLayer
|
| 38 |
+
from transformers.modeling_outputs import (
|
| 39 |
+
BaseModelOutput,
|
| 40 |
+
BaseModelOutputWithPast,
|
| 41 |
+
CausalLMOutputWithPast,
|
| 42 |
+
)
|
| 43 |
+
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
| 44 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
| 45 |
+
from transformers.processing_utils import Unpack
|
| 46 |
+
from transformers.utils import TransformersKwargs
|
| 47 |
+
from transformers.utils import generic as _generic
|
| 48 |
+
from transformers.utils.deprecation import deprecate_kwarg
|
| 49 |
+
|
| 50 |
+
try:
|
| 51 |
+
from .configuration_jais_asr import (
|
| 52 |
+
JaisASRConfig,
|
| 53 |
+
K2HorizonConfig,
|
| 54 |
+
Qwen3OmniMoeAudioEncoderConfig,
|
| 55 |
+
)
|
| 56 |
+
except ImportError:
|
| 57 |
+
if __package__:
|
| 58 |
+
raise
|
| 59 |
+
from configuration_jais_asr import (
|
| 60 |
+
JaisASRConfig,
|
| 61 |
+
K2HorizonConfig,
|
| 62 |
+
Qwen3OmniMoeAudioEncoderConfig,
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
# Shared attention utility (identical in the upstream-derived components).
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 69 |
+
"""
|
| 70 |
+
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
| 71 |
+
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
| 72 |
+
"""
|
| 73 |
+
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
| 74 |
+
if n_rep == 1:
|
| 75 |
+
return hidden_states
|
| 76 |
+
hidden_states = hidden_states[:, :, None, :, :].expand(
|
| 77 |
+
batch, num_key_value_heads, n_rep, slen, head_dim
|
| 78 |
+
)
|
| 79 |
+
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
# Qwen3-Omni audio encoder: adapted from Hugging Face Transformers v5.16.1.
|
| 83 |
+
# https://github.com/huggingface/transformers/tree/v5.16.1/src/transformers/models/qwen3_omni_moe
|
| 84 |
+
# Retains Jais-ASR's packed chunking, generalized lengths and positional-buffer fixes.
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def materialize_sinusoids_(module) -> None:
|
| 88 |
+
"""(Re)compute a SinusoidsPositionEmbedding's non-persistent buffer IN PLACE.
|
| 89 |
+
|
| 90 |
+
The sinusoid table is registered persistent=False, so it is NOT in the checkpoint.
|
| 91 |
+
transformers >= 5.x builds the model on the meta device (discarding what __init__
|
| 92 |
+
computed) and re-materializes modules through _init_weights, so both _init_weights
|
| 93 |
+
hooks in this repo call this function; without it the buffer would be uninitialized
|
| 94 |
+
memory on those versions.
|
| 95 |
+
"""
|
| 96 |
+
log_timescale_increment = np.log(module.max_timescale) / (module.channels // 2 - 1)
|
| 97 |
+
inv_timescales = torch.exp(
|
| 98 |
+
-log_timescale_increment * torch.arange(module.channels // 2).float()
|
| 99 |
+
)
|
| 100 |
+
scaled_time = (
|
| 101 |
+
torch.arange(module.length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
|
| 102 |
+
)
|
| 103 |
+
emb = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
|
| 104 |
+
with torch.no_grad():
|
| 105 |
+
module.positional_embedding.copy_(emb.to(module.positional_embedding.dtype))
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
class Qwen3OmniMoeAudioPreTrainedModel(PreTrainedModel):
|
| 109 |
+
"""Base class carrying the loading / attention flags for the audio tower."""
|
| 110 |
+
|
| 111 |
+
config_class = Qwen3OmniMoeAudioEncoderConfig
|
| 112 |
+
base_model_prefix = "model"
|
| 113 |
+
supports_gradient_checkpointing = True
|
| 114 |
+
_no_split_modules = ["Qwen3OmniMoeAudioEncoderLayer"]
|
| 115 |
+
_skip_keys_device_placement = "past_key_values"
|
| 116 |
+
_supports_flash_attn = True
|
| 117 |
+
_supports_sdpa = True
|
| 118 |
+
_can_compile_fullgraph = False
|
| 119 |
+
_supports_attention_backend = True
|
| 120 |
+
|
| 121 |
+
def _init_weights(self, module):
|
| 122 |
+
super()._init_weights(module)
|
| 123 |
+
if isinstance(module, SinusoidsPositionEmbedding):
|
| 124 |
+
materialize_sinusoids_(module)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def qwen_audio_eager_attention_forward(
|
| 128 |
+
module: nn.Module,
|
| 129 |
+
query: torch.Tensor,
|
| 130 |
+
key: torch.Tensor,
|
| 131 |
+
value: torch.Tensor,
|
| 132 |
+
attention_mask: Optional[torch.Tensor],
|
| 133 |
+
scaling: float,
|
| 134 |
+
dropout: float = 0.0,
|
| 135 |
+
**kwargs,
|
| 136 |
+
):
|
| 137 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 138 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 139 |
+
|
| 140 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 141 |
+
if attention_mask is not None:
|
| 142 |
+
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
| 143 |
+
attn_weights = attn_weights + causal_mask
|
| 144 |
+
|
| 145 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
|
| 146 |
+
query.dtype
|
| 147 |
+
)
|
| 148 |
+
attn_weights = nn.functional.dropout(
|
| 149 |
+
attn_weights, p=dropout, training=module.training
|
| 150 |
+
)
|
| 151 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 152 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 153 |
+
|
| 154 |
+
return attn_output, attn_weights
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def _get_feat_extract_output_lengths(input_lengths, chunk_size=100):
|
| 158 |
+
"""
|
| 159 |
+
Computes the output length of the convolutional layers and the output length of the audio encoder
|
| 160 |
+
"""
|
| 161 |
+
|
| 162 |
+
chunk_output_size = chunk_size
|
| 163 |
+
for _ in range(3): # three stride-2 convolutions on each (padded) chunk
|
| 164 |
+
chunk_output_size = (chunk_output_size - 1) // 2 + 1
|
| 165 |
+
input_lengths_leave = input_lengths % chunk_size
|
| 166 |
+
feat_lengths = (input_lengths_leave - 1) // 2 + 1
|
| 167 |
+
output_lengths = (
|
| 168 |
+
((feat_lengths - 1) // 2 + 1 - 1) // 2
|
| 169 |
+
+ 1
|
| 170 |
+
+ (input_lengths // chunk_size) * chunk_output_size
|
| 171 |
+
)
|
| 172 |
+
return output_lengths
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
class SinusoidsPositionEmbedding(nn.Module):
|
| 176 |
+
def __init__(self, length, channels, max_timescale=10000):
|
| 177 |
+
super().__init__()
|
| 178 |
+
self.length = length
|
| 179 |
+
self.channels = channels
|
| 180 |
+
self.max_timescale = max_timescale
|
| 181 |
+
if channels % 2 != 0:
|
| 182 |
+
raise ValueError("SinusoidsPositionEmbedding needs even channels input")
|
| 183 |
+
log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
|
| 184 |
+
inv_timescales = torch.exp(
|
| 185 |
+
-log_timescale_increment * torch.arange(channels // 2).float()
|
| 186 |
+
)
|
| 187 |
+
scaled_time = (
|
| 188 |
+
torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
|
| 189 |
+
)
|
| 190 |
+
self.register_buffer(
|
| 191 |
+
"positional_embedding",
|
| 192 |
+
torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1),
|
| 193 |
+
persistent=False,
|
| 194 |
+
)
|
| 195 |
+
|
| 196 |
+
def forward(self, seqlen: int):
|
| 197 |
+
return self.positional_embedding[:seqlen, :]
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
class Qwen3OmniMoeAudioAttention(nn.Module):
|
| 201 |
+
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
| 202 |
+
|
| 203 |
+
def __init__(self, config):
|
| 204 |
+
super().__init__()
|
| 205 |
+
self.embed_dim = config.d_model
|
| 206 |
+
self.num_heads = config.encoder_attention_heads
|
| 207 |
+
self.dropout = config.attention_dropout
|
| 208 |
+
self.head_dim = self.embed_dim // self.num_heads
|
| 209 |
+
self.num_key_value_groups = 1 # needed for eager attention
|
| 210 |
+
self.config = config
|
| 211 |
+
|
| 212 |
+
if (self.head_dim * self.num_heads) != self.embed_dim:
|
| 213 |
+
raise ValueError(
|
| 214 |
+
f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
|
| 215 |
+
f" and `num_heads`: {self.num_heads})."
|
| 216 |
+
)
|
| 217 |
+
self.scaling = self.head_dim**-0.5
|
| 218 |
+
self.attention_dropout = 0.0
|
| 219 |
+
self.is_decoder = False
|
| 220 |
+
self.is_causal = False
|
| 221 |
+
self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
|
| 222 |
+
self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
|
| 223 |
+
self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
|
| 224 |
+
self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
|
| 225 |
+
|
| 226 |
+
def forward(
|
| 227 |
+
self,
|
| 228 |
+
hidden_states: torch.Tensor,
|
| 229 |
+
cu_seqlens: torch.Tensor,
|
| 230 |
+
max_seqlen: int | None = None,
|
| 231 |
+
**kwargs,
|
| 232 |
+
) -> torch.Tensor:
|
| 233 |
+
"""Input shape: Batch x Time x Channel"""
|
| 234 |
+
|
| 235 |
+
seq_length, _ = hidden_states.size()
|
| 236 |
+
|
| 237 |
+
query_states = self.q_proj(hidden_states).reshape(
|
| 238 |
+
seq_length, self.num_heads, -1
|
| 239 |
+
)
|
| 240 |
+
key_states = self.k_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
|
| 241 |
+
value_states = self.v_proj(hidden_states).reshape(
|
| 242 |
+
seq_length, self.num_heads, -1
|
| 243 |
+
)
|
| 244 |
+
|
| 245 |
+
query_states = query_states.transpose(0, 1).unsqueeze(0)
|
| 246 |
+
key_states = key_states.transpose(0, 1).unsqueeze(0)
|
| 247 |
+
value_states = value_states.transpose(0, 1).unsqueeze(0)
|
| 248 |
+
|
| 249 |
+
attention_interface: Callable = (
|
| 250 |
+
qwen_audio_eager_attention_forward
|
| 251 |
+
if self.config._attn_implementation == "eager"
|
| 252 |
+
else ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
| 253 |
+
)
|
| 254 |
+
|
| 255 |
+
if self.config._attn_implementation in (
|
| 256 |
+
"flash_attention_2",
|
| 257 |
+
"flash_attention_3",
|
| 258 |
+
):
|
| 259 |
+
# Flash Attention: Use cu_seqlens for variable length attention
|
| 260 |
+
max_seqlen = (
|
| 261 |
+
int((cu_seqlens[1:] - cu_seqlens[:-1]).max())
|
| 262 |
+
if max_seqlen is None
|
| 263 |
+
else max_seqlen
|
| 264 |
+
)
|
| 265 |
+
attn_output, _ = attention_interface(
|
| 266 |
+
self,
|
| 267 |
+
query_states,
|
| 268 |
+
key_states,
|
| 269 |
+
value_states,
|
| 270 |
+
attention_mask=None,
|
| 271 |
+
scaling=self.scaling,
|
| 272 |
+
dropout=0.0 if not self.training else self.attention_dropout,
|
| 273 |
+
cu_seq_lens_q=cu_seqlens,
|
| 274 |
+
cu_seq_lens_k=cu_seqlens,
|
| 275 |
+
max_length_q=max_seqlen,
|
| 276 |
+
max_length_k=max_seqlen,
|
| 277 |
+
is_causal=False,
|
| 278 |
+
**kwargs,
|
| 279 |
+
)
|
| 280 |
+
else:
|
| 281 |
+
# Other implementations: Process each chunk separately
|
| 282 |
+
lengths = cu_seqlens[1:] - cu_seqlens[:-1]
|
| 283 |
+
splits = [
|
| 284 |
+
torch.split(tensor, lengths.tolist(), dim=2)
|
| 285 |
+
for tensor in (query_states, key_states, value_states)
|
| 286 |
+
]
|
| 287 |
+
attn_outputs = [
|
| 288 |
+
attention_interface(
|
| 289 |
+
self,
|
| 290 |
+
q,
|
| 291 |
+
k,
|
| 292 |
+
v,
|
| 293 |
+
attention_mask=None,
|
| 294 |
+
scaling=self.scaling,
|
| 295 |
+
dropout=0.0 if not self.training else self.attention_dropout,
|
| 296 |
+
is_causal=False,
|
| 297 |
+
**kwargs,
|
| 298 |
+
)[0]
|
| 299 |
+
for q, k, v in zip(*splits)
|
| 300 |
+
]
|
| 301 |
+
attn_output = torch.cat(attn_outputs, dim=1)
|
| 302 |
+
|
| 303 |
+
attn_output = attn_output.reshape(seq_length, -1).contiguous()
|
| 304 |
+
attn_output = self.out_proj(attn_output)
|
| 305 |
+
|
| 306 |
+
return attn_output
|
| 307 |
+
|
| 308 |
+
|
| 309 |
+
class Qwen3OmniMoeAudioEncoderLayer(GradientCheckpointingLayer):
|
| 310 |
+
def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig):
|
| 311 |
+
super().__init__()
|
| 312 |
+
self.embed_dim = config.d_model
|
| 313 |
+
self.self_attn = Qwen3OmniMoeAudioAttention(config)
|
| 314 |
+
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
|
| 315 |
+
self.dropout = config.dropout
|
| 316 |
+
self.activation_fn = ACT2FN[config.activation_function]
|
| 317 |
+
self.activation_dropout = config.activation_dropout
|
| 318 |
+
self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
|
| 319 |
+
self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
|
| 320 |
+
self.final_layer_norm = nn.LayerNorm(self.embed_dim)
|
| 321 |
+
|
| 322 |
+
def forward(
|
| 323 |
+
self,
|
| 324 |
+
hidden_states: torch.Tensor,
|
| 325 |
+
cu_seqlens: torch.Tensor,
|
| 326 |
+
attention_mask=None,
|
| 327 |
+
**kwargs,
|
| 328 |
+
) -> torch.Tensor:
|
| 329 |
+
"""
|
| 330 |
+
Args:
|
| 331 |
+
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
|
| 332 |
+
attention_mask (`torch.FloatTensor`): attention mask of size
|
| 333 |
+
`(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
|
| 334 |
+
"""
|
| 335 |
+
residual = hidden_states
|
| 336 |
+
hidden_states = self.self_attn_layer_norm(hidden_states)
|
| 337 |
+
hidden_states = self.self_attn(
|
| 338 |
+
hidden_states=hidden_states,
|
| 339 |
+
cu_seqlens=cu_seqlens,
|
| 340 |
+
**kwargs,
|
| 341 |
+
)
|
| 342 |
+
hidden_states = residual + hidden_states
|
| 343 |
+
residual = hidden_states
|
| 344 |
+
hidden_states = self.final_layer_norm(hidden_states)
|
| 345 |
+
hidden_states = self.fc1(hidden_states)
|
| 346 |
+
hidden_states = self.activation_fn(hidden_states)
|
| 347 |
+
hidden_states = self.fc2(hidden_states)
|
| 348 |
+
hidden_states = residual + hidden_states
|
| 349 |
+
|
| 350 |
+
if hidden_states.dtype == torch.float16:
|
| 351 |
+
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
| 352 |
+
hidden_states = torch.clamp(
|
| 353 |
+
hidden_states, min=-clamp_value, max=clamp_value
|
| 354 |
+
)
|
| 355 |
+
|
| 356 |
+
outputs = (hidden_states,)
|
| 357 |
+
|
| 358 |
+
return outputs
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
class Qwen3OmniMoeAudioEncoder(Qwen3OmniMoeAudioPreTrainedModel):
|
| 362 |
+
config: Qwen3OmniMoeAudioEncoderConfig
|
| 363 |
+
main_input_name = "input_features"
|
| 364 |
+
_no_split_modules = ["Qwen3OmniMoeAudioEncoderLayer"]
|
| 365 |
+
_supports_sdpa = True
|
| 366 |
+
|
| 367 |
+
def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig):
|
| 368 |
+
super().__init__(config)
|
| 369 |
+
self.dropout = config.dropout
|
| 370 |
+
|
| 371 |
+
embed_dim = config.d_model
|
| 372 |
+
self.num_mel_bins = config.num_mel_bins
|
| 373 |
+
self.max_source_positions = config.max_source_positions
|
| 374 |
+
self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
|
| 375 |
+
self.n_window = config.n_window
|
| 376 |
+
self.positional_embedding = SinusoidsPositionEmbedding(
|
| 377 |
+
self.max_source_positions, embed_dim
|
| 378 |
+
)
|
| 379 |
+
self.layers = nn.ModuleList(
|
| 380 |
+
[
|
| 381 |
+
Qwen3OmniMoeAudioEncoderLayer(config)
|
| 382 |
+
for _ in range(config.encoder_layers)
|
| 383 |
+
]
|
| 384 |
+
)
|
| 385 |
+
self.ln_post = nn.LayerNorm(config.d_model)
|
| 386 |
+
self.gradient_checkpointing = False
|
| 387 |
+
self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
|
| 388 |
+
self.conv2d2 = nn.Conv2d(
|
| 389 |
+
config.downsample_hidden_size,
|
| 390 |
+
config.downsample_hidden_size,
|
| 391 |
+
3,
|
| 392 |
+
2,
|
| 393 |
+
padding=1,
|
| 394 |
+
)
|
| 395 |
+
self.conv2d3 = nn.Conv2d(
|
| 396 |
+
config.downsample_hidden_size,
|
| 397 |
+
config.downsample_hidden_size,
|
| 398 |
+
3,
|
| 399 |
+
2,
|
| 400 |
+
padding=1,
|
| 401 |
+
)
|
| 402 |
+
self.conv_out = nn.Linear(
|
| 403 |
+
config.downsample_hidden_size
|
| 404 |
+
* ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2),
|
| 405 |
+
config.d_model,
|
| 406 |
+
bias=False,
|
| 407 |
+
)
|
| 408 |
+
self.proj1 = nn.Linear(config.d_model, config.d_model)
|
| 409 |
+
self.act = ACT2FN[config.activation_function]
|
| 410 |
+
self.proj2 = nn.Linear(config.d_model, config.output_dim)
|
| 411 |
+
self.n_window_infer = self.config.n_window_infer
|
| 412 |
+
self.conv_chunksize = self.config.conv_chunksize
|
| 413 |
+
if self.n_window_infer % (self.n_window * 2) != 0:
|
| 414 |
+
raise ValueError(
|
| 415 |
+
f"n_window_infer ({self.n_window_infer}) must be a multiple of n_window * 2 "
|
| 416 |
+
f"({self.n_window * 2}): attention windows are assembled from whole "
|
| 417 |
+
f"{self.n_window * 2}-mel-frame chunks"
|
| 418 |
+
)
|
| 419 |
+
# Initialize weights and apply final processing
|
| 420 |
+
self.post_init()
|
| 421 |
+
|
| 422 |
+
def _freeze_parameters(self):
|
| 423 |
+
for param in self.parameters():
|
| 424 |
+
param.requires_grad = False
|
| 425 |
+
self._requires_grad = False
|
| 426 |
+
|
| 427 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 428 |
+
return self.conv2d1
|
| 429 |
+
|
| 430 |
+
def set_input_embeddings(self, value: nn.Module):
|
| 431 |
+
self.conv2d1 = value
|
| 432 |
+
|
| 433 |
+
def _prepare_attention_mask(
|
| 434 |
+
self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor
|
| 435 |
+
) -> torch.Tensor:
|
| 436 |
+
# Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
|
| 437 |
+
# NOTE: the created attention masl only approximates the ragged FA2 attention by
|
| 438 |
+
# allowing bidirectional attention within `cu_seqlens` blocks, and not attending between
|
| 439 |
+
# blocks. Though it will not be a 100% match for FA2's `varlen` path
|
| 440 |
+
if self.config._attn_implementation == "flash_attention_2":
|
| 441 |
+
return None
|
| 442 |
+
|
| 443 |
+
seq_length = inputs_tensor.shape[0]
|
| 444 |
+
attention_mask = torch.full(
|
| 445 |
+
[1, 1, seq_length, seq_length],
|
| 446 |
+
torch.finfo(inputs_tensor.dtype).min,
|
| 447 |
+
device=inputs_tensor.device,
|
| 448 |
+
dtype=inputs_tensor.dtype,
|
| 449 |
+
)
|
| 450 |
+
for i in range(1, len(cu_seqlens)):
|
| 451 |
+
attention_mask[
|
| 452 |
+
...,
|
| 453 |
+
cu_seqlens[i - 1] : cu_seqlens[i],
|
| 454 |
+
cu_seqlens[i - 1] : cu_seqlens[i],
|
| 455 |
+
] = 0
|
| 456 |
+
return attention_mask
|
| 457 |
+
|
| 458 |
+
def forward(
|
| 459 |
+
self,
|
| 460 |
+
input_features,
|
| 461 |
+
feature_lens=None,
|
| 462 |
+
aftercnn_lens=None,
|
| 463 |
+
):
|
| 464 |
+
r"""
|
| 465 |
+
feature_lens (`torch.LongTensor` of shape `(batch_size,)`):
|
| 466 |
+
mel length
|
| 467 |
+
aftercnn_lens (`torch.LongTensor` of shape `(batch_size,)`):
|
| 468 |
+
mel length after cnn
|
| 469 |
+
"""
|
| 470 |
+
if feature_lens is None:
|
| 471 |
+
raise ValueError(
|
| 472 |
+
"feature_lens is required: pass the per-clip mel frame counts as a "
|
| 473 |
+
"(batch_size,) LongTensor matching the packed input_features"
|
| 474 |
+
)
|
| 475 |
+
if input_features.dim() != 2:
|
| 476 |
+
raise ValueError(
|
| 477 |
+
"input_features must be a packed 2-D tensor of shape (num_mel_bins, total_frames) "
|
| 478 |
+
f"with clips concatenated along the frame axis; got shape {tuple(input_features.shape)}"
|
| 479 |
+
)
|
| 480 |
+
if (feature_lens <= 0).any():
|
| 481 |
+
raise ValueError(
|
| 482 |
+
f"every clip must have at least one mel frame; got feature_lens={feature_lens.tolist()}"
|
| 483 |
+
)
|
| 484 |
+
chunk_size = self.n_window * 2
|
| 485 |
+
if aftercnn_lens is None:
|
| 486 |
+
aftercnn_lens = _get_feat_extract_output_lengths(feature_lens, chunk_size)
|
| 487 |
+
chunk_num = torch.ceil(feature_lens / chunk_size).long()
|
| 488 |
+
|
| 489 |
+
chunk_lengths = torch.tensor(
|
| 490 |
+
[self.n_window * 2] * chunk_num.sum(),
|
| 491 |
+
dtype=torch.long,
|
| 492 |
+
device=feature_lens.device,
|
| 493 |
+
)
|
| 494 |
+
tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
|
| 495 |
+
chunk_lengths[tail_chunk_index] = feature_lens % chunk_size
|
| 496 |
+
chunk_lengths[chunk_lengths == 0] = chunk_size
|
| 497 |
+
|
| 498 |
+
chunk_size_after_cnn = int(
|
| 499 |
+
_get_feat_extract_output_lengths(torch.tensor(chunk_size), chunk_size)
|
| 500 |
+
)
|
| 501 |
+
|
| 502 |
+
chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
|
| 503 |
+
padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True)
|
| 504 |
+
padded_feature = F.pad(
|
| 505 |
+
padded_feature, (0, 0, 0, chunk_size - padded_feature.shape[1])
|
| 506 |
+
).transpose(1, 2)
|
| 507 |
+
feature_lens_after_cnn = _get_feat_extract_output_lengths(
|
| 508 |
+
chunk_lengths, chunk_size
|
| 509 |
+
)
|
| 510 |
+
padded_mask_after_cnn = nn.utils.rnn.pad_sequence(
|
| 511 |
+
[
|
| 512 |
+
torch.ones(length, dtype=torch.bool, device=padded_feature.device)
|
| 513 |
+
for length in feature_lens_after_cnn
|
| 514 |
+
],
|
| 515 |
+
batch_first=True,
|
| 516 |
+
)
|
| 517 |
+
padded_mask_after_cnn = F.pad(
|
| 518 |
+
padded_mask_after_cnn,
|
| 519 |
+
(0, chunk_size_after_cnn - padded_mask_after_cnn.shape[-1]),
|
| 520 |
+
)
|
| 521 |
+
padded_feature = padded_feature.unsqueeze(1)
|
| 522 |
+
# Split to chunk to avoid OOM during convolution
|
| 523 |
+
padded_embeds = []
|
| 524 |
+
for chunk in padded_feature.split(self.conv_chunksize, dim=0):
|
| 525 |
+
padded_embed = F.gelu(self.conv2d1(chunk))
|
| 526 |
+
padded_embed = F.gelu(self.conv2d2(padded_embed))
|
| 527 |
+
padded_embed = F.gelu(self.conv2d3(padded_embed))
|
| 528 |
+
padded_embeds.append(padded_embed)
|
| 529 |
+
padded_embed = torch.cat(padded_embeds, dim=0)
|
| 530 |
+
b, c, f, t = padded_embed.size()
|
| 531 |
+
padded_embed = self.conv_out(
|
| 532 |
+
padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f)
|
| 533 |
+
)
|
| 534 |
+
|
| 535 |
+
positional_embedding = (
|
| 536 |
+
self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
|
| 537 |
+
.unsqueeze(0)
|
| 538 |
+
.to(padded_embed.dtype)
|
| 539 |
+
)
|
| 540 |
+
padded_embed = padded_embed + positional_embedding
|
| 541 |
+
hidden_states = padded_embed[padded_mask_after_cnn]
|
| 542 |
+
cu_chunk_lens = [0]
|
| 543 |
+
window_aftercnn = padded_mask_after_cnn.shape[-1] * (
|
| 544 |
+
self.n_window_infer // (self.n_window * 2)
|
| 545 |
+
)
|
| 546 |
+
for cnn_len in aftercnn_lens:
|
| 547 |
+
cu_chunk_lens += [window_aftercnn] * (cnn_len // window_aftercnn)
|
| 548 |
+
remainder = cnn_len % window_aftercnn
|
| 549 |
+
if remainder != 0:
|
| 550 |
+
cu_chunk_lens += [remainder]
|
| 551 |
+
cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(
|
| 552 |
+
-1, dtype=torch.int32
|
| 553 |
+
)
|
| 554 |
+
|
| 555 |
+
for encoder_layer in self.layers:
|
| 556 |
+
layer_outputs = encoder_layer(
|
| 557 |
+
hidden_states,
|
| 558 |
+
cu_seqlens,
|
| 559 |
+
)
|
| 560 |
+
|
| 561 |
+
hidden_states = layer_outputs[0]
|
| 562 |
+
|
| 563 |
+
hidden_states = self.ln_post(hidden_states)
|
| 564 |
+
hidden_states = self.proj1(hidden_states)
|
| 565 |
+
hidden_states = self.act(hidden_states)
|
| 566 |
+
hidden_states = self.proj2(hidden_states)
|
| 567 |
+
return BaseModelOutput(last_hidden_state=hidden_states)
|
| 568 |
+
|
| 569 |
+
|
| 570 |
+
# Dense K2-Horizon decoder: adapted from IFM/K2-Horizon-3.7B.
|
| 571 |
+
# Source revision: 633f52ad28b17edeabd82afc61d2d13b4c59a561.
|
| 572 |
+
# https://huggingface.co/IFM/K2-Horizon-3.7B/tree/633f52ad28b17edeabd82afc61d2d13b4c59a561
|
| 573 |
+
# Retains grouped RMSNorm, partial RoPE and Transformers 4.57/5.x compatibility.
|
| 574 |
+
|
| 575 |
+
if hasattr(_generic, "merge_with_config_defaults"):
|
| 576 |
+
from transformers.utils.output_capturing import capture_outputs as _capture_outputs
|
| 577 |
+
|
| 578 |
+
def check_model_inputs(func):
|
| 579 |
+
# Modern Transformers split config defaults and output recording.
|
| 580 |
+
return _generic.merge_with_config_defaults(_capture_outputs(func))
|
| 581 |
+
|
| 582 |
+
else:
|
| 583 |
+
# Older releases combine both responsibilities in check_model_inputs.
|
| 584 |
+
_check_model_inputs = _generic.check_model_inputs
|
| 585 |
+
if "func" in _inspect.signature(_check_model_inputs).parameters:
|
| 586 |
+
check_model_inputs = _check_model_inputs
|
| 587 |
+
else:
|
| 588 |
+
|
| 589 |
+
def check_model_inputs(func):
|
| 590 |
+
return _check_model_inputs()(func)
|
| 591 |
+
|
| 592 |
+
|
| 593 |
+
# The mask helpers changed shape across releases: 5.x renamed input_embeds ->
|
| 594 |
+
# inputs_embeds and removed cache_position. Read the installed signature and pass
|
| 595 |
+
# only what it accepts.
|
| 596 |
+
_MASK_PARAMS = set(_inspect.signature(create_causal_mask).parameters)
|
| 597 |
+
_MASK_EMBEDS_KW = "inputs_embeds" if "inputs_embeds" in _MASK_PARAMS else "input_embeds"
|
| 598 |
+
|
| 599 |
+
|
| 600 |
+
def _mask_kwargs(embeds, **kw):
|
| 601 |
+
kw[_MASK_EMBEDS_KW] = embeds
|
| 602 |
+
return {k: v for k, v in kw.items() if k in _MASK_PARAMS}
|
| 603 |
+
|
| 604 |
+
|
| 605 |
+
def rotate_half(x):
|
| 606 |
+
"""Rotates half the hidden dims of the input."""
|
| 607 |
+
x1 = x[..., : x.shape[-1] // 2]
|
| 608 |
+
x2 = x[..., x.shape[-1] // 2 :]
|
| 609 |
+
return torch.cat((-x2, x1), dim=-1)
|
| 610 |
+
|
| 611 |
+
|
| 612 |
+
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
| 613 |
+
"""Applies Rotary Position Embedding to the query and key tensors.
|
| 614 |
+
|
| 615 |
+
Args:
|
| 616 |
+
q (`torch.Tensor`): The query tensor.
|
| 617 |
+
k (`torch.Tensor`): The key tensor.
|
| 618 |
+
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
| 619 |
+
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
| 620 |
+
position_ids (`torch.Tensor`, *optional*):
|
| 621 |
+
Deprecated and unused.
|
| 622 |
+
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
| 623 |
+
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
| 624 |
+
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
| 625 |
+
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
| 626 |
+
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
| 627 |
+
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
| 628 |
+
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
| 629 |
+
Returns:
|
| 630 |
+
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
| 631 |
+
"""
|
| 632 |
+
cos = cos.unsqueeze(unsqueeze_dim)
|
| 633 |
+
sin = sin.unsqueeze(unsqueeze_dim)
|
| 634 |
+
q_embed = (q * cos) + (rotate_half(q) * sin)
|
| 635 |
+
k_embed = (k * cos) + (rotate_half(k) * sin)
|
| 636 |
+
return q_embed, k_embed
|
| 637 |
+
|
| 638 |
+
|
| 639 |
+
def _permute_interleaved(x):
|
| 640 |
+
return x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).reshape(*x.shape[:-1], -1)
|
| 641 |
+
|
| 642 |
+
|
| 643 |
+
def _permute_halved(x):
|
| 644 |
+
return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1)
|
| 645 |
+
|
| 646 |
+
|
| 647 |
+
def k2_eager_attention_forward(
|
| 648 |
+
module: nn.Module,
|
| 649 |
+
query: torch.Tensor,
|
| 650 |
+
key: torch.Tensor,
|
| 651 |
+
value: torch.Tensor,
|
| 652 |
+
attention_mask: Optional[torch.Tensor],
|
| 653 |
+
scaling: float,
|
| 654 |
+
dropout: float = 0.0,
|
| 655 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 656 |
+
):
|
| 657 |
+
key_states = repeat_kv(key, module.num_key_value_groups)
|
| 658 |
+
value_states = repeat_kv(value, module.num_key_value_groups)
|
| 659 |
+
|
| 660 |
+
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
| 661 |
+
if attention_mask is not None:
|
| 662 |
+
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
| 663 |
+
attn_weights = attn_weights + causal_mask
|
| 664 |
+
|
| 665 |
+
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
|
| 666 |
+
query.dtype
|
| 667 |
+
)
|
| 668 |
+
attn_weights = nn.functional.dropout(
|
| 669 |
+
attn_weights, p=dropout, training=module.training
|
| 670 |
+
)
|
| 671 |
+
attn_output = torch.matmul(attn_weights, value_states)
|
| 672 |
+
attn_output = attn_output.transpose(1, 2).contiguous()
|
| 673 |
+
|
| 674 |
+
return attn_output, attn_weights
|
| 675 |
+
|
| 676 |
+
|
| 677 |
+
class K2HorizonAttention(nn.Module):
|
| 678 |
+
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
| 679 |
+
|
| 680 |
+
def __init__(self, config: K2HorizonConfig, layer_idx: int):
|
| 681 |
+
super().__init__()
|
| 682 |
+
self.config = config
|
| 683 |
+
self.layer_idx = layer_idx
|
| 684 |
+
self.head_dim = getattr(
|
| 685 |
+
config, "head_dim", config.hidden_size // config.num_attention_heads
|
| 686 |
+
)
|
| 687 |
+
self.num_key_value_groups = (
|
| 688 |
+
config.num_attention_heads // config.num_key_value_heads
|
| 689 |
+
)
|
| 690 |
+
self.scaling = self.head_dim**-0.5
|
| 691 |
+
self.attention_dropout = config.attention_dropout
|
| 692 |
+
self.is_causal = True
|
| 693 |
+
|
| 694 |
+
self.rope_head_dim = (
|
| 695 |
+
self.head_dim if config.rope_head_dim is None else config.rope_head_dim
|
| 696 |
+
)
|
| 697 |
+
|
| 698 |
+
self.q_proj = nn.Linear(
|
| 699 |
+
config.hidden_size,
|
| 700 |
+
config.num_attention_heads * self.head_dim,
|
| 701 |
+
bias=config.attention_bias,
|
| 702 |
+
)
|
| 703 |
+
self.k_proj = nn.Linear(
|
| 704 |
+
config.hidden_size,
|
| 705 |
+
config.num_key_value_heads * self.head_dim,
|
| 706 |
+
bias=config.attention_bias,
|
| 707 |
+
)
|
| 708 |
+
self.v_proj = nn.Linear(
|
| 709 |
+
config.hidden_size,
|
| 710 |
+
config.num_key_value_heads * self.head_dim,
|
| 711 |
+
bias=config.attention_bias,
|
| 712 |
+
)
|
| 713 |
+
self.o_proj = nn.Linear(
|
| 714 |
+
config.num_attention_heads * self.head_dim,
|
| 715 |
+
config.hidden_size,
|
| 716 |
+
bias=config.attention_bias,
|
| 717 |
+
)
|
| 718 |
+
|
| 719 |
+
self.gate_func = config.attention_gate_func
|
| 720 |
+
if self.gate_func is not None:
|
| 721 |
+
self.gate_proj = nn.Linear(
|
| 722 |
+
config.hidden_size,
|
| 723 |
+
config.num_attention_heads * self.head_dim,
|
| 724 |
+
bias=False,
|
| 725 |
+
)
|
| 726 |
+
|
| 727 |
+
if config.query_key_norm:
|
| 728 |
+
self.q_norm = K2HorizonRMSNorm(
|
| 729 |
+
hidden_size=config.num_attention_heads * self.head_dim,
|
| 730 |
+
n_groups=config.num_attention_heads,
|
| 731 |
+
eps=config.rms_norm_eps,
|
| 732 |
+
)
|
| 733 |
+
self.k_norm = K2HorizonRMSNorm(
|
| 734 |
+
hidden_size=config.num_key_value_heads * self.head_dim,
|
| 735 |
+
n_groups=config.num_key_value_heads,
|
| 736 |
+
eps=config.rms_norm_eps,
|
| 737 |
+
)
|
| 738 |
+
|
| 739 |
+
self.sliding_window = getattr(config, "sliding_window", None)
|
| 740 |
+
|
| 741 |
+
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
| 742 |
+
def forward(
|
| 743 |
+
self,
|
| 744 |
+
hidden_states: torch.Tensor,
|
| 745 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 746 |
+
attention_mask: Optional[torch.Tensor],
|
| 747 |
+
past_key_values: Optional[Cache] = None,
|
| 748 |
+
cache_position: Optional[torch.LongTensor] = None,
|
| 749 |
+
**kwargs: Unpack[FlashAttentionKwargs],
|
| 750 |
+
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
|
| 751 |
+
input_shape = hidden_states.shape[:-1]
|
| 752 |
+
hidden_shape = (*input_shape, -1, self.head_dim)
|
| 753 |
+
|
| 754 |
+
if self.config.query_key_norm:
|
| 755 |
+
query_states = (
|
| 756 |
+
self.q_norm(self.q_proj(hidden_states))
|
| 757 |
+
.view(hidden_shape)
|
| 758 |
+
.transpose(1, 2)
|
| 759 |
+
)
|
| 760 |
+
key_states = (
|
| 761 |
+
self.k_norm(self.k_proj(hidden_states))
|
| 762 |
+
.view(hidden_shape)
|
| 763 |
+
.transpose(1, 2)
|
| 764 |
+
)
|
| 765 |
+
else:
|
| 766 |
+
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
| 767 |
+
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
| 768 |
+
|
| 769 |
+
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
| 770 |
+
|
| 771 |
+
cos, sin = position_embeddings
|
| 772 |
+
if self.rope_head_dim == self.head_dim:
|
| 773 |
+
query_states, key_states = apply_rotary_pos_emb(
|
| 774 |
+
query_states, key_states, cos, sin
|
| 775 |
+
)
|
| 776 |
+
else:
|
| 777 |
+
query_states, query_states_ = torch.split(
|
| 778 |
+
_permute_interleaved(query_states),
|
| 779 |
+
split_size_or_sections=[
|
| 780 |
+
self.rope_head_dim,
|
| 781 |
+
self.head_dim - self.rope_head_dim,
|
| 782 |
+
],
|
| 783 |
+
dim=-1,
|
| 784 |
+
)
|
| 785 |
+
|
| 786 |
+
key_states, key_states_ = torch.split(
|
| 787 |
+
_permute_interleaved(key_states),
|
| 788 |
+
split_size_or_sections=[
|
| 789 |
+
self.rope_head_dim,
|
| 790 |
+
self.head_dim - self.rope_head_dim,
|
| 791 |
+
],
|
| 792 |
+
dim=-1,
|
| 793 |
+
)
|
| 794 |
+
|
| 795 |
+
query_states, key_states = apply_rotary_pos_emb(
|
| 796 |
+
_permute_halved(query_states), _permute_halved(key_states), cos, sin
|
| 797 |
+
)
|
| 798 |
+
|
| 799 |
+
query_states = _permute_halved(
|
| 800 |
+
torch.cat([_permute_interleaved(query_states), query_states_], dim=-1)
|
| 801 |
+
)
|
| 802 |
+
key_states = _permute_halved(
|
| 803 |
+
torch.cat([_permute_interleaved(key_states), key_states_], dim=-1)
|
| 804 |
+
)
|
| 805 |
+
|
| 806 |
+
if past_key_values is not None:
|
| 807 |
+
# sin and cos are specific to RoPE models; cache_position needed for the static cache
|
| 808 |
+
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
|
| 809 |
+
key_states, value_states = past_key_values.update(
|
| 810 |
+
key_states, value_states, self.layer_idx, cache_kwargs
|
| 811 |
+
)
|
| 812 |
+
|
| 813 |
+
attention_interface: Callable = k2_eager_attention_forward
|
| 814 |
+
if self.config._attn_implementation != "eager":
|
| 815 |
+
attention_interface = ALL_ATTENTION_FUNCTIONS[
|
| 816 |
+
self.config._attn_implementation
|
| 817 |
+
]
|
| 818 |
+
|
| 819 |
+
attn_output, attn_weights = attention_interface(
|
| 820 |
+
self,
|
| 821 |
+
query_states,
|
| 822 |
+
key_states,
|
| 823 |
+
value_states,
|
| 824 |
+
attention_mask,
|
| 825 |
+
dropout=0.0 if not self.training else self.attention_dropout,
|
| 826 |
+
scaling=self.scaling,
|
| 827 |
+
sliding_window=self.sliding_window, # diff with Llama
|
| 828 |
+
**kwargs,
|
| 829 |
+
)
|
| 830 |
+
|
| 831 |
+
if self.gate_func is not None:
|
| 832 |
+
gate = self.gate_proj(hidden_states).view(input_shape + (-1, self.head_dim))
|
| 833 |
+
if self.gate_func == "silu":
|
| 834 |
+
gate = F.silu(gate)
|
| 835 |
+
else:
|
| 836 |
+
assert self.gate_func == "softplus"
|
| 837 |
+
gate = F.softplus(gate, beta=math.log(2))
|
| 838 |
+
|
| 839 |
+
attn_output = attn_output * gate
|
| 840 |
+
|
| 841 |
+
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
| 842 |
+
attn_output = self.o_proj(attn_output)
|
| 843 |
+
return attn_output, attn_weights
|
| 844 |
+
|
| 845 |
+
|
| 846 |
+
class K2HorizonMLP(nn.Module):
|
| 847 |
+
def __init__(self, config, intermediate_size=None):
|
| 848 |
+
super().__init__()
|
| 849 |
+
self.config = config
|
| 850 |
+
self.hidden_size = config.hidden_size
|
| 851 |
+
self.intermediate_size = (
|
| 852 |
+
intermediate_size
|
| 853 |
+
if intermediate_size is not None
|
| 854 |
+
else config.intermediate_size
|
| 855 |
+
)
|
| 856 |
+
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
| 857 |
+
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
| 858 |
+
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
| 859 |
+
self.act_fn = ACT2FN[config.hidden_act]
|
| 860 |
+
|
| 861 |
+
def forward(self, x):
|
| 862 |
+
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
| 863 |
+
return down_proj
|
| 864 |
+
|
| 865 |
+
|
| 866 |
+
class K2HorizonRMSNorm(nn.Module):
|
| 867 |
+
def __init__(self, hidden_size: int, n_groups: int, eps=1e-6):
|
| 868 |
+
"""
|
| 869 |
+
K2HorizonRMSNorm is equivalent to T5LayerNorm
|
| 870 |
+
"""
|
| 871 |
+
super().__init__()
|
| 872 |
+
self.n_groups = n_groups
|
| 873 |
+
self.hidden_size = hidden_size
|
| 874 |
+
assert hidden_size % n_groups == 0
|
| 875 |
+
self.weight = nn.Parameter(torch.ones(hidden_size))
|
| 876 |
+
self.variance_epsilon = eps
|
| 877 |
+
|
| 878 |
+
def forward(self, hidden_states):
|
| 879 |
+
input_dtype = hidden_states.dtype
|
| 880 |
+
hidden_states = hidden_states.to(torch.float32)
|
| 881 |
+
|
| 882 |
+
hidden_states = hidden_states.reshape(
|
| 883 |
+
*hidden_states.shape[:-1], self.n_groups, -1
|
| 884 |
+
)
|
| 885 |
+
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
| 886 |
+
|
| 887 |
+
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
| 888 |
+
hidden_states = hidden_states.reshape(*hidden_states.shape[:-2], -1)
|
| 889 |
+
hidden_states = self.weight * hidden_states
|
| 890 |
+
|
| 891 |
+
return hidden_states.to(input_dtype)
|
| 892 |
+
|
| 893 |
+
def extra_repr(self):
|
| 894 |
+
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
| 895 |
+
|
| 896 |
+
|
| 897 |
+
class K2HorizonDecoderLayer(GradientCheckpointingLayer):
|
| 898 |
+
def __init__(self, config: K2HorizonConfig, layer_idx: int):
|
| 899 |
+
super().__init__()
|
| 900 |
+
self.hidden_size = config.hidden_size
|
| 901 |
+
|
| 902 |
+
self.self_attn = K2HorizonAttention(config, layer_idx)
|
| 903 |
+
|
| 904 |
+
self.mlp = K2HorizonMLP(config, intermediate_size=config.intermediate_size)
|
| 905 |
+
|
| 906 |
+
assert config.hidden_size % config.layernorm_num_groups == 0
|
| 907 |
+
self.input_layernorm = K2HorizonRMSNorm(
|
| 908 |
+
hidden_size=config.hidden_size,
|
| 909 |
+
n_groups=config.layernorm_num_groups,
|
| 910 |
+
eps=config.rms_norm_eps,
|
| 911 |
+
)
|
| 912 |
+
self.post_attention_layernorm = K2HorizonRMSNorm(
|
| 913 |
+
hidden_size=config.hidden_size,
|
| 914 |
+
n_groups=config.layernorm_num_groups,
|
| 915 |
+
eps=config.rms_norm_eps,
|
| 916 |
+
)
|
| 917 |
+
|
| 918 |
+
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
| 919 |
+
def forward(
|
| 920 |
+
self,
|
| 921 |
+
hidden_states: torch.Tensor,
|
| 922 |
+
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
| 923 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 924 |
+
position_ids: Optional[torch.LongTensor] = None,
|
| 925 |
+
past_key_values: Optional[Cache] = None,
|
| 926 |
+
cache_position: Optional[torch.LongTensor] = None,
|
| 927 |
+
**kwargs: Unpack[FlashAttentionKwargs],
|
| 928 |
+
) -> torch.FloatTensor:
|
| 929 |
+
"""
|
| 930 |
+
Args:
|
| 931 |
+
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
|
| 932 |
+
attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
|
| 933 |
+
`(batch, sequence_length)` where padding elements are indicated by 0.
|
| 934 |
+
output_attentions (`bool`, *optional*):
|
| 935 |
+
Whether or not to return the attentions tensors of all attention layers. See `attentions` under
|
| 936 |
+
returned tensors for more detail.
|
| 937 |
+
output_router_logits (`bool`, *optional*):
|
| 938 |
+
Whether or not to return the logits of all the routers. They are useful for computing the router loss,
|
| 939 |
+
and should not be returned during inference.
|
| 940 |
+
use_cache (`bool`, *optional*):
|
| 941 |
+
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
|
| 942 |
+
(see `past_key_values`).
|
| 943 |
+
past_key_values (`Cache`, *optional*): cached past key and value projection states
|
| 944 |
+
cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
|
| 945 |
+
Indices depicting the position of the input sequence tokens in the sequence.
|
| 946 |
+
position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
|
| 947 |
+
Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
|
| 948 |
+
with `head_dim` being the embedding dimension of each attention head.
|
| 949 |
+
kwargs (`dict`, *optional*):
|
| 950 |
+
Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
|
| 951 |
+
into the model
|
| 952 |
+
"""
|
| 953 |
+
residual = hidden_states
|
| 954 |
+
|
| 955 |
+
hidden_states = self.input_layernorm(hidden_states)
|
| 956 |
+
|
| 957 |
+
# Self Attention
|
| 958 |
+
hidden_states, _ = self.self_attn(
|
| 959 |
+
hidden_states=hidden_states,
|
| 960 |
+
position_embeddings=position_embeddings,
|
| 961 |
+
attention_mask=attention_mask,
|
| 962 |
+
position_ids=position_ids,
|
| 963 |
+
past_key_values=past_key_values,
|
| 964 |
+
cache_position=cache_position,
|
| 965 |
+
**kwargs,
|
| 966 |
+
)
|
| 967 |
+
hidden_states = residual + hidden_states
|
| 968 |
+
|
| 969 |
+
# Fully Connected
|
| 970 |
+
residual = hidden_states
|
| 971 |
+
hidden_states = self.post_attention_layernorm(hidden_states)
|
| 972 |
+
hidden_states = self.mlp(hidden_states)
|
| 973 |
+
hidden_states = residual + hidden_states
|
| 974 |
+
|
| 975 |
+
return hidden_states
|
| 976 |
+
|
| 977 |
+
|
| 978 |
+
class K2HorizonRotaryEmbedding(nn.Module):
|
| 979 |
+
inv_freq: torch.Tensor # fix linting for `register_buffer`
|
| 980 |
+
|
| 981 |
+
@staticmethod
|
| 982 |
+
def compute_default_rope_parameters(config, device=None, seq_len=None):
|
| 983 |
+
# Standard RoPE inverse frequencies. Version-agnostic replacement for
|
| 984 |
+
# transformers<=4.57's ROPE_INIT_FUNCTIONS["default"] (removed in 5.x, where
|
| 985 |
+
# the default lives as a static method on each model's RotaryEmbedding).
|
| 986 |
+
base = getattr(config, "rope_theta", None)
|
| 987 |
+
rp = getattr(config, "rope_parameters", None)
|
| 988 |
+
if base is None and isinstance(rp, dict):
|
| 989 |
+
base = rp.get("rope_theta", 10000.0)
|
| 990 |
+
partial = 1.0
|
| 991 |
+
if isinstance(rp, dict):
|
| 992 |
+
partial = rp.get("partial_rotary_factor", 1.0)
|
| 993 |
+
head_dim = getattr(config, "head_dim", None) or (
|
| 994 |
+
config.hidden_size // config.num_attention_heads
|
| 995 |
+
)
|
| 996 |
+
dim = int(head_dim * partial)
|
| 997 |
+
inv_freq = 1.0 / (
|
| 998 |
+
base
|
| 999 |
+
** (
|
| 1000 |
+
torch.arange(0, dim, 2, dtype=torch.int64).to(
|
| 1001 |
+
device=device, dtype=torch.float
|
| 1002 |
+
)
|
| 1003 |
+
/ dim
|
| 1004 |
+
)
|
| 1005 |
+
)
|
| 1006 |
+
return inv_freq, 1.0
|
| 1007 |
+
|
| 1008 |
+
def __init__(self, config: K2HorizonConfig, device=None):
|
| 1009 |
+
super().__init__()
|
| 1010 |
+
# rope_type: transformers>=5 uses config.rope_parameters; <=4.57 used config.rope_scaling
|
| 1011 |
+
rope_type = "default"
|
| 1012 |
+
rp = getattr(config, "rope_parameters", None)
|
| 1013 |
+
if isinstance(rp, dict) and rp.get("rope_type"):
|
| 1014 |
+
rope_type = rp["rope_type"]
|
| 1015 |
+
elif isinstance(getattr(config, "rope_scaling", None), dict):
|
| 1016 |
+
rope_type = config.rope_scaling.get(
|
| 1017 |
+
"rope_type", config.rope_scaling.get("type", "default")
|
| 1018 |
+
)
|
| 1019 |
+
partial = getattr(config, "partial_rotary_factor", None)
|
| 1020 |
+
if partial is None and isinstance(rp, dict):
|
| 1021 |
+
partial = rp.get("partial_rotary_factor")
|
| 1022 |
+
if partial is not None and partial != 1.0:
|
| 1023 |
+
raise ValueError(
|
| 1024 |
+
f"partial_rotary_factor={partial} is not supported: K2HorizonAttention derives its "
|
| 1025 |
+
"rotary width from rope_head_dim only. Express partial RoPE via rope_head_dim instead."
|
| 1026 |
+
)
|
| 1027 |
+
self.rope_type = rope_type
|
| 1028 |
+
self.max_seq_len_cached = config.max_position_embeddings
|
| 1029 |
+
self.original_max_seq_len = config.max_position_embeddings
|
| 1030 |
+
|
| 1031 |
+
self.config = config
|
| 1032 |
+
# "default" was dropped from ROPE_INIT_FUNCTIONS in transformers 5.x
|
| 1033 |
+
if self.rope_type == "default" or self.rope_type not in ROPE_INIT_FUNCTIONS:
|
| 1034 |
+
self.rope_init_fn = self.compute_default_rope_parameters
|
| 1035 |
+
else:
|
| 1036 |
+
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
| 1037 |
+
|
| 1038 |
+
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
| 1039 |
+
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
| 1040 |
+
self.original_inv_freq = self.inv_freq
|
| 1041 |
+
|
| 1042 |
+
@torch.no_grad()
|
| 1043 |
+
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
| 1044 |
+
def forward(self, x, position_ids):
|
| 1045 |
+
inv_freq_expanded = (
|
| 1046 |
+
self.inv_freq[None, :, None]
|
| 1047 |
+
.float()
|
| 1048 |
+
.expand(position_ids.shape[0], -1, 1)
|
| 1049 |
+
.to(x.device)
|
| 1050 |
+
)
|
| 1051 |
+
position_ids_expanded = position_ids[:, None, :].float()
|
| 1052 |
+
|
| 1053 |
+
device_type = (
|
| 1054 |
+
x.device.type
|
| 1055 |
+
if isinstance(x.device.type, str) and x.device.type != "mps"
|
| 1056 |
+
else "cpu"
|
| 1057 |
+
)
|
| 1058 |
+
with torch.autocast(device_type=device_type, enabled=False): # Force float32
|
| 1059 |
+
freqs = (
|
| 1060 |
+
inv_freq_expanded.float() @ position_ids_expanded.float()
|
| 1061 |
+
).transpose(1, 2)
|
| 1062 |
+
emb = torch.cat((freqs, freqs), dim=-1)
|
| 1063 |
+
cos = emb.cos() * self.attention_scaling
|
| 1064 |
+
sin = emb.sin() * self.attention_scaling
|
| 1065 |
+
|
| 1066 |
+
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
| 1067 |
+
|
| 1068 |
+
|
| 1069 |
+
class K2HorizonPreTrainedModel(PreTrainedModel):
|
| 1070 |
+
config: K2HorizonConfig
|
| 1071 |
+
base_model_prefix = "model"
|
| 1072 |
+
supports_gradient_checkpointing = True
|
| 1073 |
+
_no_split_modules = ["K2HorizonDecoderLayer"]
|
| 1074 |
+
_skip_keys_device_placement = ["past_key_values"]
|
| 1075 |
+
_supports_flash_attn = True
|
| 1076 |
+
_supports_sdpa = True
|
| 1077 |
+
_supports_flex_attn = True
|
| 1078 |
+
_can_compile_fullgraph = False
|
| 1079 |
+
_supports_attention_backend = True
|
| 1080 |
+
_can_record_outputs = {
|
| 1081 |
+
"hidden_states": K2HorizonDecoderLayer,
|
| 1082 |
+
"attentions": K2HorizonAttention,
|
| 1083 |
+
}
|
| 1084 |
+
|
| 1085 |
+
|
| 1086 |
+
class K2HorizonModel(K2HorizonPreTrainedModel):
|
| 1087 |
+
def __init__(self, config: K2HorizonConfig):
|
| 1088 |
+
super().__init__(config)
|
| 1089 |
+
self.vocab_size = config.vocab_size
|
| 1090 |
+
|
| 1091 |
+
# no padding_idx: the configured pad token doubles as a real token (BOS/EOS family), and
|
| 1092 |
+
# nn.Embedding's padding_idx would permanently zero that row's gradient in any fine-tune.
|
| 1093 |
+
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
| 1094 |
+
self.layers = nn.ModuleList(
|
| 1095 |
+
[
|
| 1096 |
+
K2HorizonDecoderLayer(config, layer_idx)
|
| 1097 |
+
for layer_idx in range(config.num_hidden_layers)
|
| 1098 |
+
]
|
| 1099 |
+
)
|
| 1100 |
+
|
| 1101 |
+
assert config.hidden_size % config.layernorm_num_groups == 0
|
| 1102 |
+
self.norm = K2HorizonRMSNorm(
|
| 1103 |
+
hidden_size=config.hidden_size,
|
| 1104 |
+
n_groups=config.layernorm_num_groups,
|
| 1105 |
+
eps=config.rms_norm_eps,
|
| 1106 |
+
)
|
| 1107 |
+
|
| 1108 |
+
rope_config = config
|
| 1109 |
+
if config.rope_head_dim is not None and config.rope_head_dim != config.head_dim:
|
| 1110 |
+
# dedicated config for the rotary embedding: it must see head_dim == rope_head_dim
|
| 1111 |
+
# persistently, because @dynamic_rope_update re-reads it at runtime (a temporary
|
| 1112 |
+
# head_dim mutation would already be restored by then).
|
| 1113 |
+
rope_config = copy.deepcopy(config)
|
| 1114 |
+
rope_config.head_dim = config.rope_head_dim
|
| 1115 |
+
self.rotary_emb = K2HorizonRotaryEmbedding(config=rope_config)
|
| 1116 |
+
|
| 1117 |
+
self.gradient_checkpointing = False
|
| 1118 |
+
|
| 1119 |
+
# Initialize weights and apply final processing
|
| 1120 |
+
self.post_init()
|
| 1121 |
+
|
| 1122 |
+
@check_model_inputs
|
| 1123 |
+
def forward(
|
| 1124 |
+
self,
|
| 1125 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1126 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1127 |
+
position_ids: Optional[torch.LongTensor] = None,
|
| 1128 |
+
past_key_values: Optional[Cache] = None,
|
| 1129 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1130 |
+
use_cache: Optional[bool] = None,
|
| 1131 |
+
cache_position: Optional[torch.LongTensor] = None,
|
| 1132 |
+
**kwargs: Unpack[TransformersKwargs],
|
| 1133 |
+
) -> BaseModelOutputWithPast:
|
| 1134 |
+
if (input_ids is None) ^ (inputs_embeds is not None):
|
| 1135 |
+
raise ValueError(
|
| 1136 |
+
"You must specify exactly one of input_ids or inputs_embeds"
|
| 1137 |
+
)
|
| 1138 |
+
|
| 1139 |
+
if use_cache and past_key_values is None:
|
| 1140 |
+
past_key_values = DynamicCache(config=self.config)
|
| 1141 |
+
|
| 1142 |
+
if inputs_embeds is None:
|
| 1143 |
+
inputs_embeds = self.embed_tokens(input_ids)
|
| 1144 |
+
|
| 1145 |
+
if cache_position is None:
|
| 1146 |
+
past_seen_tokens = (
|
| 1147 |
+
past_key_values.get_seq_length() if past_key_values is not None else 0
|
| 1148 |
+
)
|
| 1149 |
+
cache_position = torch.arange(
|
| 1150 |
+
past_seen_tokens,
|
| 1151 |
+
past_seen_tokens + inputs_embeds.shape[1],
|
| 1152 |
+
device=inputs_embeds.device,
|
| 1153 |
+
)
|
| 1154 |
+
if position_ids is None:
|
| 1155 |
+
position_ids = cache_position.unsqueeze(0)
|
| 1156 |
+
|
| 1157 |
+
mask_function = (
|
| 1158 |
+
create_causal_mask
|
| 1159 |
+
if self.config.sliding_window is None
|
| 1160 |
+
else create_sliding_window_causal_mask
|
| 1161 |
+
)
|
| 1162 |
+
causal_mask = mask_function(
|
| 1163 |
+
**_mask_kwargs(
|
| 1164 |
+
inputs_embeds,
|
| 1165 |
+
config=self.config,
|
| 1166 |
+
attention_mask=attention_mask,
|
| 1167 |
+
cache_position=cache_position,
|
| 1168 |
+
past_key_values=past_key_values,
|
| 1169 |
+
position_ids=position_ids,
|
| 1170 |
+
)
|
| 1171 |
+
)
|
| 1172 |
+
|
| 1173 |
+
hidden_states = inputs_embeds
|
| 1174 |
+
|
| 1175 |
+
# create position embeddings to be shared across the decoder layers
|
| 1176 |
+
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
| 1177 |
+
|
| 1178 |
+
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
| 1179 |
+
hidden_states = decoder_layer(
|
| 1180 |
+
hidden_states,
|
| 1181 |
+
position_embeddings=position_embeddings,
|
| 1182 |
+
attention_mask=causal_mask,
|
| 1183 |
+
position_ids=position_ids,
|
| 1184 |
+
past_key_values=past_key_values,
|
| 1185 |
+
use_cache=use_cache,
|
| 1186 |
+
cache_position=cache_position,
|
| 1187 |
+
**kwargs,
|
| 1188 |
+
)
|
| 1189 |
+
|
| 1190 |
+
hidden_states = self.norm(hidden_states)
|
| 1191 |
+
|
| 1192 |
+
return BaseModelOutputWithPast(
|
| 1193 |
+
last_hidden_state=hidden_states,
|
| 1194 |
+
past_key_values=past_key_values,
|
| 1195 |
+
)
|
| 1196 |
+
|
| 1197 |
+
|
| 1198 |
+
# Jais-ASR composition: preserve model.audio_tower.*, model.language_model.*,
|
| 1199 |
+
# model.embed_audio.* and lm_head.weight for the existing checkpoint.
|
| 1200 |
+
|
| 1201 |
+
|
| 1202 |
+
class JaisASRProjector(nn.Module):
|
| 1203 |
+
"""Fresh audio projector: LayerNorm -> Linear -> act -> Linear (audio output_dim -> LLM hidden)."""
|
| 1204 |
+
|
| 1205 |
+
def __init__(self, config: "JaisASRConfig"):
|
| 1206 |
+
super().__init__()
|
| 1207 |
+
in_dim = config.audio_config.output_dim # audio-tower output width
|
| 1208 |
+
out_dim = config.text_config.hidden_size # 2560 (the text backbone hidden)
|
| 1209 |
+
self.norm = nn.LayerNorm(in_dim)
|
| 1210 |
+
self.linear_1 = nn.Linear(in_dim, out_dim)
|
| 1211 |
+
self.act = ACT2FN[config.projector_hidden_act]
|
| 1212 |
+
self.linear_2 = nn.Linear(out_dim, out_dim)
|
| 1213 |
+
|
| 1214 |
+
def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
|
| 1215 |
+
return self.linear_2(self.act(self.linear_1(self.norm(audio_features))))
|
| 1216 |
+
|
| 1217 |
+
|
| 1218 |
+
class JaisASRPreTrainedModel(PreTrainedModel):
|
| 1219 |
+
config_class = JaisASRConfig
|
| 1220 |
+
base_model_prefix = "model"
|
| 1221 |
+
supports_gradient_checkpointing = True
|
| 1222 |
+
_no_split_modules = ["K2HorizonDecoderLayer", "Qwen3OmniMoeAudioEncoderLayer"]
|
| 1223 |
+
_supports_flash_attn = True
|
| 1224 |
+
_supports_sdpa = True
|
| 1225 |
+
_supports_attention_backend = True
|
| 1226 |
+
|
| 1227 |
+
def _init_weights(self, module):
|
| 1228 |
+
std = getattr(self.config.text_config, "initializer_range", 0.02)
|
| 1229 |
+
if isinstance(module, nn.Linear):
|
| 1230 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 1231 |
+
if module.bias is not None:
|
| 1232 |
+
module.bias.data.zero_()
|
| 1233 |
+
elif isinstance(module, nn.LayerNorm):
|
| 1234 |
+
module.weight.data.fill_(1.0)
|
| 1235 |
+
module.bias.data.zero_()
|
| 1236 |
+
elif isinstance(module, SinusoidsPositionEmbedding):
|
| 1237 |
+
# persistent=False -> not in the checkpoint; must be recomputed after the
|
| 1238 |
+
# meta-device build or it is uninitialized memory (transformers >= 5.x)
|
| 1239 |
+
materialize_sinusoids_(module)
|
| 1240 |
+
|
| 1241 |
+
|
| 1242 |
+
class JaisASRModel(JaisASRPreTrainedModel):
|
| 1243 |
+
"""Container holding the LLM decoder + audio encoder + projector (matches saved prefix `model.`)."""
|
| 1244 |
+
|
| 1245 |
+
def __init__(self, config: "JaisASRConfig"):
|
| 1246 |
+
super().__init__(config)
|
| 1247 |
+
self.language_model = K2HorizonModel(config.text_config)
|
| 1248 |
+
# dtype=ambient default: _from_config would otherwise re-apply the sub-config's
|
| 1249 |
+
# serialized dtype (bfloat16 in config.json), yielding a mixed-dtype model under
|
| 1250 |
+
# from_config(dtype=...) overrides.
|
| 1251 |
+
self.audio_tower = Qwen3OmniMoeAudioEncoder._from_config(
|
| 1252 |
+
config.audio_config, dtype=torch.get_default_dtype()
|
| 1253 |
+
)
|
| 1254 |
+
self.embed_audio = JaisASRProjector(config)
|
| 1255 |
+
self.post_init()
|
| 1256 |
+
|
| 1257 |
+
def get_input_embeddings(self):
|
| 1258 |
+
return self.language_model.embed_tokens
|
| 1259 |
+
|
| 1260 |
+
def set_input_embeddings(self, value):
|
| 1261 |
+
self.language_model.embed_tokens = value
|
| 1262 |
+
|
| 1263 |
+
def get_audio_features(self, input_features, feature_attention_mask=None):
|
| 1264 |
+
# Unpad+pack the mel via the feature-level mask and
|
| 1265 |
+
# derive real per-clip frame lengths; audio encoder then returns a PACKED (total_frames, output_dim=2048)
|
| 1266 |
+
# last_hidden_state, which the projector maps to (total_frames, hidden=2560).
|
| 1267 |
+
dtype = self.audio_tower.dtype
|
| 1268 |
+
if feature_attention_mask is None and input_features.dim() == 3:
|
| 1269 |
+
# no mask: every clip is full length (equivalent to an all-ones mask)
|
| 1270 |
+
feature_attention_mask = torch.ones(
|
| 1271 |
+
input_features.shape[0],
|
| 1272 |
+
input_features.shape[-1],
|
| 1273 |
+
dtype=torch.long,
|
| 1274 |
+
device=input_features.device,
|
| 1275 |
+
)
|
| 1276 |
+
feature_lens = None
|
| 1277 |
+
if feature_attention_mask is not None:
|
| 1278 |
+
feature_lens = feature_attention_mask.sum(dim=1)
|
| 1279 |
+
input_features = input_features.permute(0, 2, 1)[
|
| 1280 |
+
feature_attention_mask.bool()
|
| 1281 |
+
].permute(1, 0)
|
| 1282 |
+
audio_outputs = self.audio_tower(
|
| 1283 |
+
input_features.to(dtype), feature_lens=feature_lens
|
| 1284 |
+
)
|
| 1285 |
+
return self.embed_audio(audio_outputs.last_hidden_state)
|
| 1286 |
+
|
| 1287 |
+
def _get_audio_placeholder_mask(
|
| 1288 |
+
self, input_ids, inputs_embeds, n_audio_frames, audio_lens=None
|
| 1289 |
+
):
|
| 1290 |
+
# Placeholder positions, validated against the audio-tower frame count: masked_scatter fills
|
| 1291 |
+
# True positions in sequence order, so any count mismatch would silently shift every frame.
|
| 1292 |
+
if input_ids is None:
|
| 1293 |
+
audio_embed = self.get_input_embeddings()(
|
| 1294 |
+
torch.tensor(
|
| 1295 |
+
self.config.audio_token_id,
|
| 1296 |
+
dtype=torch.long,
|
| 1297 |
+
device=inputs_embeds.device,
|
| 1298 |
+
)
|
| 1299 |
+
)
|
| 1300 |
+
special_audio_mask = (inputs_embeds == audio_embed).all(-1)
|
| 1301 |
+
else:
|
| 1302 |
+
special_audio_mask = input_ids == self.config.audio_token_id
|
| 1303 |
+
n_audio_tokens = int(special_audio_mask.sum())
|
| 1304 |
+
if n_audio_tokens != n_audio_frames:
|
| 1305 |
+
raise ValueError(
|
| 1306 |
+
f"Audio placeholder tokens ({n_audio_tokens}) != audio-tower output frames ({n_audio_frames}). "
|
| 1307 |
+
"Encode prompts with JaisASRProcessor so each clip expands to exactly its frame count."
|
| 1308 |
+
)
|
| 1309 |
+
if audio_lens is not None and special_audio_mask.dim() == 2:
|
| 1310 |
+
# Per-row check: clips fill rows in order, so each row's placeholder count must equal
|
| 1311 |
+
# the frame count of the consecutive clips it consumes; compensating per-row miscounts
|
| 1312 |
+
# would otherwise silently scatter one row's audio into another.
|
| 1313 |
+
clip_lens = [int(l) for l in audio_lens]
|
| 1314 |
+
clip = 0
|
| 1315 |
+
for row, count in enumerate(special_audio_mask.sum(dim=1).tolist()):
|
| 1316 |
+
got = 0
|
| 1317 |
+
while got < count and clip < len(clip_lens):
|
| 1318 |
+
got += clip_lens[clip]
|
| 1319 |
+
clip += 1
|
| 1320 |
+
if got != count:
|
| 1321 |
+
raise ValueError(
|
| 1322 |
+
f"Row {row} has {count} audio placeholder tokens, but the clips assigned to it "
|
| 1323 |
+
f"in order produce {got} frames. Encode prompts with JaisASRProcessor so "
|
| 1324 |
+
"each clip expands to exactly its frame count."
|
| 1325 |
+
)
|
| 1326 |
+
return (
|
| 1327 |
+
special_audio_mask.unsqueeze(-1)
|
| 1328 |
+
.expand_as(inputs_embeds)
|
| 1329 |
+
.to(inputs_embeds.device)
|
| 1330 |
+
)
|
| 1331 |
+
|
| 1332 |
+
def forward(
|
| 1333 |
+
self,
|
| 1334 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1335 |
+
input_features: Optional[torch.FloatTensor] = None,
|
| 1336 |
+
feature_attention_mask: Optional[torch.Tensor] = None,
|
| 1337 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1338 |
+
position_ids: Optional[torch.LongTensor] = None,
|
| 1339 |
+
past_key_values: Optional[Cache] = None,
|
| 1340 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1341 |
+
use_cache: Optional[bool] = None,
|
| 1342 |
+
cache_position: Optional[torch.LongTensor] = None,
|
| 1343 |
+
**kwargs,
|
| 1344 |
+
):
|
| 1345 |
+
if inputs_embeds is None:
|
| 1346 |
+
inputs_embeds = self.get_input_embeddings()(input_ids)
|
| 1347 |
+
|
| 1348 |
+
# Merge audio wherever placeholder tokens appear in the CURRENT window — not only when
|
| 1349 |
+
# cache_position starts at 0 — so continuation from a reused KV cache and chunked prefill
|
| 1350 |
+
# get their audio merged instead of silently dropped. During autoregressive decoding the
|
| 1351 |
+
# audio embeddings already live in the KV cache and the 1-token window holds no
|
| 1352 |
+
# placeholder, so generate() never re-runs the audio tower per step.
|
| 1353 |
+
if input_features is not None:
|
| 1354 |
+
if input_ids is not None:
|
| 1355 |
+
has_audio_tokens = bool((input_ids == self.config.audio_token_id).any())
|
| 1356 |
+
else:
|
| 1357 |
+
audio_embed = self.get_input_embeddings()(
|
| 1358 |
+
torch.tensor(
|
| 1359 |
+
self.config.audio_token_id,
|
| 1360 |
+
dtype=torch.long,
|
| 1361 |
+
device=inputs_embeds.device,
|
| 1362 |
+
)
|
| 1363 |
+
)
|
| 1364 |
+
has_audio_tokens = bool((inputs_embeds == audio_embed).all(-1).any())
|
| 1365 |
+
if has_audio_tokens:
|
| 1366 |
+
if feature_attention_mask is None and input_features.dim() == 3:
|
| 1367 |
+
feature_attention_mask = torch.ones(
|
| 1368 |
+
input_features.shape[0],
|
| 1369 |
+
input_features.shape[-1],
|
| 1370 |
+
dtype=torch.long,
|
| 1371 |
+
device=input_features.device,
|
| 1372 |
+
)
|
| 1373 |
+
audio_embeds = self.get_audio_features(
|
| 1374 |
+
input_features, feature_attention_mask
|
| 1375 |
+
)
|
| 1376 |
+
audio_lens = None
|
| 1377 |
+
if feature_attention_mask is not None:
|
| 1378 |
+
audio_lens = _get_feat_extract_output_lengths(
|
| 1379 |
+
feature_attention_mask.sum(dim=1),
|
| 1380 |
+
self.config.audio_config.n_window * 2,
|
| 1381 |
+
)
|
| 1382 |
+
special_audio_mask = self._get_audio_placeholder_mask(
|
| 1383 |
+
input_ids, inputs_embeds, audio_embeds.shape[0], audio_lens
|
| 1384 |
+
)
|
| 1385 |
+
audio_embeds = audio_embeds.to(
|
| 1386 |
+
inputs_embeds.device, inputs_embeds.dtype
|
| 1387 |
+
)
|
| 1388 |
+
inputs_embeds = inputs_embeds.masked_scatter(
|
| 1389 |
+
special_audio_mask, audio_embeds
|
| 1390 |
+
)
|
| 1391 |
+
|
| 1392 |
+
# the text decoder requires exactly one of input_ids / inputs_embeds -> pass input_ids=None.
|
| 1393 |
+
return self.language_model(
|
| 1394 |
+
input_ids=None,
|
| 1395 |
+
attention_mask=attention_mask,
|
| 1396 |
+
position_ids=position_ids,
|
| 1397 |
+
past_key_values=past_key_values,
|
| 1398 |
+
inputs_embeds=inputs_embeds,
|
| 1399 |
+
use_cache=use_cache,
|
| 1400 |
+
cache_position=cache_position,
|
| 1401 |
+
**kwargs,
|
| 1402 |
+
)
|
| 1403 |
+
|
| 1404 |
+
|
| 1405 |
+
class JaisASRForConditionalGeneration(JaisASRPreTrainedModel, GenerationMixin):
|
| 1406 |
+
def generate(self, *args, language=None, **kwargs):
|
| 1407 |
+
"""Generate with an optional language pin; prefill belongs in the prompt."""
|
| 1408 |
+
if language is not None:
|
| 1409 |
+
try:
|
| 1410 |
+
from .language import LanguagePin
|
| 1411 |
+
from .postprocess import has_prefilled_tag
|
| 1412 |
+
except ImportError:
|
| 1413 |
+
if __package__:
|
| 1414 |
+
raise
|
| 1415 |
+
from language import LanguagePin
|
| 1416 |
+
from postprocess import has_prefilled_tag
|
| 1417 |
+
input_ids = kwargs.get(
|
| 1418 |
+
"input_ids", kwargs.get("inputs", args[0] if args else None)
|
| 1419 |
+
)
|
| 1420 |
+
if input_ids is None or input_ids.ndim != 2:
|
| 1421 |
+
raise ValueError("language pinning requires batched input_ids")
|
| 1422 |
+
if any(has_prefilled_tag(row) for row in input_ids):
|
| 1423 |
+
raise ValueError(
|
| 1424 |
+
"language cannot be combined with a prefilled language tag"
|
| 1425 |
+
)
|
| 1426 |
+
processors = list(kwargs.get("logits_processor") or [])
|
| 1427 |
+
processors.append(
|
| 1428 |
+
LanguagePin(language).for_transformers(input_ids.shape[1])
|
| 1429 |
+
)
|
| 1430 |
+
kwargs["logits_processor"] = processors
|
| 1431 |
+
return super().generate(*args, **kwargs)
|
| 1432 |
+
|
| 1433 |
+
# the text decoder is NOT tied -> lm_head is its own weight (no _tied_weights_keys mapping).
|
| 1434 |
+
|
| 1435 |
+
def __init__(self, config: "JaisASRConfig"):
|
| 1436 |
+
super().__init__(config)
|
| 1437 |
+
self.model = JaisASRModel(config)
|
| 1438 |
+
self.vocab_size = config.text_config.vocab_size
|
| 1439 |
+
self.lm_head = nn.Linear(
|
| 1440 |
+
config.text_config.hidden_size, config.text_config.vocab_size, bias=False
|
| 1441 |
+
)
|
| 1442 |
+
self.post_init()
|
| 1443 |
+
|
| 1444 |
+
# ---- embeddings ----
|
| 1445 |
+
def get_input_embeddings(self):
|
| 1446 |
+
return self.model.get_input_embeddings()
|
| 1447 |
+
|
| 1448 |
+
def set_input_embeddings(self, value):
|
| 1449 |
+
self.model.set_input_embeddings(value)
|
| 1450 |
+
|
| 1451 |
+
def get_output_embeddings(self):
|
| 1452 |
+
return self.lm_head
|
| 1453 |
+
|
| 1454 |
+
def set_output_embeddings(self, new):
|
| 1455 |
+
self.lm_head = new
|
| 1456 |
+
|
| 1457 |
+
# ---- submodule accessors (for frameworks that address submodules) ----
|
| 1458 |
+
@property
|
| 1459 |
+
def language_model(self):
|
| 1460 |
+
return self.model.language_model
|
| 1461 |
+
|
| 1462 |
+
@property
|
| 1463 |
+
def audio_tower(self):
|
| 1464 |
+
return self.model.audio_tower
|
| 1465 |
+
|
| 1466 |
+
@property
|
| 1467 |
+
def embed_audio(self):
|
| 1468 |
+
return self.model.embed_audio
|
| 1469 |
+
|
| 1470 |
+
# ---- multimodal accessors -------------------------------------------------
|
| 1471 |
+
# vLLM's Transformers modeling backend looks these up on the TOP-LEVEL class.
|
| 1472 |
+
def get_audio_features(self, input_features, feature_attention_mask=None, **kwargs):
|
| 1473 |
+
return self.model.get_audio_features(input_features, feature_attention_mask)
|
| 1474 |
+
|
| 1475 |
+
def get_encoder(self, modality: str = "audio"):
|
| 1476 |
+
"""Return the encoder for `modality`. Audio-only model."""
|
| 1477 |
+
if modality in (None, "audio"):
|
| 1478 |
+
return self.model.audio_tower
|
| 1479 |
+
raise ValueError(f"audio-only model; unsupported modality {modality!r}")
|
| 1480 |
+
|
| 1481 |
+
def get_decoder(self):
|
| 1482 |
+
return self.model.language_model
|
| 1483 |
+
|
| 1484 |
+
def forward(
|
| 1485 |
+
self,
|
| 1486 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1487 |
+
input_features: Optional[torch.FloatTensor] = None,
|
| 1488 |
+
feature_attention_mask: Optional[torch.Tensor] = None,
|
| 1489 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1490 |
+
position_ids: Optional[torch.LongTensor] = None,
|
| 1491 |
+
past_key_values: Optional[Cache] = None,
|
| 1492 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1493 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1494 |
+
use_cache: Optional[bool] = None,
|
| 1495 |
+
cache_position: Optional[torch.LongTensor] = None,
|
| 1496 |
+
logits_to_keep: int = 0,
|
| 1497 |
+
mm_token_type_ids: Optional[torch.Tensor] = None,
|
| 1498 |
+
**kwargs,
|
| 1499 |
+
):
|
| 1500 |
+
outputs = self.model(
|
| 1501 |
+
input_ids=input_ids,
|
| 1502 |
+
input_features=input_features,
|
| 1503 |
+
feature_attention_mask=feature_attention_mask,
|
| 1504 |
+
attention_mask=attention_mask,
|
| 1505 |
+
position_ids=position_ids,
|
| 1506 |
+
past_key_values=past_key_values,
|
| 1507 |
+
inputs_embeds=inputs_embeds,
|
| 1508 |
+
use_cache=use_cache,
|
| 1509 |
+
cache_position=cache_position,
|
| 1510 |
+
**kwargs,
|
| 1511 |
+
)
|
| 1512 |
+
hidden_states = outputs.last_hidden_state
|
| 1513 |
+
slice_indices = (
|
| 1514 |
+
slice(-logits_to_keep, None)
|
| 1515 |
+
if isinstance(logits_to_keep, int)
|
| 1516 |
+
else logits_to_keep
|
| 1517 |
+
)
|
| 1518 |
+
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
| 1519 |
+
|
| 1520 |
+
loss = None
|
| 1521 |
+
if labels is not None:
|
| 1522 |
+
loss = self.loss_function(
|
| 1523 |
+
logits=logits, labels=labels, vocab_size=self.vocab_size, **kwargs
|
| 1524 |
+
)
|
| 1525 |
+
|
| 1526 |
+
return CausalLMOutputWithPast(
|
| 1527 |
+
loss=loss,
|
| 1528 |
+
logits=logits,
|
| 1529 |
+
past_key_values=outputs.past_key_values,
|
| 1530 |
+
hidden_states=outputs.hidden_states,
|
| 1531 |
+
attentions=outputs.attentions,
|
| 1532 |
+
)
|
| 1533 |
+
|
| 1534 |
+
|
| 1535 |
+
__all__ = [
|
| 1536 |
+
"JaisASRConfig",
|
| 1537 |
+
"K2HorizonConfig",
|
| 1538 |
+
"Qwen3OmniMoeAudioEncoderConfig",
|
| 1539 |
+
"Qwen3OmniMoeAudioPreTrainedModel",
|
| 1540 |
+
"SinusoidsPositionEmbedding",
|
| 1541 |
+
"Qwen3OmniMoeAudioAttention",
|
| 1542 |
+
"Qwen3OmniMoeAudioEncoderLayer",
|
| 1543 |
+
"Qwen3OmniMoeAudioEncoder",
|
| 1544 |
+
"K2HorizonAttention",
|
| 1545 |
+
"K2HorizonMLP",
|
| 1546 |
+
"K2HorizonRMSNorm",
|
| 1547 |
+
"K2HorizonDecoderLayer",
|
| 1548 |
+
"K2HorizonRotaryEmbedding",
|
| 1549 |
+
"K2HorizonPreTrainedModel",
|
| 1550 |
+
"K2HorizonModel",
|
| 1551 |
+
"JaisASRProjector",
|
| 1552 |
+
"JaisASRPreTrainedModel",
|
| 1553 |
+
"JaisASRModel",
|
| 1554 |
+
"JaisASRForConditionalGeneration",
|
| 1555 |
+
]
|
pipeline_jais_asr.py
ADDED
|
@@ -0,0 +1,187 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Jais-ASR pipeline with native language controls and audio-aware batching."""
|
| 2 |
+
import os
|
| 3 |
+
import warnings
|
| 4 |
+
from collections.abc import Iterable, Mapping
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import torch
|
| 8 |
+
from torch.utils.data import DataLoader, Dataset, IterableDataset
|
| 9 |
+
from transformers import Pipeline
|
| 10 |
+
from transformers.pipelines.pt_utils import PipelineDataset, PipelineIterator
|
| 11 |
+
|
| 12 |
+
try:
|
| 13 |
+
from .language import resolve_language_controls
|
| 14 |
+
except ImportError:
|
| 15 |
+
if __package__:
|
| 16 |
+
raise
|
| 17 |
+
from language import resolve_language_controls
|
| 18 |
+
|
| 19 |
+
_UNSET = object()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class _AudioCollator:
|
| 23 |
+
"""Match the processor's left-padded text and right-padded mel time axis."""
|
| 24 |
+
|
| 25 |
+
def __init__(self, pad_token_id):
|
| 26 |
+
self.pad_token_id = pad_token_id
|
| 27 |
+
|
| 28 |
+
def __call__(self, items):
|
| 29 |
+
batch = {}
|
| 30 |
+
for key in ("input_ids", "attention_mask", "input_features", "feature_attention_mask"):
|
| 31 |
+
values = [item[key] for item in items]
|
| 32 |
+
width = max(value.shape[-1] for value in values)
|
| 33 |
+
shape = (len(values), *values[0].shape[1:-1], width)
|
| 34 |
+
fill = self.pad_token_id if key == "input_ids" else 0
|
| 35 |
+
padded = values[0].new_full(shape, fill)
|
| 36 |
+
for row, value in enumerate(values):
|
| 37 |
+
length = value.shape[-1]
|
| 38 |
+
if key in ("input_ids", "attention_mask"):
|
| 39 |
+
padded[row, ..., width - length:] = value[0]
|
| 40 |
+
else:
|
| 41 |
+
padded[row, ..., :length] = value[0]
|
| 42 |
+
batch[key] = padded
|
| 43 |
+
return batch
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
class JaisASRTranscriptionPipeline(Pipeline):
|
| 47 |
+
_load_processor = True
|
| 48 |
+
_load_tokenizer = False
|
| 49 |
+
_load_feature_extractor = False
|
| 50 |
+
_load_image_processor = False
|
| 51 |
+
|
| 52 |
+
def __init__(self, *args, **kwargs):
|
| 53 |
+
super().__init__(*args, **kwargs)
|
| 54 |
+
if getattr(self, "processor", None) is None:
|
| 55 |
+
from transformers import AutoProcessor
|
| 56 |
+
self.processor = AutoProcessor.from_pretrained(
|
| 57 |
+
self.model.name_or_path, trust_remote_code=True)
|
| 58 |
+
|
| 59 |
+
def _sanitize_parameters(
|
| 60 |
+
self, language=_UNSET, language_prefill=_UNSET, diacritics=None,
|
| 61 |
+
code_switch=None, enable_thinking=None, generate_kwargs=None, **kwargs
|
| 62 |
+
):
|
| 63 |
+
pin, prefill = resolve_language_controls(
|
| 64 |
+
None if language is _UNSET else language,
|
| 65 |
+
None if language_prefill is _UNSET else language_prefill,
|
| 66 |
+
enable_thinking=enable_thinking)
|
| 67 |
+
pre = {key: value for key, value in dict(
|
| 68 |
+
diacritics=diacritics, code_switch=code_switch,
|
| 69 |
+
enable_thinking=enable_thinking).items() if value is not None}
|
| 70 |
+
forward = dict(generate_kwargs or {})
|
| 71 |
+
if "language" in forward or "language_prefill" in forward:
|
| 72 |
+
raise ValueError("Pass language controls directly to the pipeline")
|
| 73 |
+
if "mawjais_language_pin" in kwargs or "mawjais_language_pin" in forward:
|
| 74 |
+
raise ValueError("Unsupported legacy language option; use language")
|
| 75 |
+
forward.update(kwargs)
|
| 76 |
+
# None explicitly clears a control; omission retains its default.
|
| 77 |
+
if language is not _UNSET:
|
| 78 |
+
forward["language"] = pin
|
| 79 |
+
if language_prefill is not _UNSET:
|
| 80 |
+
pre["language_prefill"] = prefill
|
| 81 |
+
# Selecting one control replaces an inherited alternative, as upstream
|
| 82 |
+
# already does. Both explicitly non-None are rejected above.
|
| 83 |
+
if pin is not None and language_prefill is _UNSET:
|
| 84 |
+
pre["language_prefill"] = None
|
| 85 |
+
if prefill is not None and language is _UNSET:
|
| 86 |
+
forward["language"] = None
|
| 87 |
+
return pre, forward, {}
|
| 88 |
+
|
| 89 |
+
def __call__(self, inputs, *, num_workers=None, batch_size=None, **kwargs):
|
| 90 |
+
# Audio inputs need neither the generic pipeline's chat detection nor its
|
| 91 |
+
# generator lookahead (which consumes input before iteration starts).
|
| 92 |
+
num_workers = (0 if self._num_workers is None else self._num_workers) if num_workers is None else num_workers
|
| 93 |
+
batch_size = (1 if self._batch_size is None else self._batch_size) if batch_size is None else batch_size
|
| 94 |
+
if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size < 1:
|
| 95 |
+
raise ValueError("batch_size must be a positive integer")
|
| 96 |
+
if not isinstance(num_workers, int) or isinstance(num_workers, bool) or num_workers < 0:
|
| 97 |
+
raise ValueError("num_workers must be a nonnegative integer")
|
| 98 |
+
pre, forward, post = self._sanitize_parameters(**kwargs)
|
| 99 |
+
pre = {**self._preprocess_params, **pre}
|
| 100 |
+
forward = {**self._forward_params, **forward}
|
| 101 |
+
post = {**self._postprocess_params, **post}
|
| 102 |
+
# Validate the effective controls after constructor/call overrides.
|
| 103 |
+
language, prefill = resolve_language_controls(
|
| 104 |
+
forward.get("language"), pre.get("language_prefill"),
|
| 105 |
+
enable_thinking=pre.get("enable_thinking", False))
|
| 106 |
+
if "language" in forward:
|
| 107 |
+
forward["language"] = language
|
| 108 |
+
if "language_prefill" in pre:
|
| 109 |
+
pre["language_prefill"] = prefill
|
| 110 |
+
self.call_count += 1
|
| 111 |
+
if isinstance(inputs, list):
|
| 112 |
+
return list(self.get_iterator(inputs, num_workers, batch_size, pre, forward, post))
|
| 113 |
+
if isinstance(inputs, (Iterable, Dataset)) and not isinstance(
|
| 114 |
+
inputs, (str, bytes, Mapping, np.ndarray, torch.Tensor)):
|
| 115 |
+
return self.get_iterator(inputs, num_workers, batch_size, pre, forward, post)
|
| 116 |
+
return self.run_single(inputs, pre, forward, post)
|
| 117 |
+
|
| 118 |
+
def get_iterator(self, inputs, num_workers, batch_size,
|
| 119 |
+
preprocess_params, forward_params, postprocess_params):
|
| 120 |
+
if hasattr(inputs, "__len__") and hasattr(inputs, "__getitem__") and not isinstance(inputs, IterableDataset):
|
| 121 |
+
dataset = PipelineDataset(inputs, self.preprocess, preprocess_params)
|
| 122 |
+
else:
|
| 123 |
+
if num_workers > 1:
|
| 124 |
+
warnings.warn("Streaming audio uses num_workers=1 to avoid duplicate inputs.", stacklevel=2)
|
| 125 |
+
num_workers = 1
|
| 126 |
+
dataset = PipelineIterator(inputs, self.preprocess, preprocess_params)
|
| 127 |
+
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
| 128 |
+
loader = DataLoader(dataset, num_workers=num_workers, batch_size=batch_size,
|
| 129 |
+
collate_fn=_AudioCollator(self.processor.tokenizer.pad_token_id))
|
| 130 |
+
outputs = PipelineIterator(loader, self.forward, forward_params, loader_batch_size=batch_size)
|
| 131 |
+
return PipelineIterator(outputs, self.postprocess, postprocess_params)
|
| 132 |
+
|
| 133 |
+
def preprocess(
|
| 134 |
+
self,
|
| 135 |
+
inputs,
|
| 136 |
+
language_prefill=None,
|
| 137 |
+
diacritics=False,
|
| 138 |
+
code_switch=False,
|
| 139 |
+
enable_thinking=False,
|
| 140 |
+
):
|
| 141 |
+
sr = None
|
| 142 |
+
if isinstance(inputs, Mapping):
|
| 143 |
+
sr = inputs.get("sampling_rate")
|
| 144 |
+
inputs = inputs.get("array", inputs.get("raw"))
|
| 145 |
+
if isinstance(inputs, (str, bytes, os.PathLike)):
|
| 146 |
+
import io
|
| 147 |
+
import soundfile as sf
|
| 148 |
+
|
| 149 |
+
inputs, sr = sf.read(
|
| 150 |
+
io.BytesIO(inputs) if isinstance(inputs, bytes) else inputs,
|
| 151 |
+
dtype="float32",
|
| 152 |
+
)
|
| 153 |
+
wav = np.asarray(inputs)
|
| 154 |
+
prompt = self.processor.tokenizer.apply_chat_template(
|
| 155 |
+
[{"role": "user", "content": [{"type": "audio"}]}],
|
| 156 |
+
tokenize=False,
|
| 157 |
+
add_generation_prompt=True,
|
| 158 |
+
language_prefill=language_prefill,
|
| 159 |
+
diacritics=diacritics,
|
| 160 |
+
code_switch=code_switch,
|
| 161 |
+
enable_thinking=enable_thinking,
|
| 162 |
+
)
|
| 163 |
+
return self.processor(
|
| 164 |
+
text=[prompt],
|
| 165 |
+
audio=[wav],
|
| 166 |
+
return_tensors="pt",
|
| 167 |
+
**({"sampling_rate": sr} if sr is not None else {})
|
| 168 |
+
)
|
| 169 |
+
|
| 170 |
+
def _forward(self, model_inputs, **generate_kwargs):
|
| 171 |
+
# Native generate(language=...) binds its own pin to this batch's width.
|
| 172 |
+
# The pipeline must not append another pin processor.
|
| 173 |
+
out = self.model.generate(**model_inputs, **generate_kwargs)
|
| 174 |
+
if hasattr(out, "sequences"):
|
| 175 |
+
out = out.sequences
|
| 176 |
+
if out.shape[0] != model_inputs["input_ids"].shape[0]:
|
| 177 |
+
raise ValueError("The transcription pipeline requires one returned sequence per input")
|
| 178 |
+
width = model_inputs["input_ids"].shape[1]
|
| 179 |
+
return {"out": out, "width": [width] * out.shape[0]}
|
| 180 |
+
|
| 181 |
+
def postprocess(self, model_outputs):
|
| 182 |
+
width = model_outputs["width"]
|
| 183 |
+
if isinstance(width, list): # single calls and batches of one are not split
|
| 184 |
+
width = width[0]
|
| 185 |
+
t = self.processor.post_process_transcription(
|
| 186 |
+
model_outputs["out"], input_lengths=int(width))[0]
|
| 187 |
+
return {"text": t.text}
|
postprocess.py
ADDED
|
@@ -0,0 +1,404 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Reading Jais-ASR output: the leading `[...]` decision tag.
|
| 2 |
+
|
| 3 |
+
Every generation starts with a bracketed decision:
|
| 4 |
+
|
| 5 |
+
"[EG] النص هنا" -> a transcript, in Egyptian Arabic
|
| 6 |
+
"[msa] النص هنا" -> a transcript, in Modern Standard Arabic
|
| 7 |
+
"[IFM-NoTranscript]" -> a deliberate abstention: no transcribable speech
|
| 8 |
+
|
| 9 |
+
The rule is general: the transcript never begins with a bracketed span. Anything the
|
| 10 |
+
model writes inside leading `[...]` spans is decision metadata, routed to `.tag` (and
|
| 11 |
+
to the API's `reasoning` field), never into `.text`:
|
| 12 |
+
|
| 13 |
+
* a span whose body opens with 'IF' (any case) is the refusal branch: `.refused`
|
| 14 |
+
is True and `.text` is empty -- nothing after a refusal is a transcript.
|
| 15 |
+
* otherwise the first alphabetic component that case-folds onto a known tag sets
|
| 16 |
+
`.language`; a span naming no supported language leaves `.language` as None.
|
| 17 |
+
* no leading span -> `.tag` is None and `.text` is the whole output, untouched.
|
| 18 |
+
|
| 19 |
+
from postprocess import split_text
|
| 20 |
+
|
| 21 |
+
r = split_text(raw)
|
| 22 |
+
r.text # the transcript
|
| 23 |
+
r.language # canonical tag ("EG", "msa", ...) or None
|
| 24 |
+
r.refused # True when the model abstained
|
| 25 |
+
r.tag # the leading tag as emitted, e.g. "[EG]", or None
|
| 26 |
+
|
| 27 |
+
For well-formed output the exact token-id fast path (`split_token_ids`) gives the same
|
| 28 |
+
answer with no string work; the text rule is the authority otherwise.
|
| 29 |
+
|
| 30 |
+
This module is pure python: no torch, no transformers, no vllm at import time.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
from __future__ import annotations
|
| 34 |
+
|
| 35 |
+
import re
|
| 36 |
+
from dataclasses import dataclass
|
| 37 |
+
|
| 38 |
+
__all__ = [
|
| 39 |
+
"LANGUAGE_TAGS",
|
| 40 |
+
"REFUSAL_TAG",
|
| 41 |
+
"Transcription",
|
| 42 |
+
"split_text",
|
| 43 |
+
"split_token_ids",
|
| 44 |
+
"classify_tag_body",
|
| 45 |
+
"stream_split",
|
| 46 |
+
"verify_against_tokenizer",
|
| 47 |
+
]
|
| 48 |
+
|
| 49 |
+
# ---------------------------------------------------------------------------------
|
| 50 |
+
# The decision inventory. Token ids are those of this repo's tokenizer;
|
| 51 |
+
# `verify_against_tokenizer` below re-derives every one of them.
|
| 52 |
+
# ---------------------------------------------------------------------------------
|
| 53 |
+
BRACKET_OPEN_ID = 60 # '['
|
| 54 |
+
BRACKET_CLOSE_ID = 62 # ']'
|
| 55 |
+
THINK_OPEN_ID = 250029 # <ifm|think>
|
| 56 |
+
THINK_CLOSE_ID = 250030 # '</ifm|think>', the last token of every rendered prompt
|
| 57 |
+
|
| 58 |
+
REFUSAL_TAG = "[IFM-NoTranscript]"
|
| 59 |
+
REFUSAL_TOKEN_IDS = (
|
| 60 |
+
60,
|
| 61 |
+
5126,
|
| 62 |
+
46,
|
| 63 |
+
131126,
|
| 64 |
+
195862,
|
| 65 |
+
62,
|
| 66 |
+
) # '[' 'IF' 'M' '-No' 'Transcript' ']'
|
| 67 |
+
|
| 68 |
+
#: decision-token id -> tag surface. 24 tags, all distinct.
|
| 69 |
+
LANGUAGE_TAG_IDS = {
|
| 70 |
+
226606: "msa", # Modern Standard Arabic
|
| 71 |
+
272: "en",
|
| 72 |
+
2075: "fr",
|
| 73 |
+
273: "es",
|
| 74 |
+
286: "it",
|
| 75 |
+
7646: "EG",
|
| 76 |
+
8048: "SA",
|
| 77 |
+
6101: "MA",
|
| 78 |
+
29354: "IQ",
|
| 79 |
+
23129: "JO",
|
| 80 |
+
20056: "SY",
|
| 81 |
+
27526: "QA",
|
| 82 |
+
64489: "DZ",
|
| 83 |
+
23852: "LB",
|
| 84 |
+
5594: "OM",
|
| 85 |
+
28472: "YE",
|
| 86 |
+
43852: "TN",
|
| 87 |
+
17604: "MR",
|
| 88 |
+
25893: "KW",
|
| 89 |
+
9373: "SD",
|
| 90 |
+
12933: "LY",
|
| 91 |
+
5402: "AE",
|
| 92 |
+
4295: "PS",
|
| 93 |
+
41905: "BH",
|
| 94 |
+
}
|
| 95 |
+
LANGUAGE_TAGS = tuple(LANGUAGE_TAG_IDS.values())
|
| 96 |
+
|
| 97 |
+
#: longest well-formed prefix, in tokens -- the id fast path never looks further
|
| 98 |
+
MAX_TAG_TOKENS = len(REFUSAL_TOKEN_IDS)
|
| 99 |
+
|
| 100 |
+
#: a decision span is SHORT. 64 chars is far beyond any real tag; anything longer
|
| 101 |
+
#: is not tag-shaped and stays in the transcript untouched.
|
| 102 |
+
TAG_SPAN_MAX_CHARS = 64
|
| 103 |
+
#: total leading-metadata budget (a few stacked spans at most)
|
| 104 |
+
TAG_ZONE_MAX_CHARS = 160
|
| 105 |
+
|
| 106 |
+
# one bracketed span (no newline inside; the first ']' closes it); used with .match(pos)
|
| 107 |
+
_SPAN_RE = re.compile(r"[ \t]*\[([^\]\n]{0,%d})\]" % TAG_SPAN_MAX_CHARS)
|
| 108 |
+
|
| 109 |
+
_CANONICAL = {t.lower(): t for t in LANGUAGE_TAGS}
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
@dataclass(frozen=True)
|
| 113 |
+
class Transcription:
|
| 114 |
+
"""What the model decided, and what it transcribed."""
|
| 115 |
+
|
| 116 |
+
text: str
|
| 117 |
+
language: str | None = None
|
| 118 |
+
refused: bool = False
|
| 119 |
+
tag: str | None = None
|
| 120 |
+
thinking: str = ""
|
| 121 |
+
truncated: bool = False
|
| 122 |
+
|
| 123 |
+
def __bool__(self) -> bool: # `if result:` is False on an abstention
|
| 124 |
+
return not self.refused and bool(self.text)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def has_prefilled_tag(prompt_ids):
|
| 128 |
+
"""Only a complete known [language] suffix counts as a prompt prefill."""
|
| 129 |
+
return (
|
| 130 |
+
len(prompt_ids) >= 3
|
| 131 |
+
and int(prompt_ids[-3]) == BRACKET_OPEN_ID
|
| 132 |
+
and int(prompt_ids[-2]) in LANGUAGE_TAG_IDS
|
| 133 |
+
and int(prompt_ids[-1]) == BRACKET_CLOSE_ID
|
| 134 |
+
)
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def classify_tag_body(body: str):
|
| 138 |
+
"""(language | None, refused) for one span body. See the module contract."""
|
| 139 |
+
b = body.strip()
|
| 140 |
+
if b[:2].upper() == "IF":
|
| 141 |
+
return None, True
|
| 142 |
+
for comp in re.split(r"[^A-Za-z]+", b):
|
| 143 |
+
if comp and comp.lower() in _CANONICAL:
|
| 144 |
+
return _CANONICAL[comp.lower()], False
|
| 145 |
+
return None, False
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
def _leading_spans(text: str):
|
| 149 |
+
"""Parse the leading metadata zone -> (spans, refused, language, end_pos)."""
|
| 150 |
+
spans, refused, language, pos = [], False, None, 0
|
| 151 |
+
while pos <= TAG_ZONE_MAX_CHARS:
|
| 152 |
+
m = _SPAN_RE.match(text, pos)
|
| 153 |
+
if not m:
|
| 154 |
+
break
|
| 155 |
+
body = m.group(1)
|
| 156 |
+
spans.append(f"[{body}]")
|
| 157 |
+
lang, ref = classify_tag_body(body)
|
| 158 |
+
refused = refused or ref
|
| 159 |
+
if language is None and lang is not None:
|
| 160 |
+
language = lang
|
| 161 |
+
pos = m.end()
|
| 162 |
+
return spans, refused, language, pos
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
# ---------------------------------------------------------------------------------
|
| 166 |
+
# token-id path -- exact fast path for WELL-FORMED output
|
| 167 |
+
# ---------------------------------------------------------------------------------
|
| 168 |
+
def tag_span_len(token_ids) -> int:
|
| 169 |
+
"""Tokens spanned by a WELL-FORMED leading tag, else 0 (the text rule then decides)."""
|
| 170 |
+
ids = list(token_ids[:MAX_TAG_TOKENS])
|
| 171 |
+
if not ids or ids[0] != BRACKET_OPEN_ID:
|
| 172 |
+
return 0
|
| 173 |
+
if tuple(ids[: len(REFUSAL_TOKEN_IDS)]) == REFUSAL_TOKEN_IDS:
|
| 174 |
+
return len(REFUSAL_TOKEN_IDS)
|
| 175 |
+
if len(ids) >= 3 and ids[1] in LANGUAGE_TAG_IDS and ids[2] == BRACKET_CLOSE_ID:
|
| 176 |
+
return 3
|
| 177 |
+
return 0
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
def split_token_ids(token_ids, *, strip: bool = True):
|
| 181 |
+
"""-> (tag or None, refused, remaining_token_ids). Exact-form fast path only;
|
| 182 |
+
anything else returns (None, False, ids) and callers fall back to `split_text`."""
|
| 183 |
+
ids = list(token_ids)
|
| 184 |
+
n = tag_span_len(ids)
|
| 185 |
+
if n == 0:
|
| 186 |
+
return None, False, ids
|
| 187 |
+
if n == len(REFUSAL_TOKEN_IDS):
|
| 188 |
+
return REFUSAL_TAG, True, ([] if strip else ids)
|
| 189 |
+
return f"[{LANGUAGE_TAG_IDS[ids[1]]}]", False, (ids[n:] if strip else ids)
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
# ---------------------------------------------------------------------------------
|
| 193 |
+
# text path -- the authority
|
| 194 |
+
# ---------------------------------------------------------------------------------
|
| 195 |
+
def _split_think(text: str):
|
| 196 |
+
"""Generated-side think block (enable_thinking=True): the prompt opens
|
| 197 |
+
'<ifm|think>', so the generation carries thought text up to the first
|
| 198 |
+
'</ifm|think>'. Without that marker this is a no-op."""
|
| 199 |
+
head = text.lstrip()
|
| 200 |
+
if head.startswith("<ifm|think>"):
|
| 201 |
+
text = head[len("<ifm|think>") :]
|
| 202 |
+
idx = text.find("</ifm|think>")
|
| 203 |
+
if idx < 0:
|
| 204 |
+
return "", text
|
| 205 |
+
return text[:idx].strip(), text[idx + len("</ifm|think>") :]
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def split_text(text: str, *, strip: bool = True) -> Transcription:
|
| 209 |
+
"""Split a decoded generation per the general contract (module docstring)."""
|
| 210 |
+
thinking, text = _split_think(text)
|
| 211 |
+
spans, refused, language, pos = _leading_spans(text)
|
| 212 |
+
if not spans:
|
| 213 |
+
return Transcription(text=text, thinking=thinking)
|
| 214 |
+
tag = " ".join(spans)
|
| 215 |
+
if refused:
|
| 216 |
+
return Transcription(
|
| 217 |
+
text="" if strip else text, refused=True, tag=tag, thinking=thinking
|
| 218 |
+
)
|
| 219 |
+
rest = text[pos:].lstrip() if strip else text
|
| 220 |
+
return Transcription(text=rest, language=language, tag=tag, thinking=thinking)
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
def stream_split(text: str):
|
| 224 |
+
"""Streaming view of the contract over the text-so-far.
|
| 225 |
+
|
| 226 |
+
-> (reasoning_so_far, content_so_far, holding)
|
| 227 |
+
|
| 228 |
+
* while the head could still be inside the metadata zone (only whitespace yet, or
|
| 229 |
+
a short unclosed '[' span) the undecided tail is HELD (both strings frozen);
|
| 230 |
+
* a refusal-branch output puts its closed leading span(s) in reasoning; nothing
|
| 231 |
+
after a refusal is ever content;
|
| 232 |
+
* otherwise reasoning is the closed leading span(s), content starts at the first
|
| 233 |
+
non-metadata character. Both returned strings only ever grow (delta-safe).
|
| 234 |
+
"""
|
| 235 |
+
spans, refused, _lang, pos = _leading_spans(text)
|
| 236 |
+
if refused:
|
| 237 |
+
return text[:pos], "", False
|
| 238 |
+
rest = text[pos:]
|
| 239 |
+
head = rest.lstrip(" \t")
|
| 240 |
+
ws = len(rest) - len(head)
|
| 241 |
+
if not spans:
|
| 242 |
+
if head == "" and len(text) <= TAG_ZONE_MAX_CHARS:
|
| 243 |
+
return "", "", True # nothing decisive yet
|
| 244 |
+
if (
|
| 245 |
+
head.startswith("[")
|
| 246 |
+
and "]" not in head
|
| 247 |
+
and "\n" not in head
|
| 248 |
+
and len(head) <= TAG_SPAN_MAX_CHARS + 2
|
| 249 |
+
):
|
| 250 |
+
return "", "", True # a span may be forming
|
| 251 |
+
return "", text, False # no metadata: all content
|
| 252 |
+
if (
|
| 253 |
+
head.startswith("[")
|
| 254 |
+
and "]" not in head
|
| 255 |
+
and "\n" not in head
|
| 256 |
+
and len(head) <= TAG_SPAN_MAX_CHARS + 2
|
| 257 |
+
and pos + ws <= TAG_ZONE_MAX_CHARS
|
| 258 |
+
):
|
| 259 |
+
return text[:pos], "", True # another span may be forming
|
| 260 |
+
return text[:pos], text[pos + ws :], False
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
def verify_against_tokenizer(tokenizer) -> None:
|
| 264 |
+
"""Assert the id table matches this tokenizer. Cheap; call it in your tests."""
|
| 265 |
+
enc = lambda s: tokenizer(s, add_special_tokens=False).input_ids # noqa: E731
|
| 266 |
+
assert (
|
| 267 |
+
tuple(enc(REFUSAL_TAG)) == REFUSAL_TOKEN_IDS
|
| 268 |
+
), f"{REFUSAL_TAG} -> {enc(REFUSAL_TAG)}, expected {list(REFUSAL_TOKEN_IDS)}"
|
| 269 |
+
assert tuple(enc("</ifm|think>")) == (
|
| 270 |
+
THINK_CLOSE_ID,
|
| 271 |
+
), f"</ifm|think> -> {enc('</ifm|think>')}, expected [{THINK_CLOSE_ID}]"
|
| 272 |
+
for tid, tag in LANGUAGE_TAG_IDS.items():
|
| 273 |
+
ids = enc(f"[{tag}] ")
|
| 274 |
+
assert (
|
| 275 |
+
ids[0] == BRACKET_OPEN_ID and ids[1] == tid
|
| 276 |
+
), f"[{tag}] -> {ids[:3]}, expected decision id {tid} at position 1"
|
| 277 |
+
assert len(set(LANGUAGE_TAG_IDS)) == 24
|
| 278 |
+
|
| 279 |
+
|
| 280 |
+
# =================================================================================
|
| 281 |
+
# vLLM integration -- metadata becomes `message.reasoning`, transcript is `content`
|
| 282 |
+
#
|
| 283 |
+
# vllm serve <model> --trust-remote-code --reasoning-parser jais_asr_tag
|
| 284 |
+
#
|
| 285 |
+
# Registration happens on import inside a vLLM process; the vllm_plugin/ package in
|
| 286 |
+
# this repo does that automatically at server start.
|
| 287 |
+
# =================================================================================
|
| 288 |
+
try: # pragma: no cover - only meaningful inside a vLLM process
|
| 289 |
+
from vllm.reasoning.abs_reasoning_parsers import (
|
| 290 |
+
ReasoningParser,
|
| 291 |
+
ReasoningParserManager,
|
| 292 |
+
)
|
| 293 |
+
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
|
| 294 |
+
|
| 295 |
+
@ReasoningParserManager.register_module("jais_asr_tag")
|
| 296 |
+
class JaisASRTagParser(ReasoningParser):
|
| 297 |
+
"""reasoning = the leading decision metadata; content = the transcript only.
|
| 298 |
+
|
| 299 |
+
Same general contract as split_text: any leading bracketed span(s) go to
|
| 300 |
+
`reasoning`; a refusal-branch output puts its leading tag in `reasoning`
|
| 301 |
+
(the same value as `.tag`) with empty content; no leading span -> empty
|
| 302 |
+
reasoning.
|
| 303 |
+
"""
|
| 304 |
+
|
| 305 |
+
def __init__(self, tokenizer, *args, **kwargs):
|
| 306 |
+
template = kwargs.get("chat_template_kwargs") or {}
|
| 307 |
+
self.language_prefill = template.get("language_prefill")
|
| 308 |
+
self._prefill_emitted = False
|
| 309 |
+
super().__init__(tokenizer, *args, **kwargs)
|
| 310 |
+
|
| 311 |
+
# ReasoningConfig auto-initialization tokenizes these at server start;
|
| 312 |
+
# they are the prompt's think-block delimiters, already closed in every
|
| 313 |
+
# rendered prompt, so the thinking-budget machinery stays inert.
|
| 314 |
+
@property
|
| 315 |
+
def reasoning_start_str(self) -> str:
|
| 316 |
+
return "<ifm|think>"
|
| 317 |
+
|
| 318 |
+
@property
|
| 319 |
+
def reasoning_end_str(self) -> str:
|
| 320 |
+
return "</ifm|think>"
|
| 321 |
+
|
| 322 |
+
def is_reasoning_end(self, input_ids) -> bool:
|
| 323 |
+
ids = list(input_ids)
|
| 324 |
+
# vLLM also calls this with PROMPT ids (seeding per-request state) and
|
| 325 |
+
# with prompt+output ids; every rendered prompt ends with '</ifm|think>'
|
| 326 |
+
# (+ an optional prefill), so only what follows the LAST one is output.
|
| 327 |
+
for i in range(len(ids) - 1, -1, -1):
|
| 328 |
+
if ids[i] == THINK_CLOSE_ID:
|
| 329 |
+
ids = ids[i + 1 :]
|
| 330 |
+
break
|
| 331 |
+
else:
|
| 332 |
+
if ids and ids[0] == 0: # BOS-led, no anchor: a bare prompt
|
| 333 |
+
return False
|
| 334 |
+
text = self.model_tokenizer.decode(ids[:48])
|
| 335 |
+
_r, content, holding = stream_split(text)
|
| 336 |
+
return (not holding) and content != ""
|
| 337 |
+
|
| 338 |
+
def extract_content_ids(self, input_ids):
|
| 339 |
+
ids = list(input_ids)
|
| 340 |
+
text = self.model_tokenizer.decode(ids)
|
| 341 |
+
r = split_text(text)
|
| 342 |
+
if r.refused:
|
| 343 |
+
return []
|
| 344 |
+
if r.tag is None:
|
| 345 |
+
return ids
|
| 346 |
+
# the space after ']' merges into the first content token, so keep the
|
| 347 |
+
# boundary token: cut after the last prefix that fits INSIDE the tag zone
|
| 348 |
+
cut = len(text) - len(r.text)
|
| 349 |
+
keep = 0
|
| 350 |
+
for k in range(1, len(ids) + 1):
|
| 351 |
+
if len(self.model_tokenizer.decode(ids[:k])) > cut:
|
| 352 |
+
break
|
| 353 |
+
keep = k
|
| 354 |
+
return ids[keep:]
|
| 355 |
+
|
| 356 |
+
def extract_reasoning(self, model_output, request):
|
| 357 |
+
r = split_text(model_output)
|
| 358 |
+
tag = r.tag
|
| 359 |
+
if tag is None and not r.refused:
|
| 360 |
+
kw = getattr(request, "chat_template_kwargs", None) or {}
|
| 361 |
+
if kw.get("language_prefill"):
|
| 362 |
+
tag = f"[{kw['language_prefill']}]" # tag was prefilled into the prompt
|
| 363 |
+
reasoning = "\n".join(p for p in (r.thinking, tag or "") if p) or None
|
| 364 |
+
if r.refused:
|
| 365 |
+
return reasoning, ""
|
| 366 |
+
if tag is None and not r.thinking:
|
| 367 |
+
return None, model_output
|
| 368 |
+
return reasoning, r.text
|
| 369 |
+
|
| 370 |
+
def extract_reasoning_streaming(
|
| 371 |
+
self,
|
| 372 |
+
previous_text,
|
| 373 |
+
current_text,
|
| 374 |
+
delta_text,
|
| 375 |
+
previous_token_ids,
|
| 376 |
+
current_token_ids,
|
| 377 |
+
delta_token_ids,
|
| 378 |
+
):
|
| 379 |
+
pr, pc, _ = stream_split(previous_text)
|
| 380 |
+
cr, cc, _holding = stream_split(current_text)
|
| 381 |
+
dr = cr[len(pr) :] if cr.startswith(pr) else cr
|
| 382 |
+
dc = cc[len(pc) :] if cc.startswith(pc) else cc
|
| 383 |
+
if self.language_prefill and not self._prefill_emitted:
|
| 384 |
+
dr = f"[{self.language_prefill}]" + ("\n" + dr if dr else "")
|
| 385 |
+
self._prefill_emitted = True
|
| 386 |
+
if not dr and not dc:
|
| 387 |
+
return None
|
| 388 |
+
return DeltaMessage(reasoning=dr or None, content=dc or None)
|
| 389 |
+
|
| 390 |
+
__all__.append("JaisASRTagParser")
|
| 391 |
+
except (
|
| 392 |
+
Exception
|
| 393 |
+
) as exc: # vLLM absent, or an API shift -- the pure-python half still works
|
| 394 |
+
import importlib.util
|
| 395 |
+
|
| 396 |
+
if importlib.util.find_spec("vllm") is not None:
|
| 397 |
+
import logging
|
| 398 |
+
|
| 399 |
+
logging.getLogger(__name__).warning(
|
| 400 |
+
"jais_asr_tag reasoning parser NOT registered (%s: %s); "
|
| 401 |
+
"`--reasoning-parser jais_asr_tag` will fail with a KeyError",
|
| 402 |
+
type(exc).__name__,
|
| 403 |
+
exc,
|
| 404 |
+
)
|
preprocessor_config.json
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"dither": 0.0,
|
| 3 |
+
"feature_extractor_type": "WhisperFeatureExtractor",
|
| 4 |
+
"feature_size": 128,
|
| 5 |
+
"hop_length": 160,
|
| 6 |
+
"n_fft": 400,
|
| 7 |
+
"n_samples": 4800000,
|
| 8 |
+
"nb_max_frames": 30000,
|
| 9 |
+
"padding_side": "right",
|
| 10 |
+
"padding_value": 0.0,
|
| 11 |
+
"processor_class": "JaisASRProcessor",
|
| 12 |
+
"return_attention_mask": true,
|
| 13 |
+
"sampling_rate": 16000
|
| 14 |
+
}
|
processing_jais_asr.py
ADDED
|
@@ -0,0 +1,390 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Processor for Jais-ASR.
|
| 2 |
+
|
| 3 |
+
`processor(text=..., audio=..., sampling_rate=...)` takes one `<|ifm|audio|>`
|
| 4 |
+
placeholder per clip and expands it to `<audio>` + one placeholder per encoder frame
|
| 5 |
+
+ `</audio>`, so the model's `masked_scatter` lines up exactly with the frames the
|
| 6 |
+
audio tower emits.
|
| 7 |
+
|
| 8 |
+
Audio input is standardized automatically: pass `sampling_rate=` with whatever rate
|
| 9 |
+
your file has (a scalar, or one rate per clip) and it is resampled to the model's
|
| 10 |
+
16 kHz; multi-channel audio is downmixed to mono and integer PCM is converted to
|
| 11 |
+
float. When `sampling_rate` is omitted the audio is assumed to already be 16 kHz.
|
| 12 |
+
A bare 2-D numpy array is always a batch of mono clips (leading dim = batch); pass
|
| 13 |
+
a multi-channel clip inside a list to have its channels downmixed.
|
| 14 |
+
|
| 15 |
+
Audio length is unrestricted at the long end: no maximum duration, no truncation.
|
| 16 |
+
The mel is extracted per clip at full length and long audio is chunked inside the
|
| 17 |
+
encoder; each clip expands to its own frame count (~13/sec). Clips must be longer
|
| 18 |
+
than `n_fft // 2` samples (200 samples = 12.5 ms at 16 kHz), the STFT minimum.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import re
|
| 22 |
+
|
| 23 |
+
import numpy as np
|
| 24 |
+
|
| 25 |
+
import transformers
|
| 26 |
+
from transformers.feature_extraction_utils import BatchFeature
|
| 27 |
+
from transformers.processing_utils import ProcessorMixin
|
| 28 |
+
|
| 29 |
+
if tuple(int(p) for p in transformers.__version__.split(".")[:2]) < (4, 57):
|
| 30 |
+
raise ImportError(
|
| 31 |
+
"MawJais-stt requires transformers >= 4.57 (earlier versions lack the "
|
| 32 |
+
"WhisperFeatureExtractor attention-mask trim and corrupt the audio features); "
|
| 33 |
+
f"found transformers=={transformers.__version__}. "
|
| 34 |
+
"Upgrade with `pip install -U 'transformers>=4.57'`.")
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def _audio_output_lengths(input_lengths):
|
| 38 |
+
"""Encoder output-frame count from the mel-frame count.
|
| 39 |
+
|
| 40 |
+
Three stride-2 stages plus per-100-frame windowing: 100 mel frames/sec -> 13 output
|
| 41 |
+
tokens/sec. Works elementwise on a python int or a torch / numpy integer tensor.
|
| 42 |
+
"""
|
| 43 |
+
input_lengths_leave = input_lengths % 100
|
| 44 |
+
feat_lengths = (input_lengths_leave - 1) // 2 + 1
|
| 45 |
+
return ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _resample(wave, orig_sr, target_sr):
|
| 49 |
+
"""Resample one mono float32 clip with the best library installed."""
|
| 50 |
+
try:
|
| 51 |
+
import soxr
|
| 52 |
+
return np.asarray(soxr.resample(wave, orig_sr, target_sr, quality="HQ"),
|
| 53 |
+
dtype=np.float32)
|
| 54 |
+
except ImportError:
|
| 55 |
+
pass
|
| 56 |
+
try:
|
| 57 |
+
import librosa
|
| 58 |
+
return np.asarray(librosa.resample(wave, orig_sr=orig_sr, target_sr=target_sr),
|
| 59 |
+
dtype=np.float32)
|
| 60 |
+
except ImportError:
|
| 61 |
+
pass
|
| 62 |
+
try:
|
| 63 |
+
from math import gcd
|
| 64 |
+
from scipy.signal import resample_poly
|
| 65 |
+
g = gcd(int(orig_sr), int(target_sr))
|
| 66 |
+
return np.asarray(resample_poly(wave, target_sr // g, orig_sr // g),
|
| 67 |
+
dtype=np.float32)
|
| 68 |
+
except ImportError:
|
| 69 |
+
pass
|
| 70 |
+
try:
|
| 71 |
+
import torch
|
| 72 |
+
import torchaudio.functional as _taf
|
| 73 |
+
return _taf.resample(torch.from_numpy(wave), orig_sr, target_sr).numpy()
|
| 74 |
+
except ImportError:
|
| 75 |
+
raise ImportError(
|
| 76 |
+
f"this audio is {orig_sr} Hz and the model expects {target_sr} Hz; install a "
|
| 77 |
+
"resampler (`pip install soxr`, or librosa / scipy / torchaudio) or resample "
|
| 78 |
+
"the audio yourself.")
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def _standardize_audio(wave, sr, target_sr):
|
| 82 |
+
"""One clip -> float32 mono at `target_sr`.
|
| 83 |
+
|
| 84 |
+
Accepts any dtype (integer PCM is rescaled to [-1, 1]), mono or multi-channel
|
| 85 |
+
(channels are averaged), any sample rate (`sr=None` means already `target_sr`).
|
| 86 |
+
A clip that is already float mono at `target_sr` passes through untouched.
|
| 87 |
+
"""
|
| 88 |
+
a = np.asarray(wave)
|
| 89 |
+
# integer PCM -> [-1, 1] float FIRST: channel averaging below returns float and
|
| 90 |
+
# would otherwise hide the integer dtype, leaving the waveform at PCM amplitude
|
| 91 |
+
if np.issubdtype(a.dtype, np.unsignedinteger): # e.g. uint8 PCM (offset binary)
|
| 92 |
+
half = float(np.iinfo(a.dtype).max + 1) / 2.0
|
| 93 |
+
a = (a.astype(np.float32) - half) / half
|
| 94 |
+
elif np.issubdtype(a.dtype, np.integer): # e.g. int16 / int32 PCM
|
| 95 |
+
a = a.astype(np.float32) / float(np.iinfo(a.dtype).max + 1)
|
| 96 |
+
if a.ndim == 2:
|
| 97 |
+
ch_axis = 0 if a.shape[0] <= a.shape[1] else 1
|
| 98 |
+
if a.shape[ch_axis] > 8:
|
| 99 |
+
raise ValueError(
|
| 100 |
+
f"cannot tell channels from samples in audio of shape {tuple(a.shape)}; "
|
| 101 |
+
"pass each clip as 1-D mono or (channels, samples) / (samples, channels)")
|
| 102 |
+
a = a.mean(axis=ch_axis)
|
| 103 |
+
if a.ndim != 1:
|
| 104 |
+
raise ValueError(
|
| 105 |
+
f"each audio clip must be 1-D (mono) or 2-D (with a channel axis), "
|
| 106 |
+
f"got shape {tuple(a.shape)}")
|
| 107 |
+
a = a.astype(np.float32, copy=False)
|
| 108 |
+
if sr is not None and int(sr) != int(target_sr):
|
| 109 |
+
a = _resample(a, int(sr), int(target_sr))
|
| 110 |
+
return a
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
class JaisASRProcessor(ProcessorMixin):
|
| 114 |
+
attributes = ["feature_extractor", "tokenizer"]
|
| 115 |
+
tokenizer_class = "AutoTokenizer"
|
| 116 |
+
# transformers>=5 resolves the feature extractor from preprocessor_config.json
|
| 117 |
+
# (`feature_extractor_type`) and deprecates the class attribute; <5 still needs it.
|
| 118 |
+
if int(transformers.__version__.split(".")[0]) < 5:
|
| 119 |
+
feature_extractor_class = "WhisperFeatureExtractor"
|
| 120 |
+
valid_kwargs = ["chat_template", "audio_token", "boa_token", "eoa_token"]
|
| 121 |
+
|
| 122 |
+
#: fields this processor can produce (read by serving backends)
|
| 123 |
+
model_input_names = [
|
| 124 |
+
"input_ids", "attention_mask", "input_features", "feature_attention_mask",
|
| 125 |
+
]
|
| 126 |
+
|
| 127 |
+
def __init__(
|
| 128 |
+
self,
|
| 129 |
+
feature_extractor=None,
|
| 130 |
+
tokenizer=None,
|
| 131 |
+
chat_template=None,
|
| 132 |
+
audio_token="<|ifm|audio|>",
|
| 133 |
+
boa_token="<audio>",
|
| 134 |
+
eoa_token="</audio>",
|
| 135 |
+
**kwargs,
|
| 136 |
+
):
|
| 137 |
+
self.audio_token = audio_token
|
| 138 |
+
self.boa_token = boa_token # span delimiters; never scattered onto
|
| 139 |
+
self.eoa_token = eoa_token
|
| 140 |
+
self.audio_token_id = (
|
| 141 |
+
tokenizer.convert_tokens_to_ids(audio_token) if tokenizer is not None else None
|
| 142 |
+
)
|
| 143 |
+
super().__init__(feature_extractor, tokenizer, chat_template=chat_template)
|
| 144 |
+
|
| 145 |
+
# ---- ids some serving backends look for by these exact names -----------------
|
| 146 |
+
@property
|
| 147 |
+
def audio_token_ids(self):
|
| 148 |
+
return [self.audio_token_id]
|
| 149 |
+
|
| 150 |
+
@property
|
| 151 |
+
def audio_processor(self):
|
| 152 |
+
return self.feature_extractor
|
| 153 |
+
|
| 154 |
+
def __call__(
|
| 155 |
+
self,
|
| 156 |
+
text=None,
|
| 157 |
+
audio=None,
|
| 158 |
+
return_tensors=None,
|
| 159 |
+
return_mm_token_type_ids: bool = False,
|
| 160 |
+
**kwargs,
|
| 161 |
+
):
|
| 162 |
+
if text is None:
|
| 163 |
+
raise ValueError("`text` is required.")
|
| 164 |
+
# rate of the INCOMING audio; consumed here (audio is standardized to the
|
| 165 |
+
# feature extractor's rate below), never forwarded into the tokenizer
|
| 166 |
+
sr_kw = kwargs.pop("sampling_rate", None)
|
| 167 |
+
# ProcessorMixin.apply_chat_template(tokenize=True) always passes images= and
|
| 168 |
+
# videos=; they must not fall through **kwargs into the tokenizer
|
| 169 |
+
for key in ("images", "videos"):
|
| 170 |
+
if kwargs.pop(key, None) is not None:
|
| 171 |
+
raise ValueError(f"`{key}` is not supported: this is an audio-only model.")
|
| 172 |
+
if isinstance(text, str):
|
| 173 |
+
text = [text]
|
| 174 |
+
|
| 175 |
+
audio_inputs = {}
|
| 176 |
+
if audio is not None:
|
| 177 |
+
sr_target = int(self.feature_extractor.sampling_rate)
|
| 178 |
+
if isinstance(audio, np.ndarray):
|
| 179 |
+
if audio.ndim == 1:
|
| 180 |
+
audio = [audio] # one mono clip
|
| 181 |
+
else:
|
| 182 |
+
audio = list(audio) # leading dim is always the batch of clips
|
| 183 |
+
elif not isinstance(audio, (list, tuple)):
|
| 184 |
+
audio = [audio]
|
| 185 |
+
if len(audio) == 0:
|
| 186 |
+
raise ValueError("`audio` is empty; pass at least one clip or omit `audio=`.")
|
| 187 |
+
if isinstance(sr_kw, (list, tuple)):
|
| 188 |
+
if len(sr_kw) != len(audio):
|
| 189 |
+
raise ValueError(
|
| 190 |
+
f"sampling_rate has {len(sr_kw)} entries but {len(audio)} audio clips "
|
| 191 |
+
"were passed; give one rate per clip (or a single scalar).")
|
| 192 |
+
audio = [_standardize_audio(a, s, sr_target) for a, s in zip(audio, sr_kw)]
|
| 193 |
+
else:
|
| 194 |
+
audio = [_standardize_audio(a, sr_kw, sr_target) for a in audio]
|
| 195 |
+
min_samples = int(self.feature_extractor.n_fft) // 2 + 1
|
| 196 |
+
for i, a in enumerate(audio):
|
| 197 |
+
if len(a) < min_samples:
|
| 198 |
+
raise ValueError(
|
| 199 |
+
f"audio clip {i} has {len(a)} samples after standardization; the STFT "
|
| 200 |
+
f"needs at least {min_samples} samples "
|
| 201 |
+
f"({1000.0 * min_samples / sr_target:.1f} ms at {sr_target} Hz). Note a "
|
| 202 |
+
"bare 2-D array is a batch of mono clips; pass a multi-channel clip "
|
| 203 |
+
"inside a list.")
|
| 204 |
+
n_placeholders = sum(t.count(self.audio_token) for t in text)
|
| 205 |
+
if n_placeholders != len(audio):
|
| 206 |
+
raise ValueError(
|
| 207 |
+
f"text contains {n_placeholders} '{self.audio_token}' placeholders but "
|
| 208 |
+
f"{len(audio)} audio clips were passed; they must match one-to-one (in order).")
|
| 209 |
+
sr = self.feature_extractor.sampling_rate
|
| 210 |
+
# each clip runs through the feature extractor ALONE: its mel values and
|
| 211 |
+
# frame count never depend on which other clips share the batch
|
| 212 |
+
feats = [
|
| 213 |
+
np.asarray(
|
| 214 |
+
self.feature_extractor(
|
| 215 |
+
a,
|
| 216 |
+
sampling_rate=sr,
|
| 217 |
+
return_attention_mask=True,
|
| 218 |
+
padding="longest",
|
| 219 |
+
truncation=False,
|
| 220 |
+
)["input_features"]
|
| 221 |
+
)[0]
|
| 222 |
+
for a in audio
|
| 223 |
+
]
|
| 224 |
+
mel_lens = [f.shape[-1] for f in feats]
|
| 225 |
+
t_max = max(mel_lens)
|
| 226 |
+
input_features = np.zeros((len(feats), feats[0].shape[0], t_max), dtype=np.float32)
|
| 227 |
+
feature_attention_mask = np.zeros((len(feats), t_max), dtype=np.int32)
|
| 228 |
+
for i, (f, n) in enumerate(zip(feats, mel_lens)):
|
| 229 |
+
input_features[i, :, :n] = f
|
| 230 |
+
feature_attention_mask[i, :n] = 1
|
| 231 |
+
audio_inputs["input_features"] = input_features
|
| 232 |
+
audio_inputs["feature_attention_mask"] = feature_attention_mask
|
| 233 |
+
n_tokens = [int(_audio_output_lengths(int(n))) for n in mel_lens]
|
| 234 |
+
repl = iter(
|
| 235 |
+
f"{self.boa_token}{self.audio_token * n}{self.eoa_token}" for n in n_tokens
|
| 236 |
+
)
|
| 237 |
+
pattern = re.escape(self.audio_token)
|
| 238 |
+
text = [re.sub(pattern, lambda _m: next(repl), t) for t in text]
|
| 239 |
+
else:
|
| 240 |
+
n_placeholders = sum(t.count(self.audio_token) for t in text)
|
| 241 |
+
if n_placeholders:
|
| 242 |
+
raise ValueError(
|
| 243 |
+
f"text contains {n_placeholders} '{self.audio_token}' placeholders but no "
|
| 244 |
+
"audio was passed; pass `audio=` with one clip per placeholder.")
|
| 245 |
+
|
| 246 |
+
kwargs.setdefault("add_special_tokens", False)
|
| 247 |
+
text_inputs = self.tokenizer(text=text, return_tensors=return_tensors, **kwargs)
|
| 248 |
+
data = {**text_inputs, **audio_inputs}
|
| 249 |
+
# PreTrainedTokenizerFast emits `token_type_ids` by default. This model has no
|
| 250 |
+
# segment embeddings, and `generate()` rejects unknown kwargs outright, so
|
| 251 |
+
# `model.generate(**processor(...))` would raise. Drop it here so the processor's
|
| 252 |
+
# output is always directly usable.
|
| 253 |
+
data.pop("token_type_ids", None)
|
| 254 |
+
|
| 255 |
+
if return_mm_token_type_ids:
|
| 256 |
+
# 1 wherever the audio embeddings will be scattered, 0 elsewhere. Mirrors
|
| 257 |
+
# `input_ids` row for row (lists stay lists, so ragged batches work).
|
| 258 |
+
import torch
|
| 259 |
+
ids = data["input_ids"]
|
| 260 |
+
if isinstance(ids, torch.Tensor):
|
| 261 |
+
data["mm_token_type_ids"] = (ids == self.audio_token_id).long()
|
| 262 |
+
elif isinstance(ids, np.ndarray):
|
| 263 |
+
data["mm_token_type_ids"] = (ids == self.audio_token_id).astype(np.int64)
|
| 264 |
+
else:
|
| 265 |
+
data["mm_token_type_ids"] = [
|
| 266 |
+
[1 if t == self.audio_token_id else 0 for t in row] for row in ids
|
| 267 |
+
]
|
| 268 |
+
|
| 269 |
+
return BatchFeature(data=data, tensor_type=return_tensors)
|
| 270 |
+
|
| 271 |
+
def _get_num_multimodal_tokens(self, audio_lengths=None, **kwargs):
|
| 272 |
+
"""Placeholder count per clip, without running the feature extractor.
|
| 273 |
+
|
| 274 |
+
`audio_lengths` are MEL-frame counts (100/sec). Used by serving backends to
|
| 275 |
+
size the placeholder span before scheduling.
|
| 276 |
+
"""
|
| 277 |
+
if audio_lengths is None or len(audio_lengths) == 0:
|
| 278 |
+
return {"num_audio_tokens": []}
|
| 279 |
+
return {"num_audio_tokens": [int(_audio_output_lengths(int(n))) for n in audio_lengths]}
|
| 280 |
+
|
| 281 |
+
# ---- tokenizer passthroughs --------------------------------------------------
|
| 282 |
+
def batch_decode(self, *args, **kwargs):
|
| 283 |
+
return self.tokenizer.batch_decode(*args, **kwargs)
|
| 284 |
+
|
| 285 |
+
def decode(self, *args, **kwargs):
|
| 286 |
+
return self.tokenizer.decode(*args, **kwargs)
|
| 287 |
+
|
| 288 |
+
# ---- reading the model's output ----------------------------------------------
|
| 289 |
+
def post_process_transcription(self, generated, *, input_lengths=None, strip_tag=True,
|
| 290 |
+
skip_special_tokens=True):
|
| 291 |
+
"""Decode generations and split off the leading `[...]` decision tag.
|
| 292 |
+
|
| 293 |
+
`generated` is what `model.generate(...)` returned (ids, incl. the prompt — a
|
| 294 |
+
batch or a single row) or already-decoded text. Pass `input_lengths` to trim
|
| 295 |
+
the prompt first: either the padded prompt width as a single int, or the true
|
| 296 |
+
prompt length per row (leading pad tokens are then skipped automatically, so
|
| 297 |
+
per-row lengths stay correct under left padding).
|
| 298 |
+
|
| 299 |
+
Returns a list of `Transcription(text, language, refused, tag, thinking,
|
| 300 |
+
truncated)`; see postprocess.py. `.truncated` is True when a row hit the
|
| 301 |
+
generation cap (no end-of-turn token in the generated ids -- raise
|
| 302 |
+
`max_new_tokens` for long-form audio). With `strip_tag=False` the tag is
|
| 303 |
+
left in `.text`.
|
| 304 |
+
"""
|
| 305 |
+
try:
|
| 306 |
+
from .postprocess import (split_text, BRACKET_OPEN_ID, BRACKET_CLOSE_ID,
|
| 307 |
+
LANGUAGE_TAG_IDS, THINK_OPEN_ID, THINK_CLOSE_ID)
|
| 308 |
+
except ImportError:
|
| 309 |
+
if __package__:
|
| 310 |
+
raise # Preserve dependency errors raised inside sibling modules.
|
| 311 |
+
from postprocess import (split_text, BRACKET_OPEN_ID, BRACKET_CLOSE_ID,
|
| 312 |
+
LANGUAGE_TAG_IDS, THINK_OPEN_ID, THINK_CLOSE_ID)
|
| 313 |
+
|
| 314 |
+
if isinstance(generated, str):
|
| 315 |
+
return [split_text(generated, strip=strip_tag)]
|
| 316 |
+
first = generated[0] if len(generated) else None
|
| 317 |
+
if isinstance(first, str):
|
| 318 |
+
return [split_text(t, strip=strip_tag) for t in generated]
|
| 319 |
+
if getattr(generated, "ndim", None) == 1 or (
|
| 320 |
+
first is not None and not hasattr(first, "__iter__")):
|
| 321 |
+
generated = [generated] # a single row of ids -> batch of one
|
| 322 |
+
|
| 323 |
+
pad_id = self.tokenizer.pad_token_id
|
| 324 |
+
bos_id = self.tokenizer.bos_token_id
|
| 325 |
+
eos_ids = {self.tokenizer.eos_token_id,
|
| 326 |
+
self.tokenizer.convert_tokens_to_ids("<|ifm|endoftext|>")}
|
| 327 |
+
eos_ids.discard(None)
|
| 328 |
+
rows = []
|
| 329 |
+
for i, seq in enumerate(generated):
|
| 330 |
+
seq = [int(t) for t in seq]
|
| 331 |
+
prefilled = None
|
| 332 |
+
thinking_open = False
|
| 333 |
+
if input_lengths is not None:
|
| 334 |
+
if isinstance(input_lengths, int):
|
| 335 |
+
n = input_lengths # padded prompt width (pads included)
|
| 336 |
+
else:
|
| 337 |
+
n = int(input_lengths[i]) # true per-row prompt length: skip the
|
| 338 |
+
if pad_id is not None: # leading pad run first (left padding)
|
| 339 |
+
lead = 0
|
| 340 |
+
while lead < len(seq) and seq[lead] == pad_id:
|
| 341 |
+
lead += 1
|
| 342 |
+
if lead and pad_id == bos_id:
|
| 343 |
+
lead -= 1 # that last one is the prompt's own BOS
|
| 344 |
+
n += lead
|
| 345 |
+
if (n >= 3 and seq[n - 1] == BRACKET_CLOSE_ID
|
| 346 |
+
and seq[n - 3] == BRACKET_OPEN_ID
|
| 347 |
+
and seq[n - 2] in LANGUAGE_TAG_IDS):
|
| 348 |
+
prefilled = LANGUAGE_TAG_IDS[seq[n - 2]] # tag prefilled into the prompt
|
| 349 |
+
for token_id in reversed(seq[:n]):
|
| 350 |
+
if token_id == THINK_CLOSE_ID:
|
| 351 |
+
break
|
| 352 |
+
if token_id == THINK_OPEN_ID:
|
| 353 |
+
thinking_open = True
|
| 354 |
+
break
|
| 355 |
+
seq = seq[n:]
|
| 356 |
+
if seq and seq[0] == THINK_OPEN_ID:
|
| 357 |
+
thinking_open = True
|
| 358 |
+
pending_thinking = thinking_open and THINK_CLOSE_ID not in seq
|
| 359 |
+
rows.append((self.tokenizer.decode(seq, skip_special_tokens=skip_special_tokens),
|
| 360 |
+
not any(t in eos_ids for t in seq), prefilled, pending_thinking))
|
| 361 |
+
import dataclasses
|
| 362 |
+
out = []
|
| 363 |
+
for t, cut, pre, pending_thinking in rows:
|
| 364 |
+
# A capped generation still inside <ifm|think> is reasoning, not a
|
| 365 |
+
# transcript. The prompt establishes the open block when supplied.
|
| 366 |
+
r = split_text(t + "</ifm|think>" if pending_thinking else t, strip=strip_tag)
|
| 367 |
+
if r.tag is None and pre is not None:
|
| 368 |
+
r = dataclasses.replace(r, tag=f"[{pre}]", language=pre)
|
| 369 |
+
out.append(dataclasses.replace(r, truncated=cut))
|
| 370 |
+
return out
|
| 371 |
+
|
| 372 |
+
def language_pin(self, language, force_bracket=True):
|
| 373 |
+
"""A `LanguagePin` for this repo: constrain the language tag while keeping the
|
| 374 |
+
model's own transcribe-or-abstain decision. `language` is a tag, an ISO code
|
| 375 |
+
covering several tags (`"ar"`), or a list of tags.
|
| 376 |
+
|
| 377 |
+
pin = processor.language_pin("ar")
|
| 378 |
+
out = model.generate(**inputs,
|
| 379 |
+
logits_processor=[pin.for_transformers(inputs["input_ids"].shape[1])])
|
| 380 |
+
"""
|
| 381 |
+
try:
|
| 382 |
+
from .language import LanguagePin
|
| 383 |
+
except ImportError:
|
| 384 |
+
if __package__:
|
| 385 |
+
raise # Preserve dependency errors raised inside sibling modules.
|
| 386 |
+
from language import LanguagePin
|
| 387 |
+
return LanguagePin(language, force_bracket=force_bracket)
|
| 388 |
+
|
| 389 |
+
|
| 390 |
+
__all__ = ["JaisASRProcessor"]
|
processor_config.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"audio_token": "<|ifm|audio|>",
|
| 3 |
+
"auto_map": {
|
| 4 |
+
"AutoProcessor": "processing_jais_asr.JaisASRProcessor"
|
| 5 |
+
},
|
| 6 |
+
"boa_token": "<audio>",
|
| 7 |
+
"eoa_token": "</audio>",
|
| 8 |
+
"processor_class": "JaisASRProcessor"
|
| 9 |
+
}
|
special_tokens_map.json
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token": {
|
| 3 |
+
"content": "<|ifm|begin_of_text|>",
|
| 4 |
+
"lstrip": false,
|
| 5 |
+
"normalized": false,
|
| 6 |
+
"rstrip": false,
|
| 7 |
+
"single_word": false
|
| 8 |
+
},
|
| 9 |
+
"eos_token": {
|
| 10 |
+
"content": "<|ifm|im_end|>",
|
| 11 |
+
"lstrip": false,
|
| 12 |
+
"normalized": false,
|
| 13 |
+
"rstrip": false,
|
| 14 |
+
"single_word": false
|
| 15 |
+
},
|
| 16 |
+
"pad_token": {
|
| 17 |
+
"content": "<|ifm|begin_of_text|>",
|
| 18 |
+
"lstrip": false,
|
| 19 |
+
"normalized": false,
|
| 20 |
+
"rstrip": false,
|
| 21 |
+
"single_word": false
|
| 22 |
+
}
|
| 23 |
+
}
|
tokenizer.json
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7e0f10516f96686b52ae4831cd83d302fe4f9224342f937c8292d3db09903634
|
| 3 |
+
size 20637248
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"auto_map": {
|
| 3 |
+
"AutoProcessor": "processing_jais_asr.JaisASRProcessor"
|
| 4 |
+
},
|
| 5 |
+
"bos_token": "<|ifm|begin_of_text|>",
|
| 6 |
+
"clean_up_tokenization_spaces": false,
|
| 7 |
+
"eos_token": "<|ifm|im_end|>",
|
| 8 |
+
"model_max_length": 32768,
|
| 9 |
+
"pad_token": "<|ifm|begin_of_text|>",
|
| 10 |
+
"padding_side": "left",
|
| 11 |
+
"processor_class": "JaisASRProcessor",
|
| 12 |
+
"tokenizer_class": "PreTrainedTokenizerFast",
|
| 13 |
+
"chat_template": "{%- if language is defined and language is not none -%}\n {{- raise_exception(\"language is a generation-time pin. Pass it to model.generate(), the pipeline, or the API; use language_prefill here to insert a tag.\") -}}\n{%- endif -%}\n{#- ===================================================================================\n Jais-ASR — audio transcription chat template.\n\n Optional flags (keyword arguments to apply_chat_template, or\n chat_template_kwargs on a server):\n diacritics=True -> fully diacritized transcript\n code_switch=True -> keep foreign-language spans verbatim\n language_prefill=\"EG\" -> prefill the language tag (see the model card)\n An explicit system message always wins over the mode flags.\n=================================================================================== -#}\n{{- bos_token }}\n{#- ---- mode: the system message content is the switch -------------------------- -#}\n{%- set _m = namespace(modes='') -%}\n{%- if code_switch is defined and code_switch -%}\n {%- set _m.modes = '[Code-switch]' -%}\n{%- endif -%}\n{%- if diacritics is defined and diacritics -%}\n {%- set _m.modes = (_m.modes ~ ' [Diacritics]') if _m.modes else '[Diacritics]' -%}\n{%- endif -%}\n{%- set _has_system = messages and messages[0].role == 'system' -%}\n{%- if _has_system -%}\n {%- if messages[0].content is string -%}\n {%- set system_content = messages[0].content -%}\n {%- elif messages[0].content -%}\n {%- set _sys = namespace(text='') -%}\n {%- for part in messages[0].content -%}\n {%- if part.type == 'text' -%}\n {%- set _sys.text = _sys.text + part.text -%}\n {%- endif -%}\n {%- endfor -%}\n {%- set system_content = _sys.text -%}\n {%- else -%}\n {%- set system_content = '' -%}\n {%- endif -%}\n{%- else -%}\n {%- set system_content = _m.modes -%}\n{%- endif -%}\n{%- if system_content -%}\n {{- '<|ifm|im_start|>system\\n' + system_content + '<|ifm|im_end|>' }}\n{%- endif -%}\n{#- ---- turns ------------------------------------------------------------------ -#}\n{%- for message in messages -%}\n {%- if message.content is string -%}\n {%- set content = message.content -%}\n {%- elif not message.content -%}\n {%- set content = '' -%}\n {%- else -%}\n {#- content parts: {\"type\": \"audio\"} renders the audio placeholder,\n which the processor expands -#}\n {%- set _ns = namespace(text='') -%}\n {%- for part in message.content -%}\n {%- if part.type == 'audio' -%}\n {%- set _ns.text = _ns.text + '<|ifm|audio|>' -%}\n {%- elif part.type == 'text' -%}\n {%- set _ns.text = _ns.text + part.text -%}\n {%- else -%}\n {{- raise_exception(\"audio-only model: unsupported content part type '\" ~ part.type ~ \"' (use 'audio' or 'text')\") -}}\n {%- endif -%}\n {%- endfor -%}\n {%- set content = _ns.text -%}\n {%- endif -%}\n {%- if message.role == 'system' -%}\n {%- if not loop.first -%}\n {{- '<|ifm|im_start|>system\\n' + content + '<|ifm|im_end|>' }}\n {%- endif -%}\n {%- elif message.role == 'user' -%}\n {{- '<|ifm|im_start|>user\\n' + content + '<|ifm|im_end|>' }}\n {%- elif message.role == 'assistant' -%}\n {%- set thinking = message.reasoning if message.reasoning is defined and message.reasoning is string else '' -%}\n {{- '<|ifm|im_start|>assistant' }}\n {%- generation -%}\n {{- '\\n<ifm|think>\\n' + thinking + '</ifm|think>' + content }}\n {{- '<|ifm|im_end|>' -}}\n {%- endgeneration -%}\n {%- else -%}\n {{- raise_exception(\"unsupported role '\" ~ message.role ~ \"' (use system, user or assistant)\") -}}\n {%- endif -%}\n{%- endfor -%}\n{#- ---- generation prompt ------------------------------------------------------- -#}\n{%- if add_generation_prompt -%}\n {%- if enable_thinking is defined and enable_thinking -%}\n {%- if language_prefill is defined and language_prefill -%}\n {{- raise_exception(\"language_prefill= prefill is not supported with enable_thinking=True: the prefill cannot follow an open think block. Drop language_prefill= or disable thinking.\") -}}\n {%- endif -%}\n {{- '<|ifm|im_start|>assistant\\n<ifm|think>\\n' }}\n {%- else -%}\n {{- '<|ifm|im_start|>assistant\\n<ifm|think>\\n</ifm|think>' }}\n {#- language_prefill= prefills the decision tag into the prompt: a transcript is then\n guaranteed and [IFM-NoTranscript] cannot be produced. Unset = the model\n decides. To pin a language while keeping abstention, use language.LanguagePin. -#}\n {%- if language_prefill is defined and language_prefill -%}\n {%- set _tags = ['msa','en','fr','es','it','EG','SA','MA','IQ','JO','SY','QA','DZ','LB','OM','YE','TN','MR','KW','SD','LY','AE','PS','BH'] -%}\n {%- if language_prefill not in _tags -%}\n {{- raise_exception(\"language_prefill must be one of \" ~ (_tags | join(', ')) ~ \" (got '\" ~ language_prefill ~ \"'). An ISO code such as 'ar' covers several tags and cannot be prefilled; use the LanguagePin logits processor for that.\") -}}\n {%- endif -%}\n {{- '[' ~ language_prefill ~ ']' }}\n {%- endif -%}\n {%- endif -%}\n{%- endif -%}\n",
|
| 14 |
+
"model_input_names": [
|
| 15 |
+
"input_ids",
|
| 16 |
+
"attention_mask"
|
| 17 |
+
]
|
| 18 |
+
}
|