Automatic Speech Recognition
Transformers
qwen3-asr
latent-reasoning
test-time-compute
parameter-efficient
Instructions to use voidful/latentASR with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use voidful/latentASR with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("automatic-speech-recognition", model="voidful/latentASR")# pip install -U transformers accelerate # Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("voidful/latentASR", device_map="auto") - Notebooks
- Google Colab
- Kaggle
Release latentASR adapter, code, docs, and reproducibility artifacts
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +2 -33
- .gitignore +23 -0
- LICENSE +18 -0
- README.md +126 -0
- checkpoints/README.md +28 -0
- checkpoints/latentASR_adapter.pth +3 -0
- config.py +195 -0
- data.py +258 -0
- docs/HUGGINGFACE.md +38 -0
- docs/METHOD.md +77 -0
- docs/REPRODUCIBILITY.md +107 -0
- docs/RESULTS.md +54 -0
- docs/paper_latentasr_activation500_updated.tex +453 -0
- eval.py +2337 -0
- examples/transcribe_file.py +85 -0
- experiments/README.md +28 -0
- experiments/eval_paper_tbd_checkpoints.sh +158 -0
- experiments/lib/common.sh +35 -0
- experiments/run_all_training_modes.sh +128 -0
- experiments/run_english_clean_and_noise.sh +169 -0
- experiments/run_multilingual_streaming.sh +152 -0
- experiments/run_paper_tbd_retrain.sh +112 -0
- experiments/run_threshold_sweep.sh +145 -0
- hf_upload/README.md +11 -0
- hf_upload/upload_to_hf.py +54 -0
- losses.py +66 -0
- model.py +1972 -0
- model_metadata.json +19 -0
- peft_utils.py +94 -0
- pyproject.toml +36 -0
- requirements.txt +12 -0
- results/paper500_ascend_20260518/ascend_clean.json +75 -0
- results/paper500_english_clean_noise_20260518/fleurs_en_us_clean.json +75 -0
- results/paper500_english_clean_noise_20260518/fleurs_en_us_snr0db.json +75 -0
- results/paper500_english_clean_noise_20260518/minds14_en_us_clean.json +75 -0
- results/paper500_english_clean_noise_20260518/minds14_en_us_snr0db.json +75 -0
- results/paper500_english_clean_noise_20260518/showcase_report.md +33 -0
- results/paper500_english_clean_noise_20260518/voxpopuli_en_clean.json +75 -0
- results/paper500_english_clean_noise_20260518/voxpopuli_en_snr0db.json +75 -0
- results/paper500_main_eval_20260518/main_fleurs_en_us.json +75 -0
- results/paper500_main_eval_20260518/main_voxpopuli_en.json +75 -0
- results/paper500_main_eval_20260518/paper500_complete_results.md +176 -0
- results/paper500_more_testing_stress_20260525/ami_ihm_sdm_snr0_500.json +106 -0
- results/paper500_more_testing_stress_20260525/common_voice_en_snr0_500.json +75 -0
- results/paper500_more_testing_stress_20260525/extreme_asr_pony_500.json +75 -0
- results/paper500_more_testing_stress_20260525/gigaspeech_test_snr0_500.json +75 -0
- results/paper500_more_testing_stress_20260525/librispeech_clean_other_snr0_500.json +106 -0
- results/paper500_more_testing_stress_20260525/peoples_speech_test_snr0_500.json +75 -0
- results/paper500_more_testing_stress_20260525/showcase_report.md +35 -0
- results/paper500_more_testing_stress_20260525/tedlium_release1_snr0_500.json +75 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,4 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
-
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.
|
| 25 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 3 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
.gitignore
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__pycache__/
|
| 2 |
+
*.py[cod]
|
| 3 |
+
.pytest_cache/
|
| 4 |
+
.ruff_cache/
|
| 5 |
+
.mypy_cache/
|
| 6 |
+
.venv/
|
| 7 |
+
venv/
|
| 8 |
+
env/
|
| 9 |
+
logs/
|
| 10 |
+
eval_runs/
|
| 11 |
+
wandb/
|
| 12 |
+
*.log
|
| 13 |
+
core
|
| 14 |
+
*.tmp
|
| 15 |
+
|
| 16 |
+
# Keep the released adapter checkpoint, ignore accidental large local artifacts.
|
| 17 |
+
*.bin
|
| 18 |
+
*.safetensors
|
| 19 |
+
*.ckpt
|
| 20 |
+
*.pt
|
| 21 |
+
*.pth
|
| 22 |
+
!checkpoints/latentASR_adapter.pth
|
| 23 |
+
|
LICENSE
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Apache License
|
| 2 |
+
Version 2.0, January 2004
|
| 3 |
+
http://www.apache.org/licenses/
|
| 4 |
+
|
| 5 |
+
Copyright 2026 LatentASR contributors
|
| 6 |
+
|
| 7 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 8 |
+
you may not use this file except in compliance with the License.
|
| 9 |
+
You may obtain a copy of the License at
|
| 10 |
+
|
| 11 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 12 |
+
|
| 13 |
+
Unless required by applicable law or agreed to in writing, software
|
| 14 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 15 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 16 |
+
See the License for the specific language governing permissions and
|
| 17 |
+
limitations under the License.
|
| 18 |
+
|
README.md
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
library_name: transformers
|
| 4 |
+
base_model: Qwen/Qwen3-ASR-0.6B
|
| 5 |
+
tags:
|
| 6 |
+
- automatic-speech-recognition
|
| 7 |
+
- qwen3-asr
|
| 8 |
+
- latent-reasoning
|
| 9 |
+
- test-time-compute
|
| 10 |
+
- parameter-efficient
|
| 11 |
+
datasets:
|
| 12 |
+
- google/fleurs
|
| 13 |
+
- facebook/voxpopuli
|
| 14 |
+
metrics:
|
| 15 |
+
- wer
|
| 16 |
+
- cer
|
| 17 |
+
---
|
| 18 |
+
|
| 19 |
+
# latentASR
|
| 20 |
+
|
| 21 |
+
`latentASR` is the release repo for **Listen, Think, Transcribe: Continuous
|
| 22 |
+
Latent Test-Time Scaling for ASR**. The method adds a small trainable latent
|
| 23 |
+
adapter and value head on top of a frozen `Qwen/Qwen3-ASR-0.6B` backbone.
|
| 24 |
+
|
| 25 |
+
The released checkpoint is an adapter-only artifact. It does not redistribute
|
| 26 |
+
the Qwen3-ASR base weights.
|
| 27 |
+
|
| 28 |
+
## What Is Included
|
| 29 |
+
|
| 30 |
+
- Core model/training/evaluation code: `model.py`, `train.py`, `eval.py`.
|
| 31 |
+
- Paper-aligned adapter checkpoint: `checkpoints/latentASR_adapter.pth`.
|
| 32 |
+
- Experiment runners: `experiments/`.
|
| 33 |
+
- Analysis and summarization scripts: `scripts/`.
|
| 34 |
+
- Reproducibility artifacts: `results/`.
|
| 35 |
+
- Documentation: `docs/`.
|
| 36 |
+
|
| 37 |
+
## Method Summary
|
| 38 |
+
|
| 39 |
+
LatentASR inserts `N=4` latent prefix positions before transcript generation.
|
| 40 |
+
At each latent step, a frozen decoder hidden state is projected into a bounded
|
| 41 |
+
delta, gated, and added around a fixed latent-token embedding. A value head
|
| 42 |
+
predicts whether further latent computation is useful and halts dynamically.
|
| 43 |
+
|
| 44 |
+
The frozen ASR backbone is never updated in the LatentASR setting. The adapter
|
| 45 |
+
has 5,251,077 trainable parameters.
|
| 46 |
+
|
| 47 |
+
## Install
|
| 48 |
+
|
| 49 |
+
```bash
|
| 50 |
+
git clone https://huggingface.co/voidful/latentASR
|
| 51 |
+
cd latentASR
|
| 52 |
+
python -m pip install -r requirements.txt
|
| 53 |
+
```
|
| 54 |
+
|
| 55 |
+
The code expects the `qwen-asr` package, which provides
|
| 56 |
+
`Qwen3ASRModel.from_pretrained`.
|
| 57 |
+
|
| 58 |
+
## Quick Evaluation
|
| 59 |
+
|
| 60 |
+
Evaluate the released adapter on FLEURS English:
|
| 61 |
+
|
| 62 |
+
```bash
|
| 63 |
+
python eval.py \
|
| 64 |
+
--model-id Qwen/Qwen3-ASR-0.6B \
|
| 65 |
+
--dataset-name google/fleurs \
|
| 66 |
+
--configs en_us \
|
| 67 |
+
--split test \
|
| 68 |
+
--latent-ckpt checkpoints/latentASR_adapter.pth \
|
| 69 |
+
--dynamic-halt-threshold 0.0 \
|
| 70 |
+
--skip-baseline-ft \
|
| 71 |
+
--skip-prompt-tuning \
|
| 72 |
+
--skip-lora-r16 \
|
| 73 |
+
--output-json eval_runs/fleurs_en_us.json
|
| 74 |
+
```
|
| 75 |
+
|
| 76 |
+
Use `--max-samples-per-config 100` for a quick smoke test.
|
| 77 |
+
|
| 78 |
+
## Single-File Transcription
|
| 79 |
+
|
| 80 |
+
```bash
|
| 81 |
+
python examples/transcribe_file.py path/to/audio.wav \
|
| 82 |
+
--latent-ckpt checkpoints/latentASR_adapter.pth \
|
| 83 |
+
--theta 0.0
|
| 84 |
+
```
|
| 85 |
+
|
| 86 |
+
## Training
|
| 87 |
+
|
| 88 |
+
The paper-aligned minimal-data activation setting uses:
|
| 89 |
+
|
| 90 |
+
- `N_LATENT=4`
|
| 91 |
+
- `TRAIN_MAX_SAMPLES=500`
|
| 92 |
+
- effective batch size 16
|
| 93 |
+
- 10 epochs
|
| 94 |
+
- value forced-negative probability `0.3`
|
| 95 |
+
- latent input noise `0.5`
|
| 96 |
+
- latent-loop dropout `0.15`
|
| 97 |
+
|
| 98 |
+
Run the maintained sequential baseline/adapter trainer:
|
| 99 |
+
|
| 100 |
+
```bash
|
| 101 |
+
./experiments/run_all_training_modes.sh
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
Run the paper ablation checkpoint retraining suite:
|
| 105 |
+
|
| 106 |
+
```bash
|
| 107 |
+
./experiments/run_paper_tbd_retrain.sh
|
| 108 |
+
```
|
| 109 |
+
|
| 110 |
+
## Result Scope
|
| 111 |
+
|
| 112 |
+
The clean benchmark deltas are single-run point estimates. They are reported as
|
| 113 |
+
directional, not statistically established, and should be confirmed with paired
|
| 114 |
+
multi-seed testing before making strong claims.
|
| 115 |
+
|
| 116 |
+
ASCEND is included only as an overlap-confounded in-domain diagnostic because
|
| 117 |
+
ASCEND contributes 5/500 activation utterances and speaker-disjointness was not
|
| 118 |
+
enforced.
|
| 119 |
+
|
| 120 |
+
## Important Files
|
| 121 |
+
|
| 122 |
+
- `docs/METHOD.md`: architecture and loss details.
|
| 123 |
+
- `docs/REPRODUCIBILITY.md`: exact commands and evaluation workflow.
|
| 124 |
+
- `docs/RESULTS.md`: headline tables and caveats.
|
| 125 |
+
- `model_metadata.json`: released checkpoint metadata.
|
| 126 |
+
|
checkpoints/README.md
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Checkpoints
|
| 2 |
+
|
| 3 |
+
This directory contains the released LatentASR adapter checkpoint:
|
| 4 |
+
|
| 5 |
+
- `latentASR_adapter.pth`
|
| 6 |
+
|
| 7 |
+
The checkpoint stores only the lightweight LatentASR modules:
|
| 8 |
+
|
| 9 |
+
- `init_proj`
|
| 10 |
+
- `delta_proj`
|
| 11 |
+
- `step_proj`
|
| 12 |
+
- `step_embed`
|
| 13 |
+
- `log_scale`
|
| 14 |
+
- `value_head`
|
| 15 |
+
- `injection_gate`
|
| 16 |
+
|
| 17 |
+
It does not include the frozen `Qwen/Qwen3-ASR-0.6B` backbone. Evaluation and
|
| 18 |
+
inference load the base model from Hugging Face and then attach this adapter.
|
| 19 |
+
|
| 20 |
+
Checkpoint metadata:
|
| 21 |
+
|
| 22 |
+
- Source run: `activation_500_epoch10.pth`
|
| 23 |
+
- Base model: `Qwen/Qwen3-ASR-0.6B`
|
| 24 |
+
- Latent budget: `N=4`
|
| 25 |
+
- Training set size: 500 utterances
|
| 26 |
+
- Adapter parameters: 5,251,077
|
| 27 |
+
- SHA256: `f0ce39fa5e6952fced6992508f3e2b32ea8467442b545678781a0f04e64f2430`
|
| 28 |
+
|
checkpoints/latentASR_adapter.pth
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f0ce39fa5e6952fced6992508f3e2b32ea8467442b545678781a0f04e64f2430
|
| 3 |
+
size 10507457
|
config.py
ADDED
|
@@ -0,0 +1,195 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Training configuration dataclass for lr_whisper.
|
| 2 |
+
|
| 3 |
+
All hyperparameters are read from environment variables via ``TrainingConfig.from_env()``.
|
| 4 |
+
Use ``get_config()`` to obtain a process-wide singleton (lazy, never called at import time).
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import os
|
| 8 |
+
from dataclasses import dataclass, field
|
| 9 |
+
from typing import List, Optional
|
| 10 |
+
|
| 11 |
+
from utils import env_flag, normalize_train_mode
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
@dataclass
|
| 15 |
+
class TrainingConfig:
|
| 16 |
+
# Model / dataset
|
| 17 |
+
model_id: str = "Qwen/Qwen3-ASR-0.6B"
|
| 18 |
+
dataset_name: str = "openslr/librispeech_asr"
|
| 19 |
+
dataset_config: str = "clean"
|
| 20 |
+
|
| 21 |
+
# Training batch / epoch
|
| 22 |
+
batch_size: int = 4
|
| 23 |
+
grad_accum_steps: int = 4
|
| 24 |
+
num_epochs: int = 10
|
| 25 |
+
eval_samples: int = 1000
|
| 26 |
+
pretrain_eval_samples: int = 10
|
| 27 |
+
|
| 28 |
+
# Latent tokens
|
| 29 |
+
n_latent: int = 4
|
| 30 |
+
is_deq_train: bool = False
|
| 31 |
+
deq_train_iters: int = 15
|
| 32 |
+
halt_threshold: float = 0.0
|
| 33 |
+
latent_drop_prob: float = 0.15
|
| 34 |
+
latent_input_noise_std: float = 0.5
|
| 35 |
+
value_forced_neg_prob: float = 0.3
|
| 36 |
+
latent_use_bounded_delta: bool = True
|
| 37 |
+
latent_use_injection_gate: bool = True
|
| 38 |
+
latent_use_embedding_anchor: bool = True
|
| 39 |
+
train_max_samples: int = 0
|
| 40 |
+
checkpoint_prefix: str = ""
|
| 41 |
+
|
| 42 |
+
# Thought layout
|
| 43 |
+
thought_mode: str = "prefix" # "prefix" | "interleaved"
|
| 44 |
+
thought_group_size: int = 1 # words per NT token in interleaved mode
|
| 45 |
+
|
| 46 |
+
# Training mode
|
| 47 |
+
train_mode: str = "latent"
|
| 48 |
+
|
| 49 |
+
# Learning rates
|
| 50 |
+
lr_adapter: float = 1e-4
|
| 51 |
+
|
| 52 |
+
lr_scale: float = 5e-5
|
| 53 |
+
lr_baseline_ft: float = 1e-5
|
| 54 |
+
lr_prompt_tuning: float = 5e-4
|
| 55 |
+
lr_lora_r16: float = 1e-4
|
| 56 |
+
|
| 57 |
+
# Prompt text
|
| 58 |
+
user_prompt_text: str = "Transcribe the audio into text."
|
| 59 |
+
|
| 60 |
+
# Prompt-tuning hyperparameters
|
| 61 |
+
prompt_tuning_init_mode: str = "text"
|
| 62 |
+
prompt_tuning_num_virtual_tokens: int = 4
|
| 63 |
+
prompt_tuning_init_text: str = "Transcribe the audio into text."
|
| 64 |
+
|
| 65 |
+
# LoRA hyperparameters
|
| 66 |
+
lora_rank: int = 16
|
| 67 |
+
lora_alpha: int = 32
|
| 68 |
+
lora_dropout: float = 0.05
|
| 69 |
+
lora_target_modules: List[str] = field(
|
| 70 |
+
default_factory=lambda: [
|
| 71 |
+
"q_proj", "k_proj", "v_proj", "o_proj",
|
| 72 |
+
"gate_proj", "up_proj", "down_proj",
|
| 73 |
+
]
|
| 74 |
+
)
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
# Delta stabilisation / scale
|
| 78 |
+
delta_tanh_c: float = 5.0
|
| 79 |
+
scale_max: float = 3.0
|
| 80 |
+
scale_init: float = 0.2
|
| 81 |
+
|
| 82 |
+
# Regularisation weights
|
| 83 |
+
w_smooth: float = 0.0 # Disabled — replaced by w_cycle
|
| 84 |
+
w_cycle: float = 0.1 # Cycle consistency loss (vocabulary-independent) - main regularizer
|
| 85 |
+
w_value: float = 3.0
|
| 86 |
+
|
| 87 |
+
# Logging intervals
|
| 88 |
+
log_every: int = 50
|
| 89 |
+
grad_log_every: int = 50
|
| 90 |
+
|
| 91 |
+
# Derived flags (set in __post_init__)
|
| 92 |
+
use_latent_reasoning: bool = False
|
| 93 |
+
use_prompt_tuning: bool = False
|
| 94 |
+
use_lora_r16: bool = False
|
| 95 |
+
use_peft_mode: bool = False
|
| 96 |
+
|
| 97 |
+
def __post_init__(self) -> None:
|
| 98 |
+
valid_modes = {"latent", "baseline", "prompt_tuning", "lora_r16"}
|
| 99 |
+
if self.train_mode not in valid_modes:
|
| 100 |
+
raise ValueError(
|
| 101 |
+
f"Unsupported train_mode={self.train_mode!r}. "
|
| 102 |
+
f"Use one of: {sorted(valid_modes)}"
|
| 103 |
+
)
|
| 104 |
+
valid_thought_modes = {"prefix", "interleaved"}
|
| 105 |
+
if self.thought_mode not in valid_thought_modes:
|
| 106 |
+
raise ValueError(
|
| 107 |
+
f"Unsupported thought_mode={self.thought_mode!r}. "
|
| 108 |
+
f"Use one of: {sorted(valid_thought_modes)}"
|
| 109 |
+
)
|
| 110 |
+
if self.thought_group_size < 1:
|
| 111 |
+
raise ValueError(f"thought_group_size must be >= 1, got {self.thought_group_size}")
|
| 112 |
+
self.use_latent_reasoning = self.train_mode == "latent"
|
| 113 |
+
self.use_prompt_tuning = self.train_mode == "prompt_tuning"
|
| 114 |
+
self.use_lora_r16 = self.train_mode == "lora_r16"
|
| 115 |
+
self.use_peft_mode = self.use_lora_r16
|
| 116 |
+
|
| 117 |
+
@classmethod
|
| 118 |
+
def from_env(cls) -> "TrainingConfig":
|
| 119 |
+
"""Construct a ``TrainingConfig`` by reading environment variables."""
|
| 120 |
+
lora_modules_raw = os.getenv(
|
| 121 |
+
"LORA_TARGET_MODULES",
|
| 122 |
+
"q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
|
| 123 |
+
)
|
| 124 |
+
lora_target_modules = [x.strip() for x in lora_modules_raw.split(",") if x.strip()]
|
| 125 |
+
|
| 126 |
+
train_mode = normalize_train_mode(os.getenv("TRAIN_MODE", "latent"))
|
| 127 |
+
default_dataset = "openslr/librispeech_asr"
|
| 128 |
+
default_config = "clean"
|
| 129 |
+
if train_mode == "latent":
|
| 130 |
+
default_dataset = "SpeechTest/extreme_asr_pony"
|
| 131 |
+
default_config = "default"
|
| 132 |
+
|
| 133 |
+
return cls(
|
| 134 |
+
model_id=os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-0.6B"),
|
| 135 |
+
dataset_name=os.getenv("DATASET_NAME", default_dataset),
|
| 136 |
+
dataset_config=os.getenv("DATASET_CONFIG", default_config),
|
| 137 |
+
batch_size=int(os.getenv("BATCH_SIZE", "4")),
|
| 138 |
+
grad_accum_steps=int(os.getenv("GRAD_ACCUM_STEPS", "4")),
|
| 139 |
+
num_epochs=int(os.getenv("NUM_EPOCHS", "10")),
|
| 140 |
+
eval_samples=int(os.getenv("EVAL_SAMPLES", "1000")),
|
| 141 |
+
pretrain_eval_samples=int(os.getenv("PRETRAIN_EVAL_SAMPLES", "10")),
|
| 142 |
+
n_latent=int(os.getenv("N_LATENT", "4")),
|
| 143 |
+
is_deq_train=os.getenv("IS_DEQ_TRAIN", "true").strip().lower() == "true",
|
| 144 |
+
deq_train_iters=int(os.getenv("DEQ_TRAIN_ITERS", "15")),
|
| 145 |
+
halt_threshold=float(os.getenv("HALT_THRESHOLD", "0.0")),
|
| 146 |
+
latent_drop_prob=float(os.getenv("LATENT_DROP_PROB", "0.15")),
|
| 147 |
+
latent_input_noise_std=float(os.getenv("LATENT_INPUT_NOISE_STD", "0.5")),
|
| 148 |
+
value_forced_neg_prob=float(os.getenv("VALUE_FORCED_NEG_PROB", "0.3")),
|
| 149 |
+
latent_use_bounded_delta=env_flag("LATENT_USE_BOUNDED_DELTA", True),
|
| 150 |
+
latent_use_injection_gate=env_flag("LATENT_USE_INJECTION_GATE", True),
|
| 151 |
+
latent_use_embedding_anchor=env_flag("LATENT_USE_EMBEDDING_ANCHOR", True),
|
| 152 |
+
train_max_samples=int(os.getenv("TRAIN_MAX_SAMPLES", "0")),
|
| 153 |
+
checkpoint_prefix=os.getenv("CHECKPOINT_PREFIX", "").strip(),
|
| 154 |
+
thought_mode=os.getenv("THOUGHT_MODE", "prefix").strip().lower(),
|
| 155 |
+
thought_group_size=int(os.getenv("THOUGHT_GROUP_SIZE", "1")),
|
| 156 |
+
train_mode=train_mode,
|
| 157 |
+
lr_adapter=float(os.getenv("LR_ADAPTER", "1e-4")),
|
| 158 |
+
lr_scale=float(os.getenv("LR_SCALE", "5e-5")),
|
| 159 |
+
lr_baseline_ft=float(os.getenv("LR_BASELINE_FT", "1e-5")),
|
| 160 |
+
lr_prompt_tuning=float(os.getenv("LR_PROMPT_TUNING", "5e-4")),
|
| 161 |
+
lr_lora_r16=float(os.getenv("LR_LORA_R16", "1e-4")),
|
| 162 |
+
user_prompt_text=os.getenv(
|
| 163 |
+
"USER_PROMPT_TEXT", "Transcribe the audio into text."
|
| 164 |
+
).strip(),
|
| 165 |
+
prompt_tuning_init_mode=os.getenv("PROMPT_TUNING_INIT_MODE", "text").strip().lower(),
|
| 166 |
+
prompt_tuning_num_virtual_tokens=int(
|
| 167 |
+
os.getenv("PROMPT_TUNING_NUM_VIRTUAL_TOKENS", "4")
|
| 168 |
+
),
|
| 169 |
+
prompt_tuning_init_text=os.getenv(
|
| 170 |
+
"PROMPT_TUNING_INIT_TEXT", "Transcribe the audio into text."
|
| 171 |
+
).strip(),
|
| 172 |
+
lora_rank=int(os.getenv("LORA_RANK", "16")),
|
| 173 |
+
lora_alpha=int(os.getenv("LORA_ALPHA", "32")),
|
| 174 |
+
lora_dropout=float(os.getenv("LORA_DROPOUT", "0.05")),
|
| 175 |
+
lora_target_modules=lora_target_modules,
|
| 176 |
+
delta_tanh_c=float(os.getenv("DELTA_TANH_C", "5.0")),
|
| 177 |
+
scale_max=float(os.getenv("SCALE_MAX", "3.0")),
|
| 178 |
+
scale_init=float(os.getenv("SCALE_INIT", "0.2")),
|
| 179 |
+
w_smooth=float(os.getenv("W_SMOOTH", "0.0")),
|
| 180 |
+
w_cycle=float(os.getenv("W_CYCLE", "0.1")),
|
| 181 |
+
w_value=float(os.getenv("W_VALUE", "3.0")),
|
| 182 |
+
log_every=int(os.getenv("LOG_EVERY", "50")),
|
| 183 |
+
grad_log_every=int(os.getenv("GRAD_LOG_EVERY", "50")),
|
| 184 |
+
)
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
_config: Optional[TrainingConfig] = None
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def get_config() -> TrainingConfig:
|
| 191 |
+
"""Return the process-wide ``TrainingConfig`` singleton (lazy init from env)."""
|
| 192 |
+
global _config
|
| 193 |
+
if _config is None:
|
| 194 |
+
_config = TrainingConfig.from_env()
|
| 195 |
+
return _config
|
data.py
ADDED
|
@@ -0,0 +1,258 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Dataset preparation and collation for lr_whisper."""
|
| 2 |
+
|
| 3 |
+
import numpy as np
|
| 4 |
+
from typing import Any, Dict, List, Optional
|
| 5 |
+
|
| 6 |
+
import torch
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
# ---------------------------------------------------------------------------
|
| 10 |
+
# Audio token sequence helpers
|
| 11 |
+
# ---------------------------------------------------------------------------
|
| 12 |
+
|
| 13 |
+
def build_audio_token_seq(
|
| 14 |
+
audio_bos_id: int,
|
| 15 |
+
audio_token_id: int,
|
| 16 |
+
audio_eos_id: int,
|
| 17 |
+
l_audio: int,
|
| 18 |
+
device: torch.device,
|
| 19 |
+
) -> torch.Tensor:
|
| 20 |
+
"""Build ``[<|audio_start|>, <audio_pad>×L, <|audio_end|>]`` token tensor.
|
| 21 |
+
|
| 22 |
+
Args:
|
| 23 |
+
audio_bos_id: Token id for ``<|audio_start|>``.
|
| 24 |
+
audio_token_id: Token id for the audio placeholder (repeated L times).
|
| 25 |
+
audio_eos_id: Token id for ``<|audio_end|>``.
|
| 26 |
+
l_audio: Number of audio placeholder tokens.
|
| 27 |
+
device: Target device.
|
| 28 |
+
|
| 29 |
+
Returns:
|
| 30 |
+
1-D LongTensor of length ``l_audio + 2``.
|
| 31 |
+
"""
|
| 32 |
+
audio_bos = torch.tensor([audio_bos_id], dtype=torch.long, device=device)
|
| 33 |
+
audio_eos = torch.tensor([audio_eos_id], dtype=torch.long, device=device)
|
| 34 |
+
audio_pads = torch.full((l_audio,), audio_token_id, dtype=torch.long, device=device)
|
| 35 |
+
return torch.cat([audio_bos, audio_pads, audio_eos])
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _coerce_feat_attention_mask(
|
| 39 |
+
mask: Any,
|
| 40 |
+
target_len: int,
|
| 41 |
+
) -> torch.Tensor:
|
| 42 |
+
"""Coerce a feature-attention mask to a LongTensor of length *target_len*.
|
| 43 |
+
|
| 44 |
+
Pads with zeros or truncates as needed.
|
| 45 |
+
|
| 46 |
+
Args:
|
| 47 |
+
mask: A list, numpy array, or Tensor representing the mask.
|
| 48 |
+
target_len: Required length.
|
| 49 |
+
|
| 50 |
+
Returns:
|
| 51 |
+
1-D LongTensor of length ``target_len``.
|
| 52 |
+
"""
|
| 53 |
+
if not isinstance(mask, torch.Tensor):
|
| 54 |
+
mask = torch.tensor(mask, dtype=torch.long)
|
| 55 |
+
mask = mask.long()
|
| 56 |
+
cur = mask.size(-1)
|
| 57 |
+
if cur < target_len:
|
| 58 |
+
mask = torch.cat([mask, torch.zeros(target_len - cur, dtype=torch.long)])
|
| 59 |
+
elif cur > target_len:
|
| 60 |
+
mask = mask[:target_len]
|
| 61 |
+
return mask
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
# ---------------------------------------------------------------------------
|
| 65 |
+
# Dataset preparation
|
| 66 |
+
# ---------------------------------------------------------------------------
|
| 67 |
+
|
| 68 |
+
def prepare_dataset(
|
| 69 |
+
processor: Any,
|
| 70 |
+
lang_id: int,
|
| 71 |
+
transcribe_id: int,
|
| 72 |
+
nt_id: Optional[int],
|
| 73 |
+
n_latent: int,
|
| 74 |
+
batch: Dict[str, Any],
|
| 75 |
+
thought_mode: str = "prefix",
|
| 76 |
+
thought_group_size: int = 1,
|
| 77 |
+
language: str = "English",
|
| 78 |
+
) -> Dict[str, Any]:
|
| 79 |
+
"""Prepare a single sample for training.
|
| 80 |
+
|
| 81 |
+
This helper converts raw audio and text into model inputs and labels. The
|
| 82 |
+
audio waveform is passed through the processor's feature extractor and
|
| 83 |
+
encoded into log-Mel features. The transcript text is tokenised without
|
| 84 |
+
special tokens.
|
| 85 |
+
|
| 86 |
+
Labels use Qwen Chat Template format to be compatible with the pre-trained
|
| 87 |
+
model. Two layout modes are supported:
|
| 88 |
+
|
| 89 |
+
Prefix mode (default)::
|
| 90 |
+
|
| 91 |
+
<|im_start|>assistant\\n[NT*N]transcription<|im_end|>
|
| 92 |
+
|
| 93 |
+
Interleaved mode (``thought_mode="interleaved"``)::
|
| 94 |
+
|
| 95 |
+
<|im_start|>assistant\\n[NT][word1][NT][word2]...<|im_end|>
|
| 96 |
+
|
| 97 |
+
In interleaved mode, one NT token is placed before each group of
|
| 98 |
+
``thought_group_size`` word tokens. NT positions are masked in the model's
|
| 99 |
+
labels so no CE loss is computed on them.
|
| 100 |
+
|
| 101 |
+
The user turn with audio is handled as input in the forward method.
|
| 102 |
+
|
| 103 |
+
Args:
|
| 104 |
+
processor: The Qwen3-ASR processor returned by ``asr_wrapper.processor``.
|
| 105 |
+
lang_id: Integer token id for ``<|im_start|>``.
|
| 106 |
+
transcribe_id: Unused (kept for API compatibility).
|
| 107 |
+
nt_id: Integer token id for the latent "NT" token.
|
| 108 |
+
n_latent: Number of latent tokens to prepend (prefix mode only).
|
| 109 |
+
batch: A dictionary with keys ``audio`` and ``text`` as provided by
|
| 110 |
+
``load_dataset``.
|
| 111 |
+
thought_mode: ``"prefix"`` (default) or ``"interleaved"``.
|
| 112 |
+
thought_group_size: Words per NT token in interleaved mode (default 1).
|
| 113 |
+
|
| 114 |
+
Returns:
|
| 115 |
+
A dictionary containing the input features, decoder labels and
|
| 116 |
+
reference text.
|
| 117 |
+
"""
|
| 118 |
+
audio = batch["audio"]
|
| 119 |
+
feat_out = processor.feature_extractor(
|
| 120 |
+
audio["array"], sampling_rate=audio["sampling_rate"],
|
| 121 |
+
return_attention_mask=True,
|
| 122 |
+
)
|
| 123 |
+
feats = feat_out.input_features[0]
|
| 124 |
+
if hasattr(feat_out, "attention_mask") and feat_out.attention_mask is not None:
|
| 125 |
+
feat_attention_mask = feat_out.attention_mask[0]
|
| 126 |
+
else:
|
| 127 |
+
feat_attention_mask = [1] * len(feats[0]) if hasattr(feats, "__len__") else None
|
| 128 |
+
|
| 129 |
+
raw_text = batch["text"].strip()
|
| 130 |
+
if raw_text:
|
| 131 |
+
fmt_text = raw_text.capitalize()
|
| 132 |
+
# Add a period if it doesn't have one to match base model grammar
|
| 133 |
+
if not fmt_text.endswith((".", "!", "?")):
|
| 134 |
+
fmt_text += "."
|
| 135 |
+
else:
|
| 136 |
+
fmt_text = ""
|
| 137 |
+
|
| 138 |
+
text_ids = processor.tokenizer(fmt_text, add_special_tokens=False).input_ids
|
| 139 |
+
|
| 140 |
+
im_start = lang_id
|
| 141 |
+
_im_end_token = "<|im_end|>"
|
| 142 |
+
im_end = processor.tokenizer.convert_tokens_to_ids(_im_end_token)
|
| 143 |
+
if im_end is None or im_end == processor.tokenizer.unk_token_id:
|
| 144 |
+
im_end = processor.tokenizer.eos_token_id
|
| 145 |
+
|
| 146 |
+
# Original: asst_nl = processor.tokenizer.encode("assistant\n", add_special_tokens=False)
|
| 147 |
+
# if not asst_nl:
|
| 148 |
+
# raise ValueError("Tokenizer returned empty ids for 'assistant\\n'.")
|
| 149 |
+
|
| 150 |
+
# New construction for asst_nl to include language formatting
|
| 151 |
+
asst = processor.tokenizer.encode("assistant", add_special_tokens=False)
|
| 152 |
+
nl = processor.tokenizer.encode("\n", add_special_tokens=False)
|
| 153 |
+
if not asst or not nl:
|
| 154 |
+
raise ValueError("Tokenizer returned empty ids for 'assistant' or '\\n'.")
|
| 155 |
+
asr_prefix_ids = processor.tokenizer.encode(f"language {language}<asr_text>", add_special_tokens=False)
|
| 156 |
+
asst_nl_prefix = asst + nl + asr_prefix_ids
|
| 157 |
+
|
| 158 |
+
# For labels, we mask the prefix part including `asst_nl`
|
| 159 |
+
# Let's verify `nt_id` logic.
|
| 160 |
+
interleaved = (thought_mode == "interleaved")
|
| 161 |
+
if interleaved:
|
| 162 |
+
# Interleave: [NT, tok1, NT, tok2, ...] with one NT per thought_group_size words.
|
| 163 |
+
if nt_id is None or nt_id < 0:
|
| 164 |
+
raise ValueError("interleaved thought_mode requires a valid nt_id.")
|
| 165 |
+
group = max(1, int(thought_group_size))
|
| 166 |
+
# Build interleaved sequence
|
| 167 |
+
num_toks = len(text_ids)
|
| 168 |
+
num_groups = (num_toks + group - 1) // group
|
| 169 |
+
actual_nt = min(n_latent, num_groups) # Use n_latent as max NTs
|
| 170 |
+
|
| 171 |
+
interleaved_body: List[int] = []
|
| 172 |
+
for i, tok in enumerate(text_ids):
|
| 173 |
+
if i % group == 0 and (i // group) < actual_nt:
|
| 174 |
+
interleaved_body.append(nt_id)
|
| 175 |
+
interleaved_body.append(tok)
|
| 176 |
+
body = interleaved_body
|
| 177 |
+
else:
|
| 178 |
+
# Prefix mode: [NT*N, tok1, tok2, ...]
|
| 179 |
+
latent_tokens: List[int] = []
|
| 180 |
+
if n_latent > 0:
|
| 181 |
+
if nt_id is None or nt_id < 0:
|
| 182 |
+
raise ValueError("n_latent > 0 requires a valid nt_id.")
|
| 183 |
+
latent_tokens = [nt_id] * int(n_latent)
|
| 184 |
+
body = latent_tokens + text_ids
|
| 185 |
+
|
| 186 |
+
labels: List[int] = [im_start] + asst_nl_prefix + body + [im_end]
|
| 187 |
+
if labels[-1] != im_end:
|
| 188 |
+
raise ValueError("Label sequence must end with <|im_end|>.")
|
| 189 |
+
|
| 190 |
+
result: Dict[str, Any] = {
|
| 191 |
+
"input_features": feats,
|
| 192 |
+
"labels": labels,
|
| 193 |
+
"reference_text": fmt_text, # evaluate on formatted text to allow cleaner debug
|
| 194 |
+
}
|
| 195 |
+
if feat_attention_mask is not None:
|
| 196 |
+
result["feature_attention_mask"] = feat_attention_mask
|
| 197 |
+
return result
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
# ---------------------------------------------------------------------------
|
| 201 |
+
# Collation
|
| 202 |
+
# ---------------------------------------------------------------------------
|
| 203 |
+
|
| 204 |
+
class DataCollatorQwenASR:
|
| 205 |
+
"""Collate function for Qwen-ASR latent finetuning.
|
| 206 |
+
|
| 207 |
+
This collator pads the list of input feature arrays into a batch tensor and
|
| 208 |
+
pads the decoder label sequences. Padding positions in the labels are
|
| 209 |
+
replaced with ``-100`` so that they are ignored by the cross-entropy loss.
|
| 210 |
+
|
| 211 |
+
Args:
|
| 212 |
+
processor: The Qwen3-ASR processor used for tokenisation and feature
|
| 213 |
+
extraction.
|
| 214 |
+
"""
|
| 215 |
+
|
| 216 |
+
def __init__(self, processor: Any) -> None:
|
| 217 |
+
self.processor = processor
|
| 218 |
+
|
| 219 |
+
def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
|
| 220 |
+
# Pad input features
|
| 221 |
+
input_features = [{"input_features": f["input_features"]} for f in features]
|
| 222 |
+
batch_inputs = self.processor.feature_extractor.pad(
|
| 223 |
+
input_features, return_tensors="pt", return_attention_mask=True
|
| 224 |
+
)
|
| 225 |
+
|
| 226 |
+
# Build feature_attention_mask from per-sample masks saved in
|
| 227 |
+
# prepare_dataset. The collator's pad() creates an attention_mask
|
| 228 |
+
# but it's ALL 1s when every sample is pre-padded to 3000 frames.
|
| 229 |
+
if "feature_attention_mask" in features[0]:
|
| 230 |
+
max_len = batch_inputs["input_features"].size(-1)
|
| 231 |
+
masks = []
|
| 232 |
+
for f in features:
|
| 233 |
+
m = f["feature_attention_mask"]
|
| 234 |
+
if hasattr(m, "tolist"):
|
| 235 |
+
m = m if isinstance(m, list) else m.tolist()
|
| 236 |
+
m = list(m) + [0] * (max_len - len(m))
|
| 237 |
+
masks.append(m[:max_len])
|
| 238 |
+
feature_attention_mask = torch.tensor(masks, dtype=torch.long)
|
| 239 |
+
else:
|
| 240 |
+
feature_attention_mask = batch_inputs["attention_mask"]
|
| 241 |
+
|
| 242 |
+
# Pad labels using the tokenizer
|
| 243 |
+
label_features = [{"input_ids": f["labels"]} for f in features]
|
| 244 |
+
labels_batch = self.processor.tokenizer.pad(
|
| 245 |
+
label_features,
|
| 246 |
+
padding=True,
|
| 247 |
+
return_tensors="pt",
|
| 248 |
+
return_attention_mask=True,
|
| 249 |
+
)
|
| 250 |
+
labels = labels_batch["input_ids"]
|
| 251 |
+
attn = labels_batch["attention_mask"]
|
| 252 |
+
labels = labels.masked_fill(attn.eq(0), -100)
|
| 253 |
+
|
| 254 |
+
return {
|
| 255 |
+
"input_features": batch_inputs["input_features"],
|
| 256 |
+
"feature_attention_mask": feature_attention_mask,
|
| 257 |
+
"labels": labels,
|
| 258 |
+
}
|
docs/HUGGINGFACE.md
ADDED
|
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Hugging Face Release
|
| 2 |
+
|
| 3 |
+
The public model repo is intended to be:
|
| 4 |
+
|
| 5 |
+
```text
|
| 6 |
+
voidful/latentASR
|
| 7 |
+
```
|
| 8 |
+
|
| 9 |
+
The repository hosts:
|
| 10 |
+
|
| 11 |
+
- LatentASR code
|
| 12 |
+
- adapter checkpoint
|
| 13 |
+
- documentation
|
| 14 |
+
- model card
|
| 15 |
+
- reproducibility outputs
|
| 16 |
+
|
| 17 |
+
## Upload
|
| 18 |
+
|
| 19 |
+
From this project root:
|
| 20 |
+
|
| 21 |
+
```bash
|
| 22 |
+
python hf_upload/upload_to_hf.py --repo-id voidful/latentASR
|
| 23 |
+
```
|
| 24 |
+
|
| 25 |
+
The script creates the model repo if needed and uploads the current folder.
|
| 26 |
+
|
| 27 |
+
## Download Checkpoint Programmatically
|
| 28 |
+
|
| 29 |
+
```python
|
| 30 |
+
from huggingface_hub import hf_hub_download
|
| 31 |
+
|
| 32 |
+
ckpt = hf_hub_download(
|
| 33 |
+
repo_id="voidful/latentASR",
|
| 34 |
+
filename="checkpoints/latentASR_adapter.pth",
|
| 35 |
+
)
|
| 36 |
+
print(ckpt)
|
| 37 |
+
```
|
| 38 |
+
|
docs/METHOD.md
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Method
|
| 2 |
+
|
| 3 |
+
LatentASR adds continuous latent test-time scaling to a frozen ASR backbone.
|
| 4 |
+
The implementation wraps `Qwen/Qwen3-ASR-0.6B` with two lightweight modules:
|
| 5 |
+
|
| 6 |
+
- **Latent Adapter**: converts decoder hidden states into bounded latent
|
| 7 |
+
deltas.
|
| 8 |
+
- **Value Head**: predicts whether additional latent compute is useful and
|
| 9 |
+
enables dynamic halting.
|
| 10 |
+
|
| 11 |
+
## Prefix Layout
|
| 12 |
+
|
| 13 |
+
The default layout is prefix mode:
|
| 14 |
+
|
| 15 |
+
```text
|
| 16 |
+
system/user/audio prompt
|
| 17 |
+
assistant language <asr_text>
|
| 18 |
+
<|latent|> <|latent|> <|latent|> <|latent|>
|
| 19 |
+
transcript tokens
|
| 20 |
+
```
|
| 21 |
+
|
| 22 |
+
Transcript loss masks all prompt and latent positions.
|
| 23 |
+
|
| 24 |
+
## Stable Injection
|
| 25 |
+
|
| 26 |
+
Each latent delta is stabilized by three mechanisms:
|
| 27 |
+
|
| 28 |
+
1. **Bounded delta**: `delta_proj(h_k)` is L2-normalized and scaled by a
|
| 29 |
+
learned per-step scalar.
|
| 30 |
+
2. **Sigmoid gate**: a zero-initialized gate starts at `0.5` and learns how
|
| 31 |
+
much of the delta to apply.
|
| 32 |
+
3. **Fixed embedding anchor**: the injected vector is
|
| 33 |
+
`embedding(<|latent|>) + gate * delta`, keeping the input near a real token
|
| 34 |
+
embedding.
|
| 35 |
+
|
| 36 |
+
These mechanisms correspond to the ablation flags:
|
| 37 |
+
|
| 38 |
+
```bash
|
| 39 |
+
LATENT_USE_BOUNDED_DELTA=1
|
| 40 |
+
LATENT_USE_INJECTION_GATE=1
|
| 41 |
+
LATENT_USE_EMBEDDING_ANCHOR=1
|
| 42 |
+
```
|
| 43 |
+
|
| 44 |
+
## Value Head Target
|
| 45 |
+
|
| 46 |
+
For each training utterance, the value head predicts the latent-vs-baseline
|
| 47 |
+
accuracy gain. The default target is:
|
| 48 |
+
|
| 49 |
+
```text
|
| 50 |
+
0.9 * tanh(3.0 * (latent_accuracy - baseline_accuracy))
|
| 51 |
+
```
|
| 52 |
+
|
| 53 |
+
If both baseline and latent accuracies are zero for an utterance, the target
|
| 54 |
+
uses a clamped CE-difference fallback:
|
| 55 |
+
|
| 56 |
+
```text
|
| 57 |
+
0.9 * tanh(0.5 * clamp(baseline_ce - latent_ce, -2, 2))
|
| 58 |
+
```
|
| 59 |
+
|
| 60 |
+
With probability `VALUE_FORCED_NEG_PROB=0.3` per minibatch, the target is
|
| 61 |
+
flipped to `-|target|` for conservative calibration.
|
| 62 |
+
|
| 63 |
+
The current implementation supervises the value prediction at every produced
|
| 64 |
+
latent state, matching the fact that inference can halt after any step.
|
| 65 |
+
|
| 66 |
+
## Dynamic Halting
|
| 67 |
+
|
| 68 |
+
At inference time:
|
| 69 |
+
|
| 70 |
+
1. The initial value prediction gates the whole loop.
|
| 71 |
+
2. If `v_0 < theta`, all latent tokens are removed and generation falls back to
|
| 72 |
+
the frozen baseline.
|
| 73 |
+
3. After each latent step, the value head is re-evaluated.
|
| 74 |
+
4. If `v_k < theta`, unused latent tokens are removed before decoding.
|
| 75 |
+
|
| 76 |
+
The paper uses `theta=0.0` for the deployed setting.
|
| 77 |
+
|
docs/REPRODUCIBILITY.md
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Reproducibility
|
| 2 |
+
|
| 3 |
+
This document describes how to run the released code and adapter checkpoint.
|
| 4 |
+
|
| 5 |
+
## Environment
|
| 6 |
+
|
| 7 |
+
```bash
|
| 8 |
+
python -m pip install -r requirements.txt
|
| 9 |
+
```
|
| 10 |
+
|
| 11 |
+
Recommended runtime:
|
| 12 |
+
|
| 13 |
+
- Python 3.10+
|
| 14 |
+
- CUDA GPU with bfloat16 support for practical speed
|
| 15 |
+
- `qwen-asr>=0.0.6`
|
| 16 |
+
|
| 17 |
+
## Released Adapter
|
| 18 |
+
|
| 19 |
+
The released checkpoint is:
|
| 20 |
+
|
| 21 |
+
```text
|
| 22 |
+
checkpoints/latentASR_adapter.pth
|
| 23 |
+
```
|
| 24 |
+
|
| 25 |
+
It was copied from:
|
| 26 |
+
|
| 27 |
+
```text
|
| 28 |
+
activation_500_epoch10.pth
|
| 29 |
+
```
|
| 30 |
+
|
| 31 |
+
Metadata:
|
| 32 |
+
|
| 33 |
+
- Base model: `Qwen/Qwen3-ASR-0.6B`
|
| 34 |
+
- Latent budget: `N=4`
|
| 35 |
+
- Activation size: 500 utterances
|
| 36 |
+
- Adapter parameters: 5,251,077
|
| 37 |
+
- SHA256: `f0ce39fa5e6952fced6992508f3e2b32ea8467442b545678781a0f04e64f2430`
|
| 38 |
+
|
| 39 |
+
## Main Evaluation
|
| 40 |
+
|
| 41 |
+
FLEURS English:
|
| 42 |
+
|
| 43 |
+
```bash
|
| 44 |
+
python eval.py \
|
| 45 |
+
--model-id Qwen/Qwen3-ASR-0.6B \
|
| 46 |
+
--dataset-name google/fleurs \
|
| 47 |
+
--configs en_us \
|
| 48 |
+
--split test \
|
| 49 |
+
--latent-ckpt checkpoints/latentASR_adapter.pth \
|
| 50 |
+
--dynamic-halt-threshold 0.0 \
|
| 51 |
+
--max-samples-per-config 0 \
|
| 52 |
+
--skip-baseline-ft \
|
| 53 |
+
--skip-prompt-tuning \
|
| 54 |
+
--skip-lora-r16 \
|
| 55 |
+
--output-json eval_runs/fleurs_en_us.json
|
| 56 |
+
```
|
| 57 |
+
|
| 58 |
+
VoxPopuli English:
|
| 59 |
+
|
| 60 |
+
```bash
|
| 61 |
+
python eval.py \
|
| 62 |
+
--model-id Qwen/Qwen3-ASR-0.6B \
|
| 63 |
+
--dataset-name facebook/voxpopuli \
|
| 64 |
+
--configs en \
|
| 65 |
+
--split test \
|
| 66 |
+
--latent-ckpt checkpoints/latentASR_adapter.pth \
|
| 67 |
+
--dynamic-halt-threshold 0.0 \
|
| 68 |
+
--max-samples-per-config 0 \
|
| 69 |
+
--skip-baseline-ft \
|
| 70 |
+
--skip-prompt-tuning \
|
| 71 |
+
--skip-lora-r16 \
|
| 72 |
+
--output-json eval_runs/voxpopuli_en.json
|
| 73 |
+
```
|
| 74 |
+
|
| 75 |
+
## Threshold Sweep
|
| 76 |
+
|
| 77 |
+
```bash
|
| 78 |
+
LATENT_CKPT=checkpoints/latentASR_adapter.pth \
|
| 79 |
+
./experiments/run_threshold_sweep.sh
|
| 80 |
+
```
|
| 81 |
+
|
| 82 |
+
## Minimal-Data Activation Training
|
| 83 |
+
|
| 84 |
+
```bash
|
| 85 |
+
TRAIN_MODE=latent \
|
| 86 |
+
N_LATENT=4 \
|
| 87 |
+
TRAIN_MAX_SAMPLES=500 \
|
| 88 |
+
BATCH_SIZE=4 \
|
| 89 |
+
GRAD_ACCUM_STEPS=4 \
|
| 90 |
+
NUM_EPOCHS=10 \
|
| 91 |
+
python train.py
|
| 92 |
+
```
|
| 93 |
+
|
| 94 |
+
The sequential comparison runner defaults to the paper-aligned protocol:
|
| 95 |
+
|
| 96 |
+
```bash
|
| 97 |
+
./experiments/run_all_training_modes.sh
|
| 98 |
+
```
|
| 99 |
+
|
| 100 |
+
## Notes
|
| 101 |
+
|
| 102 |
+
- The clean benchmark improvements are small single-run point estimates.
|
| 103 |
+
- The ASCEND diagnostic is overlap-confounded and should not be treated as
|
| 104 |
+
speaker-independent generalization evidence.
|
| 105 |
+
- Results in `results/` are included for transparency and summarization, not as
|
| 106 |
+
a substitute for rerunning evaluation in your environment.
|
| 107 |
+
|
docs/RESULTS.md
ADDED
|
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Results
|
| 2 |
+
|
| 3 |
+
This release includes JSON and Markdown outputs under `results/`. The most
|
| 4 |
+
important caveat is that the clean benchmark gains are single-run directional
|
| 5 |
+
point estimates, not statistically established improvements.
|
| 6 |
+
|
| 7 |
+
## Clean Main Benchmarks
|
| 8 |
+
|
| 9 |
+
| Model | FLEURS WER | VoxPopuli WER |
|
| 10 |
+
| --- | ---: | ---: |
|
| 11 |
+
| Frozen Qwen3-ASR baseline | 4.900 | 9.038 |
|
| 12 |
+
| LatentASR, N=4, theta=0 | 4.776 | 8.995 |
|
| 13 |
+
|
| 14 |
+
Relative WER deltas:
|
| 15 |
+
|
| 16 |
+
- FLEURS: `-2.54%`
|
| 17 |
+
- VoxPopuli: `-0.47%`
|
| 18 |
+
|
| 19 |
+
## Robustness Diagnostics
|
| 20 |
+
|
| 21 |
+
At SNR `0` dB:
|
| 22 |
+
|
| 23 |
+
| Dataset | Baseline WER | LatentASR WER | Relative delta |
|
| 24 |
+
| --- | ---: | ---: | ---: |
|
| 25 |
+
| FLEURS | 29.81 | 29.42 | -1.32% |
|
| 26 |
+
| VoxPopuli | 19.71 | 19.51 | -1.02% |
|
| 27 |
+
|
| 28 |
+
Stress broad suite weighted aggregate:
|
| 29 |
+
|
| 30 |
+
- Baseline WER: `36.843`
|
| 31 |
+
- LatentASR WER: `36.265`
|
| 32 |
+
- Delta: `-0.578` pp (`-1.57%` relative)
|
| 33 |
+
|
| 34 |
+
## ASCEND Diagnostic
|
| 35 |
+
|
| 36 |
+
ASCEND CER:
|
| 37 |
+
|
| 38 |
+
- Baseline: `57.81`
|
| 39 |
+
- LatentASR: `48.55`
|
| 40 |
+
- Relative CER delta: `-16.02%`
|
| 41 |
+
|
| 42 |
+
This is explicitly treated as an overlap-confounded in-domain diagnostic.
|
| 43 |
+
ASCEND contributes 5/500 activation utterances and speaker-disjointness was not
|
| 44 |
+
enforced.
|
| 45 |
+
|
| 46 |
+
## Included Result Folders
|
| 47 |
+
|
| 48 |
+
- `results/paper500_main_eval_20260518/`
|
| 49 |
+
- `results/paper500_english_clean_noise_20260518/`
|
| 50 |
+
- `results/paper500_multilingual_streaming_20260518/`
|
| 51 |
+
- `results/paper500_threshold_20260518/`
|
| 52 |
+
- `results/paper500_more_testing_stress_20260525/`
|
| 53 |
+
- `results/paper500_ascend_20260518/`
|
| 54 |
+
|
docs/paper_latentasr_activation500_updated.tex
ADDED
|
@@ -0,0 +1,453 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
\documentclass[a4paper]{article}
|
| 2 |
+
\usepackage{ISCSLP2026}
|
| 3 |
+
\usepackage{iftex}
|
| 4 |
+
\ifPDFTeX\else
|
| 5 |
+
\usepackage{fontspec}
|
| 6 |
+
\setmainfont{TeX Gyre Termes}
|
| 7 |
+
\setmonofont{TeX Gyre Cursor}
|
| 8 |
+
\fi
|
| 9 |
+
\usepackage{ifthen}
|
| 10 |
+
\newboolean{blind}
|
| 11 |
+
\setboolean{blind}{true}
|
| 12 |
+
|
| 13 |
+
%==================================================================================
|
| 14 |
+
% Packages
|
| 15 |
+
\usepackage{cite}
|
| 16 |
+
\usepackage{amsmath,amssymb,amsfonts}
|
| 17 |
+
\usepackage{graphicx}
|
| 18 |
+
\usepackage{textcomp}
|
| 19 |
+
\usepackage{xcolor}
|
| 20 |
+
\usepackage{booktabs}
|
| 21 |
+
\usepackage{bm}
|
| 22 |
+
\usepackage{multirow}
|
| 23 |
+
\usepackage{tikz}
|
| 24 |
+
\usetikzlibrary{positioning, arrows.meta, fit, backgrounds, calc, decorations.pathreplacing}
|
| 25 |
+
\usepackage{pgfplots}
|
| 26 |
+
\pgfplotsset{compat=1.18}
|
| 27 |
+
\usepackage{comment}
|
| 28 |
+
|
| 29 |
+
%==================================================================================
|
| 30 |
+
% Custom commands
|
| 31 |
+
\newcommand{\method}{LatentASR}
|
| 32 |
+
\newcommand{\vdelta}{\boldsymbol{\delta}}
|
| 33 |
+
|
| 34 |
+
%==================================================================================
|
| 35 |
+
% Title and Authors (ISCSLP-style: name/address/email with blind toggle)
|
| 36 |
+
\title{Listen, Think, Transcribe: Continuous Latent Test-Time Scaling for ASR}
|
| 37 |
+
|
| 38 |
+
\name{
|
| 39 |
+
\ifthenelse{\boolean{blind}}{Anonymous to ISCSLP}
|
| 40 |
+
{Name$^1$, Co-author Name$^2$}
|
| 41 |
+
}
|
| 42 |
+
|
| 43 |
+
\address{
|
| 44 |
+
\ifthenelse{\boolean{blind}}{Anonymous to ISCSLP}
|
| 45 |
+
{
|
| 46 |
+
$^1$Author Affiliation\\
|
| 47 |
+
$^2$Co-author Affiliation
|
| 48 |
+
}
|
| 49 |
+
}
|
| 50 |
+
|
| 51 |
+
\email{
|
| 52 |
+
\ifthenelse{\boolean{blind}}{Anonymous to ISCSLP}
|
| 53 |
+
{author@university.edu, coauthor@company.com}
|
| 54 |
+
}
|
| 55 |
+
|
| 56 |
+
\begin{document}
|
| 57 |
+
|
| 58 |
+
\maketitle
|
| 59 |
+
|
| 60 |
+
\begin{abstract}
|
| 61 |
+
End-to-end ASR models transcribe in a single pass, leaving no room for the decoder to revisit hard inputs. We propose \emph{LatentASR}, a parameter-efficient method that adds continuous latent test-time scaling to a frozen ASR backbone. Two small trainable modules drive it: a Latent Adapter that iteratively refines a few latent prefix positions through bounded, stabilized updates, and a Value Head that predicts whether extra computation will help and halts the loop early. The Qwen3-ASR-0.6B backbone stays fully frozen, and we train only ${\sim}5.25$M extra parameters. We activate this loop with a deliberately small, diverse 500-utterance training set. Under this minimal-data regime, standard adaptation methods regress: full fine-tuning, LoRA, and prompt tuning each increase WER on the two clean main benchmarks. \method{} instead gives the only negative single-run point estimate on both clean benchmarks, lowering FLEURS WER from $4.900\%$ to $4.776\%$ and VoxPopuli WER from $9.038\%$ to $8.995\%$. These clean-set deltas are small and should be read as directional rather than as statistically established gains. Same-protocol diagnostics on multilingual and noisy stress conditions suggest that the latent loop preserves the frozen backbone while converting fixed per-utterance compute into input-dependent compute. Dynamic halting allocates compute sparsely on clean read speech, skipping $47.0\%$ of FLEURS and $63.0\%$ of VoxPopuli utterances at the entry gate. Our results suggest that a small, carefully chosen activation set can switch on input-dependent latent computation inside a frozen ASR model without adapting, and thereby corrupting, the model itself.
|
| 62 |
+
\end{abstract}
|
| 63 |
+
\noindent\textbf{Index Terms}: automatic speech recognition, latent test-time scaling, continuous thought, parameter-efficient finetuning
|
| 64 |
+
|
| 65 |
+
\section{Introduction}
|
| 66 |
+
|
| 67 |
+
End-to-end automatic speech recognition (ASR) models map audio directly to text in a single forward pass~\cite{radford2023robust,peng2025owsm}. This design is simple and effective. However, it forces a single left-to-right decoder to handle acoustic disambiguation, language modeling, and error correction at the same time.
|
| 68 |
+
|
| 69 |
+
Recent work has shown that allocating extra computation at inference time can substantially improve model performance. This paradigm is broadly known as test-time scaling, and includes explicit chain-of-thought reasoning~\cite{wei2022chain}, parallel sampling, and iterative refinement. Within this line, latent test-time scaling has drawn growing attention. For example, Coconut replaces discrete chain-of-thought tokens with continuous hidden states that are iteratively fed back into the model~\cite{hao2024training}. Quiet-STaR trains models to generate and exploit implicit thoughts to improve next-token prediction~\cite{zelikman2024quiet}. Pause tokens provide extra compute steps without requiring explicit intermediate text~\cite{goyal2024think}. Overall, these findings suggest that latent compute scaling is particularly beneficial on challenging inputs.
|
| 70 |
+
|
| 71 |
+
Motivated by this, we ask a simple question. Can an ASR decoder benefit from latent test-time scaling before it commits to a transcript?
|
| 72 |
+
|
| 73 |
+
Applying continuous latent scaling to ASR introduces two challenges absent in standard NLP tasks. First, ASR is a direct transcription task with no intermediate reasoning trajectory to distill, unlike Coconut and Quiet-STaR which rely on explicit chain-of-thought rationales. The model must discover how to use extra latent computation entirely unsupervised. Second, modern ASR backbones are typically kept frozen for efficiency and zero-shot generalization. Standard parameter-efficient methods (LoRA, prompt tuning, prefix tuning) follow this regime~\cite{hu2022lora,lester2021power,li2021prefix}, but recent test-time scaling techniques often fine-tune the entire backbone, which inevitably triggers catastrophic forgetting. Conversely, directly injecting unconstrained continuous states into a frozen backbone causes a geometric mismatch that pushes the decoder out of its native vocabulary distribution and leads to representation collapse.
|
| 74 |
+
|
| 75 |
+
We propose \method{}, a parameter-efficient method that adds two small trainable modules (a Latent Adapter and a Value Head) to a frozen ASR backbone. The Latent Adapter refines the embeddings of $N$ latent prefix positions through bounded recurrent updates; the Value Head predicts per-utterance utility and halts the loop early. Both modules consume only decoder hidden states; the encoder and decoder weights are never updated.
|
| 76 |
+
|
| 77 |
+
Our contributions are as follows:
|
| 78 |
+
\begin{itemize}
|
| 79 |
+
\item We introduce continuous latent test-time scaling into frozen ASR backbones. We propose a soft-gated additive injection strategy that produces bounded updates and constructs scale-invariant compute states without triggering representation collapse. Same-protocol component ablations confirm that the stabilization mechanisms are individually important, with removal costs up to tens of WER points.
|
| 80 |
+
\item We design a Value Head that predicts per-utterance utility and halts computation when extra steps will not help. At the deployed threshold, dynamic halting uses only $1.20$ average latent steps on FLEURS and $0.70$ on VoxPopuli, while skipping $47.0\%$ and $63.0\%$ of utterances respectively. An ASCEND accented/code-switched diagnostic shows heavier compute allocation. ASCEND accounts for only $5/500$ activation utterances ($1.0\%$), but because speaker-disjointness was not enforced, we still treat this row as an overlap-confounded in-domain diagnostic rather than a headline generalization result.
|
| 81 |
+
\item We demonstrate \emph{minimal-data activation}: a small training set can \emph{activate} latent test-time scaling on a frozen backbone, while being too small to \emph{adapt} (and thereby corrupt) the backbone itself. An activation-size sweep from 100 to 800 utterances selects the 500-utterance checkpoint used for the main results. This checkpoint gives negative single-run WER point estimates on FLEURS ($-0.124$~pp) and VoxPopuli ($-0.043$~pp), while standard adaptation methods of comparable scale regress; we treat these clean-set deltas as directional and leave statistical confirmation to multi-seed paired testing.
|
| 82 |
+
\end{itemize}
|
| 83 |
+
|
| 84 |
+
\section{Related Work}
|
| 85 |
+
\label{sec:related}
|
| 86 |
+
|
| 87 |
+
\textbf{Latent test-time scaling.}\quad Beyond explicit chain-of-thought reasoning~\cite{wei2022chain}, recent work injects \emph{continuous} computation into language models. Coconut~\cite{hao2024training} feeds last-layer hidden states back into the model in place of discrete tokens, allowing reasoning to proceed in a continuous embedding space. Quiet-STaR~\cite{zelikman2024quiet} trains models to generate and exploit implicit thoughts that improve next-token prediction without explicit rationale supervision. Pause tokens~\cite{goyal2024think} insert filler positions that grant the model extra forward compute without producing intermediate text. These methods primarily target language modeling and reasoning benchmarks where intermediate rationales either exist or can be distilled from teacher models. \method{} extends the latent-scaling principle to direct speech-to-text transcription, where the target sequence is the only supervision signal and the backbone must remain frozen to preserve broad acoustic generalization.
|
| 88 |
+
|
| 89 |
+
\textbf{Adaptive computation.}\quad Spending variable compute per input has a long history. Adaptive Computation Time (ACT)~\cite{graves2016adaptive} learns a halting probability per step inside an RNN, terminating when accumulated halting mass exceeds a threshold. PonderNet~\cite{banino2021pondernet} reformulates this with a probabilistic framework and explicit regularization on the expected pondering steps. Early-exit transformers~\cite{schuster2022confident,xin2020deebert} cut inference cost by exiting at intermediate layers when token-level confidence is high. \method{}'s Value Head sits in this lineage but differs along two axes. First, it operates over a small fixed budget of \emph{latent prefix} positions rather than over network depth or over output sequence length; second, its halting signal is the predicted \emph{latent-vs-baseline accuracy gain}, not a confidence proxy or a learned ponder cost. The signal is therefore directly aligned with the deployment objective (``would extra compute reduce WER?''), rather than a heuristic surrogate.
|
| 90 |
+
|
| 91 |
+
\textbf{Parameter-efficient adaptation.}\quad LoRA~\cite{hu2022lora}, prompt tuning~\cite{lester2021power}, and prefix tuning~\cite{li2021prefix} keep a backbone frozen while inserting low-rank weight updates or learnable surface prompts. As shown in Table~\ref{tab:main}, under our 500-utterance low-resource regime these gradient-based input or weight perturbations consistently \emph{increase} WER on a strong zero-shot ASR backbone. \method{} differs in two ways. First, it leaves both backbone weights and surface prompts untouched. Second, it does not minimize the input-output discrepancy; it instead introduces additional \emph{computation} along a bounded latent trajectory governed by a learned halting policy.
|
| 92 |
+
|
| 93 |
+
\textbf{Robust pre-trained ASR.}\quad Foundation ASR models such as Whisper~\cite{radford2023robust}, OWSM~\cite{peng2025owsm}, and Qwen-ASR~\cite{yang2024qwen2audio} achieve strong zero-shot transcription via massive multilingual pre-training. Their generalization, however, is brittle: even small downstream gradient updates can erase it (Section~\ref{sec:exp_setup}). To our knowledge, \method{} is the first method to add a learnable test-time refinement loop on top of such a backbone without modifying any of its parameters, yielding negative clean-set point estimates in a regime where conventional adaptation actively regresses.
|
| 94 |
+
|
| 95 |
+
\section{Method}
|
| 96 |
+
|
| 97 |
+
\subsection{Problem Setup}
|
| 98 |
+
\label{sec:setup}
|
| 99 |
+
|
| 100 |
+
We consider an end-to-end ASR pipeline with a frozen acoustic encoder and a frozen autoregressive text decoder. Given an utterance $a$, the encoder produces acoustic states $\mathbf{Z}=\mathrm{Enc}(a)$, and the decoder generates transcript tokens conditioned on $\mathbf{Z}$ and a text prefix. \method{} keeps both the encoder and the decoder fully frozen and intervenes \emph{only} on the decoder's input prefix: it inserts $N$ \emph{latent positions} between the system prompt and the transcript target. These positions are non-text continuous slots used solely as additional computation context; no transcript token is emitted at them. We use $N$ both as the number of latent positions and as the maximum number of refinement steps: refinement step $k$ updates the embedding at latent position $k$.
|
| 101 |
+
|
| 102 |
+
\subsection{Overview}
|
| 103 |
+
|
| 104 |
+
\method{} adds two small trainable modules to the frozen backbone (Fig.~\ref{fig:arch}): a \textbf{Latent Adapter} that refines the embeddings at the $N$ latent positions through a bounded recurrent loop, and a \textbf{Value Head} that decides how many of those refinement steps to actually run. Both modules consume only decoder hidden states; neither modifies the backbone or the encoder. One difficulty drives the whole design. Feeding an arbitrary continuous vector into a frozen decoder pushes its input off the distribution seen during pre-training and quickly destabilizes it. The Latent Adapter therefore does not inject free vectors; it produces small, bounded, gated updates around a fixed anchor embedding, controlled by three stabilization mechanisms that we detail in Section~\ref{sec:stability}. Sections~\ref{sec:refinement}--\ref{sec:valuehead} describe each component in turn.
|
| 105 |
+
|
| 106 |
+
\begin{figure}[t]
|
| 107 |
+
\centering
|
| 108 |
+
\includegraphics[width=\columnwidth]{latentASR.png}
|
| 109 |
+
\caption{Architecture of \method{}. The ASR backbone (blue) remains \textbf{frozen}. We add a trainable \textbf{Latent Adapter} (orange) that forms a recurrent compute loop. At step $k$, it maps the decoder hidden state at latent position $k$ to a bounded update $\vdelta_k$ that is added to the latent embedding via a sigmoid gate. A \textbf{Value Head} (diamond) reads each latent state to predict a utility score $v_k$, enabling the model to halt dynamically when further computation is unnecessary.}
|
| 110 |
+
\label{fig:arch}
|
| 111 |
+
\end{figure}
|
| 112 |
+
|
| 113 |
+
\subsection{Latent Prefix Refinement}
|
| 114 |
+
\label{sec:refinement}
|
| 115 |
+
|
| 116 |
+
The Latent Adapter computes a bounded update at each refinement step. Let $\mathbf{h}_0$ denote the decoder's last-layer hidden state at the final latent position before any refinement (the \emph{acoustic anchor}); we $L_2$-normalize it before downstream use so the loop starts from a unit-magnitude state. At refinement step $k\in\{1,\dots,N\}$, the frozen decoder reads $\mathbf{Z}$ and the current latent prefix; we extract the last-layer hidden state $\mathbf{h}_k$ at latent position $k$. The adapter, time-conditioned by a learned step embedding $\mathbf{e}_k$, projects $\mathbf{h}_k$ into the embedding space, $L_2$-normalizes the projection, and rescales by a learned scalar $s_k$ to produce a step-bounded delta:
|
| 117 |
+
\begin{equation}
|
| 118 |
+
\vdelta_k = s_k \cdot \frac{\text{DeltaProj}(\mathbf{h}_k)}{\lVert\text{DeltaProj}(\mathbf{h}_k)\rVert + \epsilon}, \quad \epsilon{=}10^{-6}.
|
| 119 |
+
\label{eq:delta}
|
| 120 |
+
\end{equation}
|
| 121 |
+
The $L_2$-normalization makes the update direction scale-invariant; the per-step scale $s_k$ controls how far each step moves.
|
| 122 |
+
|
| 123 |
+
\subsection{Stable Injection into a Frozen Decoder}
|
| 124 |
+
\label{sec:injection}
|
| 125 |
+
\label{sec:stability}
|
| 126 |
+
|
| 127 |
+
A frozen decoder only interprets inputs that resemble the token embeddings it saw during pre-training. Injecting an arbitrary continuous vector pushes the input off this embedding manifold and destabilizes the decoder. \method{} keeps the injection on-manifold by constraining three distinct properties of the update: how far each step moves the state, whether the step is applied, and where the resulting embedding sits.
|
| 128 |
+
|
| 129 |
+
\begin{itemize}
|
| 130 |
+
\item \textbf{Bounded delta.} The per-step scale $s_k$ and the $L_2$-normalization in Eq.~\ref{eq:delta} bound the magnitude of every update, so no single step can grow without limit.
|
| 131 |
+
\item \textbf{Sigmoid gate.} A learned gate decides how much of each update to apply:
|
| 132 |
+
\begin{equation}
|
| 133 |
+
g_k = \sigma\big(\mathbf{W}_g [\mathbf{e}_{\texttt{LT}}; \vdelta_k]\big),
|
| 134 |
+
\label{eq:gate}
|
| 135 |
+
\end{equation}
|
| 136 |
+
with $\mathbf{W}_g$ zero-initialized so $g_k{=}0.5$ at the start of training.
|
| 137 |
+
\item \textbf{Fixed-embedding anchor.} Every latent position is anchored to the same fixed token embedding $\mathbf{e}_{\texttt{LT}}$ and modified only by an additive residual,
|
| 138 |
+
\begin{equation}
|
| 139 |
+
\mathbf{x}_k = \mathbf{e}_{\texttt{LT}} + g_k \cdot \vdelta_k,
|
| 140 |
+
\label{eq:inject}
|
| 141 |
+
\end{equation}
|
| 142 |
+
so the injected embedding stays in a small ball around a real vocabulary embedding the decoder already interprets.
|
| 143 |
+
\end{itemize}
|
| 144 |
+
|
| 145 |
+
Together, Eqs.~\ref{eq:delta}--\ref{eq:inject} form a gated residual around a fixed anchor. Closing the gate recovers $\mathbf{e}_{\texttt{LT}}$ exactly, so the adapter only learns a small, optionally-zero correction rather than a full continuous input from scratch.
|
| 146 |
+
|
| 147 |
+
\subsection{Utility Prediction and Dynamic Halting}
|
| 148 |
+
\label{sec:valuehead}
|
| 149 |
+
|
| 150 |
+
Latent scaling reliably helps some utterances and merely costs compute on others. We therefore need a per-utterance gate that predicts, before any compute is spent, whether further latent steps will help. The Value Head plays this role.
|
| 151 |
+
|
| 152 |
+
\textbf{Design rationale.}\quad The signal we want is ``would extra latent compute help this utterance?'', which is most directly answered by comparing the latent path against the frozen baseline on the same training utterance, a self-distilled target whose value is the latent-vs-baseline accuracy gain $\Delta$ on supervised transcript positions. As the latent path improves during training, $\Delta$ can become positive on most utterances; we therefore expose a forced-negative sampling knob $p_{\text{neg}}$ that, with probability $p_{\text{neg}}$ per minibatch, flips $y \leftarrow -|y|$, providing a tunable counterweight against this drift. We default to $p_{\text{neg}}{=}0.3$. When both baseline and latent token accuracies are zero for an utterance, we replace the accuracy target for that utterance with a clamped cross-entropy-difference fallback.
|
| 153 |
+
|
| 154 |
+
\textbf{Implementation.}\quad The Value Head is a single linear layer over each latent hidden state, $v_k = \tanh(\text{ValueHead}(\mathbf{h}_k)) \in [-1, 1]$, so larger $v_k$ predicts a more negative $\Delta$WER. On non-degenerate utterances the target is the tanh-squashed accuracy gain
|
| 155 |
+
\begin{equation}
|
| 156 |
+
y = \alpha \cdot \tanh(\gamma \cdot \Delta),
|
| 157 |
+
\label{eq:value_target}
|
| 158 |
+
\end{equation}
|
| 159 |
+
with $\alpha{=}0.9$ and $\gamma{=}3.0$.
|
| 160 |
+
|
| 161 |
+
\textbf{Two-stage halting at inference.}\quad $v_0$ on the initial anchor gates the entire loop: if $v_0 < \theta$, all $N$ steps are skipped and the utterance falls back to the frozen Baseline. Otherwise, after each step $k\geq 1$, we re-evaluate $v_k$ and break if it drops below $\theta$. Latent compute is thus allocated only to utterances the Value Head predicts will benefit.
|
| 162 |
+
|
| 163 |
+
\subsection{Training}
|
| 164 |
+
\label{sec:training}
|
| 165 |
+
|
| 166 |
+
Three losses train the adapter and Value Head while the backbone stays frozen:
|
| 167 |
+
\begin{equation}
|
| 168 |
+
\mathcal{L} = \mathcal{L}_{\text{CE}} + w_{\text{cyc}}(t) \cdot \mathcal{L}_{\text{cyc}} + w_{\text{val}} \cdot \mathcal{L}_{\text{val}}.
|
| 169 |
+
\label{eq:loss_total}
|
| 170 |
+
\end{equation}
|
| 171 |
+
$\mathcal{L}_{\text{CE}}$ is cross-entropy on the transcript tokens (latent positions are masked out). $\mathcal{L}_{\text{val}}$ is the mean-squared error between the Value Head output at each produced latent state and the target in Eq.~\ref{eq:value_target}. $\mathcal{L}_{\text{cyc}}$ is a trajectory regularizer
|
| 172 |
+
\begin{equation}
|
| 173 |
+
\mathcal{L}_{\text{cyc}} = \tfrac{\beta}{N}\!\sum_{k=1}^{N}\!\lVert\mathbf{h}_k {-} \mathbf{h}_0\rVert_2^2 + \tfrac{1{-}\beta}{N{-}1}\!\sum_{k=1}^{N-1}\!\lVert\mathbf{h}_{k+1} {-} \mathbf{h}_k\rVert_2^2,
|
| 174 |
+
\label{eq:loss_cyc}
|
| 175 |
+
\end{equation}
|
| 176 |
+
that anchors the latent trajectory near $\mathbf{h}_0$ and damps abrupt step-to-step jumps, with $\beta{=}0.3$ favoring smooth refinement over rigid attachment. The cycle weight $w_{\text{cyc}}(t)$ decays linearly to zero. To prevent over-reliance on the latent path, we additionally drop the entire latent loop with probability $p_{\text{drop}}$ during training and add Gaussian noise of std $\sigma$ to the initial anchor.
|
| 177 |
+
|
| 178 |
+
\section{Experiments}
|
| 179 |
+
\subsection{Experiment Setup}
|
| 180 |
+
\label{sec:exp_setup}
|
| 181 |
+
|
| 182 |
+
\textbf{Training Setup.}\quad We use Qwen3-ASR-0.6B~\cite{yang2024qwen2audio} as the base model and keep all its parameters frozen. We train for 10 epochs with effective batch size $16$ on a single NVIDIA RTX 5090. We set $N{=}4$ and use AdamW with decoupled learning rates: $10^{-4}$ for the Latent Adapter and Value Head, and $5 \times 10^{-5}$ for step scales (initialized at $0.2$, capped at $3.0$). We use $w_{\text{cyc}}{=}0.1$ (linearly decayed), $w_{\text{val}}{=}3.0$, $p_{\text{drop}}{=}0.15$, input noise $\sigma{=}0.5$, and $p_{\text{neg}}{=}0.3$. With this configuration, the trainable components add only $\sim 5.25\text{M}$ parameters.
|
| 183 |
+
|
| 184 |
+
\textbf{Dataset: Minimal-Data Activation.}\quad We construct a 500-utterance training mixture from seven ASR corpora. The mixture is dominated by Common Voice 16.0~\cite{commonvoice:2020}, supplemented by FLEURS~\cite{fleurs2022arxiv} and VoxPopuli~\cite{wang-etal-2021-voxpopuli}, and regularized with challenging tail samples from LibriSpeech~\cite{panayotov2015librispeech}, GigaSpeech~\cite{GigaSpeech2021}, The People's Speech~\cite{DBLP:journals/corr/abs-2111-09344}, and ASCEND~\cite{lovenia2022ascend}. It is strongly multilingual, covering 100+ language codes and intentionally including fragmented English dialect tags. We refer to this principle as \emph{minimal-data activation}: the small, varied training set is enough to activate latent test-time scaling, but too small to adapt the frozen backbone off its pre-trained manifold. We evaluate on the standardized FLEURS (\texttt{en\_us}) test set and VoxPopuli (\texttt{en}). \emph{Disclosure.} The 500-utterance activation set is sampled from an 811-utterance mixture with seed 42. Its corpus counts are 350 Common Voice, 77 FLEURS, 52 VoxPopuli, 8 LibriSpeech, 5 ASCEND, 4 People's Speech, and 4 GigaSpeech utterances. Thus ASCEND contributes only $5/500$ activation utterances ($1.0\%$). However, ASCEND and FLEURS/VoxPopuli each contribute to both the activation mixture (drawn from training splits) and held-out test conditions (standardized test splits), and we did not enforce speaker-disjointness between train and test. For this reason, ASCEND is not used as a headline generalization result; it is reported only as an overlap-confounded diagnostic of accented/code-switched behavior. A leave-ASCEND-out activation run is required before making a cross-domain claim on ASCEND.
|
| 185 |
+
|
| 186 |
+
\textbf{Baselines and checkpoint selection.}\quad We benchmark \method{} against the frozen Qwen3-ASR-0.6B baseline, full fine-tuning, LoRA~\cite{hu2022lora} (rank 16), and prompt tuning~\cite{lester2021power} (4 soft tokens). All methods are trained for 10 epochs on the same 500-utterance mixture. Full fine-tuning was rerun with micro-batch 4 and gradient accumulation 4 after the larger micro-batch OOMed, preserving effective batch size 16. For the main \method{} row, we use the best 500-utterance $N{=}4$, $\theta{=}0$ checkpoint from the activation-size sweep, \texttt{activation\_500\_epoch10.pth}. Some ablation and cross-domain diagnostics below use independently retrained same-protocol checkpoints; we label these as diagnostics rather than as the canonical main estimate.
|
| 187 |
+
|
| 188 |
+
\subsection{Main Results}
|
| 189 |
+
\label{sec:main_results}
|
| 190 |
+
|
| 191 |
+
Table~\ref{tab:main} reports single-run point estimates on FLEURS (\texttt{en\_us}) and VoxPopuli (\texttt{en}). Throughout this paper, all reductions are reported as the relative change in WER (or CER) over the frozen Baseline; \textbf{a negative value indicates a WER reduction (improvement)} and a positive value indicates an increase (regression). Across both test sets, every conventional adaptation method we test increases WER under the 500-utterance training regime. LoRA regresses by $+31.97\%$ relative on FLEURS, full fine-tuning by $+13.38\%$, and prompt tuning collapses outright (FLEURS WER rises to $84.803\%$).
|
| 192 |
+
|
| 193 |
+
In contrast, \method{} is the only tested adaptation mechanism with a negative $\Delta$WER point estimate on both clean benchmarks. On FLEURS, the canonical 500-utterance checkpoint changes WER from $4.900\%$ to $4.776\%$ ($-0.124$~pp, $-2.54\%$ relative). On VoxPopuli, it changes WER from $9.038\%$ to $8.995\%$ ($-0.043$~pp, $-0.47\%$ relative). These absolute deltas are small. Because we do not yet have multi-seed runs or paired confidence intervals for the canonical checkpoint, we do not interpret Table~\ref{tab:main} as evidence of statistically significant clean-benchmark improvement. Its narrower role is to show that, in the same low-resource regime where conventional adaptation clearly regresses, latent test-time computation preserves the frozen backbone and gives a negative directional estimate.
|
| 194 |
+
|
| 195 |
+
\begin{table}[ht]
|
| 196 |
+
\caption{Single-run WER and CER point estimates on the FLEURS (\texttt{en\_us}, 647 utts) and VoxPopuli (\texttt{en}, 1{,}842 utts) test splits. \textbf{Rel. $\boldsymbol{\Delta}$WER (\%)} is computed from unrounded WER against the frozen Baseline. Negative = WER reduction (improvement); positive = WER increase (regression). The clean-set \method{} deltas are directional and not claimed statistically significant.}
|
| 197 |
+
\label{tab:main}
|
| 198 |
+
\centering
|
| 199 |
+
\resizebox{\columnwidth}{!}{
|
| 200 |
+
\begin{tabular}{l | c c c | c c c}
|
| 201 |
+
\toprule
|
| 202 |
+
\multirow{2}{*}{\textbf{Model}} & \multicolumn{3}{c|}{\textbf{FLEURS}} & \multicolumn{3}{c}{\textbf{VoxPopuli}} \\
|
| 203 |
+
\cmidrule(lr){2-4} \cmidrule(lr){5-7}
|
| 204 |
+
& \textbf{WER (\%)} & \textbf{CER (\%)} & \textbf{Rel.\ $\Delta$WER (\%)} & \textbf{WER (\%)} & \textbf{CER (\%)} & \textbf{Rel.\ $\Delta$WER (\%)} \\
|
| 205 |
+
\midrule
|
| 206 |
+
Baseline & 4.900 & 2.326 & -- & 9.038 & 5.900 & -- \\
|
| 207 |
+
\quad + Full Fine-Tuning & 5.556 & 2.481 & $+13.38$ & 9.485 & 6.176 & $+4.96$ \\
|
| 208 |
+
\quad + Prompt Tuning & 84.803 & 82.468 & $+1630.70$ & 88.507 & 87.901 & $+879.33$ \\
|
| 209 |
+
\quad + LoRA (rank 16) & 6.467 & 2.944 & $+31.97$ & 10.212 & 6.536 & $+13.00$ \\
|
| 210 |
+
\midrule
|
| 211 |
+
\quad + \method{} ($N{=}4$, $\theta{=}0$) & \textbf{4.776} & \textbf{2.238} & $\boldsymbol{-2.54}$ & \textbf{8.995} & \textbf{5.869} & $\boldsymbol{-0.47}$ \\
|
| 212 |
+
\bottomrule
|
| 213 |
+
\end{tabular}
|
| 214 |
+
}
|
| 215 |
+
\end{table}
|
| 216 |
+
|
| 217 |
+
\subsection{Multilingual and Cross-Domain Generalization}
|
| 218 |
+
\label{sec:multilingual}
|
| 219 |
+
|
| 220 |
+
To verify that \method{} preserves the broad multilingual capability of the frozen backbone, we evaluate on a 23{,}049-utterance multilingual benchmark spanning the FLEURS test sets of 30 languages. This suite is a same-protocol diagnostic and is reported to test whether the latent adapter preserves broad multilingual behavior.
|
| 221 |
+
|
| 222 |
+
\begin{table}[ht]
|
| 223 |
+
\caption{Multilingual aggregate WER on FLEURS 30 and its resource-tier partitions. Rel.\ $\Delta$WER is computed against the frozen Baseline; negative = WER reduction = improvement.}
|
| 224 |
+
\label{tab:multilingual_agg}
|
| 225 |
+
\centering
|
| 226 |
+
\resizebox{\columnwidth}{!}{
|
| 227 |
+
\begin{tabular}{l r c c c}
|
| 228 |
+
\toprule
|
| 229 |
+
\textbf{Subset} & \textbf{\#utts} & \textbf{Baseline WER (\%)} & \textbf{\method{} WER (\%)} & \textbf{Rel.\ $\Delta$WER (\%)} \\
|
| 230 |
+
\midrule
|
| 231 |
+
FLEURS 30 (all) & 23{,}049 & 32.65 & \textbf{32.57} & $-0.23$ \\
|
| 232 |
+
\quad core12 & 8{,}876 & 31.71 & \textbf{31.63} & $-0.26$ \\
|
| 233 |
+
\quad +8 & 5{,}597 & 17.48 & \textbf{17.45} & $-0.19$ \\
|
| 234 |
+
\quad +10 & 8{,}576 & 43.52 & \textbf{43.42} & $-0.23$ \\
|
| 235 |
+
\bottomrule
|
| 236 |
+
\end{tabular}
|
| 237 |
+
}
|
| 238 |
+
\end{table}
|
| 239 |
+
|
| 240 |
+
\begin{table}[ht]
|
| 241 |
+
\caption{Per-language results on character-based FLEURS languages. Negative Rel.\ $\Delta$CER = error reduction = improvement.}
|
| 242 |
+
\label{tab:per_language}
|
| 243 |
+
\centering
|
| 244 |
+
\resizebox{\columnwidth}{!}{
|
| 245 |
+
\begin{tabular}{l r c c c}
|
| 246 |
+
\toprule
|
| 247 |
+
\textbf{Language (FLEURS code)} & \textbf{\#utts} & \textbf{Baseline CER (\%)} & \textbf{\method{} CER (\%)} & \textbf{Rel.\ $\Delta$CER (\%)} \\
|
| 248 |
+
\midrule
|
| 249 |
+
Mandarin (cmn\_hans\_cn) & 945 & 47.94 & \textbf{47.58} & $-0.75$ \\
|
| 250 |
+
Cantonese (yue\_hant\_hk) & 819 & 49.20 & \textbf{48.80} & $-0.82$ \\
|
| 251 |
+
Japanese (ja\_jp) & 650 & 10.35 & 10.35 & $\phantom{-}0.00$ \\
|
| 252 |
+
Thai (th\_th) & 1{,}021 & 9.53 & \textbf{9.41} & $\boldsymbol{-1.31}$ \\
|
| 253 |
+
\bottomrule
|
| 254 |
+
\end{tabular}
|
| 255 |
+
}
|
| 256 |
+
\end{table}
|
| 257 |
+
|
| 258 |
+
\subsection{Robustness and In-Domain Diagnostics}
|
| 259 |
+
\label{sec:robustness}
|
| 260 |
+
|
| 261 |
+
We further evaluate same-protocol checkpoints under acoustic shift. Table~\ref{tab:robustness} reports additive Gaussian noise at SNR $=0$~dB on FLEURS and VoxPopuli. It also includes ASCEND, evaluated with CER following~\cite{lovenia2022ascend}, but this row has a different evidential status: ASCEND contributes only $5/500$ activation utterances ($1.0\%$), yet speaker-disjointness was not enforced. The $16.02\%$ relative CER reduction on ASCEND is therefore an overlap-confounded in-domain diagnostic, not a cross-domain or speaker-independent generalization result.
|
| 262 |
+
|
| 263 |
+
\begin{table}[ht]
|
| 264 |
+
\caption{Robustness and in-domain diagnostics. Rel.\ $\Delta$ is the relative WER (or CER) change vs.\ the Baseline; negative = error reduction = improvement. ASCEND$^\dagger$ accounts for only $5/500$ activation utterances ($1.0\%$), but remains overlap-confounded because speaker-disjointness was not enforced, so it is not used as headline generalization evidence.}
|
| 265 |
+
\label{tab:robustness}
|
| 266 |
+
\centering
|
| 267 |
+
\resizebox{\columnwidth}{!}{
|
| 268 |
+
\begin{tabular}{l l c c c c}
|
| 269 |
+
\toprule
|
| 270 |
+
\textbf{Dataset} & \textbf{Condition} & \textbf{Metric} & \textbf{Baseline (\%)} & \textbf{\method{} (\%)} & \textbf{Rel.\ $\Delta$ (\%)} \\
|
| 271 |
+
\midrule
|
| 272 |
+
FLEURS & SNR$=0$~dB & WER & 29.81 & \textbf{29.42} & $-1.32$ \\
|
| 273 |
+
VoxPopuli & SNR$=0$~dB & WER & 19.71 & \textbf{19.51} & $-1.02$ \\
|
| 274 |
+
\midrule
|
| 275 |
+
ASCEND$^\dagger$ & accented / code-switched & CER & 57.81 & 48.55 & $-16.02$ \\
|
| 276 |
+
\bottomrule
|
| 277 |
+
\end{tabular}
|
| 278 |
+
}
|
| 279 |
+
\end{table}
|
| 280 |
+
|
| 281 |
+
\begin{table}[ht]
|
| 282 |
+
\caption{Stress broad suite at SNR$=0$~dB. Weighted aggregate WER drops by $-0.58$~pp ($-1.57\%$ rel).}
|
| 283 |
+
\label{tab:stress_broad}
|
| 284 |
+
\centering
|
| 285 |
+
\resizebox{\columnwidth}{!}{
|
| 286 |
+
\begin{tabular}{l r c c c c}
|
| 287 |
+
\toprule
|
| 288 |
+
\textbf{Dataset} & \textbf{\#utts} & \textbf{Baseline WER (\%)} & \textbf{\method{} WER (\%)} & \textbf{$\Delta$WER (pp)} & \textbf{Rel.\ $\Delta$WER (\%)} \\
|
| 289 |
+
\midrule
|
| 290 |
+
AMI (IHM+SDM) $0$~dB & 996 & 77.784 & \textbf{75.518} & $\boldsymbol{-2.266}$ & $\boldsymbol{-2.91}$ \\
|
| 291 |
+
Common Voice (en) $0$~dB & 500 & 31.617 & 31.936 & $+0.319$ & $+1.01$ \\
|
| 292 |
+
GigaSpeech (test) $0$~dB & 391 & 20.939 & \textbf{20.875} & $-0.064$ & $-0.30$ \\
|
| 293 |
+
LibriSpeech (clean+other) $0$~dB & 1{,}000 & 18.471 & 18.584 & $+0.113$ & $+0.61$ \\
|
| 294 |
+
People's Speech $0$~dB & 500 & 31.687 & 31.771 & $+0.084$ & $+0.26$ \\
|
| 295 |
+
TEDLIUM (release 1) $0$~dB & 500 & 14.849 & \textbf{14.291} & $-0.558$ & $-3.76$ \\
|
| 296 |
+
\midrule
|
| 297 |
+
\textbf{Weighted aggregate} & \textbf{3{,}887} & \textbf{36.843} & \textbf{36.265} & $\boldsymbol{-0.578}$ & $\boldsymbol{-1.57}$ \\
|
| 298 |
+
\bottomrule
|
| 299 |
+
\end{tabular}
|
| 300 |
+
}
|
| 301 |
+
\end{table}
|
| 302 |
+
|
| 303 |
+
\subsection{Ablation Study}
|
| 304 |
+
\label{sec:ablation}
|
| 305 |
+
|
| 306 |
+
We ablate the core design choices of \method{}: the dynamic halting policy, the three stabilization mechanisms in Section~\ref{sec:stability}, and the size of the activation set. Unless stated otherwise, ablation tables use independently retrained same-protocol checkpoints; their purpose is to test mechanisms rather than to replace the canonical main estimate in Table~\ref{tab:main}.
|
| 307 |
+
|
| 308 |
+
\textbf{Compute allocation.}\quad Table~\ref{tab:halt_dist} reports the deployed step distribution. For FLEURS and VoxPopuli, these numbers come from the canonical \texttt{activation\_500\_epoch10.pth} checkpoint. ASCEND is included only as an overlap-confounded same-protocol diagnostic. Its $5/500$ activation share is small, but the absence of speaker-disjoint enforcement means the heavier compute allocation is useful for inspecting the halting policy on accented/code-switched speech, not for claiming independent generalization.
|
| 309 |
+
|
| 310 |
+
\begin{table}[th]
|
| 311 |
+
\caption{$N$-step distribution at the deployed setting $\theta{=}0.0$. FLEURS and VoxPopuli use the canonical 500-utterance checkpoint; ASCEND$^\dagger$ is an overlap-confounded same-protocol diagnostic.}
|
| 312 |
+
\label{tab:halt_dist}
|
| 313 |
+
\centering
|
| 314 |
+
\resizebox{\columnwidth}{!}{
|
| 315 |
+
\begin{tabular}{l c r r r r}
|
| 316 |
+
\toprule
|
| 317 |
+
\textbf{Dataset} & \textbf{Skip ($\mathbf{N=0}$)} & \textbf{N=1} & \textbf{N=2} & \textbf{N=3} & \textbf{N=4} \\
|
| 318 |
+
\midrule
|
| 319 |
+
FLEURS (clean) & 47.0\% & 24.9\% & 6.3\% & 4.8\% & 17.0\% \\
|
| 320 |
+
VoxPopuli (clean) & 63.0\% & 22.0\% & 5.0\% & 2.2\% & 7.8\% \\
|
| 321 |
+
ASCEND$^\dagger$ (accented) & 22.3\% & 30.3\% & 8.2\% & 5.1\% & 34.1\% \\
|
| 322 |
+
\bottomrule
|
| 323 |
+
\end{tabular}
|
| 324 |
+
}
|
| 325 |
+
\end{table}
|
| 326 |
+
|
| 327 |
+
\textbf{Component Ablation: Stable Injection Mechanisms.}\quad
|
| 328 |
+
Table~\ref{tab:component_ablation} selectively removes each stabilization mechanism in a same-protocol ablation suite. The ablations are independently retrained and demonstrate that the design is not redundant: removing the bounded delta is catastrophic, and removing either the anchor or the gate also causes large regressions.
|
| 329 |
+
|
| 330 |
+
\begin{table}[th]
|
| 331 |
+
\caption{Component ablation on the three stabilization mechanisms (FLEURS \texttt{en\_us}, $\theta{=}0.0$). This is an independently retrained same-protocol ablation suite; deltas are against the frozen Baseline in the same evaluation pipeline.}
|
| 332 |
+
\label{tab:component_ablation}
|
| 333 |
+
\centering
|
| 334 |
+
\resizebox{\columnwidth}{!}{
|
| 335 |
+
\begin{tabular}{l c c}
|
| 336 |
+
\toprule
|
| 337 |
+
\textbf{Variant} & \textbf{WER (\%)} & \textbf{$\Delta$WER (pp)} \\
|
| 338 |
+
\midrule
|
| 339 |
+
Full same-protocol \method{} ($N{=}4$) & \textbf{4.86} & $\boldsymbol{-0.04}$ \\
|
| 340 |
+
\quad $-$ bounded delta ($L_2$ + scale $s_k$) & 51.75 & $+46.85$ \\
|
| 341 |
+
\quad $-$ fixed-embedding anchor ($\mathbf{e}_{\texttt{LT}}$ removed) & 16.22 & $+11.33$ \\
|
| 342 |
+
\quad $-$ sigmoid gate ($g_k$ fixed at $1$) & 7.89 & $+2.99$ \\
|
| 343 |
+
\bottomrule
|
| 344 |
+
\end{tabular}
|
| 345 |
+
}
|
| 346 |
+
\end{table}
|
| 347 |
+
|
| 348 |
+
\textbf{Activation Set Scaling.}\quad
|
| 349 |
+
Table~\ref{tab:activation_sweep} sweeps the activation-set size from $100$ to $800$ utterances in steps of $100$. The 500-utterance checkpoint is selected as the canonical main checkpoint because it gives the best mean absolute WER delta across FLEURS and VoxPopuli in this sweep. The useful range is narrow: several neighboring sizes are near-neutral or mildly regressive.
|
| 350 |
+
|
| 351 |
+
\begin{table}[th]
|
| 352 |
+
\caption{Activation set scaling sweep ($\theta{=}0.0$). $\Delta$ (pp) is the absolute WER change vs.\ the Baseline; \textbf{Mean $\Delta$ (pp)} averages the two benchmarks.}
|
| 353 |
+
\label{tab:activation_sweep}
|
| 354 |
+
\centering
|
| 355 |
+
\resizebox{\columnwidth}{!}{
|
| 356 |
+
\begin{tabular}{c c c c c c}
|
| 357 |
+
\toprule
|
| 358 |
+
\textbf{Activation \#utts} & \textbf{FLEURS WER (\%)} & \textbf{$\Delta$ (pp)} & \textbf{VoxPopuli WER (\%)} & \textbf{$\Delta$ (pp)} & \textbf{Mean $\Delta$ (pp)} \\
|
| 359 |
+
\midrule
|
| 360 |
+
100 & 4.87 & $-0.03$ & 8.95 & $-0.08$ & $-0.055$ \\
|
| 361 |
+
200 & 4.91 & $+0.01$ & 9.01 & $-0.03$ & $-0.010$ \\
|
| 362 |
+
300 & 4.90 & $\phantom{-}0.00$ & 9.01 & $-0.03$ & $-0.015$ \\
|
| 363 |
+
400 & 4.91 & $+0.01$ & 9.04 & $\phantom{-}0.00$ & $+0.005$ \\
|
| 364 |
+
\textbf{500} & \textbf{4.78} & $\boldsymbol{-0.12}$ & 8.99 & $-0.04$ & $\boldsymbol{-0.084}$ \\
|
| 365 |
+
600 & 4.93 & $+0.03$ & 9.02 & $-0.02$ & $+0.005$ \\
|
| 366 |
+
700 & 4.82 & $-0.08$ & 9.04 & $\phantom{-}0.00$ & $-0.040$ \\
|
| 367 |
+
800 & 4.88 & $-0.02$ & 9.07 & $+0.04$ & $+0.010$ \\
|
| 368 |
+
\bottomrule
|
| 369 |
+
\end{tabular}
|
| 370 |
+
}
|
| 371 |
+
\end{table}
|
| 372 |
+
|
| 373 |
+
\subsection{Analysis}
|
| 374 |
+
\label{sec:analysis}
|
| 375 |
+
|
| 376 |
+
The aggregate gains are small, so we treat interpretability diagnostics carefully. The main checkpoint improves FLEURS more than VoxPopuli, while the same-protocol VoxPopuli diagnostic run shows the expected hard-tail pattern: improvements concentrate in the hardest quartile and are near-zero or mildly negative in easy bins. This supports a difficulty-conditional reading rather than a uniform accuracy-lift interpretation.
|
| 377 |
+
|
| 378 |
+
\begin{table}[ht]
|
| 379 |
+
\caption{Difficulty-binned same-protocol diagnostic on VoxPopuli (\texttt{en}) at $\theta{=}0.0$ ($1{,}842$ utterances total). Utterances are partitioned into Baseline-WER quartiles.}
|
| 380 |
+
\label{tab:difficulty_bins}
|
| 381 |
+
\centering
|
| 382 |
+
\resizebox{\columnwidth}{!}{
|
| 383 |
+
\begin{tabular}{l c c c c c}
|
| 384 |
+
\toprule
|
| 385 |
+
\textbf{Bin} & \textbf{\#utts} & \textbf{Baseline WER (\%)} & \textbf{\method{} WER (\%)} & \textbf{$\Delta$WER (pp)} & \textbf{Skip (\%)} \\
|
| 386 |
+
\midrule
|
| 387 |
+
Q1 (easiest) & 460 & 0.00 & 0.06 & $+0.06$ & 51.3 \\
|
| 388 |
+
Q2 & 461 & 3.03 & 3.03 & $\phantom{-}0.00$ & 63.1 \\
|
| 389 |
+
Q3 & 460 & 8.93 & \textbf{8.91} & $-0.02$ & 64.3 \\
|
| 390 |
+
Q4 (hardest) & 461 & 21.16 & \textbf{20.90} & $\boldsymbol{-0.25}$ & 62.0 \\
|
| 391 |
+
\bottomrule
|
| 392 |
+
\end{tabular}
|
| 393 |
+
}
|
| 394 |
+
\end{table}
|
| 395 |
+
|
| 396 |
+
\begin{table}[ht]
|
| 397 |
+
\caption{Value Head decision-quality diagnostic on VoxPopuli (\texttt{en}) at $\theta{=}0.0$. Counterfactual $\Delta$WER forces the $N{=}4$ latent path on each subset.}
|
| 398 |
+
\label{tab:gate_quality}
|
| 399 |
+
\centering
|
| 400 |
+
\resizebox{\columnwidth}{!}{
|
| 401 |
+
\begin{tabular}{l c c c}
|
| 402 |
+
\toprule
|
| 403 |
+
\textbf{Subset (at $\theta{=}0.0$)} & \textbf{\#utts} & \textbf{Actual $\Delta$WER (pp)} & \textbf{Counterfactual $\Delta$WER (pp)} \\
|
| 404 |
+
\midrule
|
| 405 |
+
Skipped ($v_0 < 0$) & 1{,}109 & $0.00$ (by construction) & $+0.01$ \\
|
| 406 |
+
Processed ($v_0 \geq 0$) & 733 & $\boldsymbol{-0.19}$ & $-0.06$ \\
|
| 407 |
+
\bottomrule
|
| 408 |
+
\end{tabular}
|
| 409 |
+
}
|
| 410 |
+
\end{table}
|
| 411 |
+
|
| 412 |
+
\begin{table}[ht]
|
| 413 |
+
\caption{Qualitative examples from the VoxPopuli hard-quartile diagnostic. References are shown in lower-case unpunctuated form.}
|
| 414 |
+
\label{tab:qualitative}
|
| 415 |
+
\centering
|
| 416 |
+
\footnotesize
|
| 417 |
+
\setlength{\tabcolsep}{4pt}
|
| 418 |
+
\renewcommand{\arraystretch}{1.1}
|
| 419 |
+
\begin{tabular}{@{}p{0.12\columnwidth} p{0.26\columnwidth} p{0.26\columnwidth} p{0.26\columnwidth}@{}}
|
| 420 |
+
\toprule
|
| 421 |
+
\textbf{Source} & \textbf{Reference} & \textbf{Baseline} & \textbf{\method{}} \\
|
| 422 |
+
\midrule
|
| 423 |
+
VoxPopuli Q4 & mr president thank you for your patience and indulgence. & Thank you, Mr. President. Thank you for your patience and indulgence. & Mr. President, thank you for your patience and indulgence. \\
|
| 424 |
+
\midrule
|
| 425 |
+
VoxPopuli Q4 & madam president i wish to focus my remarks on the european neighbourhood policy in relation to the middle east peace process. & Thank you, Madame President. Commissioner, I wish to focus my remarks on the European neighbourhood policy in relation to the Middle East peace process. & Madam President, I wish to focus my remarks on the European neighbourhood policy in relation to the Middle East peace process. \\
|
| 426 |
+
\midrule
|
| 427 |
+
VoxPopuli Q4 & mr president i would like to use this opportunity to say that i withdrew my name from this report. & Thank you, Mr. President. I would like to use the opportunity to say that I withdraw my name from this report. & Mr. President, I would like to use the opportunity to say that I withdraw my name from this report. \\
|
| 428 |
+
\bottomrule
|
| 429 |
+
\end{tabular}
|
| 430 |
+
\end{table}
|
| 431 |
+
|
| 432 |
+
\section{Discussion}
|
| 433 |
+
\label{sec:discussion}
|
| 434 |
+
|
| 435 |
+
\textbf{Statistical scope.}\quad The clean main-benchmark gains in Table~\ref{tab:main} are below one tenth of a WER point on VoxPopuli and only slightly above one tenth of a point on FLEURS. With a single seed and no paired confidence interval, they are not distinguishable from evaluation noise. We therefore frame them as directional point estimates, not as statistically established improvements. ASCEND has an additional confound: although it accounts for only $5/500$ activation utterances ($1.0\%$), speaker-disjointness was not enforced, so its large CER reduction is not used as headline evidence. The remaining support for \method{} comes from the lower-risk pattern across mechanism diagnostics, noisy stress-suite results, and input-dependent compute allocation.
|
| 436 |
+
|
| 437 |
+
\textbf{Where the gain comes from.}\quad The canonical main checkpoint reduces WER by $-0.124$~pp on FLEURS and $-0.043$~pp on VoxPopuli. The gains are therefore small in absolute terms, but the compute-allocation profile (Table~\ref{tab:halt_dist}) and the hard-tail diagnostic (Table~\ref{tab:difficulty_bins}) support a difficulty-conditional interpretation. On clean read-speech benchmarks, the Value Head skips roughly half or more of all utterances. The ASCEND row is compatible with this interpretation because it triggers more full-budget computation on accented/code-switched speech, but due to activation/test corpus overlap it should be read only as an in-domain diagnostic.
|
| 438 |
+
|
| 439 |
+
\textbf{When dynamic halting matters.}\quad Dynamic halting is primarily a compute-allocation mechanism in this regime. At $\theta{=}0.0$, the canonical checkpoint uses $1.20$ average latent steps on FLEURS and $0.70$ on VoxPopuli, well below the full $N{=}4$ budget. This matters because the gain is not uniformly distributed: the model should spend compute where the frozen decoder leaves correctable residual errors and avoid spending it on low-yield utterances.
|
| 440 |
+
|
| 441 |
+
\textbf{Minimal-data activation is narrow.}\quad The activation-set sweep (Table~\ref{tab:activation_sweep}) shows that the useful operating range is narrow. The 500-utterance checkpoint is best by mean WER delta in the sweep, but nearby settings are close and some larger settings are neutral or mildly regressive. We therefore do not claim monotonic improvement with more activation data. The claim is narrower: under conditions where standard adaptation regresses, a small activation set can produce a negative directional WER point estimate without modifying the backbone.
|
| 442 |
+
|
| 443 |
+
\textbf{Failure modes and open questions.}\quad Several individual regressions merit explicit discussion. On the stress suite (Table~\ref{tab:stress_broad}), Common Voice, LibriSpeech, and People's Speech at $0$~dB all show small regressions ($\leq +0.32$~pp). This signals that the method is not uniformly beneficial across all difficulty regimes. Future work should replicate the main results with multi-seed statistical testing, rerun all diagnostics from a single frozen canonical checkpoint, and run a leave-ASCEND-out activation with speaker-disjoint evaluation before making any ASCEND generalization claim.
|
| 444 |
+
|
| 445 |
+
\section{Conclusion}
|
| 446 |
+
\label{sec:conclusion}
|
| 447 |
+
|
| 448 |
+
We presented \method{}, a parameter-efficient method that adds continuous latent test-time scaling to a frozen ASR backbone. Two small modules, a Latent Adapter that refines a few latent prefix positions through bounded recurrent updates and a Value Head that halts the loop per utterance, add under $1\%$ trainable parameters and leave the backbone untouched. Under a deliberately small activation regime where standard adaptation methods regress, \method{} gives the only negative single-run WER point estimate among the tested methods on both clean main benchmarks, but these clean-set deltas are too small to claim statistical significance without paired multi-seed testing. Same-protocol diagnostics suggest that the method behaves like a difficulty-conditional refinement loop with a learned scheduler. The ASCEND result is intentionally excluded from headline generalization claims because ASCEND contributes to the activation mixture and speaker-disjointness was not enforced.
|
| 449 |
+
|
| 450 |
+
\bibliographystyle{IEEEtran}
|
| 451 |
+
\bibliography{refs}
|
| 452 |
+
|
| 453 |
+
\end{document}
|
eval.py
ADDED
|
@@ -0,0 +1,2337 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Evaluate latent-reasoning and baseline/PEFT ASR models.
|
| 3 |
+
|
| 4 |
+
Supported datasets:
|
| 5 |
+
- SpeechTest/common_voice_16_0
|
| 6 |
+
- SpeechTest/peoples_speech
|
| 7 |
+
- SpeechTest/librispeech_asr
|
| 8 |
+
- SpeechTest/fleurs
|
| 9 |
+
- SpeechTest/voxpopuli
|
| 10 |
+
- SpeechTest/gigaspeech
|
| 11 |
+
- SpeechTest/ASCEND
|
| 12 |
+
- SpeechTest/extreme_asr_pony
|
| 13 |
+
- openslr/librispeech_asr
|
| 14 |
+
- google/fleurs
|
| 15 |
+
- facebook/voxpopuli
|
| 16 |
+
- PolyAI/minds14
|
| 17 |
+
- LIUM/tedlium
|
| 18 |
+
- edinburghcstr/ami
|
| 19 |
+
|
| 20 |
+
Unknown HuggingFace datasets are also accepted with generic ASR defaults:
|
| 21 |
+
an `audio` column plus one of the fallback text columns.
|
| 22 |
+
|
| 23 |
+
Example:
|
| 24 |
+
python eval.py \
|
| 25 |
+
--baseline-ckpt baseline_qwen_asr_best.pth \
|
| 26 |
+
--prompt-tuning-ckpt prompt_tuning_qwen_asr_best \
|
| 27 |
+
--lora-r16-ckpt lora_r16_qwen_asr_best \
|
| 28 |
+
--latent-ckpt latent_qwen_asr_best.pth \
|
| 29 |
+
--max-samples-per-config 500 \
|
| 30 |
+
--output-json cv16_eval.json
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
import argparse
|
| 34 |
+
import gc
|
| 35 |
+
import itertools
|
| 36 |
+
import json
|
| 37 |
+
import os
|
| 38 |
+
import re
|
| 39 |
+
import unicodedata
|
| 40 |
+
from dataclasses import dataclass
|
| 41 |
+
from typing import Dict, List, Optional, Tuple
|
| 42 |
+
|
| 43 |
+
import numpy as np
|
| 44 |
+
import torch
|
| 45 |
+
from datasets import Audio, get_dataset_config_names, load_dataset
|
| 46 |
+
from jiwer import cer, wer
|
| 47 |
+
from tqdm import tqdm
|
| 48 |
+
|
| 49 |
+
from qwen_asr import Qwen3ASRModel # type: ignore
|
| 50 |
+
from model import LatentQwenASR
|
| 51 |
+
from utils import set_seed
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
COMMONVOICE_12_LANGS = [
|
| 55 |
+
"en",
|
| 56 |
+
]
|
| 57 |
+
|
| 58 |
+
FLEURS_12_LANGS = [
|
| 59 |
+
"en",
|
| 60 |
+
]
|
| 61 |
+
|
| 62 |
+
VOXPOPULI_30_LANGS = [
|
| 63 |
+
"en",
|
| 64 |
+
]
|
| 65 |
+
|
| 66 |
+
QWEN_FLEURS_CORE12_LANGS = [
|
| 67 |
+
"en",
|
| 68 |
+
"zh",
|
| 69 |
+
"yue",
|
| 70 |
+
"ar",
|
| 71 |
+
"de",
|
| 72 |
+
"es",
|
| 73 |
+
"fr",
|
| 74 |
+
"it",
|
| 75 |
+
"ja",
|
| 76 |
+
"ko",
|
| 77 |
+
"pt",
|
| 78 |
+
"ru",
|
| 79 |
+
]
|
| 80 |
+
|
| 81 |
+
QWEN_FLEURS_EXTRA8_LANGS = [
|
| 82 |
+
"hi",
|
| 83 |
+
"id",
|
| 84 |
+
"ms",
|
| 85 |
+
"nl",
|
| 86 |
+
"pl",
|
| 87 |
+
"th",
|
| 88 |
+
"tr",
|
| 89 |
+
"vi",
|
| 90 |
+
]
|
| 91 |
+
|
| 92 |
+
QWEN_FLEURS_EXTRA10_LANGS = [
|
| 93 |
+
"cs",
|
| 94 |
+
"da",
|
| 95 |
+
"el",
|
| 96 |
+
"fa",
|
| 97 |
+
"fi",
|
| 98 |
+
"fil",
|
| 99 |
+
"hu",
|
| 100 |
+
"mk",
|
| 101 |
+
"ro",
|
| 102 |
+
"sv",
|
| 103 |
+
]
|
| 104 |
+
|
| 105 |
+
QWEN_FLEURS_30_LANGS = (
|
| 106 |
+
QWEN_FLEURS_CORE12_LANGS
|
| 107 |
+
+ QWEN_FLEURS_EXTRA8_LANGS
|
| 108 |
+
+ QWEN_FLEURS_EXTRA10_LANGS
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
MLS_PUBLIC7_LANGS = ["de", "nl", "es", "fr", "it", "pl", "pt"]
|
| 112 |
+
|
| 113 |
+
COMMONVOICE_CONFIG_ALIASES = {
|
| 114 |
+
"zh": "zh-CN",
|
| 115 |
+
"zh-cn": "zh-CN",
|
| 116 |
+
"zh_cn": "zh-CN",
|
| 117 |
+
"zh_tw": "zh-TW",
|
| 118 |
+
"zh-tw": "zh-TW",
|
| 119 |
+
}
|
| 120 |
+
|
| 121 |
+
COMMONVOICE_LANGUAGE_HINTS = {
|
| 122 |
+
"en": "English",
|
| 123 |
+
}
|
| 124 |
+
|
| 125 |
+
GENERAL_LANGUAGE_HINTS = {
|
| 126 |
+
"ar": "Arabic",
|
| 127 |
+
"ar_eg": "Arabic",
|
| 128 |
+
"ar-eg": "Arabic",
|
| 129 |
+
"ara": "Arabic",
|
| 130 |
+
"arb": "Arabic",
|
| 131 |
+
"cs": "Czech",
|
| 132 |
+
"cs_cz": "Czech",
|
| 133 |
+
"cs-cz": "Czech",
|
| 134 |
+
"czech": "Czech",
|
| 135 |
+
"da": "Danish",
|
| 136 |
+
"da_dk": "Danish",
|
| 137 |
+
"da-dk": "Danish",
|
| 138 |
+
"danish": "Danish",
|
| 139 |
+
"en": "English",
|
| 140 |
+
"en_us": "English",
|
| 141 |
+
"en-us": "English",
|
| 142 |
+
"en_gb": "English",
|
| 143 |
+
"en-gb": "English",
|
| 144 |
+
"en_au": "English",
|
| 145 |
+
"en-au": "English",
|
| 146 |
+
"english": "English",
|
| 147 |
+
"fr": "French",
|
| 148 |
+
"fr_fr": "French",
|
| 149 |
+
"fr-fr": "French",
|
| 150 |
+
"french": "French",
|
| 151 |
+
"de": "German",
|
| 152 |
+
"de_de": "German",
|
| 153 |
+
"de-de": "German",
|
| 154 |
+
"german": "German",
|
| 155 |
+
"es": "Spanish",
|
| 156 |
+
"es_es": "Spanish",
|
| 157 |
+
"es-es": "Spanish",
|
| 158 |
+
"es_419": "Spanish",
|
| 159 |
+
"spanish": "Spanish",
|
| 160 |
+
"it": "Italian",
|
| 161 |
+
"it_it": "Italian",
|
| 162 |
+
"it-it": "Italian",
|
| 163 |
+
"italian": "Italian",
|
| 164 |
+
"pt": "Portuguese",
|
| 165 |
+
"pt_pt": "Portuguese",
|
| 166 |
+
"pt-pt": "Portuguese",
|
| 167 |
+
"pt_br": "Portuguese",
|
| 168 |
+
"pt-br": "Portuguese",
|
| 169 |
+
"portuguese": "Portuguese",
|
| 170 |
+
"nl": "Dutch",
|
| 171 |
+
"nl_nl": "Dutch",
|
| 172 |
+
"nl-nl": "Dutch",
|
| 173 |
+
"dutch": "Dutch",
|
| 174 |
+
"pl": "Polish",
|
| 175 |
+
"pl_pl": "Polish",
|
| 176 |
+
"pl-pl": "Polish",
|
| 177 |
+
"polish": "Polish",
|
| 178 |
+
"zh": "Chinese",
|
| 179 |
+
"zh_cn": "Chinese",
|
| 180 |
+
"zh-cn": "Chinese",
|
| 181 |
+
"cmn": "Chinese",
|
| 182 |
+
"cmn_hans_cn": "Chinese",
|
| 183 |
+
"cmn-hans-cn": "Chinese",
|
| 184 |
+
"yue": "Cantonese",
|
| 185 |
+
"yue_hant_hk": "Cantonese",
|
| 186 |
+
"yue-hant-hk": "Cantonese",
|
| 187 |
+
"hi": "Hindi",
|
| 188 |
+
"hi_in": "Hindi",
|
| 189 |
+
"hi-in": "Hindi",
|
| 190 |
+
"id": "Indonesian",
|
| 191 |
+
"id_id": "Indonesian",
|
| 192 |
+
"id-id": "Indonesian",
|
| 193 |
+
"ja": "Japanese",
|
| 194 |
+
"ja_jp": "Japanese",
|
| 195 |
+
"ja-jp": "Japanese",
|
| 196 |
+
"ko": "Korean",
|
| 197 |
+
"ko_kr": "Korean",
|
| 198 |
+
"ko-kr": "Korean",
|
| 199 |
+
"ms": "Malay",
|
| 200 |
+
"ms_my": "Malay",
|
| 201 |
+
"ms-my": "Malay",
|
| 202 |
+
"ru": "Russian",
|
| 203 |
+
"ru_ru": "Russian",
|
| 204 |
+
"ru-ru": "Russian",
|
| 205 |
+
"th": "Thai",
|
| 206 |
+
"th_th": "Thai",
|
| 207 |
+
"th-th": "Thai",
|
| 208 |
+
"tr": "Turkish",
|
| 209 |
+
"tr_tr": "Turkish",
|
| 210 |
+
"tr-tr": "Turkish",
|
| 211 |
+
"vi": "Vietnamese",
|
| 212 |
+
"vi_vn": "Vietnamese",
|
| 213 |
+
"vi-vn": "Vietnamese",
|
| 214 |
+
"el": "Greek",
|
| 215 |
+
"el_gr": "Greek",
|
| 216 |
+
"el-gr": "Greek",
|
| 217 |
+
"fa": "Persian",
|
| 218 |
+
"fa_ir": "Persian",
|
| 219 |
+
"fa-ir": "Persian",
|
| 220 |
+
"fi": "Finnish",
|
| 221 |
+
"fi_fi": "Finnish",
|
| 222 |
+
"fi-fi": "Finnish",
|
| 223 |
+
"fil": "Filipino",
|
| 224 |
+
"fil_ph": "Filipino",
|
| 225 |
+
"fil-ph": "Filipino",
|
| 226 |
+
"hu": "Hungarian",
|
| 227 |
+
"hu_hu": "Hungarian",
|
| 228 |
+
"hu-hu": "Hungarian",
|
| 229 |
+
"mk": "Macedonian",
|
| 230 |
+
"mk_mk": "Macedonian",
|
| 231 |
+
"mk-mk": "Macedonian",
|
| 232 |
+
"ro": "Romanian",
|
| 233 |
+
"ro_ro": "Romanian",
|
| 234 |
+
"ro-ro": "Romanian",
|
| 235 |
+
"sv": "Swedish",
|
| 236 |
+
"sv_se": "Swedish",
|
| 237 |
+
"sv-se": "Swedish",
|
| 238 |
+
}
|
| 239 |
+
|
| 240 |
+
FALLBACK_TEXT_COLUMNS = [
|
| 241 |
+
"sentence",
|
| 242 |
+
"text",
|
| 243 |
+
"normalized_text",
|
| 244 |
+
"raw_text",
|
| 245 |
+
"transcript",
|
| 246 |
+
"transcription",
|
| 247 |
+
"english_transcription",
|
| 248 |
+
]
|
| 249 |
+
|
| 250 |
+
# Language-code family aliases used when dataset config IDs differ from
|
| 251 |
+
# requested language IDs (e.g., zh -> cmn_hans_cn, fil -> tl, ms -> msa).
|
| 252 |
+
LANGUAGE_CODE_ALIASES = {
|
| 253 |
+
"en": ["eng"],
|
| 254 |
+
"zh": ["zho", "cmn", "cmn_hans_cn", "cmn_hans"],
|
| 255 |
+
"zh_tw": ["cmn_hant_tw", "cmn_hant", "zho_hant", "zh_hant_tw"],
|
| 256 |
+
"yue": ["yue_hant_hk", "zh_hk", "zh_yue"],
|
| 257 |
+
"ar": ["ara", "arb"],
|
| 258 |
+
"de": ["deu", "ger"],
|
| 259 |
+
"es": ["spa", "es_419"],
|
| 260 |
+
"fr": ["fra", "fre"],
|
| 261 |
+
"it": ["ita"],
|
| 262 |
+
"ja": ["jpn"],
|
| 263 |
+
"ko": ["kor"],
|
| 264 |
+
"pt": ["por"],
|
| 265 |
+
"ru": ["rus"],
|
| 266 |
+
"id": ["ind", "in"],
|
| 267 |
+
"ms": ["msa", "may", "zsm"],
|
| 268 |
+
"fil": ["tgl", "tl"],
|
| 269 |
+
"th": ["tha"],
|
| 270 |
+
"vi": ["vie"],
|
| 271 |
+
"tr": ["tur"],
|
| 272 |
+
"hi": ["hin"],
|
| 273 |
+
"nl": ["nld", "dut"],
|
| 274 |
+
"sv": ["swe"],
|
| 275 |
+
"da": ["dan"],
|
| 276 |
+
"fi": ["fin"],
|
| 277 |
+
"pl": ["pol"],
|
| 278 |
+
"cs": ["ces", "cze"],
|
| 279 |
+
"fa": ["fas", "per", "pes"],
|
| 280 |
+
"el": ["ell", "gre"],
|
| 281 |
+
"hu": ["hun"],
|
| 282 |
+
"mk": ["mkd", "mac"],
|
| 283 |
+
"ro": ["ron", "rum"],
|
| 284 |
+
}
|
| 285 |
+
|
| 286 |
+
DATASET_PRESETS = {
|
| 287 |
+
"SpeechTest/common_voice_16_0": {
|
| 288 |
+
"default_strategy": "fixed_languages",
|
| 289 |
+
"default_languages": COMMONVOICE_12_LANGS,
|
| 290 |
+
"config_aliases": COMMONVOICE_CONFIG_ALIASES,
|
| 291 |
+
"language_hints": COMMONVOICE_LANGUAGE_HINTS,
|
| 292 |
+
"preferred_text_columns": ["sentence"],
|
| 293 |
+
},
|
| 294 |
+
"SpeechTest/peoples_speech": {
|
| 295 |
+
"default_strategy": "all_available",
|
| 296 |
+
"config_aliases": {},
|
| 297 |
+
"language_hints": {"test": "English", "en": "English"},
|
| 298 |
+
"preferred_text_columns": ["text"],
|
| 299 |
+
},
|
| 300 |
+
"SpeechTest/librispeech_asr": {
|
| 301 |
+
"default_strategy": "all_available",
|
| 302 |
+
"config_aliases": {},
|
| 303 |
+
"language_hints": {
|
| 304 |
+
"clean": "English",
|
| 305 |
+
"other": "English",
|
| 306 |
+
"en": "English",
|
| 307 |
+
},
|
| 308 |
+
"preferred_text_columns": ["text"],
|
| 309 |
+
},
|
| 310 |
+
"SpeechTest/fleurs": {
|
| 311 |
+
"default_strategy": "fixed_languages",
|
| 312 |
+
"default_languages": FLEURS_12_LANGS,
|
| 313 |
+
"config_aliases": {},
|
| 314 |
+
"language_hints": COMMONVOICE_LANGUAGE_HINTS,
|
| 315 |
+
"preferred_text_columns": ["transcription", "raw_transcription", "text"],
|
| 316 |
+
},
|
| 317 |
+
"SpeechTest/voxpopuli": {
|
| 318 |
+
"default_strategy": "fixed_languages",
|
| 319 |
+
"default_languages": VOXPOPULI_30_LANGS,
|
| 320 |
+
"config_aliases": {},
|
| 321 |
+
"language_hints": {
|
| 322 |
+
"en": "English",
|
| 323 |
+
},
|
| 324 |
+
"preferred_text_columns": ["normalized_text", "raw_text", "text"],
|
| 325 |
+
},
|
| 326 |
+
"SpeechTest/gigaspeech": {
|
| 327 |
+
"default_strategy": "all_available",
|
| 328 |
+
"config_aliases": {},
|
| 329 |
+
"language_hints": {"test": "English", "en": "English"},
|
| 330 |
+
"preferred_text_columns": ["text"],
|
| 331 |
+
},
|
| 332 |
+
"SpeechTest/ASCEND": {
|
| 333 |
+
"default_strategy": "all_available",
|
| 334 |
+
"config_aliases": {},
|
| 335 |
+
"language_hints": {},
|
| 336 |
+
"preferred_text_columns": ["transcription", "text"],
|
| 337 |
+
},
|
| 338 |
+
"SpeechTest/extreme_asr_pony": {
|
| 339 |
+
"default_strategy": "all_available",
|
| 340 |
+
"config_aliases": {},
|
| 341 |
+
"language_hints": {},
|
| 342 |
+
"preferred_text_columns": ["text"],
|
| 343 |
+
},
|
| 344 |
+
"openslr/librispeech_asr": {
|
| 345 |
+
"default_strategy": "default_configs",
|
| 346 |
+
"default_configs": ["clean", "other"],
|
| 347 |
+
"config_aliases": {},
|
| 348 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 349 |
+
"preferred_text_columns": ["text"],
|
| 350 |
+
},
|
| 351 |
+
"google/fleurs": {
|
| 352 |
+
"default_strategy": "fixed_languages",
|
| 353 |
+
"default_languages": FLEURS_12_LANGS,
|
| 354 |
+
"config_aliases": {},
|
| 355 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 356 |
+
"preferred_text_columns": ["transcription", "raw_transcription", "text"],
|
| 357 |
+
},
|
| 358 |
+
"facebook/voxpopuli": {
|
| 359 |
+
"default_strategy": "fixed_languages",
|
| 360 |
+
"default_languages": VOXPOPULI_30_LANGS,
|
| 361 |
+
"config_aliases": {},
|
| 362 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 363 |
+
"preferred_text_columns": ["normalized_text", "raw_text", "text"],
|
| 364 |
+
},
|
| 365 |
+
"PolyAI/minds14": {
|
| 366 |
+
"default_strategy": "default_configs",
|
| 367 |
+
"default_configs": ["en-US"],
|
| 368 |
+
"config_aliases": {
|
| 369 |
+
"en": "en-US",
|
| 370 |
+
"en_us": "en-US",
|
| 371 |
+
"en-us": "en-US",
|
| 372 |
+
},
|
| 373 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 374 |
+
"preferred_text_columns": ["transcription", "english_transcription", "text"],
|
| 375 |
+
},
|
| 376 |
+
"facebook/multilingual_librispeech": {
|
| 377 |
+
"default_strategy": "fixed_languages",
|
| 378 |
+
"default_languages": MLS_PUBLIC7_LANGS,
|
| 379 |
+
"config_aliases": {
|
| 380 |
+
"de": "german",
|
| 381 |
+
"nl": "dutch",
|
| 382 |
+
"es": "spanish",
|
| 383 |
+
"fr": "french",
|
| 384 |
+
"it": "italian",
|
| 385 |
+
"pl": "polish",
|
| 386 |
+
"pt": "portuguese",
|
| 387 |
+
},
|
| 388 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 389 |
+
"preferred_text_columns": ["transcript", "text"],
|
| 390 |
+
},
|
| 391 |
+
"LIUM/tedlium": {
|
| 392 |
+
"default_strategy": "default_configs",
|
| 393 |
+
"default_configs": ["release1"],
|
| 394 |
+
"config_aliases": {
|
| 395 |
+
"r1": "release1",
|
| 396 |
+
"r2": "release2",
|
| 397 |
+
"r3": "release3",
|
| 398 |
+
},
|
| 399 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 400 |
+
"preferred_text_columns": ["text"],
|
| 401 |
+
},
|
| 402 |
+
"TwinkStart/tedlium": {
|
| 403 |
+
"default_strategy": "default_configs",
|
| 404 |
+
"default_configs": ["release1"],
|
| 405 |
+
"config_aliases": {
|
| 406 |
+
"r1": "release1",
|
| 407 |
+
"r2": "release2",
|
| 408 |
+
"r3": "release3",
|
| 409 |
+
},
|
| 410 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 411 |
+
"preferred_text_columns": ["text"],
|
| 412 |
+
},
|
| 413 |
+
"edinburghcstr/ami": {
|
| 414 |
+
"default_strategy": "default_configs",
|
| 415 |
+
"default_configs": ["ihm"],
|
| 416 |
+
"config_aliases": {},
|
| 417 |
+
"language_hints": GENERAL_LANGUAGE_HINTS,
|
| 418 |
+
"preferred_text_columns": ["text", "transcript", "sentence"],
|
| 419 |
+
},
|
| 420 |
+
}
|
| 421 |
+
|
| 422 |
+
|
| 423 |
+
@dataclass
|
| 424 |
+
class ModelBundle:
|
| 425 |
+
model: LatentQwenASR
|
| 426 |
+
processor: any
|
| 427 |
+
checkpoint: Dict[str, any]
|
| 428 |
+
|
| 429 |
+
|
| 430 |
+
def parse_args() -> argparse.Namespace:
|
| 431 |
+
parser = argparse.ArgumentParser(
|
| 432 |
+
description=(
|
| 433 |
+
"Evaluate five ASR settings on SpeechTest or HuggingFace ASR datasets: "
|
| 434 |
+
"base_model, baseline_ft, prompt_tuning, lora_r16, latent_reasoning."
|
| 435 |
+
)
|
| 436 |
+
)
|
| 437 |
+
parser.add_argument(
|
| 438 |
+
"--model-id",
|
| 439 |
+
type=str,
|
| 440 |
+
default=os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-0.6B"),
|
| 441 |
+
help="Base Qwen3-ASR model ID.",
|
| 442 |
+
)
|
| 443 |
+
parser.add_argument(
|
| 444 |
+
"--dataset-name",
|
| 445 |
+
type=str,
|
| 446 |
+
default="SpeechTest/common_voice_16_0",
|
| 447 |
+
help="HuggingFace dataset name.",
|
| 448 |
+
)
|
| 449 |
+
parser.add_argument(
|
| 450 |
+
"--split",
|
| 451 |
+
type=str,
|
| 452 |
+
default="test",
|
| 453 |
+
help="Dataset split to evaluate.",
|
| 454 |
+
)
|
| 455 |
+
parser.add_argument(
|
| 456 |
+
"--configs",
|
| 457 |
+
type=str,
|
| 458 |
+
default="",
|
| 459 |
+
help=(
|
| 460 |
+
"Comma-separated config names (dataset-specific). "
|
| 461 |
+
"Default: dataset preset "
|
| 462 |
+
"(fixed-language datasets use preset languages; "
|
| 463 |
+
"others use preset configs or all subsets). "
|
| 464 |
+
"Ignored when --all-configs is set."
|
| 465 |
+
),
|
| 466 |
+
)
|
| 467 |
+
parser.add_argument(
|
| 468 |
+
"--all-configs",
|
| 469 |
+
action="store_true",
|
| 470 |
+
help="Evaluate every config in the dataset.",
|
| 471 |
+
)
|
| 472 |
+
parser.add_argument(
|
| 473 |
+
"--max-configs",
|
| 474 |
+
type=int,
|
| 475 |
+
default=0,
|
| 476 |
+
help="Limit number of configs (0 = no limit).",
|
| 477 |
+
)
|
| 478 |
+
parser.add_argument(
|
| 479 |
+
"--max-samples-per-config",
|
| 480 |
+
type=int,
|
| 481 |
+
default=0,
|
| 482 |
+
help="Maximum test samples per config (0 = all).",
|
| 483 |
+
)
|
| 484 |
+
parser.add_argument(
|
| 485 |
+
"--streaming",
|
| 486 |
+
action="store_true",
|
| 487 |
+
help=(
|
| 488 |
+
"Load datasets with HuggingFace streaming=True and iterate samples "
|
| 489 |
+
"sequentially without materializing the split on disk."
|
| 490 |
+
),
|
| 491 |
+
)
|
| 492 |
+
parser.add_argument(
|
| 493 |
+
"--text-normalizer",
|
| 494 |
+
type=str,
|
| 495 |
+
default="english",
|
| 496 |
+
choices=["english", "basic", "simple"],
|
| 497 |
+
help=(
|
| 498 |
+
"Text normalization before WER/CER. Use 'basic' or 'simple' for "
|
| 499 |
+
"multilingual evaluations to avoid English-specific rewriting."
|
| 500 |
+
),
|
| 501 |
+
)
|
| 502 |
+
parser.add_argument(
|
| 503 |
+
"--max-new-tokens",
|
| 504 |
+
type=int,
|
| 505 |
+
default=128,
|
| 506 |
+
help="Max new tokens for generation.",
|
| 507 |
+
)
|
| 508 |
+
parser.add_argument(
|
| 509 |
+
"--baseline-ckpt",
|
| 510 |
+
type=str,
|
| 511 |
+
default="",
|
| 512 |
+
help="Path to baseline fine-tuned checkpoint (.pth).",
|
| 513 |
+
)
|
| 514 |
+
parser.add_argument(
|
| 515 |
+
"--prompt-tuning-ckpt",
|
| 516 |
+
type=str,
|
| 517 |
+
default="",
|
| 518 |
+
help=(
|
| 519 |
+
"Path to prompt-tuning checkpoint. "
|
| 520 |
+
"Supports .pth (wrapper state_dict) or a PEFT adapter directory."
|
| 521 |
+
),
|
| 522 |
+
)
|
| 523 |
+
parser.add_argument(
|
| 524 |
+
"--lora-r16-ckpt",
|
| 525 |
+
type=str,
|
| 526 |
+
default="",
|
| 527 |
+
help=(
|
| 528 |
+
"Path to LoRA(rank=16) checkpoint. "
|
| 529 |
+
"Supports .pth (wrapper state_dict) or a PEFT adapter directory."
|
| 530 |
+
),
|
| 531 |
+
)
|
| 532 |
+
parser.add_argument(
|
| 533 |
+
"--latent-ckpt",
|
| 534 |
+
type=str,
|
| 535 |
+
default="",
|
| 536 |
+
help="Path to latent reasoning checkpoint (.pth).",
|
| 537 |
+
)
|
| 538 |
+
parser.add_argument(
|
| 539 |
+
"--skip-base-model",
|
| 540 |
+
action="store_true",
|
| 541 |
+
help="Skip raw pretrained model inference baseline.",
|
| 542 |
+
)
|
| 543 |
+
parser.add_argument(
|
| 544 |
+
"--skip-baseline-ft",
|
| 545 |
+
action="store_true",
|
| 546 |
+
help="Skip baseline fine-tuned model evaluation.",
|
| 547 |
+
)
|
| 548 |
+
parser.add_argument(
|
| 549 |
+
"--skip-prompt-tuning",
|
| 550 |
+
action="store_true",
|
| 551 |
+
help="Skip prompt tuning model evaluation.",
|
| 552 |
+
)
|
| 553 |
+
parser.add_argument(
|
| 554 |
+
"--skip-lora-r16",
|
| 555 |
+
action="store_true",
|
| 556 |
+
help="Skip LoRA r16 model evaluation.",
|
| 557 |
+
)
|
| 558 |
+
parser.add_argument(
|
| 559 |
+
"--n-latent",
|
| 560 |
+
type=int,
|
| 561 |
+
default=-1,
|
| 562 |
+
help="Override latent token count. If < 0, use checkpoint n_latent.",
|
| 563 |
+
)
|
| 564 |
+
parser.add_argument(
|
| 565 |
+
"--num-beams",
|
| 566 |
+
type=int,
|
| 567 |
+
default=1,
|
| 568 |
+
help="Number of beams for beam search.",
|
| 569 |
+
)
|
| 570 |
+
parser.add_argument(
|
| 571 |
+
"--dynamic-halt-threshold",
|
| 572 |
+
type=float,
|
| 573 |
+
default=0.0,
|
| 574 |
+
help="Value head threshold for early halting of latent thoughts.",
|
| 575 |
+
)
|
| 576 |
+
|
| 577 |
+
parser.add_argument(
|
| 578 |
+
"--device",
|
| 579 |
+
type=str,
|
| 580 |
+
default="auto",
|
| 581 |
+
choices=["auto", "cuda", "cpu"],
|
| 582 |
+
help="Device to run inference.",
|
| 583 |
+
)
|
| 584 |
+
parser.add_argument(
|
| 585 |
+
"--dtype",
|
| 586 |
+
type=str,
|
| 587 |
+
default="auto",
|
| 588 |
+
choices=["auto", "float32", "float16", "bfloat16"],
|
| 589 |
+
help="Inference dtype.",
|
| 590 |
+
)
|
| 591 |
+
parser.add_argument(
|
| 592 |
+
"--seed",
|
| 593 |
+
type=int,
|
| 594 |
+
default=42,
|
| 595 |
+
help="Random seed.",
|
| 596 |
+
)
|
| 597 |
+
parser.add_argument(
|
| 598 |
+
"--print-samples",
|
| 599 |
+
type=int,
|
| 600 |
+
default=2,
|
| 601 |
+
help="Number of qualitative samples to print per config.",
|
| 602 |
+
)
|
| 603 |
+
parser.add_argument(
|
| 604 |
+
"--print-gate-skips",
|
| 605 |
+
action="store_true",
|
| 606 |
+
help=(
|
| 607 |
+
"Also print latent_reasoning samples skipped by the value-head gate. "
|
| 608 |
+
"Default is off to keep large sweep logs readable."
|
| 609 |
+
),
|
| 610 |
+
)
|
| 611 |
+
parser.add_argument(
|
| 612 |
+
"--output-json",
|
| 613 |
+
type=str,
|
| 614 |
+
default="",
|
| 615 |
+
help="If set, write detailed results to this JSON file.",
|
| 616 |
+
)
|
| 617 |
+
parser.add_argument(
|
| 618 |
+
"--no-language-hint",
|
| 619 |
+
action="store_true",
|
| 620 |
+
help="Disable language-specific prompting and use generic ASR prompt.",
|
| 621 |
+
)
|
| 622 |
+
parser.add_argument(
|
| 623 |
+
"--prompt-template",
|
| 624 |
+
type=str,
|
| 625 |
+
default="Transcribe the {language} audio into text.",
|
| 626 |
+
help=(
|
| 627 |
+
"User prompt template. Use {language} placeholder. "
|
| 628 |
+
"Ignored when --no-language-hint is set."
|
| 629 |
+
),
|
| 630 |
+
)
|
| 631 |
+
parser.add_argument(
|
| 632 |
+
"--snr-db",
|
| 633 |
+
type=float,
|
| 634 |
+
nargs="+",
|
| 635 |
+
default=None,
|
| 636 |
+
help=(
|
| 637 |
+
"Add Gaussian noise to audio at specified SNR levels (dB). "
|
| 638 |
+
"Can specify multiple values, e.g. --snr-db 20 10 5. "
|
| 639 |
+
"Each SNR level will be evaluated as a separate row. "
|
| 640 |
+
"If not set, no noise is added."
|
| 641 |
+
),
|
| 642 |
+
)
|
| 643 |
+
return parser.parse_args()
|
| 644 |
+
|
| 645 |
+
|
| 646 |
+
def choose_device(device_arg: str) -> str:
|
| 647 |
+
if device_arg == "auto":
|
| 648 |
+
return "cuda" if torch.cuda.is_available() else "cpu"
|
| 649 |
+
if device_arg == "cuda" and not torch.cuda.is_available():
|
| 650 |
+
raise RuntimeError("Requested --device cuda but CUDA is not available.")
|
| 651 |
+
return device_arg
|
| 652 |
+
|
| 653 |
+
|
| 654 |
+
def choose_dtype(dtype_arg: str, device: str) -> torch.dtype:
|
| 655 |
+
if dtype_arg == "float32":
|
| 656 |
+
return torch.float32
|
| 657 |
+
if dtype_arg == "float16":
|
| 658 |
+
return torch.float16 if device == "cuda" else torch.float32
|
| 659 |
+
if dtype_arg == "bfloat16":
|
| 660 |
+
return torch.bfloat16 if device == "cuda" else torch.float32
|
| 661 |
+
|
| 662 |
+
# auto
|
| 663 |
+
if device != "cuda":
|
| 664 |
+
return torch.float32
|
| 665 |
+
major = torch.cuda.get_device_capability(0)[0]
|
| 666 |
+
if major >= 8:
|
| 667 |
+
return torch.bfloat16
|
| 668 |
+
return torch.float16
|
| 669 |
+
|
| 670 |
+
|
| 671 |
+
def get_dataset_preset(dataset_name: str) -> Dict[str, any]:
|
| 672 |
+
return DATASET_PRESETS.get(dataset_name, {})
|
| 673 |
+
|
| 674 |
+
|
| 675 |
+
def resolve_dataset_name(dataset_name: str) -> str:
|
| 676 |
+
key = (dataset_name or "").strip()
|
| 677 |
+
if not key:
|
| 678 |
+
raise ValueError("dataset name is empty.")
|
| 679 |
+
|
| 680 |
+
if key in DATASET_PRESETS:
|
| 681 |
+
return key
|
| 682 |
+
|
| 683 |
+
lowered = key.lower()
|
| 684 |
+
for canonical in DATASET_PRESETS:
|
| 685 |
+
if canonical.lower() == lowered:
|
| 686 |
+
return canonical
|
| 687 |
+
|
| 688 |
+
print(
|
| 689 |
+
f"[warn] dataset '{dataset_name}' has no local preset; "
|
| 690 |
+
"using generic HuggingFace ASR defaults."
|
| 691 |
+
)
|
| 692 |
+
return key
|
| 693 |
+
|
| 694 |
+
|
| 695 |
+
def _map_config_alias(name: str, alias_map: Dict[str, str]) -> str:
|
| 696 |
+
n = name.strip()
|
| 697 |
+
if not n:
|
| 698 |
+
return n
|
| 699 |
+
lowered = n.lower()
|
| 700 |
+
if lowered in alias_map:
|
| 701 |
+
return alias_map[lowered]
|
| 702 |
+
return n
|
| 703 |
+
|
| 704 |
+
|
| 705 |
+
def resolve_language_hint(dataset_name: str, config_name: str) -> Optional[str]:
|
| 706 |
+
preset = get_dataset_preset(dataset_name)
|
| 707 |
+
hints_raw = preset.get("language_hints", {})
|
| 708 |
+
hints = {k.lower(): v for k, v in hints_raw.items()}
|
| 709 |
+
|
| 710 |
+
lowered = config_name.lower()
|
| 711 |
+
if lowered in hints:
|
| 712 |
+
return hints[lowered]
|
| 713 |
+
|
| 714 |
+
# Generic base-language fallback (e.g. "fr_ca" -> "fr").
|
| 715 |
+
base = lowered.split("-")[0].split("_")[0]
|
| 716 |
+
for key, val in hints.items():
|
| 717 |
+
if key.split("-")[0].split("_")[0] == base:
|
| 718 |
+
return val
|
| 719 |
+
return None
|
| 720 |
+
|
| 721 |
+
|
| 722 |
+
def _norm_cfg_key(v: str) -> str:
|
| 723 |
+
return (v or "").strip().lower().replace("-", "_")
|
| 724 |
+
|
| 725 |
+
|
| 726 |
+
def _base_lang(v: str) -> str:
|
| 727 |
+
return _norm_cfg_key(v).split("_")[0]
|
| 728 |
+
|
| 729 |
+
|
| 730 |
+
def _language_candidate_forms(lang_code: str) -> List[str]:
|
| 731 |
+
code = _norm_cfg_key(lang_code)
|
| 732 |
+
out = {
|
| 733 |
+
code,
|
| 734 |
+
code.replace("_", "-"),
|
| 735 |
+
}
|
| 736 |
+
base = _base_lang(code)
|
| 737 |
+
out.add(base)
|
| 738 |
+
|
| 739 |
+
# Common alias expansions for multilingual dataset configs.
|
| 740 |
+
if code == "zh":
|
| 741 |
+
out.update({"zh_cn", "zh-cn", "zh_hans", "zh-hans"})
|
| 742 |
+
if code in {"zh_tw", "zhtw"}:
|
| 743 |
+
out.update({"zh_tw", "zh-tw", "zh_hant", "zh-hant", "zh_hk", "zh-hk"})
|
| 744 |
+
if code == "yue":
|
| 745 |
+
out.update({"yue_hk", "yue-hk", "yue_hant_hk", "yue-hant-hk", "zh_yue", "zh-yue", "zh_hk", "zh-hk"})
|
| 746 |
+
if code == "fil":
|
| 747 |
+
out.update({"fil_ph", "fil-ph", "tl", "tl_ph", "tl-ph", "tgl", "tgl_ph", "tgl-ph"})
|
| 748 |
+
|
| 749 |
+
# Expand across ISO variants / dataset-specific aliases.
|
| 750 |
+
base_aliases = LANGUAGE_CODE_ALIASES.get(base, [])
|
| 751 |
+
for alias in base_aliases:
|
| 752 |
+
alias_norm = _norm_cfg_key(alias)
|
| 753 |
+
out.update(
|
| 754 |
+
{
|
| 755 |
+
alias_norm,
|
| 756 |
+
alias_norm.replace("_", "-"),
|
| 757 |
+
_base_lang(alias_norm),
|
| 758 |
+
}
|
| 759 |
+
)
|
| 760 |
+
|
| 761 |
+
return list(out)
|
| 762 |
+
|
| 763 |
+
|
| 764 |
+
def _pick_config_for_language(lang_code: str, available_configs: List[str]) -> Optional[str]:
|
| 765 |
+
if not available_configs:
|
| 766 |
+
return None
|
| 767 |
+
infos = []
|
| 768 |
+
for cfg in available_configs:
|
| 769 |
+
norm = _norm_cfg_key(cfg)
|
| 770 |
+
infos.append((cfg, norm, _base_lang(norm)))
|
| 771 |
+
|
| 772 |
+
candidates = set(_language_candidate_forms(lang_code))
|
| 773 |
+
code_norm = _norm_cfg_key(lang_code)
|
| 774 |
+
code_base = _base_lang(code_norm)
|
| 775 |
+
|
| 776 |
+
# 1) exact-ish match against candidate forms
|
| 777 |
+
exact_hits = [t for t in infos if t[1] in candidates]
|
| 778 |
+
if exact_hits:
|
| 779 |
+
# Prefer exact code, then prefix match (e.g. en_us for en), then lexical.
|
| 780 |
+
exact_hits.sort(
|
| 781 |
+
key=lambda t: (
|
| 782 |
+
0 if t[1] == code_norm else 1,
|
| 783 |
+
0 if t[1].startswith(code_base + "_") else 1,
|
| 784 |
+
t[1],
|
| 785 |
+
)
|
| 786 |
+
)
|
| 787 |
+
return exact_hits[0][0]
|
| 788 |
+
|
| 789 |
+
# 2) base-language fallback
|
| 790 |
+
base_hits = [t for t in infos if t[2] == code_base]
|
| 791 |
+
if base_hits:
|
| 792 |
+
# For Chinese variants, bias toward region when specified.
|
| 793 |
+
def _rank(tup: Tuple[str, str, str]) -> Tuple[int, int, str]:
|
| 794 |
+
norm = tup[1]
|
| 795 |
+
if code_norm in {"zh_tw", "zhtw"}:
|
| 796 |
+
zh_tw_bias = 0 if ("_tw" in norm or "hant" in norm) else 1
|
| 797 |
+
return (zh_tw_bias, 0 if norm.startswith(code_base + "_") else 1, norm)
|
| 798 |
+
if code_norm == "zh":
|
| 799 |
+
zh_cn_bias = 0 if ("_cn" in norm or "hans" in norm) else 1
|
| 800 |
+
return (zh_cn_bias, 0 if norm == "zh" else 1, norm)
|
| 801 |
+
return (0, 0 if norm.startswith(code_base + "_") else 1, norm)
|
| 802 |
+
|
| 803 |
+
base_hits.sort(key=_rank)
|
| 804 |
+
return base_hits[0][0]
|
| 805 |
+
|
| 806 |
+
return None
|
| 807 |
+
|
| 808 |
+
|
| 809 |
+
def resolve_language_configs(
|
| 810 |
+
language_codes: List[str],
|
| 811 |
+
available_configs: List[str],
|
| 812 |
+
) -> Tuple[List[str], List[str], List[Tuple[str, str]]]:
|
| 813 |
+
resolved: List[str] = []
|
| 814 |
+
missing: List[str] = []
|
| 815 |
+
mapping: List[Tuple[str, str]] = []
|
| 816 |
+
used = set()
|
| 817 |
+
for code in language_codes:
|
| 818 |
+
cfg = _pick_config_for_language(code, available_configs)
|
| 819 |
+
if cfg is None:
|
| 820 |
+
missing.append(code)
|
| 821 |
+
continue
|
| 822 |
+
mapping.append((code, cfg))
|
| 823 |
+
if cfg in used:
|
| 824 |
+
continue
|
| 825 |
+
resolved.append(cfg)
|
| 826 |
+
used.add(cfg)
|
| 827 |
+
return resolved, missing, mapping
|
| 828 |
+
|
| 829 |
+
|
| 830 |
+
def resolve_configs(dataset_name: str, config_arg: str, all_configs: bool, max_configs: int) -> List[str]:
|
| 831 |
+
available_configs = get_dataset_config_names(dataset_name, trust_remote_code=True)
|
| 832 |
+
available_set = set(available_configs)
|
| 833 |
+
preset = get_dataset_preset(dataset_name)
|
| 834 |
+
alias_raw = preset.get("config_aliases", {})
|
| 835 |
+
alias_map = {k.lower(): v for k, v in alias_raw.items()}
|
| 836 |
+
|
| 837 |
+
if all_configs:
|
| 838 |
+
configs = available_configs
|
| 839 |
+
else:
|
| 840 |
+
raw_configs = [x.strip() for x in config_arg.split(",") if x.strip()]
|
| 841 |
+
default_strategy = str(preset.get("default_strategy", "all_available"))
|
| 842 |
+
|
| 843 |
+
# Fixed language datasets: map language codes -> actual config names.
|
| 844 |
+
use_language_mapping = default_strategy == "fixed_languages"
|
| 845 |
+
if use_language_mapping:
|
| 846 |
+
if not raw_configs:
|
| 847 |
+
raw_configs = list(preset.get("default_languages", []))
|
| 848 |
+
mapped_langs = [_map_config_alias(x, alias_map) for x in raw_configs]
|
| 849 |
+
configs, missing_langs, mapping = resolve_language_configs(mapped_langs, available_configs)
|
| 850 |
+
if mapping:
|
| 851 |
+
mapping_str = ", ".join([f"{src}->{dst}" for src, dst in mapping])
|
| 852 |
+
print(f"[info] {dataset_name}: language->config mapping: {mapping_str}")
|
| 853 |
+
if missing_langs:
|
| 854 |
+
print(
|
| 855 |
+
f"[warn] {dataset_name}: skip missing language configs: {missing_langs}"
|
| 856 |
+
)
|
| 857 |
+
else:
|
| 858 |
+
if not raw_configs:
|
| 859 |
+
if default_strategy == "all_available":
|
| 860 |
+
raw_configs = available_configs.copy()
|
| 861 |
+
else:
|
| 862 |
+
raw_configs = list(preset.get("default_configs", []))
|
| 863 |
+
|
| 864 |
+
mapped = [_map_config_alias(x, alias_map) for x in raw_configs]
|
| 865 |
+
seen = set()
|
| 866 |
+
configs = []
|
| 867 |
+
for cfg in mapped:
|
| 868 |
+
if cfg and cfg not in seen:
|
| 869 |
+
configs.append(cfg)
|
| 870 |
+
seen.add(cfg)
|
| 871 |
+
|
| 872 |
+
missing = [cfg for cfg in configs if cfg not in available_set]
|
| 873 |
+
if missing:
|
| 874 |
+
raise ValueError(
|
| 875 |
+
"Unknown configs after alias mapping: "
|
| 876 |
+
f"{missing}. Available configs include: {sorted(list(available_set))[:20]} ..."
|
| 877 |
+
)
|
| 878 |
+
|
| 879 |
+
if not configs:
|
| 880 |
+
raise ValueError("No configs resolved. Provide --configs or use --all-configs.")
|
| 881 |
+
if max_configs > 0:
|
| 882 |
+
configs = configs[:max_configs]
|
| 883 |
+
return configs
|
| 884 |
+
|
| 885 |
+
|
| 886 |
+
def dataset_column_names(raw_ds: any) -> List[str]:
|
| 887 |
+
column_names = getattr(raw_ds, "column_names", None)
|
| 888 |
+
if column_names:
|
| 889 |
+
return list(column_names)
|
| 890 |
+
features = getattr(raw_ds, "features", None)
|
| 891 |
+
if features:
|
| 892 |
+
try:
|
| 893 |
+
return list(features.keys())
|
| 894 |
+
except Exception:
|
| 895 |
+
pass
|
| 896 |
+
return []
|
| 897 |
+
|
| 898 |
+
|
| 899 |
+
def resolve_text_column(dataset_name: str, column_names: List[str]) -> Optional[str]:
|
| 900 |
+
preset = get_dataset_preset(dataset_name)
|
| 901 |
+
preferred = list(preset.get("preferred_text_columns", []))
|
| 902 |
+
candidates = preferred + FALLBACK_TEXT_COLUMNS
|
| 903 |
+
seen = set()
|
| 904 |
+
deduped = []
|
| 905 |
+
for c in candidates:
|
| 906 |
+
if c not in seen:
|
| 907 |
+
deduped.append(c)
|
| 908 |
+
seen.add(c)
|
| 909 |
+
for c in deduped:
|
| 910 |
+
if c in column_names:
|
| 911 |
+
return c
|
| 912 |
+
return None
|
| 913 |
+
|
| 914 |
+
|
| 915 |
+
def resolve_im_start_id(tokenizer: any) -> int:
|
| 916 |
+
im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
| 917 |
+
if im_start_id is None or im_start_id == tokenizer.unk_token_id:
|
| 918 |
+
if tokenizer.bos_token_id is not None:
|
| 919 |
+
return int(tokenizer.bos_token_id)
|
| 920 |
+
if tokenizer.eos_token_id is not None:
|
| 921 |
+
return int(tokenizer.eos_token_id)
|
| 922 |
+
return 0
|
| 923 |
+
return int(im_start_id)
|
| 924 |
+
|
| 925 |
+
|
| 926 |
+
def ensure_front_prompt_token(
|
| 927 |
+
asr_model: any,
|
| 928 |
+
tokenizer: any,
|
| 929 |
+
token_name: str = "<|latent|>",
|
| 930 |
+
) -> int:
|
| 931 |
+
token_id = tokenizer.convert_tokens_to_ids(token_name)
|
| 932 |
+
if token_id is None or token_id == tokenizer.unk_token_id:
|
| 933 |
+
added = tokenizer.add_special_tokens({"additional_special_tokens": [token_name]})
|
| 934 |
+
if added > 0:
|
| 935 |
+
asr_model.thinker.resize_token_embeddings(len(tokenizer))
|
| 936 |
+
token_id = tokenizer.convert_tokens_to_ids(token_name)
|
| 937 |
+
if token_id is None or token_id < 0:
|
| 938 |
+
raise RuntimeError(f"Failed to resolve front prompt token id for {token_name!r}.")
|
| 939 |
+
return int(token_id)
|
| 940 |
+
|
| 941 |
+
|
| 942 |
+
def extract_soft_prompt_tensor(ckpt: Dict[str, any]) -> Optional[torch.Tensor]:
|
| 943 |
+
tensor = ckpt.get("soft_prompt_embed")
|
| 944 |
+
if torch.is_tensor(tensor):
|
| 945 |
+
return tensor
|
| 946 |
+
state = ckpt.get("model_state_dict")
|
| 947 |
+
if isinstance(state, dict):
|
| 948 |
+
tensor = state.get("soft_prompt_embed")
|
| 949 |
+
if torch.is_tensor(tensor):
|
| 950 |
+
return tensor
|
| 951 |
+
return None
|
| 952 |
+
|
| 953 |
+
|
| 954 |
+
def load_checkpoint(path: str) -> Dict[str, any]:
|
| 955 |
+
ckpt = torch.load(path, map_location="cpu")
|
| 956 |
+
if isinstance(ckpt, dict):
|
| 957 |
+
return ckpt
|
| 958 |
+
return {"model_state_dict": ckpt}
|
| 959 |
+
|
| 960 |
+
|
| 961 |
+
def is_path_file_or_dir(path: str) -> bool:
|
| 962 |
+
return os.path.isfile(path) or os.path.isdir(path)
|
| 963 |
+
|
| 964 |
+
|
| 965 |
+
def build_plain_bundle(
|
| 966 |
+
model_id: str,
|
| 967 |
+
device: str,
|
| 968 |
+
dtype: torch.dtype,
|
| 969 |
+
) -> Tuple[LatentQwenASR, any]:
|
| 970 |
+
asr_wrapper = Qwen3ASRModel.from_pretrained(
|
| 971 |
+
model_id,
|
| 972 |
+
dtype=dtype,
|
| 973 |
+
device_map=device if device == "cuda" else None,
|
| 974 |
+
)
|
| 975 |
+
asr_model = asr_wrapper.model
|
| 976 |
+
processor = asr_wrapper.processor
|
| 977 |
+
|
| 978 |
+
start_id = resolve_im_start_id(processor.tokenizer)
|
| 979 |
+
model = LatentQwenASR(
|
| 980 |
+
asr_model=asr_model,
|
| 981 |
+
processor=processor,
|
| 982 |
+
n_latent=0,
|
| 983 |
+
nt_token_id=-1,
|
| 984 |
+
lang_token_id=start_id,
|
| 985 |
+
transcribe_token_id=start_id,
|
| 986 |
+
freeze_base=False,
|
| 987 |
+
use_latent=False,
|
| 988 |
+
).to(device)
|
| 989 |
+
return model, processor
|
| 990 |
+
|
| 991 |
+
|
| 992 |
+
def maybe_attach_peft_adapter(
|
| 993 |
+
model: LatentQwenASR,
|
| 994 |
+
adapter_path: str,
|
| 995 |
+
mode_name: str,
|
| 996 |
+
expected_lora_rank: Optional[int] = None,
|
| 997 |
+
) -> None:
|
| 998 |
+
try:
|
| 999 |
+
from peft import PeftConfig, PeftModel
|
| 1000 |
+
except Exception as e:
|
| 1001 |
+
raise RuntimeError(
|
| 1002 |
+
f"[{mode_name}] Loading PEFT adapter requires `peft` package: {e}"
|
| 1003 |
+
) from e
|
| 1004 |
+
|
| 1005 |
+
peft_cfg = PeftConfig.from_pretrained(adapter_path)
|
| 1006 |
+
peft_type = str(getattr(peft_cfg, "peft_type", "unknown"))
|
| 1007 |
+
print(f"[{mode_name}] adapter detected: peft_type={peft_type} path={adapter_path}")
|
| 1008 |
+
|
| 1009 |
+
if expected_lora_rank is not None:
|
| 1010 |
+
rank = getattr(peft_cfg, "r", None)
|
| 1011 |
+
if rank is None:
|
| 1012 |
+
rank_pattern = getattr(peft_cfg, "rank_pattern", None)
|
| 1013 |
+
if isinstance(rank_pattern, dict) and rank_pattern:
|
| 1014 |
+
unique_ranks = sorted({int(v) for v in rank_pattern.values()})
|
| 1015 |
+
if len(unique_ranks) == 1:
|
| 1016 |
+
rank = unique_ranks[0]
|
| 1017 |
+
if rank is None:
|
| 1018 |
+
print(
|
| 1019 |
+
f"[warn][{mode_name}] unable to verify LoRA rank from adapter config."
|
| 1020 |
+
)
|
| 1021 |
+
elif int(rank) != int(expected_lora_rank):
|
| 1022 |
+
print(
|
| 1023 |
+
f"[warn][{mode_name}] expected LoRA rank={expected_lora_rank}, got rank={rank}."
|
| 1024 |
+
)
|
| 1025 |
+
|
| 1026 |
+
adapted = PeftModel.from_pretrained(
|
| 1027 |
+
model.thinker,
|
| 1028 |
+
adapter_path,
|
| 1029 |
+
is_trainable=False,
|
| 1030 |
+
)
|
| 1031 |
+
model.base_model.thinker = adapted
|
| 1032 |
+
model.thinker = model.base_model.thinker
|
| 1033 |
+
|
| 1034 |
+
|
| 1035 |
+
def sanitize_generation_config(model: LatentQwenASR) -> None:
|
| 1036 |
+
gen_cfg = getattr(model.thinker, "generation_config", None)
|
| 1037 |
+
if gen_cfg is None:
|
| 1038 |
+
return
|
| 1039 |
+
# Keep decoding deterministic and avoid warnings from stale sampling flags.
|
| 1040 |
+
if hasattr(gen_cfg, "do_sample"):
|
| 1041 |
+
gen_cfg.do_sample = False
|
| 1042 |
+
if hasattr(gen_cfg, "temperature"):
|
| 1043 |
+
gen_cfg.temperature = 1.0
|
| 1044 |
+
if hasattr(gen_cfg, "top_p"):
|
| 1045 |
+
gen_cfg.top_p = 1.0
|
| 1046 |
+
if hasattr(gen_cfg, "typical_p"):
|
| 1047 |
+
gen_cfg.typical_p = 1.0
|
| 1048 |
+
|
| 1049 |
+
|
| 1050 |
+
def build_base_model_bundle(
|
| 1051 |
+
model_id: str,
|
| 1052 |
+
device: str,
|
| 1053 |
+
dtype: torch.dtype,
|
| 1054 |
+
) -> ModelBundle:
|
| 1055 |
+
print(f"\n[Load base_model] model={model_id}")
|
| 1056 |
+
ckpt: Dict[str, any] = {}
|
| 1057 |
+
model, processor = build_plain_bundle(
|
| 1058 |
+
model_id=model_id,
|
| 1059 |
+
device=device,
|
| 1060 |
+
dtype=dtype,
|
| 1061 |
+
)
|
| 1062 |
+
|
| 1063 |
+
sanitize_generation_config(model)
|
| 1064 |
+
model.eval()
|
| 1065 |
+
return ModelBundle(model=model, processor=processor, checkpoint=ckpt)
|
| 1066 |
+
|
| 1067 |
+
|
| 1068 |
+
def build_baseline_ft_bundle(
|
| 1069 |
+
model_id: str,
|
| 1070 |
+
checkpoint_path: str,
|
| 1071 |
+
device: str,
|
| 1072 |
+
dtype: torch.dtype,
|
| 1073 |
+
) -> ModelBundle:
|
| 1074 |
+
print(f"\n[Load baseline_ft] model={model_id} ckpt={checkpoint_path}")
|
| 1075 |
+
ckpt = load_checkpoint(checkpoint_path)
|
| 1076 |
+
model, processor = build_plain_bundle(
|
| 1077 |
+
model_id=model_id,
|
| 1078 |
+
device=device,
|
| 1079 |
+
dtype=dtype,
|
| 1080 |
+
)
|
| 1081 |
+
|
| 1082 |
+
state = ckpt.get("model_state_dict", ckpt)
|
| 1083 |
+
missing, unexpected = model.load_state_dict(state, strict=False)
|
| 1084 |
+
if missing:
|
| 1085 |
+
print(f"[baseline_ft] missing keys: {len(missing)}")
|
| 1086 |
+
if unexpected:
|
| 1087 |
+
print(f"[baseline_ft] unexpected keys: {len(unexpected)}")
|
| 1088 |
+
|
| 1089 |
+
sanitize_generation_config(model)
|
| 1090 |
+
model.eval()
|
| 1091 |
+
return ModelBundle(model=model, processor=processor, checkpoint=ckpt)
|
| 1092 |
+
|
| 1093 |
+
|
| 1094 |
+
def build_prompt_tuning_bundle(
|
| 1095 |
+
model_id: str,
|
| 1096 |
+
checkpoint_path: str,
|
| 1097 |
+
device: str,
|
| 1098 |
+
dtype: torch.dtype,
|
| 1099 |
+
) -> ModelBundle:
|
| 1100 |
+
print(f"\n[Load prompt_tuning] model={model_id} ckpt={checkpoint_path}")
|
| 1101 |
+
if os.path.isdir(checkpoint_path):
|
| 1102 |
+
# Backward compatibility: old prompt-tuning PEFT adapters.
|
| 1103 |
+
model, processor = build_plain_bundle(
|
| 1104 |
+
model_id=model_id,
|
| 1105 |
+
device=device,
|
| 1106 |
+
dtype=dtype,
|
| 1107 |
+
)
|
| 1108 |
+
maybe_attach_peft_adapter(
|
| 1109 |
+
model=model,
|
| 1110 |
+
adapter_path=checkpoint_path,
|
| 1111 |
+
mode_name="prompt_tuning",
|
| 1112 |
+
)
|
| 1113 |
+
ckpt = {
|
| 1114 |
+
"adapter_path": checkpoint_path,
|
| 1115 |
+
"checkpoint_type": "peft_adapter",
|
| 1116 |
+
}
|
| 1117 |
+
else:
|
| 1118 |
+
ckpt = load_checkpoint(checkpoint_path)
|
| 1119 |
+
soft_prompt = extract_soft_prompt_tensor(ckpt)
|
| 1120 |
+
|
| 1121 |
+
if soft_prompt is not None:
|
| 1122 |
+
asr_wrapper = Qwen3ASRModel.from_pretrained(
|
| 1123 |
+
model_id,
|
| 1124 |
+
dtype=dtype,
|
| 1125 |
+
device_map=device if device == "cuda" else None,
|
| 1126 |
+
)
|
| 1127 |
+
asr_model = asr_wrapper.model
|
| 1128 |
+
processor = asr_wrapper.processor
|
| 1129 |
+
|
| 1130 |
+
nt_id = ensure_front_prompt_token(asr_model, processor.tokenizer)
|
| 1131 |
+
start_id = resolve_im_start_id(processor.tokenizer)
|
| 1132 |
+
n_latent = int(ckpt.get("n_latent", int(soft_prompt.size(0))))
|
| 1133 |
+
if n_latent <= 0:
|
| 1134 |
+
n_latent = int(soft_prompt.size(0))
|
| 1135 |
+
if n_latent <= 0:
|
| 1136 |
+
raise ValueError("prompt_tuning checkpoint has empty soft prompt tensor.")
|
| 1137 |
+
|
| 1138 |
+
model = LatentQwenASR(
|
| 1139 |
+
asr_model=asr_model,
|
| 1140 |
+
processor=processor,
|
| 1141 |
+
n_latent=n_latent,
|
| 1142 |
+
nt_token_id=nt_id,
|
| 1143 |
+
lang_token_id=start_id,
|
| 1144 |
+
transcribe_token_id=start_id,
|
| 1145 |
+
freeze_base=False,
|
| 1146 |
+
use_latent=False,
|
| 1147 |
+
use_soft_prompt=True,
|
| 1148 |
+
soft_prompt_init_mode="random",
|
| 1149 |
+
).to(device)
|
| 1150 |
+
|
| 1151 |
+
prompt = soft_prompt.detach()
|
| 1152 |
+
if prompt.dim() != 2:
|
| 1153 |
+
raise ValueError(
|
| 1154 |
+
f"soft_prompt_embed must be 2D [n_latent, hidden], got shape={tuple(prompt.shape)}"
|
| 1155 |
+
)
|
| 1156 |
+
if prompt.size(0) < n_latent:
|
| 1157 |
+
reps = (n_latent + prompt.size(0) - 1) // prompt.size(0)
|
| 1158 |
+
prompt = prompt.repeat((reps, 1))
|
| 1159 |
+
prompt = prompt[:n_latent]
|
| 1160 |
+
model.soft_prompt_embed.data.copy_(
|
| 1161 |
+
prompt.to(
|
| 1162 |
+
device=model.soft_prompt_embed.device,
|
| 1163 |
+
dtype=model.soft_prompt_embed.dtype,
|
| 1164 |
+
)
|
| 1165 |
+
)
|
| 1166 |
+
print(f"[prompt_tuning] loaded front soft prompt: n_latent={n_latent}")
|
| 1167 |
+
else:
|
| 1168 |
+
# Fallback for legacy `.pth` checkpoints that saved full wrapper state.
|
| 1169 |
+
model, processor = build_plain_bundle(
|
| 1170 |
+
model_id=model_id,
|
| 1171 |
+
device=device,
|
| 1172 |
+
dtype=dtype,
|
| 1173 |
+
)
|
| 1174 |
+
state = ckpt.get("model_state_dict", ckpt)
|
| 1175 |
+
missing, unexpected = model.load_state_dict(state, strict=False)
|
| 1176 |
+
if missing:
|
| 1177 |
+
print(f"[prompt_tuning] missing keys: {len(missing)}")
|
| 1178 |
+
if unexpected:
|
| 1179 |
+
print(f"[prompt_tuning] unexpected keys: {len(unexpected)}")
|
| 1180 |
+
|
| 1181 |
+
sanitize_generation_config(model)
|
| 1182 |
+
model.eval()
|
| 1183 |
+
return ModelBundle(model=model, processor=processor, checkpoint=ckpt)
|
| 1184 |
+
|
| 1185 |
+
|
| 1186 |
+
def build_lora_r16_bundle(
|
| 1187 |
+
model_id: str,
|
| 1188 |
+
checkpoint_path: str,
|
| 1189 |
+
device: str,
|
| 1190 |
+
dtype: torch.dtype,
|
| 1191 |
+
) -> ModelBundle:
|
| 1192 |
+
print(f"\n[Load lora_r16] model={model_id} ckpt={checkpoint_path}")
|
| 1193 |
+
model, processor = build_plain_bundle(
|
| 1194 |
+
model_id=model_id,
|
| 1195 |
+
device=device,
|
| 1196 |
+
dtype=dtype,
|
| 1197 |
+
)
|
| 1198 |
+
|
| 1199 |
+
if os.path.isdir(checkpoint_path):
|
| 1200 |
+
maybe_attach_peft_adapter(
|
| 1201 |
+
model=model,
|
| 1202 |
+
adapter_path=checkpoint_path,
|
| 1203 |
+
mode_name="lora_r16",
|
| 1204 |
+
expected_lora_rank=None,
|
| 1205 |
+
)
|
| 1206 |
+
ckpt = {
|
| 1207 |
+
"adapter_path": checkpoint_path,
|
| 1208 |
+
"checkpoint_type": "peft_adapter",
|
| 1209 |
+
}
|
| 1210 |
+
else:
|
| 1211 |
+
ckpt = load_checkpoint(checkpoint_path)
|
| 1212 |
+
state = ckpt.get("model_state_dict", ckpt)
|
| 1213 |
+
missing, unexpected = model.load_state_dict(state, strict=False)
|
| 1214 |
+
if missing:
|
| 1215 |
+
print(f"[lora_r16] missing keys: {len(missing)}")
|
| 1216 |
+
if unexpected:
|
| 1217 |
+
print(f"[lora_r16] unexpected keys: {len(unexpected)}")
|
| 1218 |
+
|
| 1219 |
+
sanitize_generation_config(model)
|
| 1220 |
+
model.eval()
|
| 1221 |
+
return ModelBundle(model=model, processor=processor, checkpoint=ckpt)
|
| 1222 |
+
|
| 1223 |
+
|
| 1224 |
+
def build_latent_bundle(
|
| 1225 |
+
model_id: str,
|
| 1226 |
+
checkpoint_path: str,
|
| 1227 |
+
n_latent_override: int,
|
| 1228 |
+
device: str,
|
| 1229 |
+
dtype: torch.dtype,
|
| 1230 |
+
) -> ModelBundle:
|
| 1231 |
+
print(f"\n[Load latent_reasoning] model={model_id} ckpt={checkpoint_path}")
|
| 1232 |
+
ckpt = load_checkpoint(checkpoint_path)
|
| 1233 |
+
|
| 1234 |
+
n_latent = int(ckpt.get("n_latent", 4))
|
| 1235 |
+
if n_latent_override >= 0:
|
| 1236 |
+
n_latent = int(n_latent_override)
|
| 1237 |
+
if n_latent <= 0:
|
| 1238 |
+
raise ValueError(f"Invalid latent token count: {n_latent}")
|
| 1239 |
+
|
| 1240 |
+
asr_wrapper = Qwen3ASRModel.from_pretrained(
|
| 1241 |
+
model_id,
|
| 1242 |
+
dtype=dtype,
|
| 1243 |
+
device_map=device if device == "cuda" else None,
|
| 1244 |
+
)
|
| 1245 |
+
asr_model = asr_wrapper.model
|
| 1246 |
+
processor = asr_wrapper.processor
|
| 1247 |
+
|
| 1248 |
+
token_id = ensure_front_prompt_token(asr_model, processor.tokenizer)
|
| 1249 |
+
|
| 1250 |
+
start_id = resolve_im_start_id(processor.tokenizer)
|
| 1251 |
+
model = LatentQwenASR(
|
| 1252 |
+
asr_model=asr_model,
|
| 1253 |
+
processor=processor,
|
| 1254 |
+
n_latent=n_latent,
|
| 1255 |
+
nt_token_id=int(token_id),
|
| 1256 |
+
lang_token_id=start_id,
|
| 1257 |
+
transcribe_token_id=start_id,
|
| 1258 |
+
freeze_base=False,
|
| 1259 |
+
use_latent=True,
|
| 1260 |
+
latent_use_bounded_delta=bool(ckpt.get("latent_use_bounded_delta", True)),
|
| 1261 |
+
latent_use_injection_gate=bool(ckpt.get("latent_use_injection_gate", True)),
|
| 1262 |
+
latent_use_embedding_anchor=bool(ckpt.get("latent_use_embedding_anchor", True)),
|
| 1263 |
+
).to(device)
|
| 1264 |
+
|
| 1265 |
+
# Prefer latent-adapter payload format from train.py.
|
| 1266 |
+
if "init_proj" in ckpt:
|
| 1267 |
+
model.init_proj.load_state_dict(ckpt["init_proj"], strict=True)
|
| 1268 |
+
model.delta_proj.load_state_dict(ckpt["delta_proj"], strict=True)
|
| 1269 |
+
model.step_proj.load_state_dict(ckpt["step_proj"], strict=True)
|
| 1270 |
+
if "step_embed" in ckpt:
|
| 1271 |
+
model.step_embed.data.copy_(
|
| 1272 |
+
ckpt["step_embed"].to(
|
| 1273 |
+
device=model.step_embed.device,
|
| 1274 |
+
dtype=model.step_embed.dtype,
|
| 1275 |
+
)
|
| 1276 |
+
)
|
| 1277 |
+
if "log_scale" in ckpt:
|
| 1278 |
+
model.log_scale.data.copy_(
|
| 1279 |
+
ckpt["log_scale"].to(
|
| 1280 |
+
device=model.log_scale.device,
|
| 1281 |
+
dtype=model.log_scale.dtype,
|
| 1282 |
+
)
|
| 1283 |
+
)
|
| 1284 |
+
if "value_head" in ckpt:
|
| 1285 |
+
model.value_head.load_state_dict(ckpt["value_head"], strict=True)
|
| 1286 |
+
if "injection_gate" in ckpt:
|
| 1287 |
+
model.injection_gate.load_state_dict(ckpt["injection_gate"], strict=True)
|
| 1288 |
+
else:
|
| 1289 |
+
# Fallback: checkpoint may contain full wrapper state_dict.
|
| 1290 |
+
state = ckpt.get("model_state_dict", ckpt)
|
| 1291 |
+
missing, unexpected = model.load_state_dict(state, strict=False)
|
| 1292 |
+
if missing:
|
| 1293 |
+
print(f"[latent_reasoning] missing keys: {len(missing)}")
|
| 1294 |
+
if unexpected:
|
| 1295 |
+
print(f"[latent_reasoning] unexpected keys: {len(unexpected)}")
|
| 1296 |
+
|
| 1297 |
+
sanitize_generation_config(model)
|
| 1298 |
+
model.eval()
|
| 1299 |
+
return ModelBundle(model=model, processor=processor, checkpoint=ckpt)
|
| 1300 |
+
|
| 1301 |
+
|
| 1302 |
+
_global_normalizer = None
|
| 1303 |
+
_global_normalizer_kind = "english"
|
| 1304 |
+
|
| 1305 |
+
|
| 1306 |
+
def configure_text_normalizer(kind: str) -> None:
|
| 1307 |
+
global _global_normalizer, _global_normalizer_kind
|
| 1308 |
+
kind = (kind or "english").strip().lower()
|
| 1309 |
+
if kind not in {"english", "basic", "simple"}:
|
| 1310 |
+
raise ValueError(f"Unknown text normalizer: {kind}")
|
| 1311 |
+
if kind != _global_normalizer_kind:
|
| 1312 |
+
_global_normalizer = None
|
| 1313 |
+
_global_normalizer_kind = kind
|
| 1314 |
+
|
| 1315 |
+
|
| 1316 |
+
def simple_text_normalize(text: str) -> str:
|
| 1317 |
+
text = unicodedata.normalize("NFKC", text or "")
|
| 1318 |
+
text = text.lower()
|
| 1319 |
+
text = "".join(
|
| 1320 |
+
" " if unicodedata.category(ch).startswith("P") else ch
|
| 1321 |
+
for ch in text
|
| 1322 |
+
)
|
| 1323 |
+
return re.sub(r"\s+", " ", text).strip()
|
| 1324 |
+
|
| 1325 |
+
|
| 1326 |
+
class _SimpleTextNormalizer:
|
| 1327 |
+
def __call__(self, text: str) -> str:
|
| 1328 |
+
return simple_text_normalize(text)
|
| 1329 |
+
|
| 1330 |
+
|
| 1331 |
+
def build_text_normalizer(kind: str):
|
| 1332 |
+
if kind == "simple":
|
| 1333 |
+
return _SimpleTextNormalizer()
|
| 1334 |
+
|
| 1335 |
+
if kind == "basic":
|
| 1336 |
+
try:
|
| 1337 |
+
from transformers.models.whisper.english_normalizer import BasicTextNormalizer
|
| 1338 |
+
return BasicTextNormalizer()
|
| 1339 |
+
except Exception:
|
| 1340 |
+
return _SimpleTextNormalizer()
|
| 1341 |
+
|
| 1342 |
+
try:
|
| 1343 |
+
# We use the standard Whisper English normalizer for English-heavy legacy runs.
|
| 1344 |
+
from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
|
| 1345 |
+
import json, urllib.request, os
|
| 1346 |
+
|
| 1347 |
+
mapping_path = os.path.join(os.path.dirname(__file__), "english_normalizer.json")
|
| 1348 |
+
if not os.path.exists(mapping_path):
|
| 1349 |
+
url = "https://raw.githubusercontent.com/huggingface/transformers/main/src/transformers/models/whisper/english_normalizer.json"
|
| 1350 |
+
req = urllib.request.Request(url)
|
| 1351 |
+
with urllib.request.urlopen(req) as response:
|
| 1352 |
+
data = json.loads(response.read().decode("utf-8"))
|
| 1353 |
+
with open(mapping_path, "w") as f:
|
| 1354 |
+
json.dump(data, f)
|
| 1355 |
+
else:
|
| 1356 |
+
with open(mapping_path, "r") as f:
|
| 1357 |
+
data = json.load(f)
|
| 1358 |
+
return EnglishTextNormalizer(data)
|
| 1359 |
+
except Exception:
|
| 1360 |
+
try:
|
| 1361 |
+
from transformers.models.whisper.english_normalizer import BasicTextNormalizer
|
| 1362 |
+
return BasicTextNormalizer()
|
| 1363 |
+
except Exception:
|
| 1364 |
+
return _SimpleTextNormalizer()
|
| 1365 |
+
|
| 1366 |
+
def normalize_text(text: str) -> str:
|
| 1367 |
+
global _global_normalizer
|
| 1368 |
+
if _global_normalizer is None:
|
| 1369 |
+
_global_normalizer = build_text_normalizer(_global_normalizer_kind)
|
| 1370 |
+
normalized = _global_normalizer(text)
|
| 1371 |
+
|
| 1372 |
+
if isinstance(normalized, str):
|
| 1373 |
+
if _global_normalizer_kind == "english":
|
| 1374 |
+
# LibriSpeech specific normalizations: expand common contractions
|
| 1375 |
+
# that Whisper leaves as "m" or "re".
|
| 1376 |
+
contractions = {
|
| 1377 |
+
r"\bi m\b": "i am",
|
| 1378 |
+
r"\byou re\b": "you are",
|
| 1379 |
+
r"\bwe re\b": "we are",
|
| 1380 |
+
r"\bthey re\b": "they are",
|
| 1381 |
+
r"\bhe s\b": "he is",
|
| 1382 |
+
r"\bshe s\b": "she is",
|
| 1383 |
+
r"\bit s\b": "it is",
|
| 1384 |
+
r"\bthat s\b": "that is",
|
| 1385 |
+
r"\bwho s\b": "who is",
|
| 1386 |
+
r"\bwhat s\b": "what is",
|
| 1387 |
+
r"\bthere s\b": "there is",
|
| 1388 |
+
r"\bhere s\b": "here is",
|
| 1389 |
+
r"\bwhere s\b": "where is",
|
| 1390 |
+
r"\blet s\b": "let us",
|
| 1391 |
+
r"\bhow s\b": "how is",
|
| 1392 |
+
# Standard possessive stripping (boy's -> boy s -> boys)
|
| 1393 |
+
r"\b s\b": "s",
|
| 1394 |
+
}
|
| 1395 |
+
for pat, repl in contractions.items():
|
| 1396 |
+
normalized = re.sub(pat, repl, normalized)
|
| 1397 |
+
|
| 1398 |
+
return re.sub(r"\s+", " ", normalized).strip()
|
| 1399 |
+
return ""
|
| 1400 |
+
|
| 1401 |
+
def clean_prediction(text: str) -> str:
|
| 1402 |
+
text = text or ""
|
| 1403 |
+
text = re.sub(r"language\s+\w+<asr_text>", "", text, flags=re.IGNORECASE)
|
| 1404 |
+
if "<asr_text>" in text:
|
| 1405 |
+
text = text.split("<asr_text>", 1)[-1]
|
| 1406 |
+
return text.strip()
|
| 1407 |
+
|
| 1408 |
+
|
| 1409 |
+
@torch.no_grad()
|
| 1410 |
+
def add_noise_at_snr(waveform: np.ndarray, snr_db: float) -> np.ndarray:
|
| 1411 |
+
"""Add Gaussian noise to a waveform at a target Signal-to-Noise Ratio (dB).
|
| 1412 |
+
|
| 1413 |
+
SNR_dB = 10 * log10(P_signal / P_noise)
|
| 1414 |
+
=> P_noise = P_signal / 10^(SNR_dB / 10)
|
| 1415 |
+
=> std_noise = sqrt(P_noise)
|
| 1416 |
+
"""
|
| 1417 |
+
signal_power = np.mean(waveform ** 2)
|
| 1418 |
+
if signal_power < 1e-10: # silence guard
|
| 1419 |
+
return waveform
|
| 1420 |
+
noise_power = signal_power / (10.0 ** (snr_db / 10.0))
|
| 1421 |
+
noise = np.random.normal(0, np.sqrt(noise_power), waveform.shape)
|
| 1422 |
+
return (waveform + noise).astype(waveform.dtype)
|
| 1423 |
+
|
| 1424 |
+
|
| 1425 |
+
def evaluate_config(
|
| 1426 |
+
model: LatentQwenASR,
|
| 1427 |
+
processor: any,
|
| 1428 |
+
dataset_name: str,
|
| 1429 |
+
config_name: str,
|
| 1430 |
+
split: str,
|
| 1431 |
+
use_baseline: bool,
|
| 1432 |
+
mode_name: str,
|
| 1433 |
+
max_samples: int,
|
| 1434 |
+
max_new_tokens: int,
|
| 1435 |
+
print_samples: int,
|
| 1436 |
+
language_hint: Optional[str],
|
| 1437 |
+
prompt_template: str,
|
| 1438 |
+
disable_language_hint: bool,
|
| 1439 |
+
num_beams: int = 1,
|
| 1440 |
+
dynamic_halt_threshold: float = 0.0,
|
| 1441 |
+
snr_db: Optional[float] = None,
|
| 1442 |
+
print_gate_skips: bool = False,
|
| 1443 |
+
streaming: bool = False,
|
| 1444 |
+
) -> Dict[str, any]:
|
| 1445 |
+
raw_ds = load_dataset(
|
| 1446 |
+
dataset_name,
|
| 1447 |
+
config_name,
|
| 1448 |
+
split=split,
|
| 1449 |
+
trust_remote_code=True,
|
| 1450 |
+
streaming=streaming,
|
| 1451 |
+
)
|
| 1452 |
+
known_total: Optional[int]
|
| 1453 |
+
try:
|
| 1454 |
+
known_total = len(raw_ds)
|
| 1455 |
+
except TypeError:
|
| 1456 |
+
known_total = None
|
| 1457 |
+
total = known_total if known_total is not None else 0
|
| 1458 |
+
if max_samples > 0:
|
| 1459 |
+
total = min(total, max_samples) if known_total is not None else max_samples
|
| 1460 |
+
column_names = dataset_column_names(raw_ds)
|
| 1461 |
+
text_column = resolve_text_column(dataset_name, column_names)
|
| 1462 |
+
audio_column = "audio" if "audio" in column_names else None
|
| 1463 |
+
|
| 1464 |
+
if audio_column is None and not streaming:
|
| 1465 |
+
raise ValueError(
|
| 1466 |
+
f"Dataset {dataset_name}/{config_name} does not have an 'audio' column. "
|
| 1467 |
+
f"Columns: {column_names}"
|
| 1468 |
+
)
|
| 1469 |
+
target_sampling_rate = int(getattr(processor.feature_extractor, "sampling_rate", 0) or 0)
|
| 1470 |
+
if target_sampling_rate > 0 and audio_column is not None:
|
| 1471 |
+
raw_ds = raw_ds.cast_column(audio_column, Audio(sampling_rate=target_sampling_rate))
|
| 1472 |
+
|
| 1473 |
+
preds: List[str] = []
|
| 1474 |
+
refs: List[str] = []
|
| 1475 |
+
shown = 0
|
| 1476 |
+
skipped = 0
|
| 1477 |
+
reasoned = 0
|
| 1478 |
+
|
| 1479 |
+
# Value Head aggregate statistics collectors
|
| 1480 |
+
all_v_preds: List[List[float]] = [] # per-sample list of v_pred values
|
| 1481 |
+
all_deq_iters: List[int] = [] # per-sample thinking steps taken
|
| 1482 |
+
all_gates: List[float] = [] # per-sample gate values
|
| 1483 |
+
all_scales: List[List[float]] = [] # per-sample scale values
|
| 1484 |
+
n_skipped_by_gate = 0 # count of N=0 (gate skip)
|
| 1485 |
+
|
| 1486 |
+
effective_language_hint = None if disable_language_hint else language_hint
|
| 1487 |
+
if effective_language_hint:
|
| 1488 |
+
try:
|
| 1489 |
+
prompt_text = prompt_template.format(language=effective_language_hint)
|
| 1490 |
+
except Exception:
|
| 1491 |
+
prompt_text = f"Transcribe the {effective_language_hint} audio into text."
|
| 1492 |
+
else:
|
| 1493 |
+
prompt_text = "Transcribe the audio into text."
|
| 1494 |
+
|
| 1495 |
+
sample_iter = iter(raw_ds)
|
| 1496 |
+
if max_samples > 0:
|
| 1497 |
+
sample_iter = itertools.islice(sample_iter, max_samples)
|
| 1498 |
+
pbar = tqdm(sample_iter, total=total if total > 0 else None, desc=f"{config_name}:{mode_name}", leave=False)
|
| 1499 |
+
seen = 0
|
| 1500 |
+
for i, sample in enumerate(pbar):
|
| 1501 |
+
seen += 1
|
| 1502 |
+
|
| 1503 |
+
if (text_column is None or audio_column is None) and isinstance(sample, dict):
|
| 1504 |
+
sample_columns = list(sample.keys())
|
| 1505 |
+
if text_column is None:
|
| 1506 |
+
text_column = resolve_text_column(dataset_name, sample_columns)
|
| 1507 |
+
if audio_column is None and "audio" in sample_columns:
|
| 1508 |
+
audio_column = "audio"
|
| 1509 |
+
|
| 1510 |
+
if audio_column is None:
|
| 1511 |
+
skipped += 1
|
| 1512 |
+
continue
|
| 1513 |
+
|
| 1514 |
+
if text_column is not None:
|
| 1515 |
+
ref = sample.get(text_column, "")
|
| 1516 |
+
else:
|
| 1517 |
+
ref = ""
|
| 1518 |
+
for cand in FALLBACK_TEXT_COLUMNS:
|
| 1519 |
+
v = sample.get(cand, "")
|
| 1520 |
+
if isinstance(v, str) and v.strip():
|
| 1521 |
+
ref = v
|
| 1522 |
+
break
|
| 1523 |
+
if not isinstance(ref, str) or not ref.strip():
|
| 1524 |
+
skipped += 1
|
| 1525 |
+
continue
|
| 1526 |
+
|
| 1527 |
+
audio = sample.get(audio_column)
|
| 1528 |
+
if not isinstance(audio, dict) or "array" not in audio or "sampling_rate" not in audio:
|
| 1529 |
+
skipped += 1
|
| 1530 |
+
continue
|
| 1531 |
+
|
| 1532 |
+
audio_array = np.array(audio["array"], dtype=np.float64)
|
| 1533 |
+
if snr_db is not None:
|
| 1534 |
+
audio_array = add_noise_at_snr(audio_array, snr_db)
|
| 1535 |
+
|
| 1536 |
+
target_dtype = model.thinker.dtype if hasattr(model.thinker, "dtype") else torch.float32
|
| 1537 |
+
feat_out = processor.feature_extractor(
|
| 1538 |
+
audio_array,
|
| 1539 |
+
sampling_rate=audio["sampling_rate"],
|
| 1540 |
+
return_attention_mask=True,
|
| 1541 |
+
)
|
| 1542 |
+
feats = torch.tensor(feat_out.input_features[0], dtype=target_dtype, device=model.base_model.device).unsqueeze(0)
|
| 1543 |
+
n_frames = feats.size(-1)
|
| 1544 |
+
|
| 1545 |
+
if getattr(feat_out, "attention_mask", None) is not None:
|
| 1546 |
+
# Coerce mask to match the Mel frame length (same logic as data.py)
|
| 1547 |
+
raw_mask = feat_out.attention_mask[0]
|
| 1548 |
+
if not isinstance(raw_mask, (list, torch.Tensor)):
|
| 1549 |
+
raw_mask = list(raw_mask)
|
| 1550 |
+
if isinstance(raw_mask, torch.Tensor):
|
| 1551 |
+
raw_mask = raw_mask.long()
|
| 1552 |
+
else:
|
| 1553 |
+
raw_mask = torch.tensor(raw_mask, dtype=torch.long)
|
| 1554 |
+
# Pad or truncate to n_frames
|
| 1555 |
+
if raw_mask.size(-1) < n_frames:
|
| 1556 |
+
raw_mask = torch.cat([raw_mask, torch.zeros(n_frames - raw_mask.size(-1), dtype=torch.long)])
|
| 1557 |
+
else:
|
| 1558 |
+
raw_mask = raw_mask[:n_frames]
|
| 1559 |
+
feature_attention_mask = raw_mask.to(device=model.base_model.device).unsqueeze(0)
|
| 1560 |
+
else:
|
| 1561 |
+
feature_attention_mask = torch.ones((1, n_frames), dtype=torch.long, device=model.base_model.device)
|
| 1562 |
+
|
| 1563 |
+
# Guard: if effective feature length is too short for audio tower, skip sample
|
| 1564 |
+
effective_len = feature_attention_mask.sum().item()
|
| 1565 |
+
if effective_len < 10:
|
| 1566 |
+
skipped += 1
|
| 1567 |
+
continue
|
| 1568 |
+
|
| 1569 |
+
gen_kwargs = {
|
| 1570 |
+
"feature_attention_mask": feature_attention_mask,
|
| 1571 |
+
"max_new_tokens": max_new_tokens,
|
| 1572 |
+
"use_baseline": use_baseline,
|
| 1573 |
+
"return_thoughts": False,
|
| 1574 |
+
"return_stats": True,
|
| 1575 |
+
"do_sample": False,
|
| 1576 |
+
"eos_token_id": [151645, 151643],
|
| 1577 |
+
"num_beams": num_beams,
|
| 1578 |
+
"language_hint": effective_language_hint,
|
| 1579 |
+
"prompt_text": prompt_text,
|
| 1580 |
+
"dynamic_halt_threshold": dynamic_halt_threshold,
|
| 1581 |
+
}
|
| 1582 |
+
|
| 1583 |
+
gen_output = model.generate(feats, **gen_kwargs)
|
| 1584 |
+
if isinstance(gen_output, tuple):
|
| 1585 |
+
gen_ids = gen_output[0]
|
| 1586 |
+
stats = gen_output[1] if isinstance(gen_output[1], dict) else (gen_output[2] if len(gen_output)>2 and isinstance(gen_output[2], dict) else {})
|
| 1587 |
+
else:
|
| 1588 |
+
gen_ids = gen_output
|
| 1589 |
+
stats = {}
|
| 1590 |
+
|
| 1591 |
+
ids = gen_ids[0]
|
| 1592 |
+
eos_id = processor.tokenizer.eos_token_id
|
| 1593 |
+
if eos_id is not None and (ids == eos_id).any():
|
| 1594 |
+
eos_pos = (ids == eos_id).nonzero(as_tuple=True)[0][0]
|
| 1595 |
+
ids = ids[:eos_pos]
|
| 1596 |
+
|
| 1597 |
+
pred_raw = processor.tokenizer.decode(ids, skip_special_tokens=True)
|
| 1598 |
+
pred = clean_prediction(pred_raw)
|
| 1599 |
+
|
| 1600 |
+
pred_norm = normalize_text(pred)
|
| 1601 |
+
ref_norm = normalize_text(ref)
|
| 1602 |
+
if not ref_norm:
|
| 1603 |
+
skipped += 1
|
| 1604 |
+
continue
|
| 1605 |
+
|
| 1606 |
+
preds.append(pred_norm)
|
| 1607 |
+
refs.append(ref_norm)
|
| 1608 |
+
|
| 1609 |
+
# Always print if it was skipped (N=0 from Gate) OR if under the print limit
|
| 1610 |
+
is_skipped_lr = stats.get("skipped", False)
|
| 1611 |
+
|
| 1612 |
+
if mode_name == "latent_reasoning":
|
| 1613 |
+
if not is_skipped_lr:
|
| 1614 |
+
reasoned += 1
|
| 1615 |
+
|
| 1616 |
+
if shown < print_samples or (print_gate_skips and is_skipped_lr):
|
| 1617 |
+
if not is_skipped_lr:
|
| 1618 |
+
shown += 1
|
| 1619 |
+
print(f"\n[{config_name} | {mode_name} | sample {shown if not is_skipped_lr else 'SKIPPED'}]")
|
| 1620 |
+
print(f" Ref : {ref.strip()}")
|
| 1621 |
+
print(f" Pred: {pred}")
|
| 1622 |
+
if "v_preds" in stats and stats["v_preds"] is not None:
|
| 1623 |
+
v_vals = stats["v_preds"][0].tolist()
|
| 1624 |
+
v_str = " | ".join([f"{v:.4f}" for v in v_vals])
|
| 1625 |
+
iters_t = stats.get("deq_iters", None)
|
| 1626 |
+
deq_iters = int(iters_t.item()) if iters_t is not None and hasattr(iters_t, "item") else -1
|
| 1627 |
+
if is_skipped_lr:
|
| 1628 |
+
print(f" v_preds: [{v_str}] -> SKIPPED (N=0)")
|
| 1629 |
+
else:
|
| 1630 |
+
print(f" v_preds: [{v_str}] -> {deq_iters} step(s)")
|
| 1631 |
+
|
| 1632 |
+
# --- Collect Value Head stats for aggregate summary ---
|
| 1633 |
+
if mode_name == "latent_reasoning":
|
| 1634 |
+
iters_t = stats.get("deq_iters", None)
|
| 1635 |
+
deq_n = int(iters_t.item()) if iters_t is not None and hasattr(iters_t, "item") else 0
|
| 1636 |
+
all_deq_iters.append(deq_n)
|
| 1637 |
+
if is_skipped_lr:
|
| 1638 |
+
n_skipped_by_gate += 1
|
| 1639 |
+
if "v_preds" in stats and stats["v_preds"] is not None:
|
| 1640 |
+
all_v_preds.append(stats["v_preds"][0].tolist())
|
| 1641 |
+
if "gate" in stats and stats["gate"] is not None:
|
| 1642 |
+
g = stats["gate"]
|
| 1643 |
+
all_gates.append(float(g[0].item()) if hasattr(g, "item") or g.dim() > 0 else float(g))
|
| 1644 |
+
if "scales" in stats and stats["scales"] is not None:
|
| 1645 |
+
s = stats["scales"]
|
| 1646 |
+
all_scales.append(s.tolist() if hasattr(s, "tolist") else list(s))
|
| 1647 |
+
|
| 1648 |
+
# --- Print Value Head aggregate summary ---
|
| 1649 |
+
if mode_name == "latent_reasoning" and all_deq_iters:
|
| 1650 |
+
import collections
|
| 1651 |
+
n_total_eval = len(all_deq_iters)
|
| 1652 |
+
n_dist = collections.Counter(all_deq_iters)
|
| 1653 |
+
print(f"\n{'='*60}")
|
| 1654 |
+
print(f"[Value-Head Summary] {config_name} | {mode_name}")
|
| 1655 |
+
print(f"{'='*60}")
|
| 1656 |
+
print(f" Total samples evaluated : {n_total_eval}")
|
| 1657 |
+
print(f" Skipped by gate (N=0) : {n_skipped_by_gate} ({100*n_skipped_by_gate/max(1,n_total_eval):.1f}%)")
|
| 1658 |
+
print(f" Reasoned (N>0) : {n_total_eval - n_skipped_by_gate} ({100*(n_total_eval - n_skipped_by_gate)/max(1,n_total_eval):.1f}%)")
|
| 1659 |
+
print(f" --- N-step distribution ---")
|
| 1660 |
+
for k in sorted(n_dist.keys()):
|
| 1661 |
+
pct = 100 * n_dist[k] / n_total_eval
|
| 1662 |
+
bar = '█' * int(pct / 2)
|
| 1663 |
+
print(f" N={k}: {n_dist[k]:5d} ({pct:5.1f}%) {bar}")
|
| 1664 |
+
if all_v_preds:
|
| 1665 |
+
flat_v = [v for vs in all_v_preds for v in vs]
|
| 1666 |
+
if flat_v:
|
| 1667 |
+
import numpy as _np
|
| 1668 |
+
v_arr = _np.array(flat_v)
|
| 1669 |
+
print(f" --- v_pred statistics (all steps, all samples) ---")
|
| 1670 |
+
print(f" count : {len(v_arr)}")
|
| 1671 |
+
print(f" mean : {v_arr.mean():.4f}")
|
| 1672 |
+
print(f" std : {v_arr.std():.4f}")
|
| 1673 |
+
print(f" min : {v_arr.min():.4f}")
|
| 1674 |
+
print(f" max : {v_arr.max():.4f}")
|
| 1675 |
+
print(f" median: {float(_np.median(v_arr)):.4f}")
|
| 1676 |
+
# Histogram buckets
|
| 1677 |
+
edges = [-1.0, -0.8, -0.5, -0.2, 0.0, 0.2, 0.5, 0.8, 1.0]
|
| 1678 |
+
print(f" --- v_pred histogram ---")
|
| 1679 |
+
for j in range(len(edges) - 1):
|
| 1680 |
+
lo, hi = edges[j], edges[j + 1]
|
| 1681 |
+
cnt = int(((v_arr >= lo) & (v_arr < hi)).sum())
|
| 1682 |
+
pct = 100 * cnt / len(v_arr)
|
| 1683 |
+
bar = '▓' * int(pct / 2)
|
| 1684 |
+
print(f" [{lo:+.1f}, {hi:+.1f}): {cnt:5d} ({pct:5.1f}%) {bar}")
|
| 1685 |
+
# Also count == 1.0
|
| 1686 |
+
cnt_one = int((v_arr >= 0.8).sum())
|
| 1687 |
+
# First-step v_pred (initial gate signal)
|
| 1688 |
+
first_v = [vs[0] for vs in all_v_preds if vs]
|
| 1689 |
+
if first_v:
|
| 1690 |
+
fv = _np.array(first_v)
|
| 1691 |
+
print(f" --- v_pred[0] (initial gate signal) ---")
|
| 1692 |
+
print(f" mean : {fv.mean():.4f}")
|
| 1693 |
+
print(f" std : {fv.std():.4f}")
|
| 1694 |
+
print(f" < 0 : {int((fv < 0).sum()):5d} ({100*(fv < 0).mean():.1f}%)")
|
| 1695 |
+
print(f" >= 0 : {int((fv >= 0).sum()):5d} ({100*(fv >= 0).mean():.1f}%)")
|
| 1696 |
+
if all_gates:
|
| 1697 |
+
import numpy as _np
|
| 1698 |
+
g_arr = _np.array(all_gates)
|
| 1699 |
+
print(f" --- Gate confidence (mapped to [0,1]) ---")
|
| 1700 |
+
print(f" mean : {g_arr.mean():.4f}")
|
| 1701 |
+
print(f" std : {g_arr.std():.4f}")
|
| 1702 |
+
print(f" min : {g_arr.min():.4f}")
|
| 1703 |
+
print(f" max : {g_arr.max():.4f}")
|
| 1704 |
+
if all_scales:
|
| 1705 |
+
import numpy as _np
|
| 1706 |
+
s_arr = _np.array(all_scales)
|
| 1707 |
+
if s_arr.ndim == 2:
|
| 1708 |
+
print(f" --- Delta scales per step ---")
|
| 1709 |
+
for si in range(s_arr.shape[1]):
|
| 1710 |
+
col = s_arr[:, si]
|
| 1711 |
+
print(f" step {si}: mean={col.mean():.4f} std={col.std():.4f}")
|
| 1712 |
+
elif s_arr.ndim == 1:
|
| 1713 |
+
print(f" --- Delta scales ---")
|
| 1714 |
+
print(f" mean={s_arr.mean():.4f} std={s_arr.std():.4f}")
|
| 1715 |
+
print(f"{'='*60}\n")
|
| 1716 |
+
|
| 1717 |
+
if not refs:
|
| 1718 |
+
return {
|
| 1719 |
+
"config": config_name,
|
| 1720 |
+
"mode": mode_name,
|
| 1721 |
+
"language_hint": effective_language_hint,
|
| 1722 |
+
"text_column": text_column,
|
| 1723 |
+
"samples_total": total if total > 0 else seen,
|
| 1724 |
+
"samples_used": 0,
|
| 1725 |
+
"samples_skipped": skipped,
|
| 1726 |
+
"samples_reasoned": reasoned,
|
| 1727 |
+
"wer": 1.0,
|
| 1728 |
+
"cer": 1.0,
|
| 1729 |
+
}
|
| 1730 |
+
|
| 1731 |
+
return {
|
| 1732 |
+
"config": config_name,
|
| 1733 |
+
"mode": mode_name,
|
| 1734 |
+
"language_hint": effective_language_hint,
|
| 1735 |
+
"text_column": text_column,
|
| 1736 |
+
"samples_total": total if total > 0 else seen,
|
| 1737 |
+
"samples_used": len(refs),
|
| 1738 |
+
"samples_skipped": skipped,
|
| 1739 |
+
"samples_reasoned": reasoned,
|
| 1740 |
+
"wer": float(wer(refs, preds)),
|
| 1741 |
+
"cer": float(cer(refs, preds)),
|
| 1742 |
+
}
|
| 1743 |
+
|
| 1744 |
+
|
| 1745 |
+
def weighted_avg(rows: List[Dict[str, any]], metric_key: str) -> Optional[float]:
|
| 1746 |
+
numer = 0.0
|
| 1747 |
+
denom = 0
|
| 1748 |
+
for row in rows:
|
| 1749 |
+
n = int(row.get("samples_used", 0))
|
| 1750 |
+
if n <= 0:
|
| 1751 |
+
continue
|
| 1752 |
+
m = row.get(metric_key)
|
| 1753 |
+
if m is None:
|
| 1754 |
+
continue
|
| 1755 |
+
numer += float(m) * n
|
| 1756 |
+
denom += n
|
| 1757 |
+
if denom == 0:
|
| 1758 |
+
return None
|
| 1759 |
+
return numer / float(denom)
|
| 1760 |
+
|
| 1761 |
+
|
| 1762 |
+
def print_summary_table(rows: List[Dict[str, any]]) -> None:
|
| 1763 |
+
header = (
|
| 1764 |
+
f"{'Config':<12} {'N':>7} "
|
| 1765 |
+
f"{'Raw-WER':>10} {'FT-WER':>10} {'PT-WER':>10} {'LoRA-WER':>10} {'Lat-WER':>10} "
|
| 1766 |
+
f"{'Raw-CER':>10} {'FT-CER':>10} {'PT-CER':>10} {'LoRA-CER':>10} {'Lat-CER':>10}"
|
| 1767 |
+
)
|
| 1768 |
+
print("\n" + header)
|
| 1769 |
+
print("-" * len(header))
|
| 1770 |
+
|
| 1771 |
+
for row in rows:
|
| 1772 |
+
cfg = row["config"]
|
| 1773 |
+
n = row.get("samples_used", 0)
|
| 1774 |
+
raw_wer = row.get("base_model_wer")
|
| 1775 |
+
ft_wer = row.get("baseline_ft_wer")
|
| 1776 |
+
pt_wer = row.get("prompt_tuning_wer")
|
| 1777 |
+
lora_wer = row.get("lora_r16_wer")
|
| 1778 |
+
lat_wer = row.get("latent_reasoning_wer")
|
| 1779 |
+
raw_cer = row.get("base_model_cer")
|
| 1780 |
+
ft_cer = row.get("baseline_ft_cer")
|
| 1781 |
+
pt_cer = row.get("prompt_tuning_cer")
|
| 1782 |
+
lora_cer = row.get("lora_r16_cer")
|
| 1783 |
+
lat_cer = row.get("latent_reasoning_cer")
|
| 1784 |
+
|
| 1785 |
+
print(
|
| 1786 |
+
f"{cfg:<12} {n:>7d} "
|
| 1787 |
+
f"{(f'{raw_wer:.4f}' if raw_wer is not None else '-'):>10} "
|
| 1788 |
+
f"{(f'{ft_wer:.4f}' if ft_wer is not None else '-'):>10} "
|
| 1789 |
+
f"{(f'{pt_wer:.4f}' if pt_wer is not None else '-'):>10} "
|
| 1790 |
+
f"{(f'{lora_wer:.4f}' if lora_wer is not None else '-'):>10} "
|
| 1791 |
+
f"{(f'{lat_wer:.4f}' if lat_wer is not None else '-'):>10} "
|
| 1792 |
+
f"{(f'{raw_cer:.4f}' if raw_cer is not None else '-'):>10} "
|
| 1793 |
+
f"{(f'{ft_cer:.4f}' if ft_cer is not None else '-'):>10} "
|
| 1794 |
+
f"{(f'{pt_cer:.4f}' if pt_cer is not None else '-'):>10} "
|
| 1795 |
+
f"{(f'{lora_cer:.4f}' if lora_cer is not None else '-'):>10} "
|
| 1796 |
+
f"{(f'{lat_cer:.4f}' if lat_cer is not None else '-'):>10}"
|
| 1797 |
+
)
|
| 1798 |
+
|
| 1799 |
+
|
| 1800 |
+
def release_model(bundle: Optional[ModelBundle]) -> None:
|
| 1801 |
+
if bundle is None:
|
| 1802 |
+
return
|
| 1803 |
+
del bundle.model
|
| 1804 |
+
gc.collect()
|
| 1805 |
+
if torch.cuda.is_available():
|
| 1806 |
+
torch.cuda.empty_cache()
|
| 1807 |
+
|
| 1808 |
+
|
| 1809 |
+
def main() -> None:
|
| 1810 |
+
args = parse_args()
|
| 1811 |
+
set_seed(args.seed)
|
| 1812 |
+
configure_text_normalizer(args.text_normalizer)
|
| 1813 |
+
|
| 1814 |
+
args.baseline_ckpt = (args.baseline_ckpt or "").strip()
|
| 1815 |
+
args.prompt_tuning_ckpt = (args.prompt_tuning_ckpt or "").strip()
|
| 1816 |
+
args.lora_r16_ckpt = (args.lora_r16_ckpt or "").strip()
|
| 1817 |
+
args.latent_ckpt = (args.latent_ckpt or "").strip()
|
| 1818 |
+
run_base_model = not bool(args.skip_base_model)
|
| 1819 |
+
run_baseline_ft = bool(args.baseline_ckpt) and not bool(args.skip_baseline_ft)
|
| 1820 |
+
run_prompt_tuning = bool(args.prompt_tuning_ckpt) and not bool(args.skip_prompt_tuning)
|
| 1821 |
+
run_lora_r16 = bool(args.lora_r16_ckpt) and not bool(args.skip_lora_r16)
|
| 1822 |
+
run_latent_reasoning = bool(args.latent_ckpt)
|
| 1823 |
+
if not (
|
| 1824 |
+
run_base_model
|
| 1825 |
+
or run_baseline_ft
|
| 1826 |
+
or run_prompt_tuning
|
| 1827 |
+
or run_lora_r16
|
| 1828 |
+
or run_latent_reasoning
|
| 1829 |
+
):
|
| 1830 |
+
raise ValueError(
|
| 1831 |
+
"No evaluation mode enabled. Provide at least one of --baseline-ckpt, "
|
| 1832 |
+
"--prompt-tuning-ckpt, --lora-r16-ckpt, --latent-ckpt, "
|
| 1833 |
+
"or remove --skip-base-model."
|
| 1834 |
+
)
|
| 1835 |
+
if run_baseline_ft and not os.path.isfile(args.baseline_ckpt):
|
| 1836 |
+
raise FileNotFoundError(
|
| 1837 |
+
f"--baseline-ckpt expects a .pth file, got: {args.baseline_ckpt}"
|
| 1838 |
+
)
|
| 1839 |
+
if run_latent_reasoning and not os.path.isfile(args.latent_ckpt):
|
| 1840 |
+
raise FileNotFoundError(
|
| 1841 |
+
f"--latent-ckpt expects a .pth file, got: {args.latent_ckpt}"
|
| 1842 |
+
)
|
| 1843 |
+
if run_prompt_tuning and not is_path_file_or_dir(args.prompt_tuning_ckpt):
|
| 1844 |
+
raise FileNotFoundError(
|
| 1845 |
+
f"--prompt-tuning-ckpt path not found: {args.prompt_tuning_ckpt}"
|
| 1846 |
+
)
|
| 1847 |
+
if run_lora_r16 and not is_path_file_or_dir(args.lora_r16_ckpt):
|
| 1848 |
+
raise FileNotFoundError(
|
| 1849 |
+
f"--lora-r16-ckpt path not found: {args.lora_r16_ckpt}"
|
| 1850 |
+
)
|
| 1851 |
+
|
| 1852 |
+
dataset_name = resolve_dataset_name(args.dataset_name)
|
| 1853 |
+
device = choose_device(args.device)
|
| 1854 |
+
dtype = choose_dtype(args.dtype, device)
|
| 1855 |
+
configs = resolve_configs(dataset_name, args.configs, args.all_configs, args.max_configs)
|
| 1856 |
+
|
| 1857 |
+
print("=== Evaluation Setup ===")
|
| 1858 |
+
print(f"Device: {device}")
|
| 1859 |
+
print(f"Dtype : {dtype}")
|
| 1860 |
+
print(f"Model : {args.model_id}")
|
| 1861 |
+
print(f"Data : {dataset_name} ({args.split})")
|
| 1862 |
+
print(f"Configs ({len(configs)}): {configs}")
|
| 1863 |
+
print(f"Max samples/config: {'all' if args.max_samples_per_config <= 0 else args.max_samples_per_config}")
|
| 1864 |
+
print(f"Streaming: {'yes' if args.streaming else 'no'}")
|
| 1865 |
+
print(f"Text normalizer: {args.text_normalizer}")
|
| 1866 |
+
print(f"Run base_model: {'yes' if run_base_model else 'no'}")
|
| 1867 |
+
print(f"Run baseline_ft: {'yes' if run_baseline_ft else 'no'}")
|
| 1868 |
+
print(f"Run prompt_tuning: {'yes' if run_prompt_tuning else 'no'}")
|
| 1869 |
+
print(f"Run lora_r16: {'yes' if run_lora_r16 else 'no'}")
|
| 1870 |
+
print(f"Run latent_reasoning: {'yes' if run_latent_reasoning else 'no'}")
|
| 1871 |
+
print(f"Language hint: {'off' if args.no_language_hint else 'on'}")
|
| 1872 |
+
if not args.no_language_hint:
|
| 1873 |
+
print(f"Prompt template: {args.prompt_template}")
|
| 1874 |
+
print(f"Num beams: {args.num_beams}")
|
| 1875 |
+
print(f"Print gate skips: {'yes' if args.print_gate_skips else 'no'}")
|
| 1876 |
+
|
| 1877 |
+
base_model_rows: Dict[str, Dict[str, any]] = {}
|
| 1878 |
+
baseline_ft_rows: Dict[str, Dict[str, any]] = {}
|
| 1879 |
+
prompt_tuning_rows: Dict[str, Dict[str, any]] = {}
|
| 1880 |
+
lora_r16_rows: Dict[str, Dict[str, any]] = {}
|
| 1881 |
+
latent_reasoning_rows: Dict[str, Dict[str, any]] = {}
|
| 1882 |
+
|
| 1883 |
+
base_model_bundle: Optional[ModelBundle] = None
|
| 1884 |
+
baseline_ft_bundle: Optional[ModelBundle] = None
|
| 1885 |
+
prompt_tuning_bundle: Optional[ModelBundle] = None
|
| 1886 |
+
lora_r16_bundle: Optional[ModelBundle] = None
|
| 1887 |
+
latent_reasoning_bundle: Optional[ModelBundle] = None
|
| 1888 |
+
|
| 1889 |
+
if run_base_model:
|
| 1890 |
+
try:
|
| 1891 |
+
base_model_bundle = build_base_model_bundle(
|
| 1892 |
+
model_id=args.model_id,
|
| 1893 |
+
device=device,
|
| 1894 |
+
dtype=dtype,
|
| 1895 |
+
)
|
| 1896 |
+
for cfg in configs:
|
| 1897 |
+
language_hint = resolve_language_hint(dataset_name, cfg)
|
| 1898 |
+
out = evaluate_config(
|
| 1899 |
+
model=base_model_bundle.model,
|
| 1900 |
+
processor=base_model_bundle.processor,
|
| 1901 |
+
dataset_name=dataset_name,
|
| 1902 |
+
config_name=cfg,
|
| 1903 |
+
split=args.split,
|
| 1904 |
+
use_baseline=True,
|
| 1905 |
+
mode_name="base_model",
|
| 1906 |
+
max_samples=args.max_samples_per_config,
|
| 1907 |
+
max_new_tokens=args.max_new_tokens,
|
| 1908 |
+
print_samples=args.print_samples,
|
| 1909 |
+
language_hint=language_hint,
|
| 1910 |
+
prompt_template=args.prompt_template,
|
| 1911 |
+
disable_language_hint=args.no_language_hint,
|
| 1912 |
+
num_beams=args.num_beams,
|
| 1913 |
+
dynamic_halt_threshold=args.dynamic_halt_threshold,
|
| 1914 |
+
snr_db=args.snr_db[0] if args.snr_db else None,
|
| 1915 |
+
print_gate_skips=args.print_gate_skips,
|
| 1916 |
+
streaming=args.streaming,
|
| 1917 |
+
)
|
| 1918 |
+
base_model_rows[cfg] = out
|
| 1919 |
+
print(
|
| 1920 |
+
f"[BaseModel][{cfg}] "
|
| 1921 |
+
f"lang={out.get('language_hint') or 'auto'} "
|
| 1922 |
+
f"text_col={out.get('text_column') or '-'} "
|
| 1923 |
+
f"WER={out['wer']:.4f} CER={out['cer']:.4f} "
|
| 1924 |
+
f"n={out['samples_used']} skip={out['samples_skipped']}"
|
| 1925 |
+
)
|
| 1926 |
+
finally:
|
| 1927 |
+
release_model(base_model_bundle)
|
| 1928 |
+
else:
|
| 1929 |
+
print("\n[Skip] base_model evaluation skipped (--skip-base-model).")
|
| 1930 |
+
|
| 1931 |
+
if run_baseline_ft:
|
| 1932 |
+
try:
|
| 1933 |
+
baseline_ft_bundle = build_baseline_ft_bundle(
|
| 1934 |
+
model_id=args.model_id,
|
| 1935 |
+
checkpoint_path=args.baseline_ckpt,
|
| 1936 |
+
device=device,
|
| 1937 |
+
dtype=dtype,
|
| 1938 |
+
)
|
| 1939 |
+
for cfg in configs:
|
| 1940 |
+
language_hint = resolve_language_hint(dataset_name, cfg)
|
| 1941 |
+
out = evaluate_config(
|
| 1942 |
+
model=baseline_ft_bundle.model,
|
| 1943 |
+
processor=baseline_ft_bundle.processor,
|
| 1944 |
+
dataset_name=dataset_name,
|
| 1945 |
+
config_name=cfg,
|
| 1946 |
+
split=args.split,
|
| 1947 |
+
use_baseline=True,
|
| 1948 |
+
mode_name="baseline_ft",
|
| 1949 |
+
max_samples=args.max_samples_per_config,
|
| 1950 |
+
max_new_tokens=args.max_new_tokens,
|
| 1951 |
+
print_samples=args.print_samples,
|
| 1952 |
+
language_hint=language_hint,
|
| 1953 |
+
prompt_template=args.prompt_template,
|
| 1954 |
+
disable_language_hint=args.no_language_hint,
|
| 1955 |
+
num_beams=args.num_beams,
|
| 1956 |
+
dynamic_halt_threshold=args.dynamic_halt_threshold,
|
| 1957 |
+
snr_db=args.snr_db[0] if args.snr_db else None,
|
| 1958 |
+
print_gate_skips=args.print_gate_skips,
|
| 1959 |
+
streaming=args.streaming,
|
| 1960 |
+
)
|
| 1961 |
+
baseline_ft_rows[cfg] = out
|
| 1962 |
+
print(
|
| 1963 |
+
f"[BaselineFT][{cfg}] "
|
| 1964 |
+
f"lang={out.get('language_hint') or 'auto'} "
|
| 1965 |
+
f"text_col={out.get('text_column') or '-'} "
|
| 1966 |
+
f"WER={out['wer']:.4f} CER={out['cer']:.4f} "
|
| 1967 |
+
f"n={out['samples_used']} skip={out['samples_skipped']}"
|
| 1968 |
+
)
|
| 1969 |
+
finally:
|
| 1970 |
+
release_model(baseline_ft_bundle)
|
| 1971 |
+
else:
|
| 1972 |
+
print("\n[Skip] baseline_ft evaluation skipped (no --baseline-ckpt).")
|
| 1973 |
+
|
| 1974 |
+
if run_prompt_tuning:
|
| 1975 |
+
try:
|
| 1976 |
+
prompt_tuning_bundle = build_prompt_tuning_bundle(
|
| 1977 |
+
model_id=args.model_id,
|
| 1978 |
+
checkpoint_path=args.prompt_tuning_ckpt,
|
| 1979 |
+
device=device,
|
| 1980 |
+
dtype=dtype,
|
| 1981 |
+
)
|
| 1982 |
+
for cfg in configs:
|
| 1983 |
+
language_hint = resolve_language_hint(dataset_name, cfg)
|
| 1984 |
+
out = evaluate_config(
|
| 1985 |
+
model=prompt_tuning_bundle.model,
|
| 1986 |
+
processor=prompt_tuning_bundle.processor,
|
| 1987 |
+
dataset_name=dataset_name,
|
| 1988 |
+
config_name=cfg,
|
| 1989 |
+
split=args.split,
|
| 1990 |
+
use_baseline=False,
|
| 1991 |
+
mode_name="prompt_tuning",
|
| 1992 |
+
max_samples=args.max_samples_per_config,
|
| 1993 |
+
max_new_tokens=args.max_new_tokens,
|
| 1994 |
+
print_samples=args.print_samples,
|
| 1995 |
+
language_hint=language_hint,
|
| 1996 |
+
prompt_template=args.prompt_template,
|
| 1997 |
+
disable_language_hint=args.no_language_hint,
|
| 1998 |
+
num_beams=args.num_beams,
|
| 1999 |
+
dynamic_halt_threshold=args.dynamic_halt_threshold,
|
| 2000 |
+
snr_db=args.snr_db[0] if args.snr_db else None,
|
| 2001 |
+
print_gate_skips=args.print_gate_skips,
|
| 2002 |
+
streaming=args.streaming,
|
| 2003 |
+
)
|
| 2004 |
+
prompt_tuning_rows[cfg] = out
|
| 2005 |
+
print(
|
| 2006 |
+
f"[PromptTuning][{cfg}] "
|
| 2007 |
+
f"text_col={out.get('text_column') or '-'} "
|
| 2008 |
+
f"WER={out['wer']:.4f} CER={out['cer']:.4f} "
|
| 2009 |
+
f"n={out['samples_used']} skip={out['samples_skipped']}"
|
| 2010 |
+
)
|
| 2011 |
+
finally:
|
| 2012 |
+
release_model(prompt_tuning_bundle)
|
| 2013 |
+
else:
|
| 2014 |
+
print("\n[Skip] prompt_tuning evaluation skipped (no --prompt-tuning-ckpt).")
|
| 2015 |
+
|
| 2016 |
+
if run_lora_r16:
|
| 2017 |
+
try:
|
| 2018 |
+
lora_r16_bundle = build_lora_r16_bundle(
|
| 2019 |
+
model_id=args.model_id,
|
| 2020 |
+
checkpoint_path=args.lora_r16_ckpt,
|
| 2021 |
+
device=device,
|
| 2022 |
+
dtype=dtype,
|
| 2023 |
+
)
|
| 2024 |
+
for cfg in configs:
|
| 2025 |
+
language_hint = resolve_language_hint(dataset_name, cfg)
|
| 2026 |
+
out = evaluate_config(
|
| 2027 |
+
model=lora_r16_bundle.model,
|
| 2028 |
+
processor=lora_r16_bundle.processor,
|
| 2029 |
+
dataset_name=dataset_name,
|
| 2030 |
+
config_name=cfg,
|
| 2031 |
+
split=args.split,
|
| 2032 |
+
use_baseline=True,
|
| 2033 |
+
mode_name="lora_r16",
|
| 2034 |
+
max_samples=args.max_samples_per_config,
|
| 2035 |
+
max_new_tokens=args.max_new_tokens,
|
| 2036 |
+
print_samples=args.print_samples,
|
| 2037 |
+
language_hint=language_hint,
|
| 2038 |
+
prompt_template=args.prompt_template,
|
| 2039 |
+
disable_language_hint=args.no_language_hint,
|
| 2040 |
+
num_beams=args.num_beams,
|
| 2041 |
+
dynamic_halt_threshold=args.dynamic_halt_threshold,
|
| 2042 |
+
snr_db=args.snr_db[0] if args.snr_db else None,
|
| 2043 |
+
print_gate_skips=args.print_gate_skips,
|
| 2044 |
+
streaming=args.streaming,
|
| 2045 |
+
)
|
| 2046 |
+
lora_r16_rows[cfg] = out
|
| 2047 |
+
print(
|
| 2048 |
+
f"[LoRA-R16][{cfg}] "
|
| 2049 |
+
f"text_col={out.get('text_column') or '-'} "
|
| 2050 |
+
f"WER={out['wer']:.4f} CER={out['cer']:.4f} "
|
| 2051 |
+
f"n={out['samples_used']} skip={out['samples_skipped']}"
|
| 2052 |
+
)
|
| 2053 |
+
finally:
|
| 2054 |
+
release_model(lora_r16_bundle)
|
| 2055 |
+
else:
|
| 2056 |
+
print("\n[Skip] lora_r16 evaluation skipped (no --lora-r16-ckpt).")
|
| 2057 |
+
|
| 2058 |
+
if run_latent_reasoning:
|
| 2059 |
+
try:
|
| 2060 |
+
latent_reasoning_bundle = build_latent_bundle(
|
| 2061 |
+
model_id=args.model_id,
|
| 2062 |
+
checkpoint_path=args.latent_ckpt,
|
| 2063 |
+
n_latent_override=args.n_latent,
|
| 2064 |
+
device=device,
|
| 2065 |
+
dtype=dtype,
|
| 2066 |
+
)
|
| 2067 |
+
for cfg in configs:
|
| 2068 |
+
language_hint = resolve_language_hint(dataset_name, cfg)
|
| 2069 |
+
out = evaluate_config(
|
| 2070 |
+
model=latent_reasoning_bundle.model,
|
| 2071 |
+
processor=latent_reasoning_bundle.processor,
|
| 2072 |
+
dataset_name=dataset_name,
|
| 2073 |
+
config_name=cfg,
|
| 2074 |
+
split=args.split,
|
| 2075 |
+
use_baseline=False,
|
| 2076 |
+
mode_name="latent_reasoning",
|
| 2077 |
+
max_samples=args.max_samples_per_config,
|
| 2078 |
+
max_new_tokens=args.max_new_tokens,
|
| 2079 |
+
print_samples=args.print_samples,
|
| 2080 |
+
language_hint=language_hint,
|
| 2081 |
+
prompt_template=args.prompt_template,
|
| 2082 |
+
disable_language_hint=args.no_language_hint,
|
| 2083 |
+
num_beams=args.num_beams,
|
| 2084 |
+
dynamic_halt_threshold=args.dynamic_halt_threshold,
|
| 2085 |
+
snr_db=args.snr_db[0] if args.snr_db else None,
|
| 2086 |
+
print_gate_skips=args.print_gate_skips,
|
| 2087 |
+
streaming=args.streaming,
|
| 2088 |
+
)
|
| 2089 |
+
latent_reasoning_rows[cfg] = out
|
| 2090 |
+
n_used = out['samples_used']
|
| 2091 |
+
n_reasoned = out.get('samples_reasoned', 0)
|
| 2092 |
+
reason_pct = (n_reasoned / n_used * 100) if n_used > 0 else 0.0
|
| 2093 |
+
print(
|
| 2094 |
+
f"[LatentReasoning][{cfg}] "
|
| 2095 |
+
f"lang={out.get('language_hint') or 'auto'} "
|
| 2096 |
+
f"text_col={out.get('text_column') or '-'} "
|
| 2097 |
+
f"WER={out['wer']:.4f} CER={out['cer']:.4f} "
|
| 2098 |
+
f"n={n_used} skip={out['samples_skipped']} "
|
| 2099 |
+
f"reasoned={n_reasoned}/{n_used} ({reason_pct:.1f}%)"
|
| 2100 |
+
)
|
| 2101 |
+
finally:
|
| 2102 |
+
release_model(latent_reasoning_bundle)
|
| 2103 |
+
else:
|
| 2104 |
+
print("\n[Skip] latent_reasoning evaluation skipped (no --latent-ckpt).")
|
| 2105 |
+
|
| 2106 |
+
merged_rows: List[Dict[str, any]] = []
|
| 2107 |
+
for cfg in configs:
|
| 2108 |
+
base = base_model_rows.get(cfg, {})
|
| 2109 |
+
ft = baseline_ft_rows.get(cfg, {})
|
| 2110 |
+
prompt = prompt_tuning_rows.get(cfg, {})
|
| 2111 |
+
lora = lora_r16_rows.get(cfg, {})
|
| 2112 |
+
lat = latent_reasoning_rows.get(cfg, {})
|
| 2113 |
+
text_col = (
|
| 2114 |
+
base.get("text_column")
|
| 2115 |
+
or ft.get("text_column")
|
| 2116 |
+
or prompt.get("text_column")
|
| 2117 |
+
or lora.get("text_column")
|
| 2118 |
+
or lat.get("text_column")
|
| 2119 |
+
)
|
| 2120 |
+
row = {
|
| 2121 |
+
"config": cfg,
|
| 2122 |
+
"samples_total": int(
|
| 2123 |
+
max(
|
| 2124 |
+
base.get("samples_total", 0),
|
| 2125 |
+
ft.get("samples_total", 0),
|
| 2126 |
+
prompt.get("samples_total", 0),
|
| 2127 |
+
lora.get("samples_total", 0),
|
| 2128 |
+
lat.get("samples_total", 0),
|
| 2129 |
+
)
|
| 2130 |
+
),
|
| 2131 |
+
"samples_used": int(
|
| 2132 |
+
max(
|
| 2133 |
+
base.get("samples_used", 0),
|
| 2134 |
+
ft.get("samples_used", 0),
|
| 2135 |
+
prompt.get("samples_used", 0),
|
| 2136 |
+
lora.get("samples_used", 0),
|
| 2137 |
+
lat.get("samples_used", 0),
|
| 2138 |
+
)
|
| 2139 |
+
),
|
| 2140 |
+
"text_column": text_col,
|
| 2141 |
+
"base_model_wer": base.get("wer"),
|
| 2142 |
+
"base_model_cer": base.get("cer"),
|
| 2143 |
+
"baseline_ft_wer": ft.get("wer"),
|
| 2144 |
+
"baseline_ft_cer": ft.get("cer"),
|
| 2145 |
+
"prompt_tuning_wer": prompt.get("wer"),
|
| 2146 |
+
"prompt_tuning_cer": prompt.get("cer"),
|
| 2147 |
+
"lora_r16_wer": lora.get("wer"),
|
| 2148 |
+
"lora_r16_cer": lora.get("cer"),
|
| 2149 |
+
"latent_reasoning_wer": lat.get("wer"),
|
| 2150 |
+
"latent_reasoning_cer": lat.get("cer"),
|
| 2151 |
+
"baseline_ft_vs_base_model_wer": None
|
| 2152 |
+
if base.get("wer") is None or ft.get("wer") is None
|
| 2153 |
+
else float(base["wer"] - ft["wer"]),
|
| 2154 |
+
"baseline_ft_vs_base_model_cer": None
|
| 2155 |
+
if base.get("cer") is None or ft.get("cer") is None
|
| 2156 |
+
else float(base["cer"] - ft["cer"]),
|
| 2157 |
+
"prompt_tuning_vs_base_model_wer": None
|
| 2158 |
+
if base.get("wer") is None or prompt.get("wer") is None
|
| 2159 |
+
else float(base["wer"] - prompt["wer"]),
|
| 2160 |
+
"prompt_tuning_vs_base_model_cer": None
|
| 2161 |
+
if base.get("cer") is None or prompt.get("cer") is None
|
| 2162 |
+
else float(base["cer"] - prompt["cer"]),
|
| 2163 |
+
"lora_r16_vs_base_model_wer": None
|
| 2164 |
+
if base.get("wer") is None or lora.get("wer") is None
|
| 2165 |
+
else float(base["wer"] - lora["wer"]),
|
| 2166 |
+
"lora_r16_vs_base_model_cer": None
|
| 2167 |
+
if base.get("cer") is None or lora.get("cer") is None
|
| 2168 |
+
else float(base["cer"] - lora["cer"]),
|
| 2169 |
+
"latent_reasoning_vs_base_model_wer": None
|
| 2170 |
+
if base.get("wer") is None or lat.get("wer") is None
|
| 2171 |
+
else float(base["wer"] - lat["wer"]),
|
| 2172 |
+
"latent_reasoning_vs_base_model_cer": None
|
| 2173 |
+
if base.get("cer") is None or lat.get("cer") is None
|
| 2174 |
+
else float(base["cer"] - lat["cer"]),
|
| 2175 |
+
"latent_reasoning_vs_baseline_ft_wer": None
|
| 2176 |
+
if ft.get("wer") is None or lat.get("wer") is None
|
| 2177 |
+
else float(ft["wer"] - lat["wer"]),
|
| 2178 |
+
"latent_reasoning_vs_baseline_ft_cer": None
|
| 2179 |
+
if ft.get("cer") is None or lat.get("cer") is None
|
| 2180 |
+
else float(ft["cer"] - lat["cer"]),
|
| 2181 |
+
"latent_reasoning_vs_prompt_tuning_wer": None
|
| 2182 |
+
if prompt.get("wer") is None or lat.get("wer") is None
|
| 2183 |
+
else float(prompt["wer"] - lat["wer"]),
|
| 2184 |
+
"latent_reasoning_vs_prompt_tuning_cer": None
|
| 2185 |
+
if prompt.get("cer") is None or lat.get("cer") is None
|
| 2186 |
+
else float(prompt["cer"] - lat["cer"]),
|
| 2187 |
+
"latent_reasoning_vs_lora_r16_wer": None
|
| 2188 |
+
if lora.get("wer") is None or lat.get("wer") is None
|
| 2189 |
+
else float(lora["wer"] - lat["wer"]),
|
| 2190 |
+
"latent_reasoning_vs_lora_r16_cer": None
|
| 2191 |
+
if lora.get("cer") is None or lat.get("cer") is None
|
| 2192 |
+
else float(lora["cer"] - lat["cer"]),
|
| 2193 |
+
}
|
| 2194 |
+
merged_rows.append(row)
|
| 2195 |
+
|
| 2196 |
+
print_summary_table(merged_rows)
|
| 2197 |
+
|
| 2198 |
+
base_list = [base_model_rows[c] for c in configs if c in base_model_rows]
|
| 2199 |
+
ft_list = [baseline_ft_rows[c] for c in configs if c in baseline_ft_rows]
|
| 2200 |
+
prompt_list = [prompt_tuning_rows[c] for c in configs if c in prompt_tuning_rows]
|
| 2201 |
+
lora_list = [lora_r16_rows[c] for c in configs if c in lora_r16_rows]
|
| 2202 |
+
lat_list = [latent_reasoning_rows[c] for c in configs if c in latent_reasoning_rows]
|
| 2203 |
+
summary = {
|
| 2204 |
+
"base_model_weighted_wer": weighted_avg(base_list, "wer"),
|
| 2205 |
+
"base_model_weighted_cer": weighted_avg(base_list, "cer"),
|
| 2206 |
+
"baseline_ft_weighted_wer": weighted_avg(ft_list, "wer"),
|
| 2207 |
+
"baseline_ft_weighted_cer": weighted_avg(ft_list, "cer"),
|
| 2208 |
+
"prompt_tuning_weighted_wer": weighted_avg(prompt_list, "wer"),
|
| 2209 |
+
"prompt_tuning_weighted_cer": weighted_avg(prompt_list, "cer"),
|
| 2210 |
+
"lora_r16_weighted_wer": weighted_avg(lora_list, "wer"),
|
| 2211 |
+
"lora_r16_weighted_cer": weighted_avg(lora_list, "cer"),
|
| 2212 |
+
"latent_reasoning_weighted_wer": weighted_avg(lat_list, "wer"),
|
| 2213 |
+
"latent_reasoning_weighted_cer": weighted_avg(lat_list, "cer"),
|
| 2214 |
+
}
|
| 2215 |
+
if summary["base_model_weighted_wer"] is not None and summary["baseline_ft_weighted_wer"] is not None:
|
| 2216 |
+
summary["baseline_ft_vs_base_model_weighted_wer"] = (
|
| 2217 |
+
summary["base_model_weighted_wer"] - summary["baseline_ft_weighted_wer"]
|
| 2218 |
+
)
|
| 2219 |
+
else:
|
| 2220 |
+
summary["baseline_ft_vs_base_model_weighted_wer"] = None
|
| 2221 |
+
if summary["base_model_weighted_cer"] is not None and summary["baseline_ft_weighted_cer"] is not None:
|
| 2222 |
+
summary["baseline_ft_vs_base_model_weighted_cer"] = (
|
| 2223 |
+
summary["base_model_weighted_cer"] - summary["baseline_ft_weighted_cer"]
|
| 2224 |
+
)
|
| 2225 |
+
else:
|
| 2226 |
+
summary["baseline_ft_vs_base_model_weighted_cer"] = None
|
| 2227 |
+
|
| 2228 |
+
if summary["base_model_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None:
|
| 2229 |
+
summary["latent_reasoning_vs_base_model_weighted_wer"] = (
|
| 2230 |
+
summary["base_model_weighted_wer"] - summary["latent_reasoning_weighted_wer"]
|
| 2231 |
+
)
|
| 2232 |
+
else:
|
| 2233 |
+
summary["latent_reasoning_vs_base_model_weighted_wer"] = None
|
| 2234 |
+
if summary["base_model_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None:
|
| 2235 |
+
summary["latent_reasoning_vs_base_model_weighted_cer"] = (
|
| 2236 |
+
summary["base_model_weighted_cer"] - summary["latent_reasoning_weighted_cer"]
|
| 2237 |
+
)
|
| 2238 |
+
else:
|
| 2239 |
+
summary["latent_reasoning_vs_base_model_weighted_cer"] = None
|
| 2240 |
+
|
| 2241 |
+
if summary["baseline_ft_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None:
|
| 2242 |
+
summary["latent_reasoning_vs_baseline_ft_weighted_wer"] = (
|
| 2243 |
+
summary["baseline_ft_weighted_wer"] - summary["latent_reasoning_weighted_wer"]
|
| 2244 |
+
)
|
| 2245 |
+
else:
|
| 2246 |
+
summary["latent_reasoning_vs_baseline_ft_weighted_wer"] = None
|
| 2247 |
+
if summary["baseline_ft_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None:
|
| 2248 |
+
summary["latent_reasoning_vs_baseline_ft_weighted_cer"] = (
|
| 2249 |
+
summary["baseline_ft_weighted_cer"] - summary["latent_reasoning_weighted_cer"]
|
| 2250 |
+
)
|
| 2251 |
+
else:
|
| 2252 |
+
summary["latent_reasoning_vs_baseline_ft_weighted_cer"] = None
|
| 2253 |
+
|
| 2254 |
+
if summary["base_model_weighted_wer"] is not None and summary["prompt_tuning_weighted_wer"] is not None:
|
| 2255 |
+
summary["prompt_tuning_vs_base_model_weighted_wer"] = (
|
| 2256 |
+
summary["base_model_weighted_wer"] - summary["prompt_tuning_weighted_wer"]
|
| 2257 |
+
)
|
| 2258 |
+
else:
|
| 2259 |
+
summary["prompt_tuning_vs_base_model_weighted_wer"] = None
|
| 2260 |
+
if summary["base_model_weighted_cer"] is not None and summary["prompt_tuning_weighted_cer"] is not None:
|
| 2261 |
+
summary["prompt_tuning_vs_base_model_weighted_cer"] = (
|
| 2262 |
+
summary["base_model_weighted_cer"] - summary["prompt_tuning_weighted_cer"]
|
| 2263 |
+
)
|
| 2264 |
+
else:
|
| 2265 |
+
summary["prompt_tuning_vs_base_model_weighted_cer"] = None
|
| 2266 |
+
|
| 2267 |
+
if summary["base_model_weighted_wer"] is not None and summary["lora_r16_weighted_wer"] is not None:
|
| 2268 |
+
summary["lora_r16_vs_base_model_weighted_wer"] = (
|
| 2269 |
+
summary["base_model_weighted_wer"] - summary["lora_r16_weighted_wer"]
|
| 2270 |
+
)
|
| 2271 |
+
else:
|
| 2272 |
+
summary["lora_r16_vs_base_model_weighted_wer"] = None
|
| 2273 |
+
if summary["base_model_weighted_cer"] is not None and summary["lora_r16_weighted_cer"] is not None:
|
| 2274 |
+
summary["lora_r16_vs_base_model_weighted_cer"] = (
|
| 2275 |
+
summary["base_model_weighted_cer"] - summary["lora_r16_weighted_cer"]
|
| 2276 |
+
)
|
| 2277 |
+
else:
|
| 2278 |
+
summary["lora_r16_vs_base_model_weighted_cer"] = None
|
| 2279 |
+
|
| 2280 |
+
if summary["prompt_tuning_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None:
|
| 2281 |
+
summary["latent_reasoning_vs_prompt_tuning_weighted_wer"] = (
|
| 2282 |
+
summary["prompt_tuning_weighted_wer"] - summary["latent_reasoning_weighted_wer"]
|
| 2283 |
+
)
|
| 2284 |
+
else:
|
| 2285 |
+
summary["latent_reasoning_vs_prompt_tuning_weighted_wer"] = None
|
| 2286 |
+
if summary["prompt_tuning_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None:
|
| 2287 |
+
summary["latent_reasoning_vs_prompt_tuning_weighted_cer"] = (
|
| 2288 |
+
summary["prompt_tuning_weighted_cer"] - summary["latent_reasoning_weighted_cer"]
|
| 2289 |
+
)
|
| 2290 |
+
else:
|
| 2291 |
+
summary["latent_reasoning_vs_prompt_tuning_weighted_cer"] = None
|
| 2292 |
+
|
| 2293 |
+
if summary["lora_r16_weighted_wer"] is not None and summary["latent_reasoning_weighted_wer"] is not None:
|
| 2294 |
+
summary["latent_reasoning_vs_lora_r16_weighted_wer"] = (
|
| 2295 |
+
summary["lora_r16_weighted_wer"] - summary["latent_reasoning_weighted_wer"]
|
| 2296 |
+
)
|
| 2297 |
+
else:
|
| 2298 |
+
summary["latent_reasoning_vs_lora_r16_weighted_wer"] = None
|
| 2299 |
+
if summary["lora_r16_weighted_cer"] is not None and summary["latent_reasoning_weighted_cer"] is not None:
|
| 2300 |
+
summary["latent_reasoning_vs_lora_r16_weighted_cer"] = (
|
| 2301 |
+
summary["lora_r16_weighted_cer"] - summary["latent_reasoning_weighted_cer"]
|
| 2302 |
+
)
|
| 2303 |
+
else:
|
| 2304 |
+
summary["latent_reasoning_vs_lora_r16_weighted_cer"] = None
|
| 2305 |
+
|
| 2306 |
+
print("\n=== Weighted Summary ===")
|
| 2307 |
+
for k, v in summary.items():
|
| 2308 |
+
if v is None:
|
| 2309 |
+
print(f"{k}: -")
|
| 2310 |
+
else:
|
| 2311 |
+
print(f"{k}: {v:.6f}")
|
| 2312 |
+
|
| 2313 |
+
if args.output_json:
|
| 2314 |
+
payload = {
|
| 2315 |
+
"model_id": args.model_id,
|
| 2316 |
+
"dataset_name": dataset_name,
|
| 2317 |
+
"dataset_name_input": args.dataset_name,
|
| 2318 |
+
"split": args.split,
|
| 2319 |
+
"configs": configs,
|
| 2320 |
+
"streaming": args.streaming,
|
| 2321 |
+
"text_normalizer": args.text_normalizer,
|
| 2322 |
+
"base_model_enabled": run_base_model,
|
| 2323 |
+
"baseline_checkpoint": args.baseline_ckpt,
|
| 2324 |
+
"prompt_tuning_checkpoint": args.prompt_tuning_ckpt,
|
| 2325 |
+
"lora_r16_checkpoint": args.lora_r16_ckpt,
|
| 2326 |
+
"latent_reasoning_checkpoint": args.latent_ckpt,
|
| 2327 |
+
"latent_checkpoint": args.latent_ckpt,
|
| 2328 |
+
"rows": merged_rows,
|
| 2329 |
+
"summary": summary,
|
| 2330 |
+
}
|
| 2331 |
+
with open(args.output_json, "w", encoding="utf-8") as f:
|
| 2332 |
+
json.dump(payload, f, ensure_ascii=False, indent=2)
|
| 2333 |
+
print(f"\nSaved JSON: {args.output_json}")
|
| 2334 |
+
|
| 2335 |
+
|
| 2336 |
+
if __name__ == "__main__":
|
| 2337 |
+
main()
|
examples/transcribe_file.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Transcribe one audio file with the released LatentASR adapter."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import argparse
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
from datasets import Audio
|
| 9 |
+
|
| 10 |
+
from eval import build_latent_bundle, choose_device, choose_dtype, clean_prediction
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def parse_args() -> argparse.Namespace:
|
| 14 |
+
parser = argparse.ArgumentParser(description="Transcribe one audio file with LatentASR.")
|
| 15 |
+
parser.add_argument("audio", help="Path to an audio file.")
|
| 16 |
+
parser.add_argument("--model-id", default="Qwen/Qwen3-ASR-0.6B")
|
| 17 |
+
parser.add_argument("--latent-ckpt", default="checkpoints/latentASR_adapter.pth")
|
| 18 |
+
parser.add_argument("--theta", type=float, default=0.0, help="Dynamic halting threshold.")
|
| 19 |
+
parser.add_argument("--max-new-tokens", type=int, default=128)
|
| 20 |
+
parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
|
| 21 |
+
parser.add_argument("--dtype", default="auto", choices=["auto", "float32", "float16", "bfloat16"])
|
| 22 |
+
parser.add_argument("--language", default="English")
|
| 23 |
+
return parser.parse_args()
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
@torch.no_grad()
|
| 27 |
+
def main() -> None:
|
| 28 |
+
args = parse_args()
|
| 29 |
+
device = choose_device(args.device)
|
| 30 |
+
dtype = choose_dtype(args.dtype, device)
|
| 31 |
+
bundle = build_latent_bundle(
|
| 32 |
+
model_id=args.model_id,
|
| 33 |
+
checkpoint_path=args.latent_ckpt,
|
| 34 |
+
n_latent_override=-1,
|
| 35 |
+
device=device,
|
| 36 |
+
dtype=dtype,
|
| 37 |
+
)
|
| 38 |
+
model = bundle.model
|
| 39 |
+
processor = bundle.processor
|
| 40 |
+
|
| 41 |
+
target_sr = int(getattr(processor.feature_extractor, "sampling_rate", 16000) or 16000)
|
| 42 |
+
audio = Audio(sampling_rate=target_sr).decode_example({"path": args.audio})
|
| 43 |
+
feat_out = processor.feature_extractor(
|
| 44 |
+
audio["array"],
|
| 45 |
+
sampling_rate=audio["sampling_rate"],
|
| 46 |
+
return_attention_mask=True,
|
| 47 |
+
)
|
| 48 |
+
target_dtype = model.thinker.dtype if hasattr(model.thinker, "dtype") else torch.float32
|
| 49 |
+
feats = torch.tensor(
|
| 50 |
+
feat_out.input_features[0],
|
| 51 |
+
dtype=target_dtype,
|
| 52 |
+
device=model.base_model.device,
|
| 53 |
+
).unsqueeze(0)
|
| 54 |
+
n_frames = feats.size(-1)
|
| 55 |
+
if getattr(feat_out, "attention_mask", None) is not None:
|
| 56 |
+
mask = torch.tensor(feat_out.attention_mask[0], dtype=torch.long)
|
| 57 |
+
if mask.size(-1) < n_frames:
|
| 58 |
+
mask = torch.cat([mask, torch.zeros(n_frames - mask.size(-1), dtype=torch.long)])
|
| 59 |
+
else:
|
| 60 |
+
mask = mask[:n_frames]
|
| 61 |
+
feature_attention_mask = mask.to(device=model.base_model.device).unsqueeze(0)
|
| 62 |
+
else:
|
| 63 |
+
feature_attention_mask = torch.ones((1, n_frames), dtype=torch.long, device=model.base_model.device)
|
| 64 |
+
|
| 65 |
+
gen_ids, stats = model.generate(
|
| 66 |
+
feats,
|
| 67 |
+
feature_attention_mask=feature_attention_mask,
|
| 68 |
+
max_new_tokens=args.max_new_tokens,
|
| 69 |
+
use_baseline=False,
|
| 70 |
+
return_stats=True,
|
| 71 |
+
do_sample=False,
|
| 72 |
+
eos_token_id=model.stop_ids,
|
| 73 |
+
num_beams=1,
|
| 74 |
+
language_hint=args.language,
|
| 75 |
+
dynamic_halt_threshold=args.theta,
|
| 76 |
+
)
|
| 77 |
+
text = clean_prediction(processor.tokenizer.decode(gen_ids[0], skip_special_tokens=True))
|
| 78 |
+
steps = int(stats.get("deq_iters", torch.tensor(0)).item()) if isinstance(stats, dict) else -1
|
| 79 |
+
print(text)
|
| 80 |
+
print(f"[latent_steps={steps} theta={args.theta}]")
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
if __name__ == "__main__":
|
| 84 |
+
main()
|
| 85 |
+
|
experiments/README.md
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Experiment Runners
|
| 2 |
+
|
| 3 |
+
This directory contains the maintained entrypoints for LatentASR experiments.
|
| 4 |
+
All runners default to Qwen/Qwen3-ASR-0.6B and `./latent_qwen_asr_best.pth`.
|
| 5 |
+
The checkpoint itself is intentionally ignored by Git.
|
| 6 |
+
|
| 7 |
+
## Runners
|
| 8 |
+
|
| 9 |
+
- `run_english_clean_and_noise.sh`: English clean suite plus optional SNR noise sweeps.
|
| 10 |
+
- `run_multilingual_streaming.sh`: FLEURS 30 and MLS public7 streaming evaluation.
|
| 11 |
+
- `run_threshold_sweep.sh`: Value-head halting threshold sweep.
|
| 12 |
+
- `run_all_training_modes.sh`: Sequential baseline, prompt tuning, LoRA, and LatentASR training.
|
| 13 |
+
- `run_legacy_generalization.sh`: older SpeechTest generalization/noise runner kept for reproducibility.
|
| 14 |
+
|
| 15 |
+
Root-level scripts with the old names are thin compatibility wrappers around
|
| 16 |
+
these maintained files.
|
| 17 |
+
|
| 18 |
+
## Common Environment Variables
|
| 19 |
+
|
| 20 |
+
- `LATENT_CKPT`: latent adapter checkpoint path.
|
| 21 |
+
- `MODEL_ID`: HuggingFace model ID.
|
| 22 |
+
- `PYTHON_BIN`: Python executable override.
|
| 23 |
+
- `MAX_SAMPLES_PER_CONFIG`: per-config cap; `0` means full split.
|
| 24 |
+
- `OUT_DIR`: output directory.
|
| 25 |
+
- `RESUME=1`: skip JSON outputs that already exist.
|
| 26 |
+
|
| 27 |
+
Full experimental JSON/log outputs are written under `eval_runs/` and are not
|
| 28 |
+
committed to Git.
|
experiments/eval_paper_tbd_checkpoints.sh
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
set -euo pipefail
|
| 3 |
+
|
| 4 |
+
# Evaluate checkpoints produced by run_paper_tbd_retrain.sh for the TBD paper
|
| 5 |
+
# tables. This script intentionally does not override --n-latent; eval.py reads
|
| 6 |
+
# the latent budget from each checkpoint, which is required for the N sweep.
|
| 7 |
+
|
| 8 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
| 9 |
+
source "${SCRIPT_DIR}/lib/common.sh"
|
| 10 |
+
ROOT_DIR="$(latent_asr_repo_root "${SCRIPT_DIR}")"
|
| 11 |
+
cd "${ROOT_DIR}"
|
| 12 |
+
|
| 13 |
+
PYTHON_BIN="$(latent_asr_python_bin "${PYTHON_BIN:-}")"
|
| 14 |
+
MODEL_ID="${MODEL_ID:-Qwen/Qwen3-ASR-0.6B}"
|
| 15 |
+
CKPT_ROOT="${CKPT_ROOT:-${ROOT_DIR}/eval_runs/paper_tbd_retrain_20260518/checkpoints}"
|
| 16 |
+
BASELINE_DIR="${BASELINE_DIR:-${ROOT_DIR}/eval_runs/hf_asr_showcase_full_20260503_152506}"
|
| 17 |
+
N4_CKPT="${N4_CKPT:-}"
|
| 18 |
+
|
| 19 |
+
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
| 20 |
+
OUT_DIR="${OUT_DIR:-${ROOT_DIR}/eval_runs/paper_tbd_eval_${TIMESTAMP}}"
|
| 21 |
+
LOG_DIR="${OUT_DIR}/logs"
|
| 22 |
+
mkdir -p "${LOG_DIR}"
|
| 23 |
+
|
| 24 |
+
MAX_SAMPLES_PER_CONFIG="${MAX_SAMPLES_PER_CONFIG:-0}"
|
| 25 |
+
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"
|
| 26 |
+
PRINT_SAMPLES="${PRINT_SAMPLES:-0}"
|
| 27 |
+
STREAMING="${STREAMING:-1}"
|
| 28 |
+
RESUME="${RESUME:-1}"
|
| 29 |
+
|
| 30 |
+
# Optional filters are regexes over the variant/dataset/theta tags.
|
| 31 |
+
VARIANT_FILTER="${VARIANT_FILTER:-}"
|
| 32 |
+
DATASET_FILTER="${DATASET_FILTER:-}"
|
| 33 |
+
THETA_FILTER="${THETA_FILTER:-}"
|
| 34 |
+
|
| 35 |
+
run_eval() {
|
| 36 |
+
local variant="$1"
|
| 37 |
+
local ckpt="$2"
|
| 38 |
+
local dataset_tag="$3"
|
| 39 |
+
local dataset_name="$4"
|
| 40 |
+
local configs="$5"
|
| 41 |
+
local theta_label="$6"
|
| 42 |
+
local theta_value="$7"
|
| 43 |
+
|
| 44 |
+
if [[ -n "${VARIANT_FILTER}" && ! "${variant}" =~ ${VARIANT_FILTER} ]]; then
|
| 45 |
+
return 0
|
| 46 |
+
fi
|
| 47 |
+
if [[ -n "${DATASET_FILTER}" && ! "${dataset_tag}" =~ ${DATASET_FILTER} ]]; then
|
| 48 |
+
return 0
|
| 49 |
+
fi
|
| 50 |
+
if [[ -n "${THETA_FILTER}" && ! "${theta_label}" =~ ${THETA_FILTER} ]]; then
|
| 51 |
+
return 0
|
| 52 |
+
fi
|
| 53 |
+
|
| 54 |
+
latent_asr_require_file "${ckpt}" "checkpoint for ${variant}"
|
| 55 |
+
|
| 56 |
+
local json_out="${OUT_DIR}/${variant}_${dataset_tag}_theta_${theta_label}.json"
|
| 57 |
+
local log_out="${LOG_DIR}/${variant}_${dataset_tag}_theta_${theta_label}.log"
|
| 58 |
+
if [[ "${RESUME}" == "1" && -f "${json_out}" ]]; then
|
| 59 |
+
echo "[skip] existing ${json_out}"
|
| 60 |
+
return 0
|
| 61 |
+
fi
|
| 62 |
+
|
| 63 |
+
local args=(
|
| 64 |
+
--model-id "${MODEL_ID}"
|
| 65 |
+
--dataset-name "${dataset_name}"
|
| 66 |
+
--configs "${configs}"
|
| 67 |
+
--split test
|
| 68 |
+
--output-json "${json_out}"
|
| 69 |
+
--max-samples-per-config "${MAX_SAMPLES_PER_CONFIG}"
|
| 70 |
+
--max-new-tokens "${MAX_NEW_TOKENS}"
|
| 71 |
+
--latent-ckpt "${ckpt}"
|
| 72 |
+
--num-beams 1
|
| 73 |
+
--dynamic-halt-threshold "${theta_value}"
|
| 74 |
+
--print-samples "${PRINT_SAMPLES}"
|
| 75 |
+
--skip-base-model
|
| 76 |
+
--skip-baseline-ft
|
| 77 |
+
--skip-prompt-tuning
|
| 78 |
+
--skip-lora-r16
|
| 79 |
+
)
|
| 80 |
+
if [[ "${STREAMING}" == "1" ]]; then
|
| 81 |
+
args+=(--streaming)
|
| 82 |
+
fi
|
| 83 |
+
|
| 84 |
+
echo "------------------------------------------------------------"
|
| 85 |
+
echo "Variant : ${variant}"
|
| 86 |
+
echo "Dataset : ${dataset_tag} (${dataset_name}/${configs})"
|
| 87 |
+
echo "Theta : ${theta_label} (${theta_value})"
|
| 88 |
+
echo "Checkpoint: ${ckpt}"
|
| 89 |
+
echo "JSON : ${json_out}"
|
| 90 |
+
echo "LOG : ${log_out}"
|
| 91 |
+
"${PYTHON_BIN}" eval.py "${args[@]}" 2>&1 | tee "${log_out}"
|
| 92 |
+
}
|
| 93 |
+
|
| 94 |
+
ckpt_epoch10() {
|
| 95 |
+
local variant="$1"
|
| 96 |
+
printf '%s/%s/%s_epoch10.pth' "${CKPT_ROOT}" "${variant}" "${variant}"
|
| 97 |
+
}
|
| 98 |
+
if [[ -z "${N4_CKPT}" ]]; then
|
| 99 |
+
N4_CKPT="$(ckpt_epoch10 n4)"
|
| 100 |
+
fi
|
| 101 |
+
|
| 102 |
+
echo "============================================================"
|
| 103 |
+
echo "LatentASR TBD Paper Evaluation"
|
| 104 |
+
echo "============================================================"
|
| 105 |
+
echo "Root dir : ${ROOT_DIR}"
|
| 106 |
+
echo "Checkpoint dir : ${CKPT_ROOT}"
|
| 107 |
+
echo "Output dir : ${OUT_DIR}"
|
| 108 |
+
echo "Baseline dir : ${BASELINE_DIR}"
|
| 109 |
+
echo "N=4 ckpt : ${N4_CKPT}"
|
| 110 |
+
echo "Python : ${PYTHON_BIN}"
|
| 111 |
+
echo "Streaming : ${STREAMING}"
|
| 112 |
+
echo "Max/config : ${MAX_SAMPLES_PER_CONFIG}"
|
| 113 |
+
echo "Variant filter : ${VARIANT_FILTER:-(none)}"
|
| 114 |
+
echo "Dataset filter : ${DATASET_FILTER:-(none)}"
|
| 115 |
+
echo "Theta filter : ${THETA_FILTER:-(none)}"
|
| 116 |
+
echo "============================================================"
|
| 117 |
+
|
| 118 |
+
# Component ablation: FLEURS only, deployed threshold.
|
| 119 |
+
for variant in component_no_bounded component_no_gate component_no_anchor; do
|
| 120 |
+
run_eval "${variant}" "$(ckpt_epoch10 "${variant}")" \
|
| 121 |
+
fleurs google/fleurs en_us zero 0.0
|
| 122 |
+
done
|
| 123 |
+
|
| 124 |
+
# N sweep: FLEURS and VoxPopuli, deployed threshold.
|
| 125 |
+
for variant in n1 n2 n8; do
|
| 126 |
+
ckpt="$(ckpt_epoch10 "${variant}")"
|
| 127 |
+
run_eval "${variant}" "${ckpt}" fleurs google/fleurs en_us zero 0.0
|
| 128 |
+
run_eval "${variant}" "${ckpt}" voxpopuli facebook/voxpopuli en zero 0.0
|
| 129 |
+
done
|
| 130 |
+
run_eval n4 "${N4_CKPT}" fleurs google/fleurs en_us zero 0.0
|
| 131 |
+
run_eval n4 "${N4_CKPT}" voxpopuli facebook/voxpopuli en zero 0.0
|
| 132 |
+
|
| 133 |
+
# Forced-negative ablation: FLEURS threshold sweep.
|
| 134 |
+
for spec in full:-2.0 neg0p2:-0.2 zero:0.0 pos0p2:0.2 pos0p5:0.5; do
|
| 135 |
+
IFS=':' read -r theta_label theta_value <<< "${spec}"
|
| 136 |
+
run_eval pneg0 "$(ckpt_epoch10 pneg0)" \
|
| 137 |
+
fleurs google/fleurs en_us "${theta_label}" "${theta_value}"
|
| 138 |
+
done
|
| 139 |
+
|
| 140 |
+
# Activation set scaling.
|
| 141 |
+
for variant in \
|
| 142 |
+
activation_100 \
|
| 143 |
+
activation_200 \
|
| 144 |
+
activation_300 \
|
| 145 |
+
activation_400 \
|
| 146 |
+
activation_500 \
|
| 147 |
+
activation_600 \
|
| 148 |
+
activation_700 \
|
| 149 |
+
activation_800
|
| 150 |
+
do
|
| 151 |
+
ckpt="$(ckpt_epoch10 "${variant}")"
|
| 152 |
+
run_eval "${variant}" "${ckpt}" fleurs google/fleurs en_us zero 0.0
|
| 153 |
+
run_eval "${variant}" "${ckpt}" voxpopuli facebook/voxpopuli en zero 0.0
|
| 154 |
+
done
|
| 155 |
+
|
| 156 |
+
"${PYTHON_BIN}" scripts/summarize_paper_tbd.py "${OUT_DIR}" --baseline-dir "${BASELINE_DIR}"
|
| 157 |
+
|
| 158 |
+
echo "Outputs: ${OUT_DIR}"
|
experiments/lib/common.sh
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
|
| 3 |
+
# Shared shell helpers for LatentASR experiment runners.
|
| 4 |
+
# This file is sourced by scripts under experiments/.
|
| 5 |
+
|
| 6 |
+
latent_asr_repo_root() {
|
| 7 |
+
local script_dir="$1"
|
| 8 |
+
cd "${script_dir}/.." && pwd
|
| 9 |
+
}
|
| 10 |
+
|
| 11 |
+
latent_asr_python_bin() {
|
| 12 |
+
local requested="${1:-}"
|
| 13 |
+
if [[ -n "${requested}" ]]; then
|
| 14 |
+
printf '%s\n' "${requested}"
|
| 15 |
+
elif [[ -x "/user_data/miniconda3/envs/py311/bin/python" ]]; then
|
| 16 |
+
printf '%s\n' "/user_data/miniconda3/envs/py311/bin/python"
|
| 17 |
+
else
|
| 18 |
+
printf '%s\n' "python"
|
| 19 |
+
fi
|
| 20 |
+
}
|
| 21 |
+
|
| 22 |
+
latent_asr_require_file() {
|
| 23 |
+
local path="$1"
|
| 24 |
+
local label="$2"
|
| 25 |
+
if [[ ! -f "${path}" ]]; then
|
| 26 |
+
echo "[error] ${label} not found: ${path}" >&2
|
| 27 |
+
exit 1
|
| 28 |
+
fi
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
latent_asr_print_kv() {
|
| 32 |
+
local key="$1"
|
| 33 |
+
local value="$2"
|
| 34 |
+
printf '%-16s: %s\n' "${key}" "${value}"
|
| 35 |
+
}
|
experiments/run_all_training_modes.sh
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
# ==============================================================================
|
| 3 |
+
# run_all_modes.sh — 依序跑 LoRA、Prompt Tuning、Baseline Fine-tune、Latent Reasoning
|
| 4 |
+
#
|
| 5 |
+
# 每個 mode 獨立訓練,checkpoint 存在各自的目錄。
|
| 6 |
+
# 用法:
|
| 7 |
+
# bash run_all_modes.sh # 用預設設定
|
| 8 |
+
# NUM_EPOCHS=5 bash run_all_modes.sh # 自訂 epoch 數
|
| 9 |
+
# ==============================================================================
|
| 10 |
+
set -euo pipefail
|
| 11 |
+
|
| 12 |
+
# ---------- 共用設定(可透過環境變數覆蓋)----------
|
| 13 |
+
export MODEL_ID="${MODEL_ID:-Qwen/Qwen3-ASR-0.6B}"
|
| 14 |
+
export DATASET_NAME="${DATASET_NAME:-SpeechTest/extreme_asr_pony}"
|
| 15 |
+
export DATASET_CONFIG="${DATASET_CONFIG:-default}"
|
| 16 |
+
export BATCH_SIZE="${BATCH_SIZE:-4}"
|
| 17 |
+
export GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-4}"
|
| 18 |
+
export NUM_EPOCHS="${NUM_EPOCHS:-10}"
|
| 19 |
+
export EVAL_SAMPLES="${EVAL_SAMPLES:-1000}"
|
| 20 |
+
export PRETRAIN_EVAL_SAMPLES="${PRETRAIN_EVAL_SAMPLES:-10}"
|
| 21 |
+
export TRAIN_MAX_SAMPLES="${TRAIN_MAX_SAMPLES:-500}"
|
| 22 |
+
|
| 23 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
| 24 |
+
source "${SCRIPT_DIR}/lib/common.sh"
|
| 25 |
+
ROOT_DIR="$(latent_asr_repo_root "${SCRIPT_DIR}")"
|
| 26 |
+
cd "${ROOT_DIR}"
|
| 27 |
+
PYTHON_BIN="$(latent_asr_python_bin "${PYTHON_BIN:-}")"
|
| 28 |
+
TRAIN_PY="${ROOT_DIR}/train.py"
|
| 29 |
+
|
| 30 |
+
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
| 31 |
+
LOG_DIR="${ROOT_DIR}/logs/${TIMESTAMP}"
|
| 32 |
+
mkdir -p "${LOG_DIR}"
|
| 33 |
+
|
| 34 |
+
echo "=============================================="
|
| 35 |
+
echo " Sequential Training — ${TIMESTAMP}"
|
| 36 |
+
echo " Model: ${MODEL_ID}"
|
| 37 |
+
echo " Dataset: ${DATASET_NAME} / ${DATASET_CONFIG}"
|
| 38 |
+
echo " Epochs: ${NUM_EPOCHS}"
|
| 39 |
+
echo " Train max: ${TRAIN_MAX_SAMPLES}"
|
| 40 |
+
echo " Logs: ${LOG_DIR}"
|
| 41 |
+
echo "=============================================="
|
| 42 |
+
|
| 43 |
+
# ---------- 1. LoRA (rank=16) ----------
|
| 44 |
+
echo ""
|
| 45 |
+
echo ">>> [1/4] Training: lora_r16"
|
| 46 |
+
echo "----------------------------------------------"
|
| 47 |
+
export TRAIN_MODE=lora_r16
|
| 48 |
+
LOG_FILE="${LOG_DIR}/lora_r16.log"
|
| 49 |
+
|
| 50 |
+
"${PYTHON_BIN}" "${TRAIN_PY}" 2>&1 | tee "${LOG_FILE}"
|
| 51 |
+
LORA_EXIT=$?
|
| 52 |
+
|
| 53 |
+
if [ ${LORA_EXIT} -ne 0 ]; then
|
| 54 |
+
echo "[FAIL] lora_r16 exited with code ${LORA_EXIT}"
|
| 55 |
+
else
|
| 56 |
+
echo "[DONE] lora_r16 completed successfully"
|
| 57 |
+
fi
|
| 58 |
+
|
| 59 |
+
# ---------- 2. Prompt Tuning ----------
|
| 60 |
+
echo ""
|
| 61 |
+
echo ">>> [2/4] Training: prompt_tuning"
|
| 62 |
+
echo "----------------------------------------------"
|
| 63 |
+
export TRAIN_MODE=prompt_tuning
|
| 64 |
+
LOG_FILE="${LOG_DIR}/prompt_tuning.log"
|
| 65 |
+
|
| 66 |
+
"${PYTHON_BIN}" "${TRAIN_PY}" 2>&1 | tee "${LOG_FILE}"
|
| 67 |
+
PT_EXIT=$?
|
| 68 |
+
|
| 69 |
+
if [ ${PT_EXIT} -ne 0 ]; then
|
| 70 |
+
echo "[FAIL] prompt_tuning exited with code ${PT_EXIT}"
|
| 71 |
+
else
|
| 72 |
+
echo "[DONE] prompt_tuning completed successfully"
|
| 73 |
+
fi
|
| 74 |
+
|
| 75 |
+
# ---------- 3. Baseline Fine-tune ----------
|
| 76 |
+
echo ""
|
| 77 |
+
echo ">>> [3/4] Training: baseline (full fine-tune)"
|
| 78 |
+
echo "----------------------------------------------"
|
| 79 |
+
export TRAIN_MODE=baseline
|
| 80 |
+
LOG_FILE="${LOG_DIR}/baseline.log"
|
| 81 |
+
|
| 82 |
+
"${PYTHON_BIN}" "${TRAIN_PY}" 2>&1 | tee "${LOG_FILE}"
|
| 83 |
+
BL_EXIT=$?
|
| 84 |
+
|
| 85 |
+
if [ ${BL_EXIT} -ne 0 ]; then
|
| 86 |
+
echo "[FAIL] baseline exited with code ${BL_EXIT}"
|
| 87 |
+
else
|
| 88 |
+
echo "[DONE] baseline completed successfully"
|
| 89 |
+
fi
|
| 90 |
+
|
| 91 |
+
# ---------- 4. Latent Reasoning ----------
|
| 92 |
+
echo ""
|
| 93 |
+
echo ">>> [4/4] Training: latent (latent reasoning)"
|
| 94 |
+
echo "----------------------------------------------"
|
| 95 |
+
export TRAIN_MODE=latent
|
| 96 |
+
LOG_FILE="${LOG_DIR}/latent.log"
|
| 97 |
+
|
| 98 |
+
"${PYTHON_BIN}" "${TRAIN_PY}" 2>&1 | tee "${LOG_FILE}"
|
| 99 |
+
LR_EXIT=$?
|
| 100 |
+
|
| 101 |
+
if [ ${LR_EXIT} -ne 0 ]; then
|
| 102 |
+
echo "[FAIL] latent exited with code ${LR_EXIT}"
|
| 103 |
+
else
|
| 104 |
+
echo "[DONE] latent completed successfully"
|
| 105 |
+
fi
|
| 106 |
+
|
| 107 |
+
# ---------- Summary ----------
|
| 108 |
+
echo ""
|
| 109 |
+
echo "=============================================="
|
| 110 |
+
echo " Training Summary"
|
| 111 |
+
echo "=============================================="
|
| 112 |
+
echo " lora_r16: exit=${LORA_EXIT:-skipped}"
|
| 113 |
+
echo " prompt_tuning: exit=${PT_EXIT:-skipped}"
|
| 114 |
+
echo " baseline: exit=${BL_EXIT}"
|
| 115 |
+
echo " latent: exit=${LR_EXIT}"
|
| 116 |
+
echo " Logs: ${LOG_DIR}/"
|
| 117 |
+
echo "=============================================="
|
| 118 |
+
|
| 119 |
+
# Print best WER from each log if available
|
| 120 |
+
for mode in lora_r16 prompt_tuning baseline latent; do
|
| 121 |
+
logf="${LOG_DIR}/${mode}.log"
|
| 122 |
+
if [ -f "${logf}" ]; then
|
| 123 |
+
best=$(grep -i "best.*wer" "${logf}" | tail -1 || true)
|
| 124 |
+
if [ -n "${best}" ]; then
|
| 125 |
+
echo " ${mode}: ${best}"
|
| 126 |
+
fi
|
| 127 |
+
fi
|
| 128 |
+
done
|
experiments/run_english_clean_and_noise.sh
ADDED
|
@@ -0,0 +1,169 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
set -euo pipefail
|
| 3 |
+
|
| 4 |
+
# Run a small-to-medium HuggingFace ASR showcase for latent reasoning.
|
| 5 |
+
#
|
| 6 |
+
# Goal:
|
| 7 |
+
# 1. Start with small, recognizable HF ASR datasets.
|
| 8 |
+
# 2. Compare raw base_model vs latent_reasoning only.
|
| 9 |
+
# 3. Include optional noisy SNR stress cases where LR benefits are easier to see.
|
| 10 |
+
#
|
| 11 |
+
# Fast discovery:
|
| 12 |
+
# ./run_lr_hf_asr_showcase.sh
|
| 13 |
+
#
|
| 14 |
+
# Full confirmation after finding promising cases:
|
| 15 |
+
# MAX_SAMPLES_PER_CONFIG=0 SNR_DB_LEVELS="5 0" ./run_lr_hf_asr_showcase.sh
|
| 16 |
+
#
|
| 17 |
+
# Filter cases:
|
| 18 |
+
# CASE_FILTER='fleurs|tedlium' ./run_lr_hf_asr_showcase.sh
|
| 19 |
+
|
| 20 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
| 21 |
+
source "${SCRIPT_DIR}/lib/common.sh"
|
| 22 |
+
ROOT_DIR="$(latent_asr_repo_root "${SCRIPT_DIR}")"
|
| 23 |
+
cd "${ROOT_DIR}"
|
| 24 |
+
|
| 25 |
+
LATENT_CKPT="${LATENT_CKPT:-./latent_qwen_asr_best.pth}"
|
| 26 |
+
MODEL_ID="${MODEL_ID:-Qwen/Qwen3-ASR-0.6B}"
|
| 27 |
+
PYTHON_BIN="$(latent_asr_python_bin "${PYTHON_BIN:-}")"
|
| 28 |
+
DYNAMIC_HALT_THRESHOLD="${DYNAMIC_HALT_THRESHOLD:-0.0}"
|
| 29 |
+
MAX_SAMPLES_PER_CONFIG="${MAX_SAMPLES_PER_CONFIG:-500}"
|
| 30 |
+
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"
|
| 31 |
+
PRINT_SAMPLES="${PRINT_SAMPLES:-0}"
|
| 32 |
+
RUN_CLEAN="${RUN_CLEAN:-1}"
|
| 33 |
+
SNR_DB_LEVELS="${SNR_DB_LEVELS-10 5 0}"
|
| 34 |
+
CASE_FILTER="${CASE_FILTER:-}"
|
| 35 |
+
RESUME="${RESUME:-1}"
|
| 36 |
+
SNR_LEVEL_ARRAY=()
|
| 37 |
+
if [[ -n "${SNR_DB_LEVELS}" ]]; then
|
| 38 |
+
read -r -a SNR_LEVEL_ARRAY <<< "${SNR_DB_LEVELS}"
|
| 39 |
+
fi
|
| 40 |
+
|
| 41 |
+
latent_asr_require_file "${LATENT_CKPT}" "latent checkpoint"
|
| 42 |
+
|
| 43 |
+
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
| 44 |
+
OUT_DIR="${OUT_DIR:-${ROOT_DIR}/eval_runs/hf_asr_showcase_${TIMESTAMP}}"
|
| 45 |
+
LOG_DIR="${OUT_DIR}/logs"
|
| 46 |
+
mkdir -p "${LOG_DIR}"
|
| 47 |
+
|
| 48 |
+
# Format: tag|dataset|configs|split
|
| 49 |
+
# Ordered roughly small/quick first, then broader canonical benchmarks.
|
| 50 |
+
HF_ASR_CASES=(
|
| 51 |
+
"fleurs_en_us|google/fleurs|en_us|test"
|
| 52 |
+
"minds14_en_us|PolyAI/minds14|en-US|train"
|
| 53 |
+
"tedlium_release1|TwinkStart/tedlium|release1|test"
|
| 54 |
+
"voxpopuli_en|facebook/voxpopuli|en|test"
|
| 55 |
+
"librispeech_clean_other|openslr/librispeech_asr|clean,other|test"
|
| 56 |
+
)
|
| 57 |
+
|
| 58 |
+
EXTRA_ARGS=("$@")
|
| 59 |
+
|
| 60 |
+
run_case() {
|
| 61 |
+
local tag="$1"
|
| 62 |
+
local dataset="$2"
|
| 63 |
+
local configs="$3"
|
| 64 |
+
local split="$4"
|
| 65 |
+
local condition="$5"
|
| 66 |
+
local snr="$6"
|
| 67 |
+
|
| 68 |
+
local json_out="${OUT_DIR}/${tag}_${condition}.json"
|
| 69 |
+
local log_out="${LOG_DIR}/${tag}_${condition}.log"
|
| 70 |
+
|
| 71 |
+
if [[ "${RESUME}" == "1" && -f "${json_out}" ]]; then
|
| 72 |
+
echo "[skip] existing ${json_out}"
|
| 73 |
+
return 0
|
| 74 |
+
fi
|
| 75 |
+
|
| 76 |
+
local args=(
|
| 77 |
+
--model-id "${MODEL_ID}"
|
| 78 |
+
--dataset-name "${dataset}"
|
| 79 |
+
--configs "${configs}"
|
| 80 |
+
--split "${split}"
|
| 81 |
+
--output-json "${json_out}"
|
| 82 |
+
--max-samples-per-config "${MAX_SAMPLES_PER_CONFIG}"
|
| 83 |
+
--max-new-tokens "${MAX_NEW_TOKENS}"
|
| 84 |
+
--latent-ckpt "${LATENT_CKPT}"
|
| 85 |
+
--n-latent 4
|
| 86 |
+
--num-beams 1
|
| 87 |
+
--dynamic-halt-threshold "${DYNAMIC_HALT_THRESHOLD}"
|
| 88 |
+
--print-samples "${PRINT_SAMPLES}"
|
| 89 |
+
--skip-baseline-ft
|
| 90 |
+
--skip-prompt-tuning
|
| 91 |
+
--skip-lora-r16
|
| 92 |
+
)
|
| 93 |
+
|
| 94 |
+
if [[ -n "${snr}" ]]; then
|
| 95 |
+
args+=(--snr-db "${snr}")
|
| 96 |
+
fi
|
| 97 |
+
|
| 98 |
+
echo "------------------------------------------------------------"
|
| 99 |
+
echo "Case : ${tag}"
|
| 100 |
+
echo "Dataset : ${dataset}"
|
| 101 |
+
echo "Configs : ${configs}"
|
| 102 |
+
echo "Split : ${split}"
|
| 103 |
+
echo "Condition : ${condition}"
|
| 104 |
+
echo "JSON : ${json_out}"
|
| 105 |
+
echo "LOG : ${log_out}"
|
| 106 |
+
|
| 107 |
+
"${PYTHON_BIN}" eval.py "${args[@]}" "${EXTRA_ARGS[@]}" 2>&1 | tee "${log_out}"
|
| 108 |
+
}
|
| 109 |
+
|
| 110 |
+
echo "============================================================"
|
| 111 |
+
echo "Latent Reasoning HuggingFace ASR Showcase"
|
| 112 |
+
echo "============================================================"
|
| 113 |
+
echo "Root dir : ${ROOT_DIR}"
|
| 114 |
+
echo "Output dir : ${OUT_DIR}"
|
| 115 |
+
echo "Python : ${PYTHON_BIN}"
|
| 116 |
+
echo "Model : ${MODEL_ID}"
|
| 117 |
+
echo "Latent ckpt : ${LATENT_CKPT}"
|
| 118 |
+
echo "Halt threshold : ${DYNAMIC_HALT_THRESHOLD}"
|
| 119 |
+
echo "Max/config : ${MAX_SAMPLES_PER_CONFIG}"
|
| 120 |
+
echo "Clean pass : ${RUN_CLEAN}"
|
| 121 |
+
echo "SNR levels : ${SNR_DB_LEVELS:-(none)}"
|
| 122 |
+
echo "Case filter : ${CASE_FILTER:-(none)}"
|
| 123 |
+
echo "Extra args : ${EXTRA_ARGS[*]:-(none)}"
|
| 124 |
+
echo "============================================================"
|
| 125 |
+
|
| 126 |
+
SUCCESS=()
|
| 127 |
+
FAILED=()
|
| 128 |
+
|
| 129 |
+
for entry in "${HF_ASR_CASES[@]}"; do
|
| 130 |
+
IFS='|' read -r TAG DATASET CONFIGS SPLIT <<< "${entry}"
|
| 131 |
+
|
| 132 |
+
if [[ -n "${CASE_FILTER}" && ! "${TAG}" =~ ${CASE_FILTER} ]]; then
|
| 133 |
+
continue
|
| 134 |
+
fi
|
| 135 |
+
|
| 136 |
+
if [[ "${RUN_CLEAN}" == "1" ]]; then
|
| 137 |
+
if run_case "${TAG}" "${DATASET}" "${CONFIGS}" "${SPLIT}" "clean" ""; then
|
| 138 |
+
SUCCESS+=("${TAG}:clean")
|
| 139 |
+
else
|
| 140 |
+
FAILED+=("${TAG}:clean")
|
| 141 |
+
fi
|
| 142 |
+
fi
|
| 143 |
+
|
| 144 |
+
for SNR in "${SNR_LEVEL_ARRAY[@]}"; do
|
| 145 |
+
CONDITION="snr${SNR}db"
|
| 146 |
+
if run_case "${TAG}" "${DATASET}" "${CONFIGS}" "${SPLIT}" "${CONDITION}" "${SNR}"; then
|
| 147 |
+
SUCCESS+=("${TAG}:${CONDITION}")
|
| 148 |
+
else
|
| 149 |
+
FAILED+=("${TAG}:${CONDITION}")
|
| 150 |
+
fi
|
| 151 |
+
done
|
| 152 |
+
done
|
| 153 |
+
|
| 154 |
+
echo ""
|
| 155 |
+
echo "============================================================"
|
| 156 |
+
echo "Showcase Summary"
|
| 157 |
+
echo "============================================================"
|
| 158 |
+
echo "Succeeded (${#SUCCESS[@]}):"
|
| 159 |
+
for item in "${SUCCESS[@]}"; do
|
| 160 |
+
echo " - ${item}"
|
| 161 |
+
done
|
| 162 |
+
echo "Failed (${#FAILED[@]}):"
|
| 163 |
+
for item in "${FAILED[@]}"; do
|
| 164 |
+
echo " - ${item}"
|
| 165 |
+
done
|
| 166 |
+
|
| 167 |
+
"${PYTHON_BIN}" scripts/summarize_lr_showcase.py "${OUT_DIR}"
|
| 168 |
+
echo "Outputs: ${OUT_DIR}"
|
| 169 |
+
echo "============================================================"
|
experiments/run_multilingual_streaming.sh
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
set -euo pipefail
|
| 3 |
+
|
| 4 |
+
# Multilingual ASR evaluation aligned with Qwen3-ASR public benchmark families.
|
| 5 |
+
#
|
| 6 |
+
# This runner always passes --streaming so HuggingFace datasets are read lazily
|
| 7 |
+
# instead of materializing full audio splits under the local cache.
|
| 8 |
+
#
|
| 9 |
+
# Default run:
|
| 10 |
+
# ./run_lr_multilingual_asr_streaming.sh
|
| 11 |
+
#
|
| 12 |
+
# Run every configured case:
|
| 13 |
+
# CASE_FILTER='' ./run_lr_multilingual_asr_streaming.sh
|
| 14 |
+
#
|
| 15 |
+
# Quick smoke test:
|
| 16 |
+
# MAX_SAMPLES_PER_CONFIG=20 CASE_FILTER='fleurs_core12' ./run_lr_multilingual_asr_streaming.sh
|
| 17 |
+
|
| 18 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
| 19 |
+
source "${SCRIPT_DIR}/lib/common.sh"
|
| 20 |
+
ROOT_DIR="$(latent_asr_repo_root "${SCRIPT_DIR}")"
|
| 21 |
+
cd "${ROOT_DIR}"
|
| 22 |
+
|
| 23 |
+
LATENT_CKPT="${LATENT_CKPT:-./latent_qwen_asr_best.pth}"
|
| 24 |
+
MODEL_ID="${MODEL_ID:-Qwen/Qwen3-ASR-0.6B}"
|
| 25 |
+
PYTHON_BIN="$(latent_asr_python_bin "${PYTHON_BIN:-}")"
|
| 26 |
+
|
| 27 |
+
DYNAMIC_HALT_THRESHOLD="${DYNAMIC_HALT_THRESHOLD:-0.0}"
|
| 28 |
+
MAX_SAMPLES_PER_CONFIG="${MAX_SAMPLES_PER_CONFIG:-0}"
|
| 29 |
+
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"
|
| 30 |
+
PRINT_SAMPLES="${PRINT_SAMPLES:-0}"
|
| 31 |
+
TEXT_NORMALIZER="${TEXT_NORMALIZER:-basic}"
|
| 32 |
+
RESUME="${RESUME:-1}"
|
| 33 |
+
|
| 34 |
+
# Default to the 30-language FLEURS grouping from the Qwen3-ASR card.
|
| 35 |
+
# Set CASE_FILTER='' to also run the public MLS case below.
|
| 36 |
+
CASE_FILTER="${CASE_FILTER:-fleurs_core12|fleurs_extra8|fleurs_extra10}"
|
| 37 |
+
|
| 38 |
+
latent_asr_require_file "${LATENT_CKPT}" "latent checkpoint"
|
| 39 |
+
|
| 40 |
+
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
| 41 |
+
OUT_DIR="${OUT_DIR:-${ROOT_DIR}/eval_runs/multilingual_asr_streaming_${TIMESTAMP}}"
|
| 42 |
+
LOG_DIR="${OUT_DIR}/logs"
|
| 43 |
+
mkdir -p "${LOG_DIR}"
|
| 44 |
+
|
| 45 |
+
# Format: tag|dataset|configs|split
|
| 46 |
+
# Qwen3-ASR public multilingual ASR benchmark families include FLEURS,
|
| 47 |
+
# CommonVoice, MLS, and MLC-SLM. FLEURS is the broadest open HF target here.
|
| 48 |
+
# HuggingFace facebook/multilingual_librispeech currently exposes seven
|
| 49 |
+
# non-English MLS configs; English is covered separately by LibriSpeech runs.
|
| 50 |
+
MULTILINGUAL_CASES=(
|
| 51 |
+
"fleurs_core12|google/fleurs|en,zh,yue,ar,de,es,fr,it,ja,ko,pt,ru|test"
|
| 52 |
+
"fleurs_extra8|google/fleurs|hi,id,ms,nl,pl,th,tr,vi|test"
|
| 53 |
+
"fleurs_extra10|google/fleurs|cs,da,el,fa,fi,fil,hu,mk,ro,sv|test"
|
| 54 |
+
"mls_public7|facebook/multilingual_librispeech|de,nl,es,fr,it,pl,pt|test"
|
| 55 |
+
)
|
| 56 |
+
|
| 57 |
+
EXTRA_ARGS=("$@")
|
| 58 |
+
|
| 59 |
+
run_case() {
|
| 60 |
+
local tag="$1"
|
| 61 |
+
local dataset="$2"
|
| 62 |
+
local configs="$3"
|
| 63 |
+
local split="$4"
|
| 64 |
+
|
| 65 |
+
local json_out="${OUT_DIR}/${tag}_streaming_clean.json"
|
| 66 |
+
local log_out="${LOG_DIR}/${tag}_streaming_clean.log"
|
| 67 |
+
|
| 68 |
+
if [[ "${RESUME}" == "1" && -f "${json_out}" ]]; then
|
| 69 |
+
echo "[skip] existing ${json_out}"
|
| 70 |
+
return 0
|
| 71 |
+
fi
|
| 72 |
+
|
| 73 |
+
local args=(
|
| 74 |
+
--model-id "${MODEL_ID}"
|
| 75 |
+
--dataset-name "${dataset}"
|
| 76 |
+
--configs "${configs}"
|
| 77 |
+
--split "${split}"
|
| 78 |
+
--output-json "${json_out}"
|
| 79 |
+
--max-samples-per-config "${MAX_SAMPLES_PER_CONFIG}"
|
| 80 |
+
--max-new-tokens "${MAX_NEW_TOKENS}"
|
| 81 |
+
--latent-ckpt "${LATENT_CKPT}"
|
| 82 |
+
--n-latent 4
|
| 83 |
+
--num-beams 1
|
| 84 |
+
--dynamic-halt-threshold "${DYNAMIC_HALT_THRESHOLD}"
|
| 85 |
+
--print-samples "${PRINT_SAMPLES}"
|
| 86 |
+
--text-normalizer "${TEXT_NORMALIZER}"
|
| 87 |
+
--streaming
|
| 88 |
+
--skip-baseline-ft
|
| 89 |
+
--skip-prompt-tuning
|
| 90 |
+
--skip-lora-r16
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
echo "------------------------------------------------------------"
|
| 94 |
+
echo "Case : ${tag}"
|
| 95 |
+
echo "Dataset : ${dataset}"
|
| 96 |
+
echo "Configs : ${configs}"
|
| 97 |
+
echo "Split : ${split}"
|
| 98 |
+
echo "Streaming : yes"
|
| 99 |
+
echo "JSON : ${json_out}"
|
| 100 |
+
echo "LOG : ${log_out}"
|
| 101 |
+
|
| 102 |
+
"${PYTHON_BIN}" eval.py "${args[@]}" "${EXTRA_ARGS[@]}" 2>&1 | tee "${log_out}"
|
| 103 |
+
}
|
| 104 |
+
|
| 105 |
+
echo "============================================================"
|
| 106 |
+
echo "Latent Reasoning Multilingual ASR Streaming"
|
| 107 |
+
echo "============================================================"
|
| 108 |
+
echo "Root dir : ${ROOT_DIR}"
|
| 109 |
+
echo "Output dir : ${OUT_DIR}"
|
| 110 |
+
echo "Python : ${PYTHON_BIN}"
|
| 111 |
+
echo "Model : ${MODEL_ID}"
|
| 112 |
+
echo "Latent ckpt : ${LATENT_CKPT}"
|
| 113 |
+
echo "Halt threshold : ${DYNAMIC_HALT_THRESHOLD}"
|
| 114 |
+
echo "Max/config : ${MAX_SAMPLES_PER_CONFIG} (0 = full split)"
|
| 115 |
+
echo "Text normalizer : ${TEXT_NORMALIZER}"
|
| 116 |
+
echo "Case filter : ${CASE_FILTER:-(none)}"
|
| 117 |
+
echo "Extra args : ${EXTRA_ARGS[*]:-(none)}"
|
| 118 |
+
echo "============================================================"
|
| 119 |
+
|
| 120 |
+
SUCCESS=()
|
| 121 |
+
FAILED=()
|
| 122 |
+
|
| 123 |
+
for entry in "${MULTILINGUAL_CASES[@]}"; do
|
| 124 |
+
IFS='|' read -r TAG DATASET CONFIGS SPLIT <<< "${entry}"
|
| 125 |
+
|
| 126 |
+
if [[ -n "${CASE_FILTER}" && ! "${TAG}" =~ ${CASE_FILTER} ]]; then
|
| 127 |
+
continue
|
| 128 |
+
fi
|
| 129 |
+
|
| 130 |
+
if run_case "${TAG}" "${DATASET}" "${CONFIGS}" "${SPLIT}"; then
|
| 131 |
+
SUCCESS+=("${TAG}:streaming_clean")
|
| 132 |
+
else
|
| 133 |
+
FAILED+=("${TAG}:streaming_clean")
|
| 134 |
+
fi
|
| 135 |
+
done
|
| 136 |
+
|
| 137 |
+
echo ""
|
| 138 |
+
echo "============================================================"
|
| 139 |
+
echo "Multilingual Streaming Summary"
|
| 140 |
+
echo "============================================================"
|
| 141 |
+
echo "Succeeded (${#SUCCESS[@]}):"
|
| 142 |
+
for item in "${SUCCESS[@]}"; do
|
| 143 |
+
echo " - ${item}"
|
| 144 |
+
done
|
| 145 |
+
echo "Failed (${#FAILED[@]}):"
|
| 146 |
+
for item in "${FAILED[@]}"; do
|
| 147 |
+
echo " - ${item}"
|
| 148 |
+
done
|
| 149 |
+
|
| 150 |
+
"${PYTHON_BIN}" scripts/summarize_lr_showcase.py "${OUT_DIR}"
|
| 151 |
+
echo "Outputs: ${OUT_DIR}"
|
| 152 |
+
echo "============================================================"
|
experiments/run_paper_tbd_retrain.sh
ADDED
|
@@ -0,0 +1,112 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
set -euo pipefail
|
| 3 |
+
|
| 4 |
+
# Retrain LatentASR variants needed by the paper TBD ablations.
|
| 5 |
+
# This runner intentionally saves per-epoch checkpoints. Downstream evaluation
|
| 6 |
+
# should use *_epoch10.pth for same-schedule comparisons unless noted otherwise.
|
| 7 |
+
|
| 8 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
| 9 |
+
source "${SCRIPT_DIR}/lib/common.sh"
|
| 10 |
+
ROOT_DIR="$(latent_asr_repo_root "${SCRIPT_DIR}")"
|
| 11 |
+
cd "${ROOT_DIR}"
|
| 12 |
+
|
| 13 |
+
PYTHON_BIN="$(latent_asr_python_bin "${PYTHON_BIN:-}")"
|
| 14 |
+
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
| 15 |
+
OUT_DIR="${OUT_DIR:-${ROOT_DIR}/eval_runs/paper_tbd_retrain_${TIMESTAMP}}"
|
| 16 |
+
LOG_DIR="${OUT_DIR}/logs"
|
| 17 |
+
CKPT_DIR="${OUT_DIR}/checkpoints"
|
| 18 |
+
mkdir -p "${LOG_DIR}" "${CKPT_DIR}"
|
| 19 |
+
|
| 20 |
+
NUM_EPOCHS="${NUM_EPOCHS:-10}"
|
| 21 |
+
BATCH_SIZE="${BATCH_SIZE:-16}"
|
| 22 |
+
GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-1}"
|
| 23 |
+
EVAL_SAMPLES="${EVAL_SAMPLES:-1}"
|
| 24 |
+
PRETRAIN_EVAL_SAMPLES="${PRETRAIN_EVAL_SAMPLES:-0}"
|
| 25 |
+
VARIANT_FILTER="${VARIANT_FILTER:-}"
|
| 26 |
+
DEFAULT_TRAIN_MAX="${DEFAULT_TRAIN_MAX:-0}"
|
| 27 |
+
|
| 28 |
+
# name|n_latent|train_max|p_neg|bounded|gate|anchor
|
| 29 |
+
VARIANTS=(
|
| 30 |
+
"component_no_bounded|4|${DEFAULT_TRAIN_MAX}|0.3|0|1|1"
|
| 31 |
+
"component_no_gate|4|${DEFAULT_TRAIN_MAX}|0.3|1|0|1"
|
| 32 |
+
"component_no_anchor|4|${DEFAULT_TRAIN_MAX}|0.3|1|1|0"
|
| 33 |
+
"n1|1|${DEFAULT_TRAIN_MAX}|0.3|1|1|1"
|
| 34 |
+
"n2|2|${DEFAULT_TRAIN_MAX}|0.3|1|1|1"
|
| 35 |
+
"n4|4|${DEFAULT_TRAIN_MAX}|0.3|1|1|1"
|
| 36 |
+
"n8|8|${DEFAULT_TRAIN_MAX}|0.3|1|1|1"
|
| 37 |
+
"pneg0|4|${DEFAULT_TRAIN_MAX}|0.0|1|1|1"
|
| 38 |
+
"activation_100|4|100|0.3|1|1|1"
|
| 39 |
+
"activation_200|4|200|0.3|1|1|1"
|
| 40 |
+
"activation_300|4|300|0.3|1|1|1"
|
| 41 |
+
"activation_400|4|400|0.3|1|1|1"
|
| 42 |
+
"activation_500|4|500|0.3|1|1|1"
|
| 43 |
+
"activation_600|4|600|0.3|1|1|1"
|
| 44 |
+
"activation_700|4|700|0.3|1|1|1"
|
| 45 |
+
"activation_800|4|800|0.3|1|1|1"
|
| 46 |
+
)
|
| 47 |
+
|
| 48 |
+
echo "============================================================"
|
| 49 |
+
echo "LatentASR paper TBD retraining"
|
| 50 |
+
echo "============================================================"
|
| 51 |
+
latent_asr_print_kv "Root" "${ROOT_DIR}"
|
| 52 |
+
latent_asr_print_kv "Output" "${OUT_DIR}"
|
| 53 |
+
latent_asr_print_kv "Python" "${PYTHON_BIN}"
|
| 54 |
+
latent_asr_print_kv "Epochs" "${NUM_EPOCHS}"
|
| 55 |
+
latent_asr_print_kv "Batch" "${BATCH_SIZE}"
|
| 56 |
+
latent_asr_print_kv "Grad accum" "${GRAD_ACCUM_STEPS}"
|
| 57 |
+
latent_asr_print_kv "Eval samples" "${EVAL_SAMPLES}"
|
| 58 |
+
latent_asr_print_kv "Variant filter" "${VARIANT_FILTER:-(none)}"
|
| 59 |
+
latent_asr_print_kv "Default train max" "${DEFAULT_TRAIN_MAX} (0 = full 811-sample activation set)"
|
| 60 |
+
echo "============================================================"
|
| 61 |
+
|
| 62 |
+
for spec in "${VARIANTS[@]}"; do
|
| 63 |
+
IFS='|' read -r name n_latent train_max p_neg bounded gate anchor <<< "${spec}"
|
| 64 |
+
if [[ -n "${VARIANT_FILTER}" && ! "${name}" =~ ${VARIANT_FILTER} ]]; then
|
| 65 |
+
continue
|
| 66 |
+
fi
|
| 67 |
+
|
| 68 |
+
variant_dir="${CKPT_DIR}/${name}"
|
| 69 |
+
mkdir -p "${variant_dir}"
|
| 70 |
+
prefix="${variant_dir}/${name}"
|
| 71 |
+
final_ckpt="${prefix}_epoch${NUM_EPOCHS}.pth"
|
| 72 |
+
log_path="${LOG_DIR}/${name}.log"
|
| 73 |
+
|
| 74 |
+
if [[ -f "${final_ckpt}" ]]; then
|
| 75 |
+
echo "[skip] ${name}: found ${final_ckpt}"
|
| 76 |
+
continue
|
| 77 |
+
fi
|
| 78 |
+
|
| 79 |
+
echo ""
|
| 80 |
+
echo "------------------------------------------------------------"
|
| 81 |
+
echo "Variant : ${name}"
|
| 82 |
+
echo "N : ${n_latent}"
|
| 83 |
+
echo "train_max : ${train_max:-0} (0 = full 811-sample activation set)"
|
| 84 |
+
echo "p_neg : ${p_neg}"
|
| 85 |
+
echo "bounded : ${bounded}"
|
| 86 |
+
echo "gate : ${gate}"
|
| 87 |
+
echo "anchor : ${anchor}"
|
| 88 |
+
echo "log : ${log_path}"
|
| 89 |
+
echo "------------------------------------------------------------"
|
| 90 |
+
|
| 91 |
+
env \
|
| 92 |
+
PYTHONUNBUFFERED=1 \
|
| 93 |
+
TOKENIZERS_PARALLELISM=false \
|
| 94 |
+
TRAIN_MODE=latent \
|
| 95 |
+
N_LATENT="${n_latent}" \
|
| 96 |
+
TRAIN_MAX_SAMPLES="${train_max}" \
|
| 97 |
+
VALUE_FORCED_NEG_PROB="${p_neg}" \
|
| 98 |
+
LATENT_USE_BOUNDED_DELTA="${bounded}" \
|
| 99 |
+
LATENT_USE_INJECTION_GATE="${gate}" \
|
| 100 |
+
LATENT_USE_EMBEDDING_ANCHOR="${anchor}" \
|
| 101 |
+
NUM_EPOCHS="${NUM_EPOCHS}" \
|
| 102 |
+
BATCH_SIZE="${BATCH_SIZE}" \
|
| 103 |
+
GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS}" \
|
| 104 |
+
EVAL_SAMPLES="${EVAL_SAMPLES}" \
|
| 105 |
+
PRETRAIN_EVAL_SAMPLES="${PRETRAIN_EVAL_SAMPLES}" \
|
| 106 |
+
LOG_EVERY=25 \
|
| 107 |
+
GRAD_LOG_EVERY=100 \
|
| 108 |
+
CHECKPOINT_PREFIX="${prefix}" \
|
| 109 |
+
"${PYTHON_BIN}" train.py 2>&1 | tee "${log_path}"
|
| 110 |
+
done
|
| 111 |
+
|
| 112 |
+
echo "Outputs: ${OUT_DIR}"
|
experiments/run_threshold_sweep.sh
ADDED
|
@@ -0,0 +1,145 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
set -euo pipefail
|
| 3 |
+
|
| 4 |
+
# Threshold sweep for the LatentASR value-head halting policy.
|
| 5 |
+
# Clean audio only, full splits by default, and HuggingFace streaming enabled
|
| 6 |
+
# so audio is not materialized under the workspace.
|
| 7 |
+
|
| 8 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
| 9 |
+
source "${SCRIPT_DIR}/lib/common.sh"
|
| 10 |
+
ROOT_DIR="$(latent_asr_repo_root "${SCRIPT_DIR}")"
|
| 11 |
+
cd "${ROOT_DIR}"
|
| 12 |
+
|
| 13 |
+
LATENT_CKPT="${LATENT_CKPT:-./latent_qwen_asr_best.pth}"
|
| 14 |
+
MODEL_ID="${MODEL_ID:-Qwen/Qwen3-ASR-0.6B}"
|
| 15 |
+
PYTHON_BIN="$(latent_asr_python_bin "${PYTHON_BIN:-}")"
|
| 16 |
+
|
| 17 |
+
MAX_SAMPLES_PER_CONFIG="${MAX_SAMPLES_PER_CONFIG:-0}"
|
| 18 |
+
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"
|
| 19 |
+
PRINT_SAMPLES="${PRINT_SAMPLES:-0}"
|
| 20 |
+
STREAMING="${STREAMING:-1}"
|
| 21 |
+
RUN_BASE="${RUN_BASE:-0}"
|
| 22 |
+
CASE_FILTER="${CASE_FILTER:-fleurs_en_us|voxpopuli_en}"
|
| 23 |
+
THRESHOLD_SPECS="${THRESHOLD_SPECS:-full:-2.0 neg0p2:-0.2 zero:0.0 pos0p2:0.2 pos0p5:0.5}"
|
| 24 |
+
RESUME="${RESUME:-1}"
|
| 25 |
+
|
| 26 |
+
latent_asr_require_file "${LATENT_CKPT}" "latent checkpoint"
|
| 27 |
+
|
| 28 |
+
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
| 29 |
+
OUT_DIR="${OUT_DIR:-${ROOT_DIR}/eval_runs/threshold_sweep_${TIMESTAMP}}"
|
| 30 |
+
LOG_DIR="${OUT_DIR}/logs"
|
| 31 |
+
mkdir -p "${LOG_DIR}"
|
| 32 |
+
|
| 33 |
+
# Format: tag|dataset|configs|split
|
| 34 |
+
CASES=(
|
| 35 |
+
"fleurs_en_us|google/fleurs|en_us|test"
|
| 36 |
+
"voxpopuli_en|facebook/voxpopuli|en|test"
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
EXTRA_ARGS=("$@")
|
| 40 |
+
|
| 41 |
+
run_case() {
|
| 42 |
+
local theta_label="$1"
|
| 43 |
+
local theta_value="$2"
|
| 44 |
+
local tag="$3"
|
| 45 |
+
local dataset="$4"
|
| 46 |
+
local configs="$5"
|
| 47 |
+
local split="$6"
|
| 48 |
+
|
| 49 |
+
local json_out="${OUT_DIR}/${tag}_theta_${theta_label}.json"
|
| 50 |
+
local log_out="${LOG_DIR}/${tag}_theta_${theta_label}.log"
|
| 51 |
+
|
| 52 |
+
if [[ "${RESUME}" == "1" && -f "${json_out}" ]]; then
|
| 53 |
+
echo "[skip] existing ${json_out}"
|
| 54 |
+
return 0
|
| 55 |
+
fi
|
| 56 |
+
|
| 57 |
+
local args=(
|
| 58 |
+
--model-id "${MODEL_ID}"
|
| 59 |
+
--dataset-name "${dataset}"
|
| 60 |
+
--configs "${configs}"
|
| 61 |
+
--split "${split}"
|
| 62 |
+
--output-json "${json_out}"
|
| 63 |
+
--max-samples-per-config "${MAX_SAMPLES_PER_CONFIG}"
|
| 64 |
+
--max-new-tokens "${MAX_NEW_TOKENS}"
|
| 65 |
+
--latent-ckpt "${LATENT_CKPT}"
|
| 66 |
+
--n-latent 4
|
| 67 |
+
--num-beams 1
|
| 68 |
+
--dynamic-halt-threshold "${theta_value}"
|
| 69 |
+
--print-samples "${PRINT_SAMPLES}"
|
| 70 |
+
--skip-baseline-ft
|
| 71 |
+
--skip-prompt-tuning
|
| 72 |
+
--skip-lora-r16
|
| 73 |
+
)
|
| 74 |
+
|
| 75 |
+
if [[ "${RUN_BASE}" != "1" ]]; then
|
| 76 |
+
args+=(--skip-base-model)
|
| 77 |
+
fi
|
| 78 |
+
if [[ "${STREAMING}" == "1" ]]; then
|
| 79 |
+
args+=(--streaming)
|
| 80 |
+
fi
|
| 81 |
+
|
| 82 |
+
echo "------------------------------------------------------------"
|
| 83 |
+
echo "Case : ${tag}"
|
| 84 |
+
echo "Dataset : ${dataset}"
|
| 85 |
+
echo "Configs : ${configs}"
|
| 86 |
+
echo "Split : ${split}"
|
| 87 |
+
echo "Theta : ${theta_label} (${theta_value})"
|
| 88 |
+
echo "Streaming : ${STREAMING}"
|
| 89 |
+
echo "Run base : ${RUN_BASE}"
|
| 90 |
+
echo "JSON : ${json_out}"
|
| 91 |
+
echo "LOG : ${log_out}"
|
| 92 |
+
|
| 93 |
+
"${PYTHON_BIN}" eval.py "${args[@]}" "${EXTRA_ARGS[@]}" 2>&1 | tee "${log_out}"
|
| 94 |
+
}
|
| 95 |
+
|
| 96 |
+
echo "============================================================"
|
| 97 |
+
echo "LatentASR Threshold Sweep"
|
| 98 |
+
echo "============================================================"
|
| 99 |
+
echo "Root dir : ${ROOT_DIR}"
|
| 100 |
+
echo "Output dir : ${OUT_DIR}"
|
| 101 |
+
echo "Python : ${PYTHON_BIN}"
|
| 102 |
+
echo "Model : ${MODEL_ID}"
|
| 103 |
+
echo "Latent ckpt : ${LATENT_CKPT}"
|
| 104 |
+
echo "Max/config : ${MAX_SAMPLES_PER_CONFIG} (0 = full split)"
|
| 105 |
+
echo "Streaming : ${STREAMING}"
|
| 106 |
+
echo "Run base : ${RUN_BASE}"
|
| 107 |
+
echo "Case filter : ${CASE_FILTER:-(none)}"
|
| 108 |
+
echo "Thresholds : ${THRESHOLD_SPECS}"
|
| 109 |
+
echo "Extra args : ${EXTRA_ARGS[*]:-(none)}"
|
| 110 |
+
echo "============================================================"
|
| 111 |
+
|
| 112 |
+
SUCCESS=()
|
| 113 |
+
FAILED=()
|
| 114 |
+
|
| 115 |
+
for spec in ${THRESHOLD_SPECS}; do
|
| 116 |
+
IFS=':' read -r THETA_LABEL THETA_VALUE <<< "${spec}"
|
| 117 |
+
for entry in "${CASES[@]}"; do
|
| 118 |
+
IFS='|' read -r TAG DATASET CONFIGS SPLIT <<< "${entry}"
|
| 119 |
+
if [[ -n "${CASE_FILTER}" && ! "${TAG}" =~ ${CASE_FILTER} ]]; then
|
| 120 |
+
continue
|
| 121 |
+
fi
|
| 122 |
+
if run_case "${THETA_LABEL}" "${THETA_VALUE}" "${TAG}" "${DATASET}" "${CONFIGS}" "${SPLIT}"; then
|
| 123 |
+
SUCCESS+=("${TAG}:${THETA_LABEL}")
|
| 124 |
+
else
|
| 125 |
+
FAILED+=("${TAG}:${THETA_LABEL}")
|
| 126 |
+
fi
|
| 127 |
+
done
|
| 128 |
+
done
|
| 129 |
+
|
| 130 |
+
"${PYTHON_BIN}" scripts/summarize_threshold_sweep.py "${OUT_DIR}"
|
| 131 |
+
|
| 132 |
+
echo ""
|
| 133 |
+
echo "============================================================"
|
| 134 |
+
echo "Threshold Sweep Summary"
|
| 135 |
+
echo "============================================================"
|
| 136 |
+
echo "Succeeded (${#SUCCESS[@]}):"
|
| 137 |
+
for item in "${SUCCESS[@]}"; do
|
| 138 |
+
echo " - ${item}"
|
| 139 |
+
done
|
| 140 |
+
echo "Failed (${#FAILED[@]}):"
|
| 141 |
+
for item in "${FAILED[@]}"; do
|
| 142 |
+
echo " - ${item}"
|
| 143 |
+
done
|
| 144 |
+
echo "Outputs: ${OUT_DIR}"
|
| 145 |
+
echo "============================================================"
|
hf_upload/README.md
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# HF Upload Helper
|
| 2 |
+
|
| 3 |
+
Run from the repository root:
|
| 4 |
+
|
| 5 |
+
```bash
|
| 6 |
+
python hf_upload/upload_to_hf.py --repo-id voidful/latentASR
|
| 7 |
+
```
|
| 8 |
+
|
| 9 |
+
The script uploads the entire project folder, including the released adapter
|
| 10 |
+
checkpoint and documentation, to a Hugging Face model repository.
|
| 11 |
+
|
hf_upload/upload_to_hf.py
ADDED
|
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Upload the latentASR project and adapter checkpoint to Hugging Face Hub."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import argparse
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
from huggingface_hub import HfApi
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def parse_args() -> argparse.Namespace:
|
| 12 |
+
parser = argparse.ArgumentParser(description="Upload latentASR to Hugging Face Hub.")
|
| 13 |
+
parser.add_argument("--repo-id", default="voidful/latentASR")
|
| 14 |
+
parser.add_argument("--private", action="store_true")
|
| 15 |
+
parser.add_argument("--revision", default=None)
|
| 16 |
+
parser.add_argument(
|
| 17 |
+
"--commit-message",
|
| 18 |
+
default="Release latentASR adapter, code, docs, and reproducibility artifacts",
|
| 19 |
+
)
|
| 20 |
+
return parser.parse_args()
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def main() -> None:
|
| 24 |
+
args = parse_args()
|
| 25 |
+
project_root = Path(__file__).resolve().parents[1]
|
| 26 |
+
api = HfApi()
|
| 27 |
+
api.create_repo(
|
| 28 |
+
repo_id=args.repo_id,
|
| 29 |
+
repo_type="model",
|
| 30 |
+
private=args.private,
|
| 31 |
+
exist_ok=True,
|
| 32 |
+
)
|
| 33 |
+
api.upload_folder(
|
| 34 |
+
repo_id=args.repo_id,
|
| 35 |
+
repo_type="model",
|
| 36 |
+
folder_path=str(project_root),
|
| 37 |
+
revision=args.revision,
|
| 38 |
+
commit_message=args.commit_message,
|
| 39 |
+
ignore_patterns=[
|
| 40 |
+
".git/*",
|
| 41 |
+
"__pycache__/*",
|
| 42 |
+
"*.pyc",
|
| 43 |
+
".pytest_cache/*",
|
| 44 |
+
"eval_runs/*",
|
| 45 |
+
"logs/*",
|
| 46 |
+
"wandb/*",
|
| 47 |
+
],
|
| 48 |
+
)
|
| 49 |
+
print(f"Uploaded {project_root} to https://huggingface.co/{args.repo_id}")
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
if __name__ == "__main__":
|
| 53 |
+
main()
|
| 54 |
+
|
losses.py
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Loss functions for latent-reasoning ASR training."""
|
| 2 |
+
|
| 3 |
+
from typing import Any
|
| 4 |
+
|
| 5 |
+
import torch
|
| 6 |
+
import torch.nn as nn
|
| 7 |
+
import torch.nn.functional as F
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
def _fmt(x: torch.Tensor) -> str:
|
| 11 |
+
"""Format a one-dimensional tensor for logging."""
|
| 12 |
+
return " | ".join([f"d{k}:{x[k].item():.3f}" for k in range(x.numel())])
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def contrastive_loss(view_a: torch.Tensor, view_b: torch.Tensor, temperature: float) -> torch.Tensor:
|
| 16 |
+
"""InfoNCE contrastive loss over batch (symmetric)."""
|
| 17 |
+
assert view_a.dim() == 2 and view_b.dim() == 2, "contrastive_loss expects (B, D) inputs"
|
| 18 |
+
B = view_a.size(0)
|
| 19 |
+
if B <= 1:
|
| 20 |
+
return view_a.new_tensor(0.0)
|
| 21 |
+
z1 = F.normalize(view_a, dim=-1)
|
| 22 |
+
z2 = F.normalize(view_b, dim=-1)
|
| 23 |
+
logits = torch.matmul(z1, z2.t()) / float(temperature)
|
| 24 |
+
labels = torch.arange(B, device=view_a.device)
|
| 25 |
+
loss_a = F.cross_entropy(logits, labels)
|
| 26 |
+
loss_b = F.cross_entropy(logits.t(), labels)
|
| 27 |
+
return 0.5 * (loss_a + loss_b)
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def trajectory_regularization_loss(thoughts: torch.Tensor, alpha: float = 0.5) -> torch.Tensor:
|
| 31 |
+
"""Trajectory regularization loss.
|
| 32 |
+
|
| 33 |
+
Encodes the intuition that the latent reasoning trajectory should:
|
| 34 |
+
1. Stay anchored to the initial acoustic state h_0 (semantic grounding)
|
| 35 |
+
2. Evolve smoothly between consecutive steps (temporal coherence)
|
| 36 |
+
|
| 37 |
+
Args:
|
| 38 |
+
thoughts: shape (B, N, D), where thoughts[:, 0, :] is the initial state h_0
|
| 39 |
+
alpha: weight balancing anchor vs smoothness (default: 0.5)
|
| 40 |
+
|
| 41 |
+
Returns:
|
| 42 |
+
Scalar loss combining anchor and smoothness terms
|
| 43 |
+
"""
|
| 44 |
+
N = thoughts.size(1)
|
| 45 |
+
# Guard: empty trajectory (e.g. n_latent=0) – return 0.
|
| 46 |
+
if N == 0:
|
| 47 |
+
return thoughts.new_tensor(0.0)
|
| 48 |
+
|
| 49 |
+
# Extract initial acoustic anchor h_0
|
| 50 |
+
h_0 = thoughts[:, 0:1, :] # shape: (B, 1, D)
|
| 51 |
+
|
| 52 |
+
# 1. Acoustic Anchor Loss (distance of all states from h_0)
|
| 53 |
+
anchor_loss = (thoughts - h_0).pow(2).sum(dim=-1).mean()
|
| 54 |
+
|
| 55 |
+
# 2. Step-wise Smoothness Loss (distance between adjacent states).
|
| 56 |
+
# When N==1 (e.g. DEQ single fixed-point), there are no consecutive pairs,
|
| 57 |
+
# so smooth_loss is zero – avoids mean() on an empty tensor returning NaN.
|
| 58 |
+
if N > 1:
|
| 59 |
+
diffs = thoughts[:, 1:, :] - thoughts[:, :-1, :] # shape: (B, N-1, D)
|
| 60 |
+
smooth_loss = diffs.pow(2).sum(dim=-1).mean()
|
| 61 |
+
else:
|
| 62 |
+
smooth_loss = thoughts.new_tensor(0.0)
|
| 63 |
+
|
| 64 |
+
# Combine final regularization loss
|
| 65 |
+
total_trajectory_loss = alpha * anchor_loss + (1.0 - alpha) * smooth_loss
|
| 66 |
+
return total_trajectory_loss
|
model.py
ADDED
|
@@ -0,0 +1,1972 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""LatentQwenASR model for lr_whisper.
|
| 2 |
+
|
| 3 |
+
The ``LatentQwenASR`` class wraps a frozen Qwen3-ASR base model and adds a
|
| 4 |
+
small controller network that emits per-step delta vectors (latent tokens).
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import math
|
| 8 |
+
import re
|
| 9 |
+
import string
|
| 10 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 11 |
+
|
| 12 |
+
import numpy as np
|
| 13 |
+
import torch
|
| 14 |
+
import torch.nn as nn
|
| 15 |
+
import torch.nn.functional as F
|
| 16 |
+
from tqdm import tqdm
|
| 17 |
+
from jiwer import wer
|
| 18 |
+
|
| 19 |
+
from data import build_audio_token_seq, _coerce_feat_attention_mask
|
| 20 |
+
|
| 21 |
+
def _detach_kv(kv_tuple: Any) -> Any:
|
| 22 |
+
"""Detach past_key_values to prevent excessive BPTT exploding gradients."""
|
| 23 |
+
if kv_tuple is None:
|
| 24 |
+
return None
|
| 25 |
+
# If it's a raw tuple (older transformers standard)
|
| 26 |
+
if isinstance(kv_tuple, tuple):
|
| 27 |
+
return tuple(tuple(t.detach() for t in layer) for layer in kv_tuple)
|
| 28 |
+
|
| 29 |
+
# Transformers >= 4.38 uses Cache objects like DynamicCache
|
| 30 |
+
if hasattr(kv_tuple, "to_legacy_cache"):
|
| 31 |
+
legacy_tuple = kv_tuple.to_legacy_cache()
|
| 32 |
+
detached_tuple = tuple(tuple(t.detach() for t in layer) for layer in legacy_tuple)
|
| 33 |
+
return type(kv_tuple).from_legacy_cache(detached_tuple)
|
| 34 |
+
|
| 35 |
+
# Fallback to cloning attributes if the specific class doesn't match
|
| 36 |
+
import copy
|
| 37 |
+
new_cache = copy.copy(kv_tuple)
|
| 38 |
+
if hasattr(new_cache, "key_cache"):
|
| 39 |
+
new_cache.key_cache = [k.detach() for k in getattr(kv_tuple, "key_cache", [])]
|
| 40 |
+
new_cache.value_cache = [v.detach() for v in getattr(kv_tuple, "value_cache", [])]
|
| 41 |
+
return new_cache
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
class LatentQwenASR(nn.Module):
|
| 45 |
+
"""Wrapper around Qwen3-ASR that supports latent prompt injection.
|
| 46 |
+
|
| 47 |
+
The class keeps the underlying ASR model frozen and introduces a small
|
| 48 |
+
controller network that emits delta vectors for each latent token. These
|
| 49 |
+
deltas are normalised and scaled, then added to the token embeddings at
|
| 50 |
+
the NT-token positions. During training, the deltas are inserted into the
|
| 51 |
+
decoder embeddings at the appropriate positions; during generation, they
|
| 52 |
+
are applied just before decoding.
|
| 53 |
+
|
| 54 |
+
Design note - float32 parameters
|
| 55 |
+
-----------------------------------
|
| 56 |
+
``log_scale`` is kept in **float32** even when the rest of the model is in bfloat16
|
| 57 |
+
to preserve gradient precision for the per-step scales.
|
| 58 |
+
"""
|
| 59 |
+
|
| 60 |
+
@staticmethod
|
| 61 |
+
def _resolve_text_model(thinker: nn.Module) -> nn.Module:
|
| 62 |
+
"""Resolve the inner text decoder module across wrapped backends (e.g. PEFT)."""
|
| 63 |
+
queue: List[object] = [thinker]
|
| 64 |
+
seen: set = set()
|
| 65 |
+
while queue:
|
| 66 |
+
node = queue.pop(0)
|
| 67 |
+
if node is None:
|
| 68 |
+
continue
|
| 69 |
+
node_id = id(node)
|
| 70 |
+
if node_id in seen:
|
| 71 |
+
continue
|
| 72 |
+
seen.add(node_id)
|
| 73 |
+
emb = getattr(node, "embed_tokens", None)
|
| 74 |
+
if isinstance(emb, nn.Module):
|
| 75 |
+
return node
|
| 76 |
+
for attr in ("model", "base_model", "module"):
|
| 77 |
+
child = getattr(node, attr, None)
|
| 78 |
+
if child is not None and child is not node:
|
| 79 |
+
queue.append(child)
|
| 80 |
+
raise AttributeError(
|
| 81 |
+
"Unable to resolve text decoder with `embed_tokens` from thinker module."
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
@staticmethod
|
| 85 |
+
def _resolve_embed_tokens(text_model: nn.Module) -> nn.Module:
|
| 86 |
+
emb = getattr(text_model, "embed_tokens", None)
|
| 87 |
+
if isinstance(emb, nn.Module):
|
| 88 |
+
return emb
|
| 89 |
+
if hasattr(text_model, "get_input_embeddings"):
|
| 90 |
+
emb = text_model.get_input_embeddings()
|
| 91 |
+
if isinstance(emb, nn.Module):
|
| 92 |
+
return emb
|
| 93 |
+
raise AttributeError("Unable to resolve input embedding module.")
|
| 94 |
+
|
| 95 |
+
def __init__(
|
| 96 |
+
self,
|
| 97 |
+
asr_model: nn.Module,
|
| 98 |
+
processor: Any,
|
| 99 |
+
n_latent: int,
|
| 100 |
+
nt_token_id: int,
|
| 101 |
+
lang_token_id: int,
|
| 102 |
+
transcribe_token_id: int,
|
| 103 |
+
freeze_base: bool = True,
|
| 104 |
+
use_latent: bool = True,
|
| 105 |
+
use_soft_prompt: bool = False,
|
| 106 |
+
soft_prompt_init_mode: str = "text",
|
| 107 |
+
soft_prompt_init_text: str = "",
|
| 108 |
+
user_prompt_text: str = "Transcribe the audio into text.",
|
| 109 |
+
|
| 110 |
+
delta_tanh_c: float = 5.0,
|
| 111 |
+
scale_max: float = 3.0,
|
| 112 |
+
scale_init: float = 0.2,
|
| 113 |
+
thought_mode: str = "prefix",
|
| 114 |
+
thought_group_size: int = 1,
|
| 115 |
+
halt_threshold: float = 0.0,
|
| 116 |
+
latent_drop_prob: float = 0.0,
|
| 117 |
+
latent_input_noise_std: float = 0.0,
|
| 118 |
+
latent_use_bounded_delta: bool = True,
|
| 119 |
+
latent_use_injection_gate: bool = True,
|
| 120 |
+
latent_use_embedding_anchor: bool = True,
|
| 121 |
+
freeze_audio_stack: bool = True,
|
| 122 |
+
) -> None:
|
| 123 |
+
super().__init__()
|
| 124 |
+
self.base_model = asr_model
|
| 125 |
+
self.processor = processor
|
| 126 |
+
self.config = getattr(asr_model, "config", None)
|
| 127 |
+
self.n_latent = int(n_latent)
|
| 128 |
+
self.use_latent = bool(use_latent and self.n_latent > 0)
|
| 129 |
+
self.nt_token_id = int(nt_token_id)
|
| 130 |
+
self.use_soft_prompt = bool(
|
| 131 |
+
use_soft_prompt
|
| 132 |
+
and (not self.use_latent)
|
| 133 |
+
and self.n_latent > 0
|
| 134 |
+
and self.nt_token_id >= 0
|
| 135 |
+
)
|
| 136 |
+
self.soft_prompt_init_mode = (soft_prompt_init_mode or "text").strip().lower()
|
| 137 |
+
self.soft_prompt_init_text = soft_prompt_init_text or ""
|
| 138 |
+
self.lang_token_id = int(lang_token_id)
|
| 139 |
+
self.transcribe_token_id = int(transcribe_token_id)
|
| 140 |
+
|
| 141 |
+
# Hyperparams stored from constructor args (originally read from env globals).
|
| 142 |
+
|
| 143 |
+
self.delta_tanh_c = float(delta_tanh_c)
|
| 144 |
+
self.halt_threshold = float(halt_threshold)
|
| 145 |
+
self.latent_drop_prob = float(latent_drop_prob)
|
| 146 |
+
self.latent_input_noise_std = float(latent_input_noise_std)
|
| 147 |
+
self.latent_use_bounded_delta = bool(latent_use_bounded_delta)
|
| 148 |
+
self.latent_use_injection_gate = bool(latent_use_injection_gate)
|
| 149 |
+
self.latent_use_embedding_anchor = bool(latent_use_embedding_anchor)
|
| 150 |
+
self.freeze_audio_stack = bool(freeze_audio_stack)
|
| 151 |
+
|
| 152 |
+
# Qwen3-ASR audio token ID
|
| 153 |
+
_audio_tok = None
|
| 154 |
+
if self.config and hasattr(self.config, "audio_token_id") and self.config.audio_token_id is not None:
|
| 155 |
+
_audio_tok = int(self.config.audio_token_id)
|
| 156 |
+
if _audio_tok is None:
|
| 157 |
+
_thinker_cfg = getattr(self.base_model, "config", None) or getattr(
|
| 158 |
+
getattr(self.base_model, "thinker", None), "config", None
|
| 159 |
+
)
|
| 160 |
+
if _thinker_cfg is not None and hasattr(_thinker_cfg, "audio_token_id") and _thinker_cfg.audio_token_id is not None:
|
| 161 |
+
_audio_tok = int(_thinker_cfg.audio_token_id)
|
| 162 |
+
if _audio_tok is None:
|
| 163 |
+
for _name in ("<|AUDIO|>", "<|audio|>", "<|audio_content|>", "<|audio_pad|>"):
|
| 164 |
+
target_tok = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
|
| 165 |
+
if hasattr(target_tok, "convert_tokens_to_ids"):
|
| 166 |
+
_try = target_tok.convert_tokens_to_ids(_name)
|
| 167 |
+
unk_id = getattr(target_tok, "unk_token_id", None)
|
| 168 |
+
if _try is not None and _try != unk_id:
|
| 169 |
+
_audio_tok = int(_try)
|
| 170 |
+
break
|
| 171 |
+
if _audio_tok is None:
|
| 172 |
+
_audio_tok = getattr(self.processor, "audio_token_id", None)
|
| 173 |
+
if _audio_tok is None:
|
| 174 |
+
_audio_tok = 151646 # last-resort fallback
|
| 175 |
+
self.audio_token_id = _audio_tok
|
| 176 |
+
|
| 177 |
+
target_tok = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
|
| 178 |
+
_decoded = target_tok.convert_ids_to_tokens(self.audio_token_id) if hasattr(target_tok, "convert_ids_to_tokens") else str(self.audio_token_id)
|
| 179 |
+
print(f"[init] audio_token_id={self.audio_token_id} decodes_to={_decoded!r}")
|
| 180 |
+
|
| 181 |
+
# Freeze all parameters in the underlying model for latent-adapter training.
|
| 182 |
+
if freeze_base:
|
| 183 |
+
for p in self.base_model.parameters():
|
| 184 |
+
p.requires_grad = False
|
| 185 |
+
|
| 186 |
+
if self.freeze_audio_stack:
|
| 187 |
+
audio_tower = getattr(self.base_model, "thinker", self.base_model)
|
| 188 |
+
audio_tower = getattr(audio_tower, "audio_tower", None)
|
| 189 |
+
if audio_tower is not None:
|
| 190 |
+
n_frozen = 0
|
| 191 |
+
for p in audio_tower.parameters():
|
| 192 |
+
if p.requires_grad:
|
| 193 |
+
p.requires_grad = False
|
| 194 |
+
n_frozen += 1
|
| 195 |
+
if n_frozen > 0:
|
| 196 |
+
print(f"[init] Froze {n_frozen} audio_tower parameters.")
|
| 197 |
+
# Also freeze multi_modal_projector (audio-to-text projection).
|
| 198 |
+
projector = getattr(getattr(self.base_model, "thinker", self.base_model), "multi_modal_projector", None)
|
| 199 |
+
if projector is not None:
|
| 200 |
+
n_frozen_proj = 0
|
| 201 |
+
for p in projector.parameters():
|
| 202 |
+
if p.requires_grad:
|
| 203 |
+
p.requires_grad = False
|
| 204 |
+
n_frozen_proj += 1
|
| 205 |
+
if n_frozen_proj > 0:
|
| 206 |
+
print(f"[init] Froze {n_frozen_proj} multi_modal_projector parameters.")
|
| 207 |
+
|
| 208 |
+
self.thinker = self.base_model.thinker
|
| 209 |
+
self.text_model = self._resolve_text_model(self.thinker)
|
| 210 |
+
self.embed_tokens = self._resolve_embed_tokens(self.text_model)
|
| 211 |
+
|
| 212 |
+
if self.config and hasattr(self.config, "text_config"):
|
| 213 |
+
d = self.config.text_config.hidden_size
|
| 214 |
+
elif self.config and hasattr(self.config, "hidden_size"):
|
| 215 |
+
d = self.config.hidden_size
|
| 216 |
+
else:
|
| 217 |
+
if hasattr(self.embed_tokens, "embedding_dim"):
|
| 218 |
+
d = int(self.embed_tokens.embedding_dim)
|
| 219 |
+
else:
|
| 220 |
+
d = int(self.embed_tokens.weight.size(-1))
|
| 221 |
+
|
| 222 |
+
print(f"[LatentQwenASR] Using hidden_size={d}")
|
| 223 |
+
|
| 224 |
+
target_dtype = self.thinker.dtype if hasattr(self.thinker, "dtype") else torch.float32
|
| 225 |
+
|
| 226 |
+
self.init_proj = nn.Linear(d, d, bias=True).to(dtype=target_dtype)
|
| 227 |
+
self.delta_proj = nn.Linear(d, d, bias=True).to(dtype=target_dtype)
|
| 228 |
+
|
| 229 |
+
self.step_embed = nn.Parameter(torch.zeros(self.n_latent, d, dtype=target_dtype))
|
| 230 |
+
nn.init.normal_(self.step_embed, mean=0.0, std=0.02)
|
| 231 |
+
self.step_proj = nn.Linear(d, d, bias=False).to(dtype=target_dtype)
|
| 232 |
+
self.scale_max = float(scale_max)
|
| 233 |
+
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
scale_init_val = float(scale_init)
|
| 237 |
+
if scale_init_val <= 0:
|
| 238 |
+
scale_init_val = 1.0
|
| 239 |
+
scale_init_unconstrained = math.log(math.expm1(scale_init_val))
|
| 240 |
+
self.log_scale = nn.Parameter(
|
| 241 |
+
torch.full((self.n_latent,), float(scale_init_unconstrained), dtype=torch.float32)
|
| 242 |
+
)
|
| 243 |
+
|
| 244 |
+
# Stability: LayerNorm for the recurrent thought loop
|
| 245 |
+
self.thought_ln = nn.LayerNorm(d, elementwise_affine=False).to(dtype=target_dtype)
|
| 246 |
+
|
| 247 |
+
# Thought Quality Analyzer (Value Head) predicting expected CE error
|
| 248 |
+
self.value_head = nn.Linear(d, 1).to(dtype=target_dtype)
|
| 249 |
+
|
| 250 |
+
# Soft Gated Injection (Scheme 4)
|
| 251 |
+
self.injection_gate = nn.Sequential(
|
| 252 |
+
nn.Linear(d * 2, d),
|
| 253 |
+
nn.Sigmoid()
|
| 254 |
+
).to(dtype=target_dtype)
|
| 255 |
+
# Initialize gate weights to 0 so sigmoid(0) = 0.5 (neutral start)
|
| 256 |
+
nn.init.zeros_(self.injection_gate[0].weight)
|
| 257 |
+
nn.init.zeros_(self.injection_gate[0].bias)
|
| 258 |
+
self._hidden_size = d
|
| 259 |
+
|
| 260 |
+
# Front-token prompt tuning
|
| 261 |
+
self.soft_prompt_embed = nn.Parameter(
|
| 262 |
+
torch.zeros(self.n_latent, d, dtype=target_dtype),
|
| 263 |
+
requires_grad=self.use_soft_prompt,
|
| 264 |
+
)
|
| 265 |
+
if self.use_soft_prompt:
|
| 266 |
+
init_prompt = None
|
| 267 |
+
if self.soft_prompt_init_mode == "text" and self.soft_prompt_init_text.strip():
|
| 268 |
+
init_ids = self.processor.tokenizer.encode(
|
| 269 |
+
self.soft_prompt_init_text.strip(),
|
| 270 |
+
add_special_tokens=False,
|
| 271 |
+
)
|
| 272 |
+
if init_ids:
|
| 273 |
+
with torch.no_grad():
|
| 274 |
+
emb_device = self.embed_tokens.weight.device
|
| 275 |
+
init_ids_t = torch.tensor(init_ids, dtype=torch.long, device=emb_device)
|
| 276 |
+
base = self.embed_tokens(init_ids_t).detach()
|
| 277 |
+
if base.numel() > 0:
|
| 278 |
+
if base.size(0) < self.n_latent:
|
| 279 |
+
reps = (self.n_latent + base.size(0) - 1) // base.size(0)
|
| 280 |
+
base = base.repeat((reps, 1))
|
| 281 |
+
init_prompt = base[: self.n_latent]
|
| 282 |
+
if init_prompt is None:
|
| 283 |
+
init_prompt = torch.empty((self.n_latent, d), dtype=target_dtype)
|
| 284 |
+
nn.init.normal_(init_prompt, mean=0.0, std=0.02)
|
| 285 |
+
self.soft_prompt_embed.data.copy_(
|
| 286 |
+
init_prompt.to(device=self.soft_prompt_embed.device, dtype=self.soft_prompt_embed.dtype)
|
| 287 |
+
)
|
| 288 |
+
else:
|
| 289 |
+
self.soft_prompt_embed.requires_grad = False
|
| 290 |
+
|
| 291 |
+
# Cast trainable modules to backbone dtype
|
| 292 |
+
target_dtype = self.base_model.dtype if hasattr(self.base_model, "dtype") else torch.float32
|
| 293 |
+
self.init_proj.to(target_dtype)
|
| 294 |
+
self.delta_proj.to(target_dtype)
|
| 295 |
+
self.step_embed.data = self.step_embed.data.to(target_dtype)
|
| 296 |
+
self.step_proj.to(target_dtype)
|
| 297 |
+
self.soft_prompt_embed.data = self.soft_prompt_embed.data.to(target_dtype)
|
| 298 |
+
# log_scale intentionally stays float32 for per-step scale gradient precision.
|
| 299 |
+
|
| 300 |
+
# Qwen3-ASR Special Audio Tokens
|
| 301 |
+
self.audio_bos_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|audio_start|>")
|
| 302 |
+
self.audio_eos_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|audio_end|>")
|
| 303 |
+
if self.audio_bos_token_id is None:
|
| 304 |
+
self.audio_bos_token_id = self.processor.tokenizer.bos_token_id
|
| 305 |
+
if self.audio_eos_token_id is None:
|
| 306 |
+
self.audio_eos_token_id = self.processor.tokenizer.eos_token_id
|
| 307 |
+
|
| 308 |
+
# Chat-template token pieces
|
| 309 |
+
self.im_start_id = self.processor.tokenizer.convert_tokens_to_ids("<|im_start|>")
|
| 310 |
+
if self.im_start_id is None or self.im_start_id == self.processor.tokenizer.unk_token_id:
|
| 311 |
+
self.im_start_id = self.lang_token_id
|
| 312 |
+
_im_end_id = self.processor.tokenizer.convert_tokens_to_ids("<|im_end|>")
|
| 313 |
+
if _im_end_id is None or _im_end_id == self.processor.tokenizer.unk_token_id:
|
| 314 |
+
_im_end_id = self.processor.tokenizer.eos_token_id
|
| 315 |
+
self.im_end_id = _im_end_id
|
| 316 |
+
# Unified stop-token set for generation AND eval truncation.
|
| 317 |
+
_stop_ids = {int(self.im_end_id)}
|
| 318 |
+
_eot_id = self.processor.tokenizer.convert_tokens_to_ids("<|endoftext|>")
|
| 319 |
+
if _eot_id is not None and _eot_id != self.processor.tokenizer.unk_token_id:
|
| 320 |
+
_stop_ids.add(int(_eot_id))
|
| 321 |
+
if self.processor.tokenizer.eos_token_id is not None:
|
| 322 |
+
_stop_ids.add(int(self.processor.tokenizer.eos_token_id))
|
| 323 |
+
self.stop_ids = sorted(_stop_ids)
|
| 324 |
+
self.user_nl_ids = self.processor.tokenizer.encode("user\n", add_special_tokens=False)
|
| 325 |
+
self.asst_nl_ids = self.processor.tokenizer.encode("assistant\n", add_special_tokens=False)
|
| 326 |
+
self.asr_prefix_ids = self.processor.tokenizer.encode("language English<asr_text>", add_special_tokens=False)
|
| 327 |
+
self.nl_ids = self.processor.tokenizer.encode("\n", add_special_tokens=False)
|
| 328 |
+
self.user_prompt_text = user_prompt_text
|
| 329 |
+
if self.user_prompt_text:
|
| 330 |
+
self.user_prompt_ids = self.processor.tokenizer.encode(
|
| 331 |
+
self.user_prompt_text, add_special_tokens=False
|
| 332 |
+
)
|
| 333 |
+
else:
|
| 334 |
+
self.user_prompt_ids = []
|
| 335 |
+
self._warned_generate_fallback = False
|
| 336 |
+
self.thought_mode = thought_mode.strip().lower()
|
| 337 |
+
self.thought_group_size = max(1, int(thought_group_size))
|
| 338 |
+
print(
|
| 339 |
+
"[latent-ablation] "
|
| 340 |
+
f"bounded_delta={self.latent_use_bounded_delta} "
|
| 341 |
+
f"injection_gate={self.latent_use_injection_gate} "
|
| 342 |
+
f"embedding_anchor={self.latent_use_embedding_anchor}"
|
| 343 |
+
)
|
| 344 |
+
|
| 345 |
+
# In non-latent modes, keep latent controller out of optimization.
|
| 346 |
+
if not self.use_latent:
|
| 347 |
+
for p in self.init_proj.parameters():
|
| 348 |
+
p.requires_grad = False
|
| 349 |
+
for p in self.delta_proj.parameters():
|
| 350 |
+
p.requires_grad = False
|
| 351 |
+
for p in self.step_proj.parameters():
|
| 352 |
+
p.requires_grad = False
|
| 353 |
+
for p in self.thought_ln.parameters():
|
| 354 |
+
p.requires_grad = False
|
| 355 |
+
self.step_embed.requires_grad = False
|
| 356 |
+
|
| 357 |
+
self.log_scale.requires_grad = False
|
| 358 |
+
if not self.use_soft_prompt:
|
| 359 |
+
self.soft_prompt_embed.requires_grad = False
|
| 360 |
+
|
| 361 |
+
# ------------------------------------------------------------------
|
| 362 |
+
# Internal helpers
|
| 363 |
+
# ------------------------------------------------------------------
|
| 364 |
+
|
| 365 |
+
def _build_chat_template_tensors(
|
| 366 |
+
self,
|
| 367 |
+
device: torch.device,
|
| 368 |
+
language: str = "English",
|
| 369 |
+
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 370 |
+
"""Build shared chat-template pieces for train/eval consistency."""
|
| 371 |
+
system_nl_ids = self.processor.tokenizer.encode("system\n", add_special_tokens=False)
|
| 372 |
+
system_turn = torch.tensor(
|
| 373 |
+
[self.im_start_id] + system_nl_ids + [self.im_end_id] + self.nl_ids,
|
| 374 |
+
dtype=torch.long,
|
| 375 |
+
device=device,
|
| 376 |
+
)
|
| 377 |
+
user_prefix = torch.tensor(
|
| 378 |
+
[self.im_start_id] + self.user_nl_ids,
|
| 379 |
+
dtype=torch.long,
|
| 380 |
+
device=device,
|
| 381 |
+
)
|
| 382 |
+
if self.user_prompt_ids:
|
| 383 |
+
user_prompt = torch.tensor(self.user_prompt_ids, dtype=torch.long, device=device)
|
| 384 |
+
else:
|
| 385 |
+
user_prompt = torch.empty((0,), dtype=torch.long, device=device)
|
| 386 |
+
user_suffix = torch.tensor(
|
| 387 |
+
[self.im_end_id] + self.nl_ids,
|
| 388 |
+
dtype=torch.long,
|
| 389 |
+
device=device,
|
| 390 |
+
)
|
| 391 |
+
asr_prefix_ids = self.processor.tokenizer.encode(f"language {language}<asr_text>", add_special_tokens=False)
|
| 392 |
+
assistant_prefix = torch.tensor(
|
| 393 |
+
[self.im_start_id] + self.asst_nl_ids + asr_prefix_ids,
|
| 394 |
+
dtype=torch.long,
|
| 395 |
+
device=device,
|
| 396 |
+
)
|
| 397 |
+
return system_turn, user_prefix, user_prompt, user_suffix, assistant_prefix
|
| 398 |
+
|
| 399 |
+
def _build_full_sequence(
|
| 400 |
+
self,
|
| 401 |
+
audio_toks: torch.Tensor,
|
| 402 |
+
text_all: torch.Tensor,
|
| 403 |
+
system_turn: torch.Tensor,
|
| 404 |
+
user_prefix: torch.Tensor,
|
| 405 |
+
user_prompt: torch.Tensor,
|
| 406 |
+
user_suffix: torch.Tensor,
|
| 407 |
+
label_pad: int = -100,
|
| 408 |
+
n_extra_mask: int = 0,
|
| 409 |
+
mask_all_nt: bool = False,
|
| 410 |
+
) -> Tuple[torch.Tensor, torch.Tensor]:
|
| 411 |
+
"""Build full input_ids and aligned labels for one sample.
|
| 412 |
+
|
| 413 |
+
The sequence layout is::
|
| 414 |
+
|
| 415 |
+
[system_turn, user_prefix, audio_toks, user_prompt,
|
| 416 |
+
user_suffix, text_all]
|
| 417 |
+
|
| 418 |
+
Labels clone input_ids, then mask everything up to (and including)
|
| 419 |
+
the assistant control tokens, plus an optional ``n_extra_mask``
|
| 420 |
+
additional positions (e.g. NT tokens in latent prefix mode).
|
| 421 |
+
|
| 422 |
+
When ``mask_all_nt=True``, all NT-token positions anywhere in the
|
| 423 |
+
sequence are additionally masked. Use this for interleaved mode where
|
| 424 |
+
NT tokens are scattered throughout the assistant turn.
|
| 425 |
+
|
| 426 |
+
Args:
|
| 427 |
+
audio_toks: 1-D token tensor for the audio segment.
|
| 428 |
+
text_all: 1-D token tensor for the target text (assistant turn).
|
| 429 |
+
system_turn: Shared system-turn token tensor.
|
| 430 |
+
user_prefix: Shared user-prefix token tensor.
|
| 431 |
+
user_prompt: Shared user-prompt token tensor.
|
| 432 |
+
user_suffix: Shared user-suffix token tensor.
|
| 433 |
+
label_pad: Padding value for ignored label positions.
|
| 434 |
+
n_extra_mask: Number of additional positions after the assistant
|
| 435 |
+
prefix to mask (set to ``self.n_latent`` for latent prefix
|
| 436 |
+
mode, 0 for native/baseline/interleaved).
|
| 437 |
+
mask_all_nt: If True, additionally mask every NT-token position
|
| 438 |
+
in the full sequence (for interleaved thought mode).
|
| 439 |
+
|
| 440 |
+
Returns:
|
| 441 |
+
``(full_ids, full_labels)`` - 1-D LongTensors.
|
| 442 |
+
"""
|
| 443 |
+
full_ids = torch.cat(
|
| 444 |
+
[system_turn, user_prefix, audio_toks, user_prompt, user_suffix, text_all]
|
| 445 |
+
)
|
| 446 |
+
n_prefix = (
|
| 447 |
+
system_turn.size(0)
|
| 448 |
+
+ user_prefix.size(0)
|
| 449 |
+
+ audio_toks.size(0)
|
| 450 |
+
+ user_prompt.size(0)
|
| 451 |
+
+ user_suffix.size(0)
|
| 452 |
+
)
|
| 453 |
+
# Mask: <|im_start|> + "assistant\n" + n_extra_mask (e.g. NT tokens)
|
| 454 |
+
n_asst_control = 1 + len(self.asst_nl_ids) + n_extra_mask
|
| 455 |
+
n_total_mask = n_prefix + n_asst_control
|
| 456 |
+
full_labels = full_ids.clone()
|
| 457 |
+
full_labels[:n_total_mask] = label_pad
|
| 458 |
+
if mask_all_nt and self.nt_token_id >= 0:
|
| 459 |
+
full_labels[full_ids == self.nt_token_id] = label_pad
|
| 460 |
+
return full_ids, full_labels
|
| 461 |
+
|
| 462 |
+
def _inject_latent_deltas(
|
| 463 |
+
self,
|
| 464 |
+
inputs_embeds: torch.Tensor,
|
| 465 |
+
input_ids: torch.LongTensor,
|
| 466 |
+
deltas: torch.Tensor,
|
| 467 |
+
) -> torch.Tensor:
|
| 468 |
+
"""Inject latent delta embeddings at NT-token positions.
|
| 469 |
+
|
| 470 |
+
Replaces the NT-token positions in *inputs_embeds* with::
|
| 471 |
+
|
| 472 |
+
base_nt_embedding + delta
|
| 473 |
+
|
| 474 |
+
Deltas are always injected at full strength. The Value Head controls
|
| 475 |
+
hard routing (N=0 skip / early halt) instead of soft scaling.
|
| 476 |
+
"""
|
| 477 |
+
B = inputs_embeds.size(0)
|
| 478 |
+
K = deltas.size(1)
|
| 479 |
+
device = inputs_embeds.device
|
| 480 |
+
nt_mask = (input_ids == self.nt_token_id)
|
| 481 |
+
nt_counts = nt_mask.sum(dim=1)
|
| 482 |
+
|
| 483 |
+
max_nt = int(nt_counts.max().item())
|
| 484 |
+
if K == 1 and max_nt > 1:
|
| 485 |
+
deltas = deltas.expand(B, max_nt, -1)
|
| 486 |
+
K = max_nt
|
| 487 |
+
|
| 488 |
+
assert (nt_counts == K).all(), (
|
| 489 |
+
f"NT token count mismatch! {nt_counts} vs expected {K}"
|
| 490 |
+
)
|
| 491 |
+
nt_token_ids = torch.full(
|
| 492 |
+
(B, K), self.nt_token_id, dtype=torch.long, device=device
|
| 493 |
+
)
|
| 494 |
+
base_nt = self.embed_tokens(nt_token_ids).to(dtype=inputs_embeds.dtype)
|
| 495 |
+
delta_scaled = deltas.float().to(dtype=inputs_embeds.dtype)
|
| 496 |
+
|
| 497 |
+
if self.latent_use_injection_gate:
|
| 498 |
+
gate_input = torch.cat([base_nt, delta_scaled], dim=-1)
|
| 499 |
+
gating_weight = self.injection_gate(gate_input)
|
| 500 |
+
else:
|
| 501 |
+
gating_weight = torch.ones_like(delta_scaled)
|
| 502 |
+
gated_delta = delta_scaled * gating_weight
|
| 503 |
+
if self.latent_use_embedding_anchor:
|
| 504 |
+
nt_embeds = base_nt + gated_delta
|
| 505 |
+
else:
|
| 506 |
+
nt_embeds = gated_delta
|
| 507 |
+
batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
|
| 508 |
+
result = inputs_embeds.clone()
|
| 509 |
+
result[batch_idx, seq_idx, :] = nt_embeds.to(dtype=inputs_embeds.dtype).view(-1, self._hidden_size)
|
| 510 |
+
return result
|
| 511 |
+
|
| 512 |
+
def _bounded_delta_from_raw(
|
| 513 |
+
self,
|
| 514 |
+
delta_raw: torch.Tensor,
|
| 515 |
+
step_idx: int,
|
| 516 |
+
scales: Optional[torch.Tensor] = None,
|
| 517 |
+
) -> torch.Tensor:
|
| 518 |
+
"""Convert a raw projected state into the configured latent delta."""
|
| 519 |
+
if not self.latent_use_bounded_delta:
|
| 520 |
+
return delta_raw
|
| 521 |
+
if scales is None:
|
| 522 |
+
scales = self.step_scales()
|
| 523 |
+
delta_raw_f32 = delta_raw.float()
|
| 524 |
+
norm_val = delta_raw_f32.norm(dim=-1, keepdim=True)
|
| 525 |
+
delta_dir = (delta_raw_f32 / (norm_val + 1e-6)).to(dtype=delta_raw.dtype)
|
| 526 |
+
scale = scales[min(int(step_idx), scales.size(0) - 1)]
|
| 527 |
+
return delta_dir * scale
|
| 528 |
+
|
| 529 |
+
def _compute_prefix_state(
|
| 530 |
+
self,
|
| 531 |
+
inputs_embeds: torch.Tensor,
|
| 532 |
+
input_ids: torch.LongTensor,
|
| 533 |
+
strict: bool = True,
|
| 534 |
+
language: str = "English",
|
| 535 |
+
) -> Tuple[torch.Tensor, object, torch.LongTensor]:
|
| 536 |
+
"""Compute prefix state and KV cache up to the first NT token."""
|
| 537 |
+
B, _, _ = inputs_embeds.shape
|
| 538 |
+
device = inputs_embeds.device
|
| 539 |
+
text_model = self.text_model
|
| 540 |
+
|
| 541 |
+
nt_mask = (input_ids == self.nt_token_id)
|
| 542 |
+
nt_counts = nt_mask.sum(dim=1)
|
| 543 |
+
if strict:
|
| 544 |
+
assert (nt_counts == self.n_latent).all(), (
|
| 545 |
+
f"NT token count mismatch! {nt_counts} vs {self.n_latent}"
|
| 546 |
+
)
|
| 547 |
+
|
| 548 |
+
first_nt = nt_mask.float().argmax(dim=1)
|
| 549 |
+
if strict:
|
| 550 |
+
asr_prefix_ids = self.processor.tokenizer.encode(
|
| 551 |
+
f"language {language}<asr_text>", add_special_tokens=False
|
| 552 |
+
)
|
| 553 |
+
expected = torch.tensor(
|
| 554 |
+
[self.im_start_id] + self.asst_nl_ids + asr_prefix_ids, dtype=torch.long, device=device
|
| 555 |
+
)
|
| 556 |
+
exp_len = expected.numel()
|
| 557 |
+
for i in range(B):
|
| 558 |
+
idx = int(first_nt[i].item())
|
| 559 |
+
start = idx - exp_len
|
| 560 |
+
assert start >= 0, (
|
| 561 |
+
f"Prompt too short to contain assistant prefix (idx={idx}, exp_len={exp_len})."
|
| 562 |
+
)
|
| 563 |
+
actual = input_ids[i, start:idx]
|
| 564 |
+
assert torch.equal(actual, expected), (
|
| 565 |
+
"Assistant prefix mismatch before NT. "
|
| 566 |
+
f"expected={expected.tolist()} actual={actual.tolist()}"
|
| 567 |
+
)
|
| 568 |
+
|
| 569 |
+
prefix_lens = first_nt
|
| 570 |
+
if strict:
|
| 571 |
+
assert (prefix_lens > 0).all(), f"Invalid prefix length(s): {prefix_lens}"
|
| 572 |
+
|
| 573 |
+
max_prefix = int(prefix_lens.max().item())
|
| 574 |
+
prefix_embeds = inputs_embeds[:, :max_prefix, :]
|
| 575 |
+
prefix_attn = (
|
| 576 |
+
torch.arange(max_prefix, device=device).unsqueeze(0) < prefix_lens.unsqueeze(1)
|
| 577 |
+
).long()
|
| 578 |
+
if prefix_attn.numel() > 0:
|
| 579 |
+
prefix_embeds = prefix_embeds * prefix_attn.unsqueeze(-1)
|
| 580 |
+
|
| 581 |
+
with torch.no_grad():
|
| 582 |
+
out = text_model(
|
| 583 |
+
inputs_embeds=prefix_embeds,
|
| 584 |
+
attention_mask=prefix_attn,
|
| 585 |
+
use_cache=True,
|
| 586 |
+
return_dict=True,
|
| 587 |
+
)
|
| 588 |
+
last_positions = torch.clamp(prefix_lens - 1, min=0)
|
| 589 |
+
batch_idx = torch.arange(B, device=device)
|
| 590 |
+
last = out.last_hidden_state[batch_idx, last_positions].detach()
|
| 591 |
+
return last, out.past_key_values, prefix_attn
|
| 592 |
+
|
| 593 |
+
def _embed_scale(self) -> float:
|
| 594 |
+
# Qwen3-ASR does NOT scale embeddings by sqrt(d_model).
|
| 595 |
+
return 1.0
|
| 596 |
+
|
| 597 |
+
|
| 598 |
+
def step_scales(self) -> torch.Tensor:
|
| 599 |
+
s = F.softplus(self.log_scale)
|
| 600 |
+
s = torch.clamp(s, max=self.scale_max)
|
| 601 |
+
return s[:self.n_latent]
|
| 602 |
+
|
| 603 |
+
def forced_decoder_ids_latent(self) -> List[Tuple[int, int]]:
|
| 604 |
+
ids: List[Tuple[int, int]] = [(0, self.im_start_id)]
|
| 605 |
+
cur = 1
|
| 606 |
+
for tid in self.asst_nl_ids:
|
| 607 |
+
ids.append((cur, tid))
|
| 608 |
+
cur += 1
|
| 609 |
+
for i in range(self.n_latent):
|
| 610 |
+
ids.append((cur + i, self.nt_token_id))
|
| 611 |
+
return ids
|
| 612 |
+
|
| 613 |
+
def forced_decoder_ids_baseline(self) -> List[Tuple[int, int]]:
|
| 614 |
+
ids: List[Tuple[int, int]] = [(0, self.im_start_id)]
|
| 615 |
+
cur = 1
|
| 616 |
+
for tid in self.asst_nl_ids:
|
| 617 |
+
ids.append((cur, tid))
|
| 618 |
+
cur += 1
|
| 619 |
+
return ids
|
| 620 |
+
|
| 621 |
+
def _shift_right(self, labels: torch.LongTensor) -> torch.LongTensor:
|
| 622 |
+
pad_id = self.processor.tokenizer.eos_token_id
|
| 623 |
+
decoder_input_ids = labels.new_full(labels.shape, pad_id)
|
| 624 |
+
decoder_input_ids[:, 0] = (
|
| 625 |
+
self.processor.tokenizer.bos_token_id
|
| 626 |
+
if self.processor.tokenizer.bos_token_id is not None
|
| 627 |
+
else self.lang_token_id
|
| 628 |
+
)
|
| 629 |
+
shifted = labels[:, :-1].clone()
|
| 630 |
+
shifted = shifted.masked_fill(shifted == -100, pad_id)
|
| 631 |
+
decoder_input_ids[:, 1:] = shifted
|
| 632 |
+
return decoder_input_ids
|
| 633 |
+
|
| 634 |
+
def _encode_audio(
|
| 635 |
+
self,
|
| 636 |
+
input_features: torch.FloatTensor,
|
| 637 |
+
feature_attention_mask: torch.LongTensor,
|
| 638 |
+
) -> Tuple[torch.Tensor, List[int]]:
|
| 639 |
+
"""Manually encode audio features to get their exact lengths per sample."""
|
| 640 |
+
if feature_attention_mask is not None:
|
| 641 |
+
feature_lens = torch.sum(feature_attention_mask, dim=1)
|
| 642 |
+
else:
|
| 643 |
+
feature_lens = torch.full(
|
| 644 |
+
(input_features.size(0),), input_features.size(2),
|
| 645 |
+
device=input_features.device, dtype=torch.long,
|
| 646 |
+
)
|
| 647 |
+
|
| 648 |
+
audio_features = []
|
| 649 |
+
lengths = []
|
| 650 |
+
audio_tower = self.thinker.audio_tower
|
| 651 |
+
projector = getattr(self.thinker, "multi_modal_projector", None)
|
| 652 |
+
|
| 653 |
+
for i in range(len(input_features)):
|
| 654 |
+
feat = input_features[i]
|
| 655 |
+
l = feature_lens[i]
|
| 656 |
+
inp = feat[:, :l]
|
| 657 |
+
length_tensor = l.unsqueeze(0)
|
| 658 |
+
out = audio_tower(inp, feature_lens=length_tensor).last_hidden_state
|
| 659 |
+
if projector is not None:
|
| 660 |
+
out = projector(out)
|
| 661 |
+
out = out.squeeze(0)
|
| 662 |
+
audio_features.append(out)
|
| 663 |
+
lengths.append(out.size(0))
|
| 664 |
+
|
| 665 |
+
return torch.cat(audio_features, dim=0), lengths
|
| 666 |
+
|
| 667 |
+
def _get_audio_lengths(
|
| 668 |
+
self,
|
| 669 |
+
input_features: torch.FloatTensor,
|
| 670 |
+
feature_attention_mask: torch.LongTensor,
|
| 671 |
+
) -> List[int]:
|
| 672 |
+
"""Compute audio output lengths using the exact same formula as the thinker."""
|
| 673 |
+
if feature_attention_mask is not None:
|
| 674 |
+
feature_lens = torch.sum(feature_attention_mask, dim=1)
|
| 675 |
+
else:
|
| 676 |
+
feature_lens = torch.full(
|
| 677 |
+
(input_features.size(0),), input_features.size(2),
|
| 678 |
+
device=input_features.device, dtype=torch.long,
|
| 679 |
+
)
|
| 680 |
+
|
| 681 |
+
input_lengths_leave = feature_lens % 100
|
| 682 |
+
feat_lengths = (input_lengths_leave - 1) // 2 + 1
|
| 683 |
+
output_lengths = (
|
| 684 |
+
((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1
|
| 685 |
+
+ (feature_lens // 100) * 13
|
| 686 |
+
)
|
| 687 |
+
return output_lengths.tolist()
|
| 688 |
+
|
| 689 |
+
def _get_native_audio_embeds(
|
| 690 |
+
self,
|
| 691 |
+
input_ids: torch.LongTensor,
|
| 692 |
+
input_features: torch.FloatTensor,
|
| 693 |
+
feature_attention_mask: torch.LongTensor,
|
| 694 |
+
attention_mask: torch.LongTensor,
|
| 695 |
+
) -> torch.Tensor:
|
| 696 |
+
"""Capture the thinker's native audio-fused embeddings via a forward hook."""
|
| 697 |
+
captured: Dict[str, torch.Tensor] = {}
|
| 698 |
+
|
| 699 |
+
def _hook(module: nn.Module, args: Any, kwargs: Any) -> None:
|
| 700 |
+
embeds = kwargs.get("inputs_embeds", None)
|
| 701 |
+
if embeds is None and len(args) > 0:
|
| 702 |
+
embeds = args[0]
|
| 703 |
+
if embeds is not None:
|
| 704 |
+
captured["embeds"] = embeds
|
| 705 |
+
return None
|
| 706 |
+
|
| 707 |
+
hook = self.text_model.register_forward_pre_hook(_hook, with_kwargs=True)
|
| 708 |
+
try:
|
| 709 |
+
with torch.no_grad():
|
| 710 |
+
self.thinker(
|
| 711 |
+
input_ids=input_ids,
|
| 712 |
+
attention_mask=attention_mask,
|
| 713 |
+
input_features=input_features,
|
| 714 |
+
feature_attention_mask=feature_attention_mask,
|
| 715 |
+
labels=None,
|
| 716 |
+
return_dict=True,
|
| 717 |
+
)
|
| 718 |
+
finally:
|
| 719 |
+
hook.remove()
|
| 720 |
+
|
| 721 |
+
if "embeds" not in captured:
|
| 722 |
+
raise RuntimeError(
|
| 723 |
+
"_get_native_audio_embeds: forward hook did not capture embeddings. "
|
| 724 |
+
"The thinker's internal forward may use a different code path."
|
| 725 |
+
)
|
| 726 |
+
return captured["embeds"].detach().clone()
|
| 727 |
+
|
| 728 |
+
def _compute_continuous_thoughts(
|
| 729 |
+
self,
|
| 730 |
+
initial_state: torch.Tensor,
|
| 731 |
+
prefix_past_key_values: Optional[object] = None,
|
| 732 |
+
prefix_attention_mask: Optional[torch.LongTensor] = None,
|
| 733 |
+
halt_threshold: Optional[float] = None,
|
| 734 |
+
) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
|
| 735 |
+
"""Compute recurrent latent states + per-step deltas using DEQ fixed-point iteration."""
|
| 736 |
+
B = initial_state.size(0)
|
| 737 |
+
D = initial_state.size(-1)
|
| 738 |
+
device = initial_state.device
|
| 739 |
+
text_model = self.text_model
|
| 740 |
+
|
| 741 |
+
initial_norm = initial_state / (initial_state.norm(dim=-1, keepdim=True) + 1e-6)
|
| 742 |
+
h = self.init_proj(initial_norm)
|
| 743 |
+
|
| 744 |
+
# Remove random corruption as teacher forcing artificially lowers CE
|
| 745 |
+
_corrupted = False
|
| 746 |
+
|
| 747 |
+
scales = self.step_scales()
|
| 748 |
+
thought_embeds: List[torch.Tensor] = []
|
| 749 |
+
state_embeds: List[torch.Tensor] = []
|
| 750 |
+
raw_norms: List[torch.Tensor] = []
|
| 751 |
+
scaled_norms: List[torch.Tensor] = []
|
| 752 |
+
v_preds: List[torch.Tensor] = []
|
| 753 |
+
past_kv = prefix_past_key_values
|
| 754 |
+
running_attn = prefix_attention_mask
|
| 755 |
+
|
| 756 |
+
# ---- Dynamic Causal Thinking Loop ----
|
| 757 |
+
# Instead of fixed-point iteration (which doesn't converge on frozen LLMs),
|
| 758 |
+
# we generate a sequence of thought tokens.
|
| 759 |
+
max_iter = self.n_latent
|
| 760 |
+
if halt_threshold is None:
|
| 761 |
+
halt_threshold = getattr(self, "halt_threshold", 0.0)
|
| 762 |
+
|
| 763 |
+
current_embed = h
|
| 764 |
+
|
| 765 |
+
# ---- N=0 Skip Check (Tanh) ----
|
| 766 |
+
# v_pred ∈ [-1,1]: expected delta CE (positive means LR helps).
|
| 767 |
+
# If v_pred < threshold → model is confident LR is NOT needed → skip.
|
| 768 |
+
if not self.training:
|
| 769 |
+
v_init = torch.tanh(self.value_head(h.unsqueeze(1)).squeeze(-1).squeeze(-1)) # (B,)
|
| 770 |
+
if (v_init < halt_threshold).all():
|
| 771 |
+
# Return empty deltas → generate() will remove all NT tokens
|
| 772 |
+
empty_deltas = torch.zeros(B, 0, D, device=device, dtype=h.dtype)
|
| 773 |
+
empty_states = torch.zeros(B, 0, D, device=device, dtype=h.dtype)
|
| 774 |
+
stats = {
|
| 775 |
+
"predicted_value": v_init.detach(),
|
| 776 |
+
"deq_iters": torch.tensor(0.0, device=device),
|
| 777 |
+
"raw_norm_mean": torch.zeros(0, device=device),
|
| 778 |
+
"raw_norm_std": torch.zeros(0, device=device),
|
| 779 |
+
"scaled_norm_mean": torch.zeros(0, device=device),
|
| 780 |
+
"scaled_norm_std": torch.zeros(0, device=device),
|
| 781 |
+
"cos_mean": torch.zeros(0, device=device),
|
| 782 |
+
"cos_std": torch.zeros(0, device=device),
|
| 783 |
+
"scales": scales.detach(),
|
| 784 |
+
"diff_norm": torch.tensor(0.0, device=device),
|
| 785 |
+
"step_cos": torch.tensor(0.0, device=device),
|
| 786 |
+
"v_preds": v_init.unsqueeze(1).detach(),
|
| 787 |
+
"gate": v_init.unsqueeze(1).detach(),
|
| 788 |
+
"skipped": True,
|
| 789 |
+
}
|
| 790 |
+
return empty_deltas, empty_states, stats
|
| 791 |
+
|
| 792 |
+
with torch.no_grad():
|
| 793 |
+
for k in range(max_iter):
|
| 794 |
+
# The input for step k is step_embed[k] projected, added to the normalized state
|
| 795 |
+
step_idx = min(k, self.step_embed.size(0) - 1)
|
| 796 |
+
step = self.step_proj(self.step_embed[step_idx]).view(1, 1, -1).expand(B, 1, -1)
|
| 797 |
+
|
| 798 |
+
h_norm = self.thought_ln(current_embed)
|
| 799 |
+
h_input = h_norm.unsqueeze(1) + step
|
| 800 |
+
|
| 801 |
+
step_attn = None
|
| 802 |
+
if running_attn is not None:
|
| 803 |
+
# Causal: attention mask GROWS by 1 at each step
|
| 804 |
+
one = torch.ones((B, 1), dtype=running_attn.dtype, device=device)
|
| 805 |
+
step_attn = torch.cat([running_attn, one], dim=1)
|
| 806 |
+
|
| 807 |
+
out = text_model(
|
| 808 |
+
inputs_embeds=h_input,
|
| 809 |
+
past_key_values=past_kv,
|
| 810 |
+
attention_mask=step_attn,
|
| 811 |
+
use_cache=True,
|
| 812 |
+
return_dict=True,
|
| 813 |
+
)
|
| 814 |
+
|
| 815 |
+
# Next state
|
| 816 |
+
h_next = out.last_hidden_state.squeeze(1)
|
| 817 |
+
|
| 818 |
+
# Compute delta to inject into this token's output representation
|
| 819 |
+
delta_raw = self.delta_proj(h_next)
|
| 820 |
+
delta = self._bounded_delta_from_raw(delta_raw, step_idx, scales)
|
| 821 |
+
|
| 822 |
+
# Early Halt Check (Value Head, tanh) — BEFORE committing delta
|
| 823 |
+
v_raw = self.value_head(h_next.unsqueeze(1)).squeeze(-1).squeeze(-1) # (B,)
|
| 824 |
+
v_pred = torch.tanh(v_raw)
|
| 825 |
+
|
| 826 |
+
# Halt when v_pred drops below threshold (model says this step is harmful)
|
| 827 |
+
if not self.training:
|
| 828 |
+
if (v_pred < halt_threshold).all():
|
| 829 |
+
break
|
| 830 |
+
|
| 831 |
+
# Only commit the delta if we did NOT halt
|
| 832 |
+
thought_embeds.append(delta)
|
| 833 |
+
state_embeds.append(h_next)
|
| 834 |
+
raw_norms.append(delta_raw.norm(dim=-1))
|
| 835 |
+
scaled_norms.append(delta.norm(dim=-1))
|
| 836 |
+
v_preds.append(v_pred)
|
| 837 |
+
|
| 838 |
+
# Update context for the *next* iteration
|
| 839 |
+
past_kv = out.past_key_values
|
| 840 |
+
running_attn = step_attn
|
| 841 |
+
current_embed = h_next
|
| 842 |
+
|
| 843 |
+
# Grab the past_kv produced exactly prior to the final step
|
| 844 |
+
with torch.no_grad():
|
| 845 |
+
# We need to re-run up to final_step_idx-1 to get the exact past_kv?
|
| 846 |
+
# Actually, no. We can just run from initial_state and re-accumulate.
|
| 847 |
+
# BUT since max_iter <= n_latent is small (e.g. 4), we can just re-run
|
| 848 |
+
# the *entire* sequence with gradients enabled!
|
| 849 |
+
pass
|
| 850 |
+
|
| 851 |
+
# To make things simple and correct for training with such a small sequence (N=4),
|
| 852 |
+
# we can just re-run the FULL causal generation with gradients enabled and overwrite.
|
| 853 |
+
# This gives exact gradients through all thinking steps (BPTT).
|
| 854 |
+
if self.training:
|
| 855 |
+
thought_embeds_grad = []
|
| 856 |
+
state_embeds_grad = []
|
| 857 |
+
current_embed_grad = h # Keep grad flow to init_proj
|
| 858 |
+
running_attn_grad = prefix_attention_mask
|
| 859 |
+
past_kv_grad = prefix_past_key_values
|
| 860 |
+
|
| 861 |
+
# Re-run for exactly the number of steps we decided to take in the no_grad pass
|
| 862 |
+
for k in range(len(thought_embeds)):
|
| 863 |
+
step_idx = min(k, self.step_embed.size(0) - 1)
|
| 864 |
+
step = self.step_proj(self.step_embed[step_idx]).view(1, 1, -1).expand(B, 1, -1)
|
| 865 |
+
|
| 866 |
+
h_norm = self.thought_ln(current_embed_grad)
|
| 867 |
+
h_input = h_norm.unsqueeze(1) + step
|
| 868 |
+
|
| 869 |
+
step_attn = None
|
| 870 |
+
if running_attn_grad is not None:
|
| 871 |
+
one = torch.ones((B, 1), dtype=running_attn_grad.dtype, device=device)
|
| 872 |
+
step_attn = torch.cat([running_attn_grad, one], dim=1)
|
| 873 |
+
|
| 874 |
+
out = text_model(
|
| 875 |
+
inputs_embeds=h_input,
|
| 876 |
+
past_key_values=past_kv_grad, # Evolve KV cache like inference!
|
| 877 |
+
attention_mask=step_attn,
|
| 878 |
+
use_cache=True,
|
| 879 |
+
return_dict=True,
|
| 880 |
+
)
|
| 881 |
+
|
| 882 |
+
h_next_g = out.last_hidden_state.squeeze(1)
|
| 883 |
+
|
| 884 |
+
delta_raw_g = self.delta_proj(h_next_g)
|
| 885 |
+
delta_g = self._bounded_delta_from_raw(delta_raw_g, step_idx, scales)
|
| 886 |
+
|
| 887 |
+
thought_embeds_grad.append(delta_g)
|
| 888 |
+
state_embeds_grad.append(h_next_g)
|
| 889 |
+
|
| 890 |
+
# Detach state so CE-loss gradients don't flow back to previous steps.
|
| 891 |
+
# Each step is independently trained on its own delta (1-step truncated BPTT).
|
| 892 |
+
past_kv_grad = out.past_key_values
|
| 893 |
+
running_attn_grad = step_attn
|
| 894 |
+
current_embed_grad = h_next_g.detach()
|
| 895 |
+
|
| 896 |
+
# Override the no_grad lists with the gradient-tracked ones
|
| 897 |
+
thought_embeds = thought_embeds_grad
|
| 898 |
+
state_embeds = state_embeds_grad
|
| 899 |
+
|
| 900 |
+
# Stack sequences — guard against empty list when all steps halted
|
| 901 |
+
if not thought_embeds:
|
| 902 |
+
# All steps rejected by Value Head → return empty (K=0)
|
| 903 |
+
deltas = torch.zeros((B, 0, D), device=device, dtype=h.dtype)
|
| 904 |
+
states = torch.zeros((B, 0, D), device=device, dtype=h.dtype)
|
| 905 |
+
# Use initial state for value prediction
|
| 906 |
+
predicted_value = torch.tanh(self.value_head(h.unsqueeze(1)))
|
| 907 |
+
stats = {
|
| 908 |
+
"predicted_value": predicted_value.squeeze(-1).squeeze(-1),
|
| 909 |
+
"deq_iters": torch.tensor(0.0, device=device),
|
| 910 |
+
"raw_norm_mean": torch.zeros(self.n_latent, device=device),
|
| 911 |
+
"raw_norm_std": torch.zeros(self.n_latent, device=device),
|
| 912 |
+
"scaled_norm_mean": torch.zeros(self.n_latent, device=device),
|
| 913 |
+
"scaled_norm_std": torch.zeros(self.n_latent, device=device),
|
| 914 |
+
"cos_mean": torch.zeros(self.n_latent, device=device),
|
| 915 |
+
"cos_std": torch.zeros(self.n_latent, device=device),
|
| 916 |
+
"scales": scales.detach(),
|
| 917 |
+
"diff_norm": torch.tensor(0.0, device=device),
|
| 918 |
+
"step_cos": torch.tensor(0.0, device=device),
|
| 919 |
+
"v_preds": None,
|
| 920 |
+
"gate": torch.zeros((B, 1), device=device),
|
| 921 |
+
"corrupted": _corrupted,
|
| 922 |
+
"skipped": True,
|
| 923 |
+
}
|
| 924 |
+
return deltas, states, stats
|
| 925 |
+
|
| 926 |
+
deltas = torch.stack(thought_embeds, dim=1) # (B, K, D)
|
| 927 |
+
states = torch.stack(state_embeds, dim=1) # (B, K, D)
|
| 928 |
+
|
| 929 |
+
# predicted_value: tanh [-1, 1] estimating LR's impact on CE at
|
| 930 |
+
# every latent state. Training supervises each step because inference
|
| 931 |
+
# may halt from any of these value predictions.
|
| 932 |
+
predicted_value = torch.tanh(self.value_head(states)).squeeze(-1) # (B, K)
|
| 933 |
+
|
| 934 |
+
# Gate: average tanh v_pred across all steps, scaled to [0, 1] for gated delta injection
|
| 935 |
+
v_preds_stacked = torch.stack(v_preds, dim=1) if v_preds else None # (B, K)
|
| 936 |
+
if v_preds_stacked is not None:
|
| 937 |
+
gate = (v_preds_stacked.mean(dim=1, keepdim=True) + 1.0) / 2.0 # (B, 1) — mean confidence mapped to [0, 1]
|
| 938 |
+
else:
|
| 939 |
+
gate = (predicted_value.squeeze(-1) + 1.0) / 2.0 # fallback, scaled to [0, 1]
|
| 940 |
+
|
| 941 |
+
deq_iters = len(thought_embeds)
|
| 942 |
+
|
| 943 |
+
raw_tensor = torch.stack(raw_norms, dim=1)
|
| 944 |
+
scaled_tensor = torch.stack(scaled_norms, dim=1)
|
| 945 |
+
scales_val = scales.detach()
|
| 946 |
+
|
| 947 |
+
with torch.no_grad():
|
| 948 |
+
flat_states = states.view(-1, D)
|
| 949 |
+
flat_states_norm = F.normalize(flat_states, dim=-1)
|
| 950 |
+
|
| 951 |
+
word_emb_weight = self.embed_tokens.weight
|
| 952 |
+
vocab_size = word_emb_weight.size(0)
|
| 953 |
+
if vocab_size > 10000:
|
| 954 |
+
idx = torch.randperm(vocab_size, device=device)[:10000]
|
| 955 |
+
word_emb_sub = word_emb_weight[idx]
|
| 956 |
+
else:
|
| 957 |
+
word_emb_sub = word_emb_weight
|
| 958 |
+
|
| 959 |
+
word_emb_norm = F.normalize(word_emb_sub, dim=-1)
|
| 960 |
+
sim_matrix = torch.matmul(flat_states_norm, word_emb_norm.t())
|
| 961 |
+
if sim_matrix.numel() > 0:
|
| 962 |
+
max_sim, _ = sim_matrix.max(dim=-1)
|
| 963 |
+
cos_mean_val = max_sim.mean()
|
| 964 |
+
cos_std_val = max_sim.std(unbiased=False)
|
| 965 |
+
else:
|
| 966 |
+
cos_mean_val = flat_states_norm.new_tensor(0.0)
|
| 967 |
+
cos_std_val = flat_states_norm.new_tensor(0.0)
|
| 968 |
+
|
| 969 |
+
if deltas.size(1) > 1:
|
| 970 |
+
diffs = deltas[:, 1:] - deltas[:, :-1]
|
| 971 |
+
diff_norm_val = diffs.norm(dim=-1).mean()
|
| 972 |
+
step_cos_val = F.cosine_similarity(deltas[:, 1:], deltas[:, :-1], dim=-1).mean()
|
| 973 |
+
else:
|
| 974 |
+
diff_norm_val = deltas.new_tensor(0.0)
|
| 975 |
+
step_cos_val = deltas.new_tensor(0.0)
|
| 976 |
+
|
| 977 |
+
stats = {
|
| 978 |
+
"predicted_value": predicted_value, # (B, K) for per-step MSE against per-sample target
|
| 979 |
+
"deq_iters": torch.tensor(deq_iters, dtype=torch.float32, device=device),
|
| 980 |
+
"raw_norm_mean": raw_tensor.mean(dim=0).detach() if raw_tensor.ndim > 1 else raw_tensor.detach(),
|
| 981 |
+
"raw_norm_std": raw_tensor.std(dim=0, unbiased=False).detach() if raw_tensor.ndim > 1 else raw_tensor.new_zeros(raw_tensor.size(-1)),
|
| 982 |
+
"scaled_norm_mean": scaled_tensor.mean(dim=0).detach() if scaled_tensor.ndim > 1 else scaled_tensor.detach(),
|
| 983 |
+
"scaled_norm_std": scaled_tensor.std(dim=0, unbiased=False).detach() if scaled_tensor.ndim > 1 else scaled_tensor.new_zeros(scaled_tensor.size(-1)),
|
| 984 |
+
"cos_mean": torch.full((self.n_latent,), cos_mean_val.item(), device=device),
|
| 985 |
+
"cos_std": torch.full((self.n_latent,), cos_std_val.item(), device=device),
|
| 986 |
+
"scales": scales_val,
|
| 987 |
+
"diff_norm": diff_norm_val.detach(),
|
| 988 |
+
"step_cos": step_cos_val.detach(),
|
| 989 |
+
"v_preds": v_preds_stacked.detach() if v_preds_stacked is not None else None,
|
| 990 |
+
"gate": gate.detach(),
|
| 991 |
+
"corrupted": _corrupted,
|
| 992 |
+
}
|
| 993 |
+
return deltas, states, stats
|
| 994 |
+
|
| 995 |
+
def _compute_causal_thoughts(
|
| 996 |
+
self,
|
| 997 |
+
inputs_embeds: torch.Tensor,
|
| 998 |
+
input_ids: torch.LongTensor,
|
| 999 |
+
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
|
| 1000 |
+
"""Two-pass causal thought injection for interleaved mode.
|
| 1001 |
+
|
| 1002 |
+
Each NT token at position ``k`` receives a thought delta computed from
|
| 1003 |
+
the hidden state at position ``k-1`` (the token immediately before it).
|
| 1004 |
+
This makes each thought *causal*: it can attend to all previous context
|
| 1005 |
+
including the word generated just before it, which aligns with ASR's
|
| 1006 |
+
left-to-right generation structure.
|
| 1007 |
+
|
| 1008 |
+
Pass 1 (``no_grad``): forward the full sequence with **base** NT
|
| 1009 |
+
embeddings to capture all hidden states ``h_all``.
|
| 1010 |
+
|
| 1011 |
+
Pass 2 (with grad): for each NT at position k, compute a delta from
|
| 1012 |
+
``h_all[:, k-1, :]`` via ``init_proj → delta_proj → normalize →
|
| 1013 |
+
scale``. Inject the modified NT embedding and return the updated
|
| 1014 |
+
``inputs_embeds``.
|
| 1015 |
+
|
| 1016 |
+
Args:
|
| 1017 |
+
inputs_embeds: ``(B, SeqLen, D)`` embedding tensor (base NT
|
| 1018 |
+
embeddings already placed at NT positions by
|
| 1019 |
+
``_get_native_audio_embeds``).
|
| 1020 |
+
input_ids: ``(B, SeqLen)`` token id tensor to locate NT positions.
|
| 1021 |
+
|
| 1022 |
+
Returns:
|
| 1023 |
+
``(modified_embeds, deltas_padded, states_padded, stats)`` where
|
| 1024 |
+
``deltas_padded`` and ``states_padded`` have shape
|
| 1025 |
+
``(B, max_nt, D)`` (zero-padded to the batch maximum NT count).
|
| 1026 |
+
"""
|
| 1027 |
+
B, SeqLen, D = inputs_embeds.shape
|
| 1028 |
+
device = inputs_embeds.device
|
| 1029 |
+
|
| 1030 |
+
nt_mask = (input_ids == self.nt_token_id) # (B, SeqLen)
|
| 1031 |
+
batch_idx, nt_pos = torch.nonzero(nt_mask, as_tuple=True)
|
| 1032 |
+
|
| 1033 |
+
if len(nt_pos) == 0:
|
| 1034 |
+
empty_d = inputs_embeds.new_zeros((B, 0, D))
|
| 1035 |
+
empty_s = inputs_embeds.new_zeros((B, 0, D))
|
| 1036 |
+
dummy = inputs_embeds.new_tensor(0.0)
|
| 1037 |
+
stats: Dict[str, torch.Tensor] = {
|
| 1038 |
+
"raw_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1039 |
+
"raw_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1040 |
+
"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1041 |
+
"scaled_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1042 |
+
"cos_mean": inputs_embeds.new_zeros((0,)),
|
| 1043 |
+
"cos_std": inputs_embeds.new_zeros((0,)),
|
| 1044 |
+
"scales": inputs_embeds.new_zeros((0,)),
|
| 1045 |
+
"diff_norm": dummy,
|
| 1046 |
+
"step_cos": dummy,
|
| 1047 |
+
}
|
| 1048 |
+
return inputs_embeds, empty_d, empty_s, stats
|
| 1049 |
+
|
| 1050 |
+
# --- Pass 1: no_grad forward to capture all hidden states ----------
|
| 1051 |
+
with torch.no_grad():
|
| 1052 |
+
out1 = self.text_model(
|
| 1053 |
+
inputs_embeds=inputs_embeds,
|
| 1054 |
+
use_cache=False,
|
| 1055 |
+
return_dict=True,
|
| 1056 |
+
)
|
| 1057 |
+
h_all = out1.last_hidden_state # (B, SeqLen, D) - detached in no_grad
|
| 1058 |
+
|
| 1059 |
+
# Context for each NT at position k is h_all[:, k-1, :]
|
| 1060 |
+
ctx_pos = (nt_pos - 1).clamp(min=0) # (N_total,)
|
| 1061 |
+
contexts = h_all[batch_idx, ctx_pos] # (N_total, D)
|
| 1062 |
+
|
| 1063 |
+
# --- Pass 2: generate deltas (with grad) ---------------------------
|
| 1064 |
+
ctx_norm = F.normalize(contexts.float(), dim=-1).to(dtype=inputs_embeds.dtype)
|
| 1065 |
+
h_proj = self.init_proj(ctx_norm) # (N_total, D)
|
| 1066 |
+
delta_raw = self.delta_proj(h_proj) # (N_total, D)
|
| 1067 |
+
if self.latent_use_bounded_delta:
|
| 1068 |
+
delta_raw_f32 = delta_raw.float()
|
| 1069 |
+
norm_val = delta_raw_f32.norm(dim=-1, keepdim=True)
|
| 1070 |
+
delta_dir = (delta_raw_f32 / (norm_val + 1e-6)).to(dtype=delta_raw.dtype)
|
| 1071 |
+
scale = self.step_scales().mean() # scalar - shared across all positions
|
| 1072 |
+
delta = delta_dir * scale # (N_total, D)
|
| 1073 |
+
else:
|
| 1074 |
+
delta = delta_raw
|
| 1075 |
+
|
| 1076 |
+
# Build NT embeddings and inject
|
| 1077 |
+
# base_nt_ids and nt_embeds are moved down to use the scaled delta
|
| 1078 |
+
|
| 1079 |
+
|
| 1080 |
+
# --- Pack into (B, max_nt, D) for loss functions -------------------
|
| 1081 |
+
nt_counts = nt_mask.sum(dim=1) # (B,)
|
| 1082 |
+
max_nt = int(nt_counts.max().item())
|
| 1083 |
+
deltas_padded = inputs_embeds.new_zeros((B, max_nt, D))
|
| 1084 |
+
states_padded = inputs_embeds.new_zeros((B, max_nt, D))
|
| 1085 |
+
for b in range(B):
|
| 1086 |
+
b_mask = (batch_idx == b)
|
| 1087 |
+
cnt = int(nt_counts[b].item())
|
| 1088 |
+
if cnt > 0:
|
| 1089 |
+
deltas_padded[b, :cnt] = delta[b_mask]
|
| 1090 |
+
states_padded[b, :cnt] = contexts[b_mask]
|
| 1091 |
+
|
| 1092 |
+
# --- Stats for logging (shape matches prefix-mode stats) ------------
|
| 1093 |
+
with torch.no_grad():
|
| 1094 |
+
delta_for_stats = delta.float()
|
| 1095 |
+
raw_norms = delta_raw.norm(dim=-1) # (N_total,)
|
| 1096 |
+
scaled_norms = delta_for_stats.norm(dim=-1)
|
| 1097 |
+
flat_ctx_norm = F.normalize(contexts, dim=-1)
|
| 1098 |
+
word_emb = self.embed_tokens.weight
|
| 1099 |
+
vocab_size = word_emb.size(0)
|
| 1100 |
+
if vocab_size > 10000:
|
| 1101 |
+
idx = torch.randperm(vocab_size, device=device)[:10000]
|
| 1102 |
+
word_emb_sub = word_emb[idx]
|
| 1103 |
+
else:
|
| 1104 |
+
word_emb_sub = word_emb
|
| 1105 |
+
word_emb_norm = F.normalize(word_emb_sub, dim=-1)
|
| 1106 |
+
sim_matrix = torch.matmul(flat_ctx_norm.to(word_emb_norm.dtype), word_emb_norm.t())
|
| 1107 |
+
max_sim, _ = sim_matrix.max(dim=-1)
|
| 1108 |
+
cos_mean_val = max_sim.mean()
|
| 1109 |
+
cos_std_val = max_sim.std()
|
| 1110 |
+
# Compute diff_norm and step_cos across padded deltas
|
| 1111 |
+
if max_nt > 1:
|
| 1112 |
+
diff = deltas_padded[:, 1:, :] - deltas_padded[:, :-1, :]
|
| 1113 |
+
diff_norm_val = diff.norm(dim=-1).mean()
|
| 1114 |
+
step_cos_val = F.cosine_similarity(
|
| 1115 |
+
deltas_padded[:, 1:, :], deltas_padded[:, :-1, :], dim=-1
|
| 1116 |
+
).mean()
|
| 1117 |
+
else:
|
| 1118 |
+
diff_norm_val = deltas_padded.new_tensor(0.0)
|
| 1119 |
+
step_cos_val = deltas_padded.new_tensor(0.0)
|
| 1120 |
+
|
| 1121 |
+
# Inject deltas into embeddings (grad-tracked for training)
|
| 1122 |
+
base_nt_ids = torch.full(
|
| 1123 |
+
(len(batch_idx),), self.nt_token_id, dtype=torch.long, device=device
|
| 1124 |
+
)
|
| 1125 |
+
base_nt = self.embed_tokens(base_nt_ids).to(dtype=inputs_embeds.dtype)
|
| 1126 |
+
delta_i = delta.to(dtype=inputs_embeds.dtype)
|
| 1127 |
+
if self.latent_use_injection_gate:
|
| 1128 |
+
gate_input = torch.cat([base_nt, delta_i], dim=-1)
|
| 1129 |
+
delta_i = delta_i * self.injection_gate(gate_input)
|
| 1130 |
+
if self.latent_use_embedding_anchor:
|
| 1131 |
+
nt_embeds = base_nt + delta_i
|
| 1132 |
+
else:
|
| 1133 |
+
nt_embeds = delta_i
|
| 1134 |
+
|
| 1135 |
+
result_embeds = inputs_embeds.clone()
|
| 1136 |
+
result_embeds[batch_idx, nt_pos] = nt_embeds
|
| 1137 |
+
|
| 1138 |
+
n_out = max_nt
|
| 1139 |
+
stats = {
|
| 1140 |
+
"raw_norm_mean": raw_norms.mean().unsqueeze(0).expand(n_out).detach(),
|
| 1141 |
+
"raw_norm_std": raw_norms.std(unbiased=False).unsqueeze(0).expand(n_out).detach(),
|
| 1142 |
+
"scaled_norm_mean": scaled_norms.mean().unsqueeze(0).expand(n_out).detach(),
|
| 1143 |
+
"scaled_norm_std": scaled_norms.std(unbiased=False).unsqueeze(0).expand(n_out).detach(),
|
| 1144 |
+
"cos_mean": torch.full((n_out,), cos_mean_val.item(), device=device),
|
| 1145 |
+
"cos_std": torch.full((n_out,), cos_std_val.item(), device=device),
|
| 1146 |
+
"scales": scale.unsqueeze(0).expand(n_out).detach(),
|
| 1147 |
+
"diff_norm": diff_norm_val.detach(),
|
| 1148 |
+
"step_cos": step_cos_val.detach(),
|
| 1149 |
+
}
|
| 1150 |
+
return result_embeds, deltas_padded, states_padded, stats
|
| 1151 |
+
|
| 1152 |
+
@torch.no_grad()
|
| 1153 |
+
def _generate_interleaved(
|
| 1154 |
+
self,
|
| 1155 |
+
base_embeds: torch.Tensor,
|
| 1156 |
+
attention_mask: torch.LongTensor,
|
| 1157 |
+
max_new_tokens: int,
|
| 1158 |
+
do_sample: bool = False,
|
| 1159 |
+
temperature: float = 1.0,
|
| 1160 |
+
) -> torch.LongTensor:
|
| 1161 |
+
"""Token-by-token generation with interleaved thought injection.
|
| 1162 |
+
|
| 1163 |
+
Before each group of ``self.thought_group_size`` generated word tokens,
|
| 1164 |
+
a thought delta is computed from the current sequence's last hidden
|
| 1165 |
+
state, and an NT token (with the modified embedding) is prepended.
|
| 1166 |
+
|
| 1167 |
+
NT tokens in the output are included as ``self.nt_token_id`` so that
|
| 1168 |
+
callers using ``skip_special_tokens=True`` in ``tokenizer.decode`` will
|
| 1169 |
+
strip them automatically (since ``<|latent|>`` is a registered special
|
| 1170 |
+
token).
|
| 1171 |
+
|
| 1172 |
+
Args:
|
| 1173 |
+
base_embeds: ``(1, SeqLen, D)`` audio-fused embeddings of the
|
| 1174 |
+
prompt prefix (no NT tokens appended yet).
|
| 1175 |
+
attention_mask: ``(1, SeqLen)`` attention mask for *base_embeds*.
|
| 1176 |
+
max_new_tokens: Maximum number of word tokens to generate.
|
| 1177 |
+
do_sample: Whether to sample instead of greedy decode.
|
| 1178 |
+
temperature: Sampling temperature (ignored if ``do_sample=False``).
|
| 1179 |
+
|
| 1180 |
+
Returns:
|
| 1181 |
+
``(1, N_generated)`` LongTensor of token ids including interleaved
|
| 1182 |
+
NT tokens and ending with the first EOS/stop token encountered.
|
| 1183 |
+
"""
|
| 1184 |
+
eos_ids: set = set(self.stop_ids)
|
| 1185 |
+
device = base_embeds.device
|
| 1186 |
+
safe_temp = float(temperature) if float(temperature) > 0 else 1.0
|
| 1187 |
+
|
| 1188 |
+
cur_embeds = base_embeds.clone()
|
| 1189 |
+
cur_attn = attention_mask.clone()
|
| 1190 |
+
generated: List[torch.Tensor] = []
|
| 1191 |
+
words_generated = 0
|
| 1192 |
+
|
| 1193 |
+
nt_id_t = torch.tensor([[self.nt_token_id]], dtype=torch.long, device=device)
|
| 1194 |
+
base_nt_emb = self.embed_tokens(nt_id_t).to(dtype=cur_embeds.dtype) # (1, 1, D)
|
| 1195 |
+
|
| 1196 |
+
def _append_thought() -> None:
|
| 1197 |
+
"""Compute a thought from the last hidden state and inject NT."""
|
| 1198 |
+
nonlocal cur_embeds, cur_attn
|
| 1199 |
+
# Get context = last hidden state of current sequence
|
| 1200 |
+
ctx_out = self.text_model(
|
| 1201 |
+
inputs_embeds=cur_embeds,
|
| 1202 |
+
use_cache=False,
|
| 1203 |
+
return_dict=True,
|
| 1204 |
+
)
|
| 1205 |
+
h_last = ctx_out.last_hidden_state[:, -1:, :] # (1, 1, D)
|
| 1206 |
+
h_norm = F.normalize(h_last.float().squeeze(1), dim=-1).to(dtype=cur_embeds.dtype)
|
| 1207 |
+
h_proj = self.init_proj(h_norm)
|
| 1208 |
+
delta_raw = self.delta_proj(h_proj)
|
| 1209 |
+
if self.latent_use_bounded_delta:
|
| 1210 |
+
delta_dir = F.normalize(delta_raw, dim=-1)
|
| 1211 |
+
scale = self.step_scales().mean()
|
| 1212 |
+
delta_scaled = (delta_dir * scale).float().to(dtype=cur_embeds.dtype)
|
| 1213 |
+
else:
|
| 1214 |
+
delta_scaled = delta_raw.float().to(dtype=cur_embeds.dtype)
|
| 1215 |
+
|
| 1216 |
+
# Use Tanh Value Head for Interleaved mode as well
|
| 1217 |
+
if hasattr(self, "value_head"):
|
| 1218 |
+
v_raw = self.value_head(h_norm) # (1, 1)
|
| 1219 |
+
v_pred = (torch.tanh(v_raw) + 1.0) / 2.0 # (1, 1) -> scaled to [0, 1]
|
| 1220 |
+
halt_thresh = float(getattr(self.config, "halt_threshold", 0.0))
|
| 1221 |
+
# Skip injecting NT if predicted value < threshold
|
| 1222 |
+
if v_pred.squeeze().item() < halt_thresh:
|
| 1223 |
+
return
|
| 1224 |
+
|
| 1225 |
+
delta_i = delta_scaled.unsqueeze(1)
|
| 1226 |
+
if self.latent_use_injection_gate:
|
| 1227 |
+
gate_input = torch.cat([base_nt_emb.expand_as(delta_i), delta_i], dim=-1)
|
| 1228 |
+
delta_i = delta_i * self.injection_gate(gate_input)
|
| 1229 |
+
if self.latent_use_embedding_anchor:
|
| 1230 |
+
nt_emb_mod = base_nt_emb + delta_i
|
| 1231 |
+
else:
|
| 1232 |
+
nt_emb_mod = delta_i
|
| 1233 |
+
cur_embeds = torch.cat([cur_embeds, nt_emb_mod], dim=1)
|
| 1234 |
+
cur_attn = torch.cat([cur_attn, cur_attn.new_ones((1, 1))], dim=1)
|
| 1235 |
+
generated.append(nt_id_t.clone())
|
| 1236 |
+
|
| 1237 |
+
for _ in range(max_new_tokens + len(generated)):
|
| 1238 |
+
# Inject thought before each new group of words
|
| 1239 |
+
if words_generated % self.thought_group_size == 0:
|
| 1240 |
+
_append_thought()
|
| 1241 |
+
|
| 1242 |
+
# Generate next word token
|
| 1243 |
+
out = self.thinker(
|
| 1244 |
+
inputs_embeds=cur_embeds,
|
| 1245 |
+
attention_mask=cur_attn,
|
| 1246 |
+
input_features=None,
|
| 1247 |
+
use_cache=False,
|
| 1248 |
+
return_dict=True,
|
| 1249 |
+
)
|
| 1250 |
+
next_logits = out.logits[:, -1, :]
|
| 1251 |
+
if do_sample:
|
| 1252 |
+
probs = F.softmax(next_logits / safe_temp, dim=-1)
|
| 1253 |
+
next_token = torch.multinomial(probs, num_samples=1)
|
| 1254 |
+
else:
|
| 1255 |
+
next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
|
| 1256 |
+
|
| 1257 |
+
tok_id = int(next_token[0, 0].item())
|
| 1258 |
+
generated.append(next_token)
|
| 1259 |
+
if tok_id in eos_ids:
|
| 1260 |
+
break
|
| 1261 |
+
|
| 1262 |
+
next_emb = self.embed_tokens(next_token)
|
| 1263 |
+
cur_embeds = torch.cat([cur_embeds, next_emb], dim=1)
|
| 1264 |
+
cur_attn = torch.cat([cur_attn, cur_attn.new_ones((1, 1))], dim=1)
|
| 1265 |
+
words_generated += 1
|
| 1266 |
+
|
| 1267 |
+
if not generated:
|
| 1268 |
+
return torch.zeros((1, 0), dtype=torch.long, device=device)
|
| 1269 |
+
return torch.cat(generated, dim=1)
|
| 1270 |
+
|
| 1271 |
+
# ------------------------------------------------------------------
|
| 1272 |
+
# Forward passes
|
| 1273 |
+
# ------------------------------------------------------------------
|
| 1274 |
+
|
| 1275 |
+
def _forward_native(
|
| 1276 |
+
self,
|
| 1277 |
+
input_features: torch.FloatTensor,
|
| 1278 |
+
labels: torch.LongTensor,
|
| 1279 |
+
feature_attention_mask: torch.LongTensor,
|
| 1280 |
+
language_hint: Optional[str] = None,
|
| 1281 |
+
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 1282 |
+
"""Native thinker forward for LoRA / baseline FT."""
|
| 1283 |
+
B = input_features.size(0)
|
| 1284 |
+
device = input_features.device
|
| 1285 |
+
|
| 1286 |
+
pad_token_id = self.processor.tokenizer.pad_token_id
|
| 1287 |
+
if pad_token_id is None:
|
| 1288 |
+
pad_token_id = self.processor.tokenizer.eos_token_id or 151643
|
| 1289 |
+
|
| 1290 |
+
input_ids_list = []
|
| 1291 |
+
labels_list = []
|
| 1292 |
+
label_pad = -100
|
| 1293 |
+
|
| 1294 |
+
audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)
|
| 1295 |
+
system_turn, user_prefix, user_prompt, user_suffix, _ = (
|
| 1296 |
+
self._build_chat_template_tensors(device, language=language_hint or "English")
|
| 1297 |
+
)
|
| 1298 |
+
|
| 1299 |
+
for i in range(B):
|
| 1300 |
+
valid_len = (labels[i] != label_pad).sum().item()
|
| 1301 |
+
text_all = labels[i, :valid_len]
|
| 1302 |
+
l_audio = audio_lengths[i]
|
| 1303 |
+
|
| 1304 |
+
audio_toks = build_audio_token_seq(
|
| 1305 |
+
self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
|
| 1306 |
+
l_audio, device,
|
| 1307 |
+
)
|
| 1308 |
+
full_ids, full_labels = self._build_full_sequence(
|
| 1309 |
+
audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
|
| 1310 |
+
label_pad=label_pad, n_extra_mask=0,
|
| 1311 |
+
)
|
| 1312 |
+
input_ids_list.append(full_ids)
|
| 1313 |
+
labels_list.append(full_labels)
|
| 1314 |
+
|
| 1315 |
+
max_len = max(x.size(0) for x in input_ids_list)
|
| 1316 |
+
input_ids_batch = torch.full((B, max_len), pad_token_id, dtype=torch.long, device=device)
|
| 1317 |
+
labels_batch = torch.full((B, max_len), label_pad, dtype=torch.long, device=device)
|
| 1318 |
+
for i in range(B):
|
| 1319 |
+
l = input_ids_list[i].size(0)
|
| 1320 |
+
input_ids_batch[i, :l] = input_ids_list[i]
|
| 1321 |
+
labels_batch[i, :l] = labels_list[i]
|
| 1322 |
+
|
| 1323 |
+
attention_mask = (input_ids_batch != pad_token_id).long()
|
| 1324 |
+
|
| 1325 |
+
out = self.thinker(
|
| 1326 |
+
input_ids=input_ids_batch,
|
| 1327 |
+
attention_mask=attention_mask,
|
| 1328 |
+
input_features=input_features,
|
| 1329 |
+
feature_attention_mask=feature_attention_mask,
|
| 1330 |
+
labels=None,
|
| 1331 |
+
return_dict=True,
|
| 1332 |
+
)
|
| 1333 |
+
logits = out.logits
|
| 1334 |
+
|
| 1335 |
+
# One-time audio diagnostic
|
| 1336 |
+
if not getattr(self, "_native_diag_done", False):
|
| 1337 |
+
self._native_diag_done = True
|
| 1338 |
+
import sys as _s
|
| 1339 |
+
_d = _s.stderr.write
|
| 1340 |
+
_d(f"\n[native-diag] thinker type: {type(self.thinker).__name__}\n")
|
| 1341 |
+
_d(f"[native-diag] has base_model (PEFT): {hasattr(self.thinker, 'base_model')}\n")
|
| 1342 |
+
n_audio_ph = (input_ids_batch[0] == self.audio_token_id).sum().item()
|
| 1343 |
+
feat_valid = (
|
| 1344 |
+
feature_attention_mask[0].sum().item()
|
| 1345 |
+
if feature_attention_mask is not None
|
| 1346 |
+
else input_features.size(2)
|
| 1347 |
+
)
|
| 1348 |
+
_d(f"[native-diag] audio_token_id={self.audio_token_id} decoded={self.processor.tokenizer.convert_ids_to_tokens(self.audio_token_id)!r}\n")
|
| 1349 |
+
_d(f"[native-diag] input_features: {list(input_features.shape)}, valid_mel_frames={feat_valid}\n")
|
| 1350 |
+
_d(f"[native-diag] input_ids shape: {list(input_ids_batch.shape)}\n")
|
| 1351 |
+
_d(f"[native-diag] input_features shape: {list(input_features.shape)}\n")
|
| 1352 |
+
_d(f"[native-diag] audio placeholder tokens in input_ids[0]: {n_audio_ph}\n")
|
| 1353 |
+
_d(f"[native-diag] audio_bos/eos in ids[0]: bos={int((input_ids_batch[0]==self.audio_bos_token_id).sum())} eos={int((input_ids_batch[0]==self.audio_eos_token_id).sum())}\n")
|
| 1354 |
+
with torch.no_grad():
|
| 1355 |
+
try:
|
| 1356 |
+
zero_feat = torch.zeros_like(input_features[:1])
|
| 1357 |
+
out_noaud = self.thinker(
|
| 1358 |
+
input_ids=input_ids_batch[:1],
|
| 1359 |
+
attention_mask=attention_mask[:1],
|
| 1360 |
+
input_features=zero_feat,
|
| 1361 |
+
feature_attention_mask=feature_attention_mask[:1] if feature_attention_mask is not None else None,
|
| 1362 |
+
labels=None,
|
| 1363 |
+
return_dict=True,
|
| 1364 |
+
)
|
| 1365 |
+
diff = (logits[0] - out_noaud.logits[0]).abs().mean().item()
|
| 1366 |
+
_d(f"[native-diag] logit diff (real vs zero audio): {diff:.6f}\n")
|
| 1367 |
+
if diff < 0.01:
|
| 1368 |
+
_d(f"[native-diag] ⚠️ AUDIO NOT BEING USED! Logits identical with zero audio.\n")
|
| 1369 |
+
else:
|
| 1370 |
+
_d(f"[native-diag] ✓ Audio IS affecting logits (diff={diff:.4f})\n")
|
| 1371 |
+
except Exception as e:
|
| 1372 |
+
_d(f"[native-diag] zero-audio test failed: {e}\n")
|
| 1373 |
+
tok_str = self.processor.tokenizer.decode(input_ids_batch[0, :15].tolist())
|
| 1374 |
+
_d(f"[native-diag] input_ids[0][:15] decoded: {repr(tok_str)}\n")
|
| 1375 |
+
_s.stderr.flush()
|
| 1376 |
+
|
| 1377 |
+
# Pre-shift labels for loss
|
| 1378 |
+
labels_for_loss = torch.full_like(labels_batch, label_pad)
|
| 1379 |
+
labels_for_loss[:, :-1] = labels_batch[:, 1:]
|
| 1380 |
+
|
| 1381 |
+
D = self._hidden_size
|
| 1382 |
+
dummy = logits.new_tensor(0.0)
|
| 1383 |
+
stats = {
|
| 1384 |
+
"raw_norm_mean": logits.new_zeros((0,)),
|
| 1385 |
+
"raw_norm_std": logits.new_zeros((0,)),
|
| 1386 |
+
"scaled_norm_mean": logits.new_zeros((0,)),
|
| 1387 |
+
"scaled_norm_std": logits.new_zeros((0,)),
|
| 1388 |
+
"cos_mean": logits.new_zeros((0,)),
|
| 1389 |
+
"cos_std": logits.new_zeros((0,)),
|
| 1390 |
+
"scales": logits.new_zeros((0,)),
|
| 1391 |
+
"diff_norm": dummy,
|
| 1392 |
+
"step_cos": dummy,
|
| 1393 |
+
}
|
| 1394 |
+
deltas = logits.new_zeros((B, 0, D))
|
| 1395 |
+
states = logits.new_zeros((B, 0, D))
|
| 1396 |
+
initial_state = logits.new_zeros((B, D))
|
| 1397 |
+
|
| 1398 |
+
return logits, stats, deltas, states, labels_for_loss, initial_state
|
| 1399 |
+
|
| 1400 |
+
def forward(
|
| 1401 |
+
self,
|
| 1402 |
+
input_features: torch.FloatTensor,
|
| 1403 |
+
labels: torch.LongTensor,
|
| 1404 |
+
feature_attention_mask: Optional[torch.LongTensor] = None,
|
| 1405 |
+
global_step: int = 0,
|
| 1406 |
+
language_hint: Optional[str] = None,
|
| 1407 |
+
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 1408 |
+
"""Forward pass for training."""
|
| 1409 |
+
B = input_features.size(0)
|
| 1410 |
+
device = input_features.device
|
| 1411 |
+
|
| 1412 |
+
if feature_attention_mask is None:
|
| 1413 |
+
feature_attention_mask = torch.ones(
|
| 1414 |
+
(B, input_features.size(2)), dtype=torch.long, device=device
|
| 1415 |
+
)
|
| 1416 |
+
|
| 1417 |
+
# Fast path: use thinker's native forward for LoRA / baseline FT.
|
| 1418 |
+
if not self.use_latent and not self.use_soft_prompt:
|
| 1419 |
+
return self._forward_native(input_features, labels, feature_attention_mask, language_hint=language_hint)
|
| 1420 |
+
|
| 1421 |
+
audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)
|
| 1422 |
+
|
| 1423 |
+
new_input_ids_list = []
|
| 1424 |
+
new_labels_list = []
|
| 1425 |
+
|
| 1426 |
+
label_pad = -100
|
| 1427 |
+
pad_token_id = self.processor.tokenizer.pad_token_id
|
| 1428 |
+
if pad_token_id is None:
|
| 1429 |
+
pad_token_id = self.processor.tokenizer.eos_token_id or 151643
|
| 1430 |
+
|
| 1431 |
+
system_turn, user_prefix, user_prompt, user_suffix, _ = (
|
| 1432 |
+
self._build_chat_template_tensors(device, language=language_hint or "English")
|
| 1433 |
+
)
|
| 1434 |
+
|
| 1435 |
+
interleaved_mode = self.use_latent and self.thought_mode == "interleaved"
|
| 1436 |
+
|
| 1437 |
+
# +++ LATENT DROP +++
|
| 1438 |
+
drop_latent_this_batch = False
|
| 1439 |
+
if self.training and self.use_latent and not interleaved_mode:
|
| 1440 |
+
drop_prob = self.latent_drop_prob
|
| 1441 |
+
if drop_prob > 0.0 and torch.rand(1).item() < drop_prob:
|
| 1442 |
+
drop_latent_this_batch = True
|
| 1443 |
+
|
| 1444 |
+
actual_n_latent = 0 if drop_latent_this_batch else self.n_latent
|
| 1445 |
+
|
| 1446 |
+
for i in range(B):
|
| 1447 |
+
valid_len = (labels[i] != label_pad).sum().item()
|
| 1448 |
+
text_all = labels[i, :valid_len]
|
| 1449 |
+
if drop_latent_this_batch:
|
| 1450 |
+
text_all = text_all[text_all != self.nt_token_id]
|
| 1451 |
+
|
| 1452 |
+
l_audio = audio_lengths[i]
|
| 1453 |
+
|
| 1454 |
+
audio_toks = build_audio_token_seq(
|
| 1455 |
+
self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
|
| 1456 |
+
l_audio, device,
|
| 1457 |
+
)
|
| 1458 |
+
if interleaved_mode:
|
| 1459 |
+
# NT tokens are scattered throughout text_all; mask all NT positions.
|
| 1460 |
+
full_ids, full_labels = self._build_full_sequence(
|
| 1461 |
+
audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
|
| 1462 |
+
label_pad=label_pad, n_extra_mask=0, mask_all_nt=True,
|
| 1463 |
+
)
|
| 1464 |
+
else:
|
| 1465 |
+
# Prefix mode: all NT tokens appear contiguously after asst_nl.
|
| 1466 |
+
full_ids, full_labels = self._build_full_sequence(
|
| 1467 |
+
audio_toks, text_all, system_turn, user_prefix, user_prompt, user_suffix,
|
| 1468 |
+
label_pad=label_pad, n_extra_mask=actual_n_latent,
|
| 1469 |
+
)
|
| 1470 |
+
new_input_ids_list.append(full_ids)
|
| 1471 |
+
new_labels_list.append(full_labels)
|
| 1472 |
+
|
| 1473 |
+
max_len = max(x.size(0) for x in new_input_ids_list)
|
| 1474 |
+
new_input_ids = torch.full((B, max_len), pad_token_id, dtype=torch.long, device=device)
|
| 1475 |
+
new_labels_raw = torch.full((B, max_len), label_pad, dtype=torch.long, device=device)
|
| 1476 |
+
|
| 1477 |
+
for i in range(B):
|
| 1478 |
+
l = new_input_ids_list[i].size(0)
|
| 1479 |
+
new_input_ids[i, :l] = new_input_ids_list[i]
|
| 1480 |
+
new_labels_raw[i, :l] = new_labels_list[i]
|
| 1481 |
+
|
| 1482 |
+
attention_mask = (new_input_ids != pad_token_id).long()
|
| 1483 |
+
|
| 1484 |
+
inputs_embeds = self._get_native_audio_embeds(
|
| 1485 |
+
input_ids=new_input_ids,
|
| 1486 |
+
input_features=input_features,
|
| 1487 |
+
feature_attention_mask=feature_attention_mask,
|
| 1488 |
+
attention_mask=attention_mask,
|
| 1489 |
+
)
|
| 1490 |
+
|
| 1491 |
+
# +++ LATENT INPUT NOISE (Value Head Regularization) +++
|
| 1492 |
+
# Save pre-LR embeddings (with noise) for fair baseline CE comparison
|
| 1493 |
+
pre_lr_embeds = None
|
| 1494 |
+
if self.training:
|
| 1495 |
+
noise_std = self.latent_input_noise_std
|
| 1496 |
+
if noise_std > 0.0:
|
| 1497 |
+
noise = torch.randn_like(inputs_embeds) * noise_std
|
| 1498 |
+
inputs_embeds = inputs_embeds + noise
|
| 1499 |
+
# Save the (possibly noisy) embeddings BEFORE LR injection
|
| 1500 |
+
if self.use_latent:
|
| 1501 |
+
pre_lr_embeds = inputs_embeds.detach().clone()
|
| 1502 |
+
|
| 1503 |
+
# Latent / soft-prompt injection
|
| 1504 |
+
if self.use_latent and interleaved_mode:
|
| 1505 |
+
# Interleaved: two-pass causal thought generation.
|
| 1506 |
+
inputs_embeds, deltas, states, stats = self._compute_causal_thoughts(
|
| 1507 |
+
inputs_embeds, new_input_ids
|
| 1508 |
+
)
|
| 1509 |
+
initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
|
| 1510 |
+
elif self.use_latent:
|
| 1511 |
+
if actual_n_latent == 0:
|
| 1512 |
+
initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
|
| 1513 |
+
deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
|
| 1514 |
+
states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
|
| 1515 |
+
stats = {
|
| 1516 |
+
"raw_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1517 |
+
"raw_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1518 |
+
"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1519 |
+
"scaled_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1520 |
+
"cos_mean": inputs_embeds.new_zeros((0,)),
|
| 1521 |
+
"cos_std": inputs_embeds.new_zeros((0,)),
|
| 1522 |
+
"scales": inputs_embeds.new_zeros((0,)),
|
| 1523 |
+
"diff_norm": inputs_embeds.new_tensor(0.0),
|
| 1524 |
+
"step_cos": inputs_embeds.new_tensor(0.0),
|
| 1525 |
+
}
|
| 1526 |
+
else:
|
| 1527 |
+
# Prefix mode: recurrent thought loop from audio-prefix state.
|
| 1528 |
+
initial_state, prefix_past_kv, prefix_attn = self._compute_prefix_state(
|
| 1529 |
+
inputs_embeds, new_input_ids, strict=True
|
| 1530 |
+
)
|
| 1531 |
+
deltas, states, stats = self._compute_continuous_thoughts(
|
| 1532 |
+
initial_state,
|
| 1533 |
+
prefix_past_key_values=prefix_past_kv,
|
| 1534 |
+
prefix_attention_mask=prefix_attn,
|
| 1535 |
+
)
|
| 1536 |
+
inputs_embeds = self._inject_latent_deltas(
|
| 1537 |
+
inputs_embeds, new_input_ids, deltas
|
| 1538 |
+
)
|
| 1539 |
+
elif self.use_soft_prompt:
|
| 1540 |
+
nt_mask = (new_input_ids == self.nt_token_id)
|
| 1541 |
+
nt_counts = nt_mask.sum(dim=1)
|
| 1542 |
+
assert (nt_counts == self.n_latent).all(), (
|
| 1543 |
+
f"NT token count mismatch! {nt_counts} vs {self.n_latent}"
|
| 1544 |
+
)
|
| 1545 |
+
prompt = self.soft_prompt_embed.to(dtype=inputs_embeds.dtype)
|
| 1546 |
+
prompt = prompt.unsqueeze(0).expand(B, -1, -1).contiguous()
|
| 1547 |
+
batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
|
| 1548 |
+
prompt_embeds = inputs_embeds.clone()
|
| 1549 |
+
prompt_embeds[batch_idx, seq_idx, :] = prompt.view(-1, self._hidden_size)
|
| 1550 |
+
inputs_embeds = prompt_embeds
|
| 1551 |
+
initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
|
| 1552 |
+
deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
|
| 1553 |
+
states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
|
| 1554 |
+
stats = {
|
| 1555 |
+
"raw_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1556 |
+
"raw_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1557 |
+
"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1558 |
+
"scaled_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1559 |
+
"cos_mean": inputs_embeds.new_zeros((0,)),
|
| 1560 |
+
"cos_std": inputs_embeds.new_zeros((0,)),
|
| 1561 |
+
"scales": inputs_embeds.new_zeros((0,)),
|
| 1562 |
+
"diff_norm": inputs_embeds.new_tensor(0.0),
|
| 1563 |
+
"step_cos": inputs_embeds.new_tensor(0.0),
|
| 1564 |
+
}
|
| 1565 |
+
else:
|
| 1566 |
+
initial_state = inputs_embeds.new_zeros((B, self._hidden_size))
|
| 1567 |
+
deltas = inputs_embeds.new_zeros((B, 0, self._hidden_size))
|
| 1568 |
+
states = inputs_embeds.new_zeros((B, 0, self._hidden_size))
|
| 1569 |
+
stats = {
|
| 1570 |
+
"raw_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1571 |
+
"raw_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1572 |
+
"scaled_norm_mean": inputs_embeds.new_zeros((0,)),
|
| 1573 |
+
"scaled_norm_std": inputs_embeds.new_zeros((0,)),
|
| 1574 |
+
"cos_mean": inputs_embeds.new_zeros((0,)),
|
| 1575 |
+
"cos_std": inputs_embeds.new_zeros((0,)),
|
| 1576 |
+
"scales": inputs_embeds.new_zeros((0,)),
|
| 1577 |
+
"diff_norm": inputs_embeds.new_tensor(0.0),
|
| 1578 |
+
"step_cos": inputs_embeds.new_tensor(0.0),
|
| 1579 |
+
}
|
| 1580 |
+
|
| 1581 |
+
out = self.thinker(
|
| 1582 |
+
inputs_embeds=inputs_embeds,
|
| 1583 |
+
input_ids=new_input_ids,
|
| 1584 |
+
attention_mask=attention_mask,
|
| 1585 |
+
input_features=None,
|
| 1586 |
+
feature_attention_mask=None,
|
| 1587 |
+
labels=None,
|
| 1588 |
+
use_cache=False,
|
| 1589 |
+
return_dict=True,
|
| 1590 |
+
)
|
| 1591 |
+
logits = out.logits
|
| 1592 |
+
|
| 1593 |
+
new_labels = torch.full_like(new_labels_raw, label_pad)
|
| 1594 |
+
new_labels[:, :-1] = new_labels_raw[:, 1:]
|
| 1595 |
+
|
| 1596 |
+
if logits.size(1) != new_labels.size(1):
|
| 1597 |
+
if logits.size(1) > new_labels.size(1):
|
| 1598 |
+
shift = logits.size(1) - new_labels.size(1)
|
| 1599 |
+
logits = logits[:, shift:, :]
|
| 1600 |
+
else:
|
| 1601 |
+
new_labels = new_labels[:, -logits.size(1):]
|
| 1602 |
+
|
| 1603 |
+
# ---- Baseline CE vs LR CE for Value Head MSE (Delta CE) ----
|
| 1604 |
+
# Compute per-sample CE for both baseline (no LR) and LR-enhanced
|
| 1605 |
+
# forward passes. The difference (baseline_ce - lr_ce) forms the
|
| 1606 |
+
# continuous target for the Value Head via tanh compression.
|
| 1607 |
+
if self.use_latent and self.training and deltas.size(1) > 0:
|
| 1608 |
+
with torch.no_grad():
|
| 1609 |
+
# Use the SAME (possibly noisy) embeddings but WITHOUT LR deltas
|
| 1610 |
+
# This ensures a fair comparison: noisy-no-LR vs noisy-with-LR
|
| 1611 |
+
if pre_lr_embeds is not None:
|
| 1612 |
+
baseline_embeds = pre_lr_embeds
|
| 1613 |
+
else:
|
| 1614 |
+
baseline_embeds = self._get_native_audio_embeds(
|
| 1615 |
+
input_ids=new_input_ids,
|
| 1616 |
+
input_features=input_features,
|
| 1617 |
+
feature_attention_mask=feature_attention_mask,
|
| 1618 |
+
attention_mask=attention_mask,
|
| 1619 |
+
)
|
| 1620 |
+
baseline_out = self.thinker(
|
| 1621 |
+
inputs_embeds=baseline_embeds,
|
| 1622 |
+
input_ids=new_input_ids,
|
| 1623 |
+
attention_mask=attention_mask,
|
| 1624 |
+
input_features=None,
|
| 1625 |
+
feature_attention_mask=None,
|
| 1626 |
+
labels=None,
|
| 1627 |
+
use_cache=False,
|
| 1628 |
+
return_dict=True,
|
| 1629 |
+
)
|
| 1630 |
+
baseline_logits = baseline_out.logits
|
| 1631 |
+
if baseline_logits.size(1) != new_labels.size(1):
|
| 1632 |
+
if baseline_logits.size(1) > new_labels.size(1):
|
| 1633 |
+
shift = baseline_logits.size(1) - new_labels.size(1)
|
| 1634 |
+
baseline_logits = baseline_logits[:, shift:, :]
|
| 1635 |
+
else:
|
| 1636 |
+
bl_labels = new_labels[:, -baseline_logits.size(1):]
|
| 1637 |
+
else:
|
| 1638 |
+
bl_labels = new_labels
|
| 1639 |
+
|
| 1640 |
+
# Per-SAMPLE CE: for each sample in batch, compute CE on valid tokens.
|
| 1641 |
+
# This is used by train.py to form the continuous delta_ce target.
|
| 1642 |
+
valid_mask = (bl_labels != -100) # (B, T)
|
| 1643 |
+
valid_per_sample = valid_mask.float().sum(dim=1).clamp(min=1) # (B,)
|
| 1644 |
+
|
| 1645 |
+
# For value head accuracy: exclude asr_prefix tokens
|
| 1646 |
+
# ("language", "English", "<asr_text>") which are trivially correct
|
| 1647 |
+
# and dilute the real transcript accuracy signal.
|
| 1648 |
+
asr_prefix_set = set(self.asr_prefix_ids)
|
| 1649 |
+
asr_prefix_mask = torch.zeros_like(bl_labels, dtype=torch.bool)
|
| 1650 |
+
for tid in asr_prefix_set:
|
| 1651 |
+
asr_prefix_mask |= (bl_labels == tid)
|
| 1652 |
+
acc_mask = valid_mask & ~asr_prefix_mask # exclude prefix tokens
|
| 1653 |
+
acc_per_sample = acc_mask.float().sum(dim=1).clamp(min=1)
|
| 1654 |
+
|
| 1655 |
+
if not hasattr(self, '_acc_debug_printed'):
|
| 1656 |
+
self._acc_debug_printed = True
|
| 1657 |
+
total_tokens = bl_labels.size(1)
|
| 1658 |
+
ce_valid = valid_per_sample.mean().item()
|
| 1659 |
+
acc_valid = acc_per_sample.mean().item()
|
| 1660 |
+
sample_acc_toks = bl_labels[0][acc_mask[0]][:20].tolist()
|
| 1661 |
+
sample_decoded = self.processor.tokenizer.decode(sample_acc_toks)
|
| 1662 |
+
print(f"\n[ACC-DEBUG] total={total_tokens} ce_valid={ce_valid:.0f} acc_valid={acc_valid:.0f}")
|
| 1663 |
+
print(f"[ACC-DEBUG] first 20 acc label ids: {sample_acc_toks}")
|
| 1664 |
+
print(f"[ACC-DEBUG] decoded: {sample_decoded}")
|
| 1665 |
+
if valid_mask.any():
|
| 1666 |
+
# Per-sample baseline CE (uses full valid_mask for loss)
|
| 1667 |
+
loss_fct_none = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none')
|
| 1668 |
+
bl_ce_per_token = loss_fct_none(
|
| 1669 |
+
baseline_logits.reshape(-1, baseline_logits.size(-1)),
|
| 1670 |
+
bl_labels.reshape(-1),
|
| 1671 |
+
).view(B, -1) # (B, T)
|
| 1672 |
+
baseline_ce_per_sample = bl_ce_per_token.sum(dim=1) / valid_per_sample # (B,)
|
| 1673 |
+
stats["baseline_ce"] = baseline_ce_per_sample # (B,)
|
| 1674 |
+
|
| 1675 |
+
# Per-sample LR CE (for fair Delta CE comparison)
|
| 1676 |
+
lr_logits_for_ce = logits
|
| 1677 |
+
if lr_logits_for_ce.size(1) != bl_labels.size(1):
|
| 1678 |
+
lr_logits_for_ce = lr_logits_for_ce[:, -bl_labels.size(1):, :]
|
| 1679 |
+
lr_ce_per_token = loss_fct_none(
|
| 1680 |
+
lr_logits_for_ce.reshape(-1, lr_logits_for_ce.size(-1)),
|
| 1681 |
+
bl_labels.reshape(-1),
|
| 1682 |
+
).view(B, -1)
|
| 1683 |
+
lr_ce_per_sample = lr_ce_per_token.sum(dim=1) / valid_per_sample
|
| 1684 |
+
stats["lr_ce"] = lr_ce_per_sample # (B,)
|
| 1685 |
+
|
| 1686 |
+
# Per-sample accuracy for VALUE HEAD (uses acc_mask, excludes prefix)
|
| 1687 |
+
baseline_preds = baseline_logits.argmax(dim=-1) # (B, T)
|
| 1688 |
+
lr_preds = logits.argmax(dim=-1) # (B, T)
|
| 1689 |
+
if lr_preds.size(1) != bl_labels.size(1):
|
| 1690 |
+
lr_preds = lr_preds[:, -bl_labels.size(1):]
|
| 1691 |
+
baseline_correct = ((baseline_preds == bl_labels) & acc_mask).float()
|
| 1692 |
+
lr_correct = ((lr_preds == bl_labels) & acc_mask).float()
|
| 1693 |
+
baseline_acc = baseline_correct.sum(dim=1) / acc_per_sample # (B,)
|
| 1694 |
+
lr_acc = lr_correct.sum(dim=1) / acc_per_sample # (B,)
|
| 1695 |
+
stats["baseline_acc"] = baseline_acc
|
| 1696 |
+
stats["lr_acc"] = lr_acc
|
| 1697 |
+
|
| 1698 |
+
# Error-conditional metrics for value head target (uses acc_mask)
|
| 1699 |
+
baseline_wrong = ((baseline_preds != bl_labels) & acc_mask) # (B, T)
|
| 1700 |
+
baseline_right = ((baseline_preds == bl_labels) & acc_mask) # (B, T)
|
| 1701 |
+
lr_fixes = ((lr_preds == bl_labels) & baseline_wrong).float().sum(dim=1) # (B,)
|
| 1702 |
+
lr_breaks = ((lr_preds != bl_labels) & baseline_right).float().sum(dim=1) # (B,)
|
| 1703 |
+
baseline_errors = baseline_wrong.float().sum(dim=1).clamp(min=1) # (B,)
|
| 1704 |
+
stats["lr_fixes"] = lr_fixes
|
| 1705 |
+
stats["lr_breaks"] = lr_breaks
|
| 1706 |
+
stats["baseline_errors"] = baseline_errors
|
| 1707 |
+
else:
|
| 1708 |
+
B_size = logits.size(0)
|
| 1709 |
+
stats["baseline_ce"] = logits.new_zeros(B_size)
|
| 1710 |
+
stats["baseline_acc"] = logits.new_ones(B_size)
|
| 1711 |
+
stats["lr_acc"] = logits.new_ones(B_size)
|
| 1712 |
+
stats["lr_fixes"] = logits.new_zeros(B_size)
|
| 1713 |
+
stats["lr_breaks"] = logits.new_zeros(B_size)
|
| 1714 |
+
stats["baseline_errors"] = logits.new_ones(B_size)
|
| 1715 |
+
|
| 1716 |
+
return logits, stats, deltas, states, new_labels, initial_state
|
| 1717 |
+
|
| 1718 |
+
@torch.no_grad()
|
| 1719 |
+
def _generate_with_forward_fallback(
|
| 1720 |
+
self,
|
| 1721 |
+
start_ids: torch.LongTensor,
|
| 1722 |
+
inputs_embeds: torch.FloatTensor,
|
| 1723 |
+
attention_mask: torch.LongTensor,
|
| 1724 |
+
max_new_tokens: int,
|
| 1725 |
+
eos_token_id: Optional[Any],
|
| 1726 |
+
do_sample: bool,
|
| 1727 |
+
temperature: float,
|
| 1728 |
+
) -> torch.LongTensor:
|
| 1729 |
+
"""Fallback decoder when ``.generate(inputs_embeds=...)`` is unsupported."""
|
| 1730 |
+
if eos_token_id is None:
|
| 1731 |
+
eos_ids: List[int] = []
|
| 1732 |
+
elif isinstance(eos_token_id, (list, tuple, set)):
|
| 1733 |
+
eos_ids = [int(x) for x in eos_token_id]
|
| 1734 |
+
else:
|
| 1735 |
+
eos_ids = [int(eos_token_id)]
|
| 1736 |
+
|
| 1737 |
+
generated_tokens: List[torch.Tensor] = []
|
| 1738 |
+
safe_temp = float(temperature) if float(temperature) > 0 else 1.0
|
| 1739 |
+
|
| 1740 |
+
cur_ids = start_ids.clone()
|
| 1741 |
+
cur_embeds = inputs_embeds.clone()
|
| 1742 |
+
cur_attn = attention_mask.clone()
|
| 1743 |
+
|
| 1744 |
+
for _ in range(int(max_new_tokens)):
|
| 1745 |
+
out = self.thinker(
|
| 1746 |
+
inputs_embeds=cur_embeds,
|
| 1747 |
+
attention_mask=cur_attn,
|
| 1748 |
+
input_features=None,
|
| 1749 |
+
use_cache=False,
|
| 1750 |
+
return_dict=True,
|
| 1751 |
+
)
|
| 1752 |
+
next_logits = out.logits[:, -1, :]
|
| 1753 |
+
|
| 1754 |
+
if do_sample:
|
| 1755 |
+
probs = F.softmax(next_logits / safe_temp, dim=-1)
|
| 1756 |
+
next_token = torch.multinomial(probs, num_samples=1)
|
| 1757 |
+
else:
|
| 1758 |
+
next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
|
| 1759 |
+
|
| 1760 |
+
generated_tokens.append(next_token)
|
| 1761 |
+
if eos_ids and int(next_token[0, 0].item()) in eos_ids:
|
| 1762 |
+
break
|
| 1763 |
+
|
| 1764 |
+
next_embed = self.embed_tokens(next_token)
|
| 1765 |
+
cur_ids = torch.cat([cur_ids, next_token], dim=1)
|
| 1766 |
+
cur_embeds = torch.cat([cur_embeds, next_embed], dim=1)
|
| 1767 |
+
cur_attn = torch.cat(
|
| 1768 |
+
[
|
| 1769 |
+
cur_attn,
|
| 1770 |
+
torch.ones(
|
| 1771 |
+
(cur_attn.size(0), 1),
|
| 1772 |
+
dtype=cur_attn.dtype,
|
| 1773 |
+
device=cur_attn.device,
|
| 1774 |
+
),
|
| 1775 |
+
],
|
| 1776 |
+
dim=1,
|
| 1777 |
+
)
|
| 1778 |
+
|
| 1779 |
+
if not generated_tokens:
|
| 1780 |
+
return start_ids.new_empty((start_ids.size(0), 0))
|
| 1781 |
+
return torch.cat(generated_tokens, dim=1)
|
| 1782 |
+
|
| 1783 |
+
@torch.no_grad()
|
| 1784 |
+
def generate(
|
| 1785 |
+
self,
|
| 1786 |
+
input_features: torch.FloatTensor,
|
| 1787 |
+
feature_attention_mask: Optional[torch.LongTensor] = None,
|
| 1788 |
+
max_new_tokens: int = 128,
|
| 1789 |
+
use_baseline: bool = False,
|
| 1790 |
+
return_thoughts: bool = False,
|
| 1791 |
+
return_stats: bool = False,
|
| 1792 |
+
language_hint: Optional[str] = None,
|
| 1793 |
+
prompt_text: Optional[str] = None,
|
| 1794 |
+
**gen_kwargs: Any,
|
| 1795 |
+
) -> Any:
|
| 1796 |
+
"""Generate a transcription with or without latent prompting."""
|
| 1797 |
+
language_hint = gen_kwargs.pop("language_hint", language_hint)
|
| 1798 |
+
prompt_text = gen_kwargs.pop("prompt_text", prompt_text)
|
| 1799 |
+
dynamic_halt_threshold = gen_kwargs.pop("dynamic_halt_threshold", getattr(self, "halt_threshold", 0.0))
|
| 1800 |
+
|
| 1801 |
+
B = input_features.size(0)
|
| 1802 |
+
assert B == 1, "Evaluation assumes batch size 1."
|
| 1803 |
+
device = input_features.device
|
| 1804 |
+
|
| 1805 |
+
if feature_attention_mask is None:
|
| 1806 |
+
feature_attention_mask = torch.ones(
|
| 1807 |
+
(B, input_features.size(2)), dtype=torch.long, device=device
|
| 1808 |
+
)
|
| 1809 |
+
audio_lengths = self._get_audio_lengths(input_features, feature_attention_mask)
|
| 1810 |
+
|
| 1811 |
+
pad_token_id = self.processor.tokenizer.pad_token_id
|
| 1812 |
+
if pad_token_id is None:
|
| 1813 |
+
pad_token_id = self.processor.tokenizer.eos_token_id or 151643
|
| 1814 |
+
|
| 1815 |
+
l_audio = audio_lengths[0]
|
| 1816 |
+
audio_toks = build_audio_token_seq(
|
| 1817 |
+
self.audio_bos_token_id, self.audio_token_id, self.audio_eos_token_id,
|
| 1818 |
+
l_audio, device,
|
| 1819 |
+
).unsqueeze(0)
|
| 1820 |
+
|
| 1821 |
+
system_turn, user_prefix, user_prompt, user_suffix, assistant_prefix = (
|
| 1822 |
+
self._build_chat_template_tensors(device, language=language_hint or "English")
|
| 1823 |
+
)
|
| 1824 |
+
prefix = torch.cat(
|
| 1825 |
+
[system_turn, user_prefix, audio_toks.squeeze(0), user_prompt, user_suffix, assistant_prefix],
|
| 1826 |
+
dim=0,
|
| 1827 |
+
).unsqueeze(0)
|
| 1828 |
+
|
| 1829 |
+
latent_active = (not use_baseline) and self.use_latent
|
| 1830 |
+
soft_prompt_active = (not use_baseline) and self.use_soft_prompt
|
| 1831 |
+
interleaved_active = latent_active and self.thought_mode == "interleaved"
|
| 1832 |
+
|
| 1833 |
+
# For prefix mode: prepend NT tokens to the prompt.
|
| 1834 |
+
# For interleaved mode: no upfront NT tokens; they are injected during generation.
|
| 1835 |
+
front_prompt_active = latent_active or soft_prompt_active
|
| 1836 |
+
if front_prompt_active and not interleaved_active:
|
| 1837 |
+
nt_toks = torch.tensor([self.nt_token_id] * self.n_latent, device=device).unsqueeze(0)
|
| 1838 |
+
start_ids = torch.cat([prefix, nt_toks], dim=1)
|
| 1839 |
+
else:
|
| 1840 |
+
start_ids = prefix
|
| 1841 |
+
|
| 1842 |
+
attention_mask = (start_ids != pad_token_id).long()
|
| 1843 |
+
inputs_embeds = self._get_native_audio_embeds(
|
| 1844 |
+
input_ids=start_ids,
|
| 1845 |
+
input_features=input_features,
|
| 1846 |
+
feature_attention_mask=feature_attention_mask,
|
| 1847 |
+
attention_mask=attention_mask,
|
| 1848 |
+
)
|
| 1849 |
+
|
| 1850 |
+
states = None
|
| 1851 |
+
stats = {}
|
| 1852 |
+
if interleaved_active:
|
| 1853 |
+
# Interleaved generate: NT tokens injected token-by-token.
|
| 1854 |
+
fallback_do_sample = bool(gen_kwargs.get("do_sample", False))
|
| 1855 |
+
fallback_temperature = float(gen_kwargs.get("temperature", 1.0))
|
| 1856 |
+
gen_ids = self._generate_interleaved(
|
| 1857 |
+
base_embeds=inputs_embeds,
|
| 1858 |
+
attention_mask=attention_mask,
|
| 1859 |
+
max_new_tokens=max_new_tokens,
|
| 1860 |
+
do_sample=fallback_do_sample,
|
| 1861 |
+
temperature=fallback_temperature,
|
| 1862 |
+
)
|
| 1863 |
+
if return_thoughts:
|
| 1864 |
+
return gen_ids, states
|
| 1865 |
+
return gen_ids
|
| 1866 |
+
|
| 1867 |
+
if latent_active:
|
| 1868 |
+
initial_state, prefix_past_kv, prefix_attn = self._compute_prefix_state(
|
| 1869 |
+
inputs_embeds,
|
| 1870 |
+
start_ids,
|
| 1871 |
+
strict=True,
|
| 1872 |
+
language=language_hint or "English",
|
| 1873 |
+
)
|
| 1874 |
+
deltas, states, stats = self._compute_continuous_thoughts(
|
| 1875 |
+
initial_state,
|
| 1876 |
+
prefix_past_key_values=prefix_past_kv,
|
| 1877 |
+
prefix_attention_mask=prefix_attn,
|
| 1878 |
+
halt_threshold=dynamic_halt_threshold,
|
| 1879 |
+
)
|
| 1880 |
+
|
| 1881 |
+
K = deltas.size(1)
|
| 1882 |
+
if K == 0:
|
| 1883 |
+
# N=0 skip: remove ALL NT tokens, completely reverting to base generation.
|
| 1884 |
+
remove_count = self.n_latent
|
| 1885 |
+
start_ids = start_ids[:, :-remove_count]
|
| 1886 |
+
inputs_embeds = inputs_embeds[:, :-remove_count, :]
|
| 1887 |
+
attention_mask = attention_mask[:, :-remove_count]
|
| 1888 |
+
elif K < self.n_latent:
|
| 1889 |
+
# N > 0 but early halted: Remove the unneeded NT tokens instead of zero-padding them.
|
| 1890 |
+
# This prevents the LLM from attending to "blank" un-updated base_nt embeddings.
|
| 1891 |
+
remove_count = self.n_latent - K
|
| 1892 |
+
start_ids = start_ids[:, :-remove_count]
|
| 1893 |
+
inputs_embeds = inputs_embeds[:, :-remove_count, :]
|
| 1894 |
+
attention_mask = attention_mask[:, :-remove_count]
|
| 1895 |
+
inputs_embeds = self._inject_latent_deltas(
|
| 1896 |
+
inputs_embeds, start_ids, deltas
|
| 1897 |
+
)
|
| 1898 |
+
else:
|
| 1899 |
+
inputs_embeds = self._inject_latent_deltas(
|
| 1900 |
+
inputs_embeds, start_ids, deltas
|
| 1901 |
+
)
|
| 1902 |
+
elif soft_prompt_active:
|
| 1903 |
+
nt_mask = (start_ids == self.nt_token_id)
|
| 1904 |
+
nt_counts = nt_mask.sum(dim=1)
|
| 1905 |
+
assert (nt_counts == self.n_latent).all(), f"Gen NT count mismatch! {nt_counts}"
|
| 1906 |
+
prompt = self.soft_prompt_embed.to(dtype=inputs_embeds.dtype)
|
| 1907 |
+
batch_idx, seq_idx = torch.nonzero(nt_mask, as_tuple=True)
|
| 1908 |
+
inputs_embeds = inputs_embeds.clone()
|
| 1909 |
+
inputs_embeds[batch_idx, seq_idx, :] = prompt.view(-1, self._hidden_size)
|
| 1910 |
+
|
| 1911 |
+
if front_prompt_active:
|
| 1912 |
+
is_nt = (start_ids == self.nt_token_id)
|
| 1913 |
+
nt_counts = is_nt.sum(dim=1)
|
| 1914 |
+
assert (nt_counts <= self.n_latent).all(), f"Gen NT count mismatch! {nt_counts}"
|
| 1915 |
+
|
| 1916 |
+
fallback_eos = gen_kwargs.get("eos_token_id", None)
|
| 1917 |
+
fallback_do_sample = bool(gen_kwargs.get("do_sample", False))
|
| 1918 |
+
fallback_temperature = float(gen_kwargs.get("temperature", 1.0))
|
| 1919 |
+
|
| 1920 |
+
try:
|
| 1921 |
+
if latent_active and K == 0:
|
| 1922 |
+
# N=0 skip -> we stripped the NT tokens from start_ids.
|
| 1923 |
+
# Do NOT pass inputs_embeds so Qwen generates organically from text + audio ids.
|
| 1924 |
+
# Must provide input_features so audio encoder handles the raw audio.
|
| 1925 |
+
gen_out = self.thinker.generate(
|
| 1926 |
+
input_ids=start_ids,
|
| 1927 |
+
attention_mask=attention_mask,
|
| 1928 |
+
input_features=input_features,
|
| 1929 |
+
feature_attention_mask=feature_attention_mask,
|
| 1930 |
+
pad_token_id=pad_token_id,
|
| 1931 |
+
max_new_tokens=max_new_tokens,
|
| 1932 |
+
**gen_kwargs,
|
| 1933 |
+
)
|
| 1934 |
+
else:
|
| 1935 |
+
gen_out = self.thinker.generate(
|
| 1936 |
+
inputs_embeds=inputs_embeds,
|
| 1937 |
+
input_ids=start_ids,
|
| 1938 |
+
attention_mask=attention_mask,
|
| 1939 |
+
input_features=None,
|
| 1940 |
+
pad_token_id=pad_token_id,
|
| 1941 |
+
max_new_tokens=max_new_tokens,
|
| 1942 |
+
**gen_kwargs,
|
| 1943 |
+
)
|
| 1944 |
+
start_len = start_ids.size(1)
|
| 1945 |
+
gen_ids = gen_out[:, start_len:]
|
| 1946 |
+
except ValueError as e:
|
| 1947 |
+
msg = str(e)
|
| 1948 |
+
if "inputs_embeds" not in msg:
|
| 1949 |
+
raise
|
| 1950 |
+
if not self._warned_generate_fallback:
|
| 1951 |
+
print(
|
| 1952 |
+
"[warn] thinker.generate does not support inputs_embeds; "
|
| 1953 |
+
"using fallback decoding loop."
|
| 1954 |
+
)
|
| 1955 |
+
self._warned_generate_fallback = True
|
| 1956 |
+
gen_ids = self._generate_with_forward_fallback(
|
| 1957 |
+
start_ids=start_ids,
|
| 1958 |
+
inputs_embeds=inputs_embeds,
|
| 1959 |
+
attention_mask=attention_mask,
|
| 1960 |
+
max_new_tokens=max_new_tokens,
|
| 1961 |
+
eos_token_id=fallback_eos,
|
| 1962 |
+
do_sample=fallback_do_sample,
|
| 1963 |
+
temperature=fallback_temperature,
|
| 1964 |
+
)
|
| 1965 |
+
if return_stats:
|
| 1966 |
+
if return_thoughts and latent_active:
|
| 1967 |
+
return gen_ids, states, stats
|
| 1968 |
+
return gen_ids, stats
|
| 1969 |
+
|
| 1970 |
+
if return_thoughts and latent_active:
|
| 1971 |
+
return gen_ids, states
|
| 1972 |
+
return gen_ids
|
model_metadata.json
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"project": "latentASR",
|
| 3 |
+
"model_repo": "voidful/latentASR",
|
| 4 |
+
"base_model": "Qwen/Qwen3-ASR-0.6B",
|
| 5 |
+
"checkpoint_file": "checkpoints/latentASR_adapter.pth",
|
| 6 |
+
"source_checkpoint": "activation_500_epoch10.pth",
|
| 7 |
+
"checkpoint_sha256": "f0ce39fa5e6952fced6992508f3e2b32ea8467442b545678781a0f04e64f2430",
|
| 8 |
+
"train_mode": "latent",
|
| 9 |
+
"n_latent": 4,
|
| 10 |
+
"train_max_samples": 500,
|
| 11 |
+
"adapter_parameters": 5251077,
|
| 12 |
+
"value_forced_neg_prob": 0.3,
|
| 13 |
+
"latent_use_bounded_delta": true,
|
| 14 |
+
"latent_use_injection_gate": true,
|
| 15 |
+
"latent_use_embedding_anchor": true,
|
| 16 |
+
"deployed_halt_threshold": 0.0,
|
| 17 |
+
"paper": "Listen, Think, Transcribe: Continuous Latent Test-Time Scaling for ASR"
|
| 18 |
+
}
|
| 19 |
+
|
peft_utils.py
ADDED
|
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""PEFT adapter utilities for lr_whisper."""
|
| 2 |
+
|
| 3 |
+
import os
|
| 4 |
+
import re
|
| 5 |
+
from typing import Any, Dict
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
import torch.nn as nn
|
| 9 |
+
|
| 10 |
+
from config import TrainingConfig
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def attach_peft_adapter(
|
| 14 |
+
asr_model: nn.Module,
|
| 15 |
+
cfg: TrainingConfig,
|
| 16 |
+
) -> Dict[str, Any]:
|
| 17 |
+
"""Attach a LoRA PEFT adapter to *asr_model.thinker*.
|
| 18 |
+
|
| 19 |
+
Args:
|
| 20 |
+
asr_model: The base ASR model (has a ``.thinker`` attribute).
|
| 21 |
+
cfg: Training configuration providing LoRA hyperparameters.
|
| 22 |
+
|
| 23 |
+
Returns:
|
| 24 |
+
A metadata dict describing the adapter configuration.
|
| 25 |
+
|
| 26 |
+
Raises:
|
| 27 |
+
RuntimeError: If the ``peft`` package is not installed.
|
| 28 |
+
ValueError: If ``cfg.train_mode`` is not ``"lora_r16"``.
|
| 29 |
+
"""
|
| 30 |
+
try:
|
| 31 |
+
from peft import LoraConfig, TaskType, get_peft_model
|
| 32 |
+
except Exception as e:
|
| 33 |
+
raise RuntimeError(
|
| 34 |
+
"PEFT mode requested but `peft` package is unavailable."
|
| 35 |
+
) from e
|
| 36 |
+
|
| 37 |
+
if cfg.train_mode == "lora_r16":
|
| 38 |
+
rank = int(cfg.lora_rank)
|
| 39 |
+
alpha = int(cfg.lora_alpha)
|
| 40 |
+
target_modules = list(cfg.lora_target_modules)
|
| 41 |
+
|
| 42 |
+
print(
|
| 43 |
+
f"[lora] rank={rank} alpha={alpha} dropout={cfg.lora_dropout} "
|
| 44 |
+
f"target_modules={target_modules}"
|
| 45 |
+
)
|
| 46 |
+
|
| 47 |
+
if any("audio" in m for m in target_modules):
|
| 48 |
+
# User explicitly listed audio modules — respect that.
|
| 49 |
+
pass
|
| 50 |
+
else:
|
| 51 |
+
# Convert plain module names into a regex that excludes audio_tower.
|
| 52 |
+
escaped = [re.escape(m) for m in target_modules]
|
| 53 |
+
target_modules = r"^(?!.*audio_tower).*\b(" + "|".join(escaped) + r")$"
|
| 54 |
+
print(f"[lora] Using regex target_modules to exclude audio_tower: {target_modules}")
|
| 55 |
+
|
| 56 |
+
peft_cfg = LoraConfig(
|
| 57 |
+
task_type=TaskType.CAUSAL_LM,
|
| 58 |
+
r=rank,
|
| 59 |
+
lora_alpha=alpha,
|
| 60 |
+
lora_dropout=cfg.lora_dropout,
|
| 61 |
+
target_modules=target_modules,
|
| 62 |
+
bias="none",
|
| 63 |
+
)
|
| 64 |
+
asr_model.thinker = get_peft_model(asr_model.thinker, peft_cfg)
|
| 65 |
+
asr_model.thinker.print_trainable_parameters()
|
| 66 |
+
return {
|
| 67 |
+
"peft_mode": "lora_r16",
|
| 68 |
+
"rank": rank,
|
| 69 |
+
"alpha": alpha,
|
| 70 |
+
"dropout": cfg.lora_dropout,
|
| 71 |
+
"target_modules": target_modules,
|
| 72 |
+
}
|
| 73 |
+
|
| 74 |
+
raise ValueError(
|
| 75 |
+
f"attach_peft_adapter called with unsupported train_mode={cfg.train_mode!r}. "
|
| 76 |
+
"Only lora_r16 uses PEFT in this trainer."
|
| 77 |
+
)
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def save_peft_adapter_checkpoint(
|
| 81 |
+
model: nn.Module,
|
| 82 |
+
ckpt_dir: str,
|
| 83 |
+
metadata: Dict[str, Any],
|
| 84 |
+
) -> None:
|
| 85 |
+
"""Save a PEFT adapter checkpoint to *ckpt_dir*.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
model: The ``LatentQwenASR`` wrapper (has ``.base_model.thinker``).
|
| 89 |
+
ckpt_dir: Directory path for the checkpoint.
|
| 90 |
+
metadata: Arbitrary metadata dict saved alongside the adapter.
|
| 91 |
+
"""
|
| 92 |
+
os.makedirs(ckpt_dir, exist_ok=True)
|
| 93 |
+
model.base_model.thinker.save_pretrained(ckpt_dir, safe_serialization=False)
|
| 94 |
+
torch.save(metadata, os.path.join(ckpt_dir, "training_state.pth"))
|
pyproject.toml
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[project]
|
| 2 |
+
name = "latentASR"
|
| 3 |
+
version = "0.1.0"
|
| 4 |
+
description = "Continuous latent test-time scaling for frozen Qwen3-ASR backbones."
|
| 5 |
+
readme = "README.md"
|
| 6 |
+
requires-python = ">=3.10"
|
| 7 |
+
license = { text = "Apache-2.0" }
|
| 8 |
+
authors = [
|
| 9 |
+
{ name = "LatentASR contributors" }
|
| 10 |
+
]
|
| 11 |
+
dependencies = [
|
| 12 |
+
"torch>=2.6",
|
| 13 |
+
"torchaudio>=2.6",
|
| 14 |
+
"transformers>=4.57",
|
| 15 |
+
"datasets>=3.6",
|
| 16 |
+
"huggingface-hub>=0.34",
|
| 17 |
+
"qwen-asr>=0.0.6",
|
| 18 |
+
"qwen-omni-utils>=0.0.8",
|
| 19 |
+
"jiwer>=4.0",
|
| 20 |
+
"peft>=0.18",
|
| 21 |
+
"numpy>=1.24",
|
| 22 |
+
"tqdm>=4.66"
|
| 23 |
+
]
|
| 24 |
+
|
| 25 |
+
[tool.setuptools]
|
| 26 |
+
py-modules = [
|
| 27 |
+
"config",
|
| 28 |
+
"data",
|
| 29 |
+
"eval",
|
| 30 |
+
"losses",
|
| 31 |
+
"model",
|
| 32 |
+
"peft_utils",
|
| 33 |
+
"train",
|
| 34 |
+
"utils"
|
| 35 |
+
]
|
| 36 |
+
|
requirements.txt
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.6
|
| 2 |
+
torchaudio>=2.6
|
| 3 |
+
transformers>=4.57
|
| 4 |
+
datasets>=3.6
|
| 5 |
+
huggingface-hub>=0.34
|
| 6 |
+
qwen-asr>=0.0.6
|
| 7 |
+
qwen-omni-utils>=0.0.8
|
| 8 |
+
jiwer>=4.0
|
| 9 |
+
peft>=0.18
|
| 10 |
+
numpy>=1.24
|
| 11 |
+
tqdm>=4.66
|
| 12 |
+
|
results/paper500_ascend_20260518/ascend_clean.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "SpeechTest/ASCEND",
|
| 4 |
+
"dataset_name_input": "SpeechTest/ASCEND",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"main"
|
| 8 |
+
],
|
| 9 |
+
"streaming": false,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "main",
|
| 20 |
+
"samples_total": 1315,
|
| 21 |
+
"samples_used": 1315,
|
| 22 |
+
"text_column": "transcription",
|
| 23 |
+
"base_model_wer": 1.387989931679252,
|
| 24 |
+
"base_model_cer": 0.5780775716694773,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 1.29809421071557,
|
| 32 |
+
"latent_reasoning_cer": 0.4854733799084558,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.08989572096368215,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.09260419176102147,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 1.387989931679252,
|
| 51 |
+
"base_model_weighted_cer": 0.5780775716694773,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 1.29809421071557,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.48547337990845574,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.08989572096368215,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.09260419176102153,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_english_clean_noise_20260518/fleurs_en_us_clean.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "google/fleurs",
|
| 4 |
+
"dataset_name_input": "google/fleurs",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en_us"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en_us",
|
| 20 |
+
"samples_total": 647,
|
| 21 |
+
"samples_used": 647,
|
| 22 |
+
"text_column": "transcription",
|
| 23 |
+
"base_model_wer": 0.04899930986887509,
|
| 24 |
+
"base_model_cer": 0.023260254681290355,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.048585231193926846,
|
| 32 |
+
"latent_reasoning_cer": 0.022687400794834767,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.0004140786749482414,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.0005728538864555877,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.04899930986887509,
|
| 51 |
+
"base_model_weighted_cer": 0.023260254681290355,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.048585231193926846,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.022687400794834767,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.0004140786749482414,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.0005728538864555877,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_english_clean_noise_20260518/fleurs_en_us_snr0db.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "google/fleurs",
|
| 4 |
+
"dataset_name_input": "google/fleurs",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en_us"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en_us",
|
| 20 |
+
"samples_total": 647,
|
| 21 |
+
"samples_used": 647,
|
| 22 |
+
"text_column": "transcription",
|
| 23 |
+
"base_model_wer": 0.2981366459627329,
|
| 24 |
+
"base_model_cer": 0.19216860999391341,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.2942028985507246,
|
| 32 |
+
"latent_reasoning_cer": 0.18863601102743732,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.003933747412008304,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.0035325989664760915,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.2981366459627329,
|
| 51 |
+
"base_model_weighted_cer": 0.19216860999391341,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.2942028985507246,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.18863601102743732,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.003933747412008304,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.0035325989664760915,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_english_clean_noise_20260518/minds14_en_us_clean.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "PolyAI/minds14",
|
| 4 |
+
"dataset_name_input": "PolyAI/minds14",
|
| 5 |
+
"split": "train",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en-US"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en-US",
|
| 20 |
+
"samples_total": 563,
|
| 21 |
+
"samples_used": 563,
|
| 22 |
+
"text_column": "transcription",
|
| 23 |
+
"base_model_wer": 0.33701067615658364,
|
| 24 |
+
"base_model_cer": 0.2950933141919606,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.33677342823250295,
|
| 32 |
+
"latent_reasoning_cer": 0.29493949958982774,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.00023724792408069684,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.00015381460213287346,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.33701067615658364,
|
| 51 |
+
"base_model_weighted_cer": 0.2950933141919606,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.33677342823250295,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.29493949958982774,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.00023724792408069684,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.00015381460213287346,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_english_clean_noise_20260518/minds14_en_us_snr0db.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "PolyAI/minds14",
|
| 4 |
+
"dataset_name_input": "PolyAI/minds14",
|
| 5 |
+
"split": "train",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en-US"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en-US",
|
| 20 |
+
"samples_total": 563,
|
| 21 |
+
"samples_used": 563,
|
| 22 |
+
"text_column": "transcription",
|
| 23 |
+
"base_model_wer": 0.402491103202847,
|
| 24 |
+
"base_model_cer": 0.3330855209187859,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.4048635824436536,
|
| 32 |
+
"latent_reasoning_cer": 0.3333675143560295,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": -0.0023724792408066353,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": -0.0002819934372436106,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.402491103202847,
|
| 51 |
+
"base_model_weighted_cer": 0.3330855209187859,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.4048635824436536,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.3333675143560295,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": -0.0023724792408066353,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": -0.0002819934372436106,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_english_clean_noise_20260518/showcase_report.md
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# LR HuggingFace ASR Showcase Report
|
| 2 |
+
|
| 3 |
+
- Generated UTC: 2026-05-18T20:13:39+00:00
|
| 4 |
+
- Output directory: `/user_data/lr_whisper/eval_runs/paper500_english_clean_noise_20260518`
|
| 5 |
+
- Delta WER is `base_model_wer - latent_reasoning_wer`; positive means LR is better.
|
| 6 |
+
|
| 7 |
+
## Best LR Wins
|
| 8 |
+
|
| 9 |
+
| Rank | Dataset | Configs | Condition | N | Base WER | LR WER | Delta WER | Relative | Delta CER | JSON |
|
| 10 |
+
|---:|---|---|---|---:|---:|---:|---:|---:|---:|---|
|
| 11 |
+
| 1 | google/fleurs | en_us | snr0db | 647 | 0.298137 | 0.294203 | 0.003934 | +1.32% | 0.003533 | `fleurs_en_us_snr0db.json` |
|
| 12 |
+
| 2 | facebook/voxpopuli | en | snr0db | 1842 | 0.197111 | 0.195108 | 0.002003 | +1.02% | -0.000284 | `voxpopuli_en_snr0db.json` |
|
| 13 |
+
| 3 | facebook/voxpopuli | en | clean | 1842 | 0.090375 | 0.089745 | 0.000630 | +0.70% | 0.000427 | `voxpopuli_en_clean.json` |
|
| 14 |
+
| 4 | google/fleurs | en_us | clean | 647 | 0.048999 | 0.048585 | 0.000414 | +0.85% | 0.000573 | `fleurs_en_us_clean.json` |
|
| 15 |
+
| 5 | PolyAI/minds14 | en-US | clean | 563 | 0.337011 | 0.336773 | 0.000237 | +0.07% | 0.000154 | `minds14_en_us_clean.json` |
|
| 16 |
+
|
| 17 |
+
## All Cases
|
| 18 |
+
|
| 19 |
+
| Dataset | Configs | Condition | N | Base WER | LR WER | Delta WER | Relative | Base CER | LR CER | Delta CER | JSON |
|
| 20 |
+
|---|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---|
|
| 21 |
+
| PolyAI/minds14 | en-US | clean | 563 | 0.337011 | 0.336773 | 0.000237 | +0.07% | 0.295093 | 0.294939 | 0.000154 | `minds14_en_us_clean.json` |
|
| 22 |
+
| facebook/voxpopuli | en | clean | 1842 | 0.090375 | 0.089745 | 0.000630 | +0.70% | 0.059004 | 0.058576 | 0.000427 | `voxpopuli_en_clean.json` |
|
| 23 |
+
| google/fleurs | en_us | clean | 647 | 0.048999 | 0.048585 | 0.000414 | +0.85% | 0.023260 | 0.022687 | 0.000573 | `fleurs_en_us_clean.json` |
|
| 24 |
+
| PolyAI/minds14 | en-US | snr0db | 563 | 0.402491 | 0.404864 | -0.002372 | -0.59% | 0.333086 | 0.333368 | -0.000282 | `minds14_en_us_snr0db.json` |
|
| 25 |
+
| facebook/voxpopuli | en | snr0db | 1842 | 0.197111 | 0.195108 | 0.002003 | +1.02% | 0.120512 | 0.120796 | -0.000284 | `voxpopuli_en_snr0db.json` |
|
| 26 |
+
| google/fleurs | en_us | snr0db | 647 | 0.298137 | 0.294203 | 0.003934 | +1.32% | 0.192169 | 0.188636 | 0.003533 | `fleurs_en_us_snr0db.json` |
|
| 27 |
+
|
| 28 |
+
## Regressions To Check
|
| 29 |
+
|
| 30 |
+
| Dataset | Configs | Condition | N | Delta WER | Relative | JSON |
|
| 31 |
+
|---|---|---|---:|---:|---:|---|
|
| 32 |
+
| PolyAI/minds14 | en-US | snr0db | 563 | -0.002372 | -0.59% | `minds14_en_us_snr0db.json` |
|
| 33 |
+
|
results/paper500_english_clean_noise_20260518/voxpopuli_en_clean.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "facebook/voxpopuli",
|
| 4 |
+
"dataset_name_input": "facebook/voxpopuli",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en",
|
| 20 |
+
"samples_total": 1842,
|
| 21 |
+
"samples_used": 1842,
|
| 22 |
+
"text_column": "normalized_text",
|
| 23 |
+
"base_model_wer": 0.09037513783558745,
|
| 24 |
+
"base_model_cer": 0.05900354572745667,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.08974503229290906,
|
| 32 |
+
"latent_reasoning_cer": 0.05857634964833996,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.0006301055426783941,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.0004271960791167109,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.09037513783558745,
|
| 51 |
+
"base_model_weighted_cer": 0.05900354572745667,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.08974503229290906,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.05857634964833996,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.0006301055426783941,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.0004271960791167109,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_english_clean_noise_20260518/voxpopuli_en_snr0db.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "facebook/voxpopuli",
|
| 4 |
+
"dataset_name_input": "facebook/voxpopuli",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en",
|
| 20 |
+
"samples_total": 1842,
|
| 21 |
+
"samples_used": 1842,
|
| 22 |
+
"text_column": "normalized_text",
|
| 23 |
+
"base_model_wer": 0.19711051601143192,
|
| 24 |
+
"base_model_cer": 0.12051201391882498,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.19510768053648986,
|
| 32 |
+
"latent_reasoning_cer": 0.1207955167713297,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.0020028354749420663,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": -0.0002835028525047245,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.1971105160114319,
|
| 51 |
+
"base_model_weighted_cer": 0.12051201391882498,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.19510768053648986,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.1207955167713297,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.0020028354749420385,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": -0.0002835028525047245,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_main_eval_20260518/main_fleurs_en_us.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "google/fleurs",
|
| 4 |
+
"dataset_name_input": "google/fleurs",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en_us"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_main_baselines_20260518/checkpoints/baseline/baseline_best.pth",
|
| 13 |
+
"prompt_tuning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_main_baselines_20260518/checkpoints/prompt_tuning/prompt_tuning_best.pth",
|
| 14 |
+
"lora_r16_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_main_baselines_20260518/checkpoints/lora_r16/lora_r16_best",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en_us",
|
| 20 |
+
"samples_total": 647,
|
| 21 |
+
"samples_used": 647,
|
| 22 |
+
"text_column": "transcription",
|
| 23 |
+
"base_model_wer": 0.04899930986887509,
|
| 24 |
+
"base_model_cer": 0.023260254681290355,
|
| 25 |
+
"baseline_ft_wer": 0.05555555555555555,
|
| 26 |
+
"baseline_ft_cer": 0.024811733957107566,
|
| 27 |
+
"prompt_tuning_wer": 0.8480331262939959,
|
| 28 |
+
"prompt_tuning_cer": 0.8246828418326551,
|
| 29 |
+
"lora_r16_wer": 0.06466528640441684,
|
| 30 |
+
"lora_r16_cer": 0.02944230287262355,
|
| 31 |
+
"latent_reasoning_wer": 0.048585231193926846,
|
| 32 |
+
"latent_reasoning_cer": 0.022687400794834767,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": -0.006556245686680465,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": -0.0015514792758172108,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": -0.7990338164251207,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": -0.8014225871513647,
|
| 37 |
+
"lora_r16_vs_base_model_wer": -0.015665976535541748,
|
| 38 |
+
"lora_r16_vs_base_model_cer": -0.006182048191333195,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.0004140786749482414,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.0005728538864555877,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": 0.006970324361628706,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": 0.0021243331622727984,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": 0.799447895100069,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": 0.8019954410378203,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": 0.01608005521048999,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": 0.0067549020777887825
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.04899930986887509,
|
| 51 |
+
"base_model_weighted_cer": 0.023260254681290355,
|
| 52 |
+
"baseline_ft_weighted_wer": 0.05555555555555555,
|
| 53 |
+
"baseline_ft_weighted_cer": 0.02481173395710757,
|
| 54 |
+
"prompt_tuning_weighted_wer": 0.848033126293996,
|
| 55 |
+
"prompt_tuning_weighted_cer": 0.8246828418326551,
|
| 56 |
+
"lora_r16_weighted_wer": 0.06466528640441684,
|
| 57 |
+
"lora_r16_weighted_cer": 0.029442302872623553,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.048585231193926846,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.022687400794834767,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": -0.006556245686680465,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": -0.0015514792758172143,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.0004140786749482414,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.0005728538864555877,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": 0.006970324361628706,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": 0.002124333162272802,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": -0.7990338164251208,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": -0.8014225871513647,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": -0.015665976535541748,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": -0.006182048191333198,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": 0.7994478951000691,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": 0.8019954410378203,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": 0.01608005521048999,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": 0.006754902077788786
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_main_eval_20260518/main_voxpopuli_en.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "facebook/voxpopuli",
|
| 4 |
+
"dataset_name_input": "facebook/voxpopuli",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_main_baselines_20260518/checkpoints/baseline/baseline_best.pth",
|
| 13 |
+
"prompt_tuning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_main_baselines_20260518/checkpoints/prompt_tuning/prompt_tuning_best.pth",
|
| 14 |
+
"lora_r16_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_main_baselines_20260518/checkpoints/lora_r16/lora_r16_best",
|
| 15 |
+
"latent_reasoning_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "/user_data/lr_whisper/eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en",
|
| 20 |
+
"samples_total": 1842,
|
| 21 |
+
"samples_used": 1842,
|
| 22 |
+
"text_column": "normalized_text",
|
| 23 |
+
"base_model_wer": 0.09037513783558745,
|
| 24 |
+
"base_model_cer": 0.05900354572745667,
|
| 25 |
+
"baseline_ft_wer": 0.09485338794248037,
|
| 26 |
+
"baseline_ft_cer": 0.061764785838838335,
|
| 27 |
+
"prompt_tuning_wer": 0.8850732497693363,
|
| 28 |
+
"prompt_tuning_cer": 0.8790103031927081,
|
| 29 |
+
"lora_r16_wer": 0.10212210545266331,
|
| 30 |
+
"lora_r16_cer": 0.06535711650413797,
|
| 31 |
+
"latent_reasoning_wer": 0.08974503229290906,
|
| 32 |
+
"latent_reasoning_cer": 0.05857634964833996,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": -0.004478250106892914,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": -0.002761240111381666,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": -0.7946981119337488,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": -0.8200067574652514,
|
| 37 |
+
"lora_r16_vs_base_model_wer": -0.01174696761707586,
|
| 38 |
+
"lora_r16_vs_base_model_cer": -0.006353570776681305,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.0006301055426783941,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.0004271960791167109,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": 0.005108355649571308,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": 0.003188436190498377,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": 0.7953282174764272,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": 0.8204339535443682,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": 0.012377073159754254,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": 0.006780766855798016
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.09037513783558745,
|
| 51 |
+
"base_model_weighted_cer": 0.05900354572745667,
|
| 52 |
+
"baseline_ft_weighted_wer": 0.09485338794248037,
|
| 53 |
+
"baseline_ft_weighted_cer": 0.061764785838838335,
|
| 54 |
+
"prompt_tuning_weighted_wer": 0.8850732497693363,
|
| 55 |
+
"prompt_tuning_weighted_cer": 0.8790103031927081,
|
| 56 |
+
"lora_r16_weighted_wer": 0.10212210545266331,
|
| 57 |
+
"lora_r16_weighted_cer": 0.06535711650413797,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.08974503229290906,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.05857634964833996,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": -0.004478250106892914,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": -0.002761240111381666,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.0006301055426783941,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.0004271960791167109,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": 0.005108355649571308,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": 0.003188436190498377,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": -0.7946981119337488,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": -0.8200067574652514,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": -0.01174696761707586,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": -0.006353570776681305,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": 0.7953282174764272,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": 0.8204339535443682,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": 0.012377073159754254,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": 0.006780766855798016
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_main_eval_20260518/paper500_complete_results.md
ADDED
|
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Paper 500-Activation Complete Results
|
| 2 |
+
|
| 3 |
+
All results below use the 500-utterance activation setting unless explicitly stated otherwise. Metrics are shown in percent. Unless noted otherwise, `Delta` is `model - baseline`, so negative means an error reduction.
|
| 4 |
+
|
| 5 |
+
Canonical paper choice: use `eval_runs/paper_tbd_retrain_20260518/checkpoints/activation_500/activation_500_epoch10.pth` for the main LatentASR result. This is the better of the two available 500-utterance/N=4/theta=0 runs by mean FLEURS/VoxPopuli WER delta, and it is the checkpoint represented by the 500 row in the activation-size sweep. The older `eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth` should be treated as an auxiliary ablation/diagnostic run unless its dependent tables are intentionally retained with that label.
|
| 6 |
+
|
| 7 |
+
## Source Artifacts
|
| 8 |
+
|
| 9 |
+
- Main baseline/adaptation JSONs: `eval_runs/paper500_main_eval_20260518/`
|
| 10 |
+
- Canonical LatentASR JSONs: `eval_runs/paper_tbd_eval_20260518/activation_500_fleurs_theta_zero.json`, `eval_runs/paper_tbd_eval_20260518/activation_500_voxpopuli_theta_zero.json`
|
| 11 |
+
- Retrained 500 checkpoints: `eval_runs/paper500_retrain_20260518/checkpoints/`
|
| 12 |
+
- Main baseline checkpoints: `eval_runs/paper500_main_baselines_20260518/checkpoints/`
|
| 13 |
+
- Threshold sweep: `eval_runs/paper500_threshold_20260518/threshold_sweep_report.md`
|
| 14 |
+
- Component/N/p_neg ablations: `eval_runs/paper500_eval_20260518/paper_tbd_results.md`
|
| 15 |
+
- Activation scaling 100-800: `eval_runs/paper_tbd_eval_20260518/activation_scaling_100_800_full_results.md`
|
| 16 |
+
- Multilingual streaming: `eval_runs/paper500_multilingual_streaming_20260518/showcase_report.md`
|
| 17 |
+
- English clean/noise suite: `eval_runs/paper500_english_clean_noise_20260518/showcase_report.md`
|
| 18 |
+
- ASCEND: `eval_runs/paper500_ascend_20260518/ascend_clean.json`
|
| 19 |
+
- Latency: `eval_runs/paper500_latency_20260518/latency.json`
|
| 20 |
+
|
| 21 |
+
## Main Results
|
| 22 |
+
|
| 23 |
+
Main baselines were retrained/evaluated with the 500-utterance activation setting. The canonical LatentASR row uses the best 500-utterance activation checkpoint from the activation-size sweep (`activation_500_epoch10.pth`). Full fine-tuning initially OOMed with micro-batch 16, so it was rerun with micro-batch 4 and gradient accumulation 4, preserving effective batch size 16.
|
| 24 |
+
|
| 25 |
+
| Model | FLEURS WER | FLEURS CER | Rel. Delta WER | VoxPopuli WER | VoxPopuli CER | Rel. Delta WER |
|
| 26 |
+
|---|---:|---:|---:|---:|---:|---:|
|
| 27 |
+
| Baseline | 4.90 | 2.33 | -- | 9.04 | 5.90 | -- |
|
| 28 |
+
| Full fine-tuning | 5.56 | 2.48 | +13.38 | 9.49 | 6.18 | +4.96 |
|
| 29 |
+
| Prompt tuning | 84.80 | 82.47 | +1630.70 | 88.51 | 87.90 | +879.33 |
|
| 30 |
+
| LoRA r16 | 6.47 | 2.94 | +31.97 | 10.21 | 6.54 | +13.00 |
|
| 31 |
+
| LatentASR N=4, theta=0 | 4.78 | 2.24 | -2.54 | 8.99 | 5.87 | -0.47 |
|
| 32 |
+
|
| 33 |
+
Prompt tuning collapses badly in the 500 setting on both FLEURS and VoxPopuli. This should be reported plainly if the table is used.
|
| 34 |
+
|
| 35 |
+
## Activation Scaling 100-800
|
| 36 |
+
|
| 37 |
+
The activation-scaling sweep was run at 100 to 800 utterances, every 100. The 500 point gives the best average absolute WER delta across FLEURS and VoxPopuli.
|
| 38 |
+
|
| 39 |
+
| Activation size | FLEURS WER | FLEURS Delta pp | VoxPopuli WER | VoxPopuli Delta pp | Mean Delta pp |
|
| 40 |
+
|---:|---:|---:|---:|---:|---:|
|
| 41 |
+
| 100 | 4.87 | -0.03 | 8.95 | -0.08 | -0.055 |
|
| 42 |
+
| 200 | 4.91 | +0.01 | 9.01 | -0.03 | -0.010 |
|
| 43 |
+
| 300 | 4.90 | +0.00 | 9.01 | -0.03 | -0.015 |
|
| 44 |
+
| 400 | 4.91 | +0.01 | 9.04 | -0.00 | +0.005 |
|
| 45 |
+
| 500 | 4.78 | -0.12 | 8.99 | -0.04 | -0.084 |
|
| 46 |
+
| 600 | 4.93 | +0.03 | 9.02 | -0.02 | +0.005 |
|
| 47 |
+
| 700 | 4.82 | -0.08 | 9.04 | +0.00 | -0.040 |
|
| 48 |
+
| 800 | 4.88 | -0.02 | 9.07 | +0.04 | +0.010 |
|
| 49 |
+
|
| 50 |
+
## Threshold Sweep
|
| 51 |
+
|
| 52 |
+
Converted from the sweep report to paper convention: Delta WER is `LatentASR - Baseline`, so negative is better.
|
| 53 |
+
|
| 54 |
+
Important consistency note: the threshold sweep below was run with `paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth`, not the canonical `activation_500_epoch10.pth`. If the main paper uses the better canonical checkpoint, this section must either be rerun with `activation_500_epoch10.pth` or explicitly labeled as an auxiliary same-protocol diagnostic run.
|
| 55 |
+
|
| 56 |
+
| Dataset | theta | Avg. steps | Skip | WER | Delta WER pp |
|
| 57 |
+
|---|---:|---:|---:|---:|---:|
|
| 58 |
+
| FLEURS | -2.0 | 4.00 | 0.0 | 4.84 | -0.062 |
|
| 59 |
+
| FLEURS | -0.2 | 1.37 | 43.7 | 4.86 | -0.041 |
|
| 60 |
+
| FLEURS | 0.0 | 1.24 | 47.0 | 4.86 | -0.041 |
|
| 61 |
+
| FLEURS | +0.2 | 0.00 | 100.0 | 4.90 | +0.000 |
|
| 62 |
+
| FLEURS | +0.5 | 0.00 | 100.0 | 4.90 | +0.000 |
|
| 63 |
+
| VoxPopuli | -2.0 | 4.00 | 0.0 | 9.02 | -0.014 |
|
| 64 |
+
| VoxPopuli | -0.2 | 0.84 | 54.9 | 8.97 | -0.070 |
|
| 65 |
+
| VoxPopuli | 0.0 | 0.71 | 60.2 | 8.97 | -0.063 |
|
| 66 |
+
| VoxPopuli | +0.2 | 0.00 | 100.0 | 9.04 | +0.000 |
|
| 67 |
+
| VoxPopuli | +0.5 | 0.00 | 100.0 | 9.04 | +0.000 |
|
| 68 |
+
|
| 69 |
+
In this 500 setting, full compute does not regress VoxPopuli; it is slightly positive. The older claim that no-halting increases VoxPopuli WER should be revised for the 500 paper version. Dynamic halting is still useful because it recovers most of the gain at much lower average steps.
|
| 70 |
+
|
| 71 |
+
## Compute Allocation at theta=0
|
| 72 |
+
|
| 73 |
+
The clean FLEURS/VoxPopuli theta=0 allocation for the canonical `activation_500_epoch10.pth` is available from the activation sweep: FLEURS average steps 1.20 with 47.0% skip; VoxPopuli average steps 0.70 with 63.0% skip. The table below additionally includes ASCEND from the older `n4_epoch10.pth` run and therefore should be rerun if strict checkpoint consistency is required.
|
| 74 |
+
|
| 75 |
+
| Dataset | N=0 | N=1 | N=2 | N=3 | N=4 |
|
| 76 |
+
|---|---:|---:|---:|---:|---:|
|
| 77 |
+
| FLEURS clean | 47.0 | 23.2 | 7.0 | 4.6 | 18.2 |
|
| 78 |
+
| VoxPopuli clean | 60.2 | 24.0 | 6.7 | 2.8 | 6.2 |
|
| 79 |
+
| ASCEND clean | 22.3 | 30.3 | 8.2 | 5.1 | 34.1 |
|
| 80 |
+
|
| 81 |
+
## Component Ablation
|
| 82 |
+
|
| 83 |
+
FLEURS, theta=0. Baseline WER for this ablation run is 4.90.
|
| 84 |
+
|
| 85 |
+
| Variant | WER | Delta WER pp |
|
| 86 |
+
|---|---:|---:|
|
| 87 |
+
| Full LatentASR N=4 | 4.86 | -0.04 |
|
| 88 |
+
| Remove bounded delta | 51.75 | +46.85 |
|
| 89 |
+
| Remove sigmoid gate | 7.89 | +2.99 |
|
| 90 |
+
| Remove fixed-embedding anchor | 16.22 | +11.33 |
|
| 91 |
+
|
| 92 |
+
This strongly supports keeping the stabilization-mechanism ablation in the paper.
|
| 93 |
+
|
| 94 |
+
## Latent Budget N Sweep
|
| 95 |
+
|
| 96 |
+
theta=0.
|
| 97 |
+
|
| 98 |
+
| N | FLEURS WER | FLEURS Delta pp | VoxPopuli WER | VoxPopuli Delta pp |
|
| 99 |
+
|---:|---:|---:|---:|---:|
|
| 100 |
+
| 1 | 4.82 | -0.08 | 8.98 | -0.06 |
|
| 101 |
+
| 2 | 4.89 | -0.01 | 9.03 | -0.00 |
|
| 102 |
+
| 4 | 4.86 | -0.04 | 8.97 | -0.06 |
|
| 103 |
+
| 8 | 4.90 | +0.00 | 9.03 | -0.00 |
|
| 104 |
+
|
| 105 |
+
N=1 is surprisingly strong in this 500 run. N=4 remains a reasonable default because it ties/bests VoxPopuli and supports deeper reasoning on the subset that needs it, but the paper should not claim N=4 is strictly best on every dataset.
|
| 106 |
+
|
| 107 |
+
## Forced-Negative Sampling
|
| 108 |
+
|
| 109 |
+
FLEURS, theta=0 unless otherwise specified.
|
| 110 |
+
|
| 111 |
+
| Setting | WER | Delta WER pp | Skip at theta=+0.2 | Skip range over sweep |
|
| 112 |
+
|---|---:|---:|---:|---:|
|
| 113 |
+
| p_neg=0.3 | 4.86 | -0.04 | 100.0 | [0, 100] |
|
| 114 |
+
| p_neg=0.0 | 4.80 | -0.10 | 100.0 | [0, 100] |
|
| 115 |
+
|
| 116 |
+
Important paper caveat: in the 500 run, removing forced-negative sampling does not collapse threshold reachability. The previous placeholder prose saying the saturated regime disappears is false for this setting and should be removed or rewritten. A safer statement is that forced-negative sampling changes value calibration/aggressiveness, while this 500 run does not show it is required for saturation.
|
| 117 |
+
|
| 118 |
+
## Multilingual Streaming
|
| 119 |
+
|
| 120 |
+
| Subset | #utts | Baseline WER | LatentASR WER | Delta WER pp | Rel. Delta WER | Baseline CER | LatentASR CER | Delta CER pp | Rel. Delta CER |
|
| 121 |
+
|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|
|
| 122 |
+
| Primary multilingual: FLEURS30 + MLS public7 | 36,982 | 24.982 | 24.941 | -0.041 | -0.163 | 9.428 | 9.411 | -0.017 | -0.178 |
|
| 123 |
+
| FLEURS 30 overall | 23,049 | 32.649 | 32.573 | -0.076 | -0.232 | 12.810 | 12.764 | -0.046 | -0.358 |
|
| 124 |
+
| FLEURS core12 | 8,876 | 31.708 | 31.627 | -0.081 | -0.255 | 12.649 | 12.564 | -0.085 | -0.671 |
|
| 125 |
+
| FLEURS +8 | 5,597 | 17.479 | 17.447 | -0.032 | -0.185 | 6.922 | 6.904 | -0.018 | -0.256 |
|
| 126 |
+
| FLEURS +10 | 8,576 | 43.523 | 43.424 | -0.099 | -0.227 | 16.819 | 16.796 | -0.024 | -0.142 |
|
| 127 |
+
| MLS public7 overall | 13,933 | 12.299 | 12.316 | +0.018 | +0.143 | 3.834 | 3.865 | +0.031 | +0.815 |
|
| 128 |
+
|
| 129 |
+
MLS public7 slightly regresses under 500. The broad multilingual aggregate remains positive because FLEURS 30 improves.
|
| 130 |
+
|
| 131 |
+
### Character-Based FLEURS Languages
|
| 132 |
+
|
| 133 |
+
| Language | #utts | Baseline CER | LatentASR CER | Rel. Delta CER |
|
| 134 |
+
|---|---:|---:|---:|---:|
|
| 135 |
+
| Mandarin (cmn_hans_cn) | 945 | 47.94 | 47.58 | -0.75 |
|
| 136 |
+
| Cantonese (yue_hant_hk) | 819 | 49.20 | 48.80 | -0.82 |
|
| 137 |
+
| Japanese (ja_jp) | 650 | 10.35 | 10.35 | +0.00 |
|
| 138 |
+
| Thai (th_th) | 1,021 | 9.53 | 9.41 | -1.31 |
|
| 139 |
+
|
| 140 |
+
## English Clean/Noise Suite and Robustness
|
| 141 |
+
|
| 142 |
+
| Dataset | Condition | Metric | Baseline | LatentASR | Delta pp | Rel. Delta |
|
| 143 |
+
|---|---|---|---:|---:|---:|---:|
|
| 144 |
+
| FLEURS | clean | WER | 4.90 | 4.78 | -0.124 | -2.54 |
|
| 145 |
+
| FLEURS | SNR=0 dB | WER | 29.81 | 29.42 | -0.393 | -1.32 |
|
| 146 |
+
| MInDS-14 | clean | WER | 33.70 | 33.68 | -0.024 | -0.07 |
|
| 147 |
+
| MInDS-14 | SNR=0 dB | WER | 40.25 | 40.49 | +0.237 | +0.59 |
|
| 148 |
+
| VoxPopuli | clean | WER | 9.04 | 8.99 | -0.043 | -0.47 |
|
| 149 |
+
| VoxPopuli | SNR=0 dB | WER | 19.71 | 19.51 | -0.200 | -1.02 |
|
| 150 |
+
| ASCEND | clean/accented | CER | 57.81 | 48.55 | -9.26 | -16.02 |
|
| 151 |
+
|
| 152 |
+
ASCEND remains the strongest robustness result. MInDS-14 SNR=0 regresses slightly and should not be summarized as uniformly robust across every noise condition.
|
| 153 |
+
|
| 154 |
+
## Latency
|
| 155 |
+
|
| 156 |
+
Measured on one NVIDIA RTX 5090, batch size 1, 100 utterances after 10 warmup utterances.
|
| 157 |
+
|
| 158 |
+
| Setting | FLEURS ms/utt | FLEURS overhead | VoxPopuli ms/utt | VoxPopuli overhead |
|
| 159 |
+
|---|---:|---:|---:|---:|
|
| 160 |
+
| Baseline | 304.31 | -- | 297.95 | -- |
|
| 161 |
+
| LatentASR theta=-2.0, full N=4 | 358.31 | +17.75 | 349.70 | +17.37 |
|
| 162 |
+
| LatentASR theta=0.0, deployed | 341.69 | +12.28 | 329.37 | +10.54 |
|
| 163 |
+
| LatentASR theta=+0.5, full skip | 321.52 | +5.66 | 314.36 | +5.51 |
|
| 164 |
+
|
| 165 |
+
The full-skip wrapper still has a fixed overhead of about 5-6%, so the paper should not claim exact zero overhead in the skip case.
|
| 166 |
+
|
| 167 |
+
## Paper Text Adjustments Required
|
| 168 |
+
|
| 169 |
+
- Replace the old main table with the canonical `activation_500_epoch10.pth` main results above.
|
| 170 |
+
- Rerun or relabel threshold, compute-allocation, multilingual, robustness, and per-sample analysis tables that currently cite `paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth`.
|
| 171 |
+
- Replace the threshold discussion: with 500, full compute does not regress VoxPopuli; dynamic halting is a compute-efficiency mechanism that preserves gains rather than the only way to avoid VoxPopuli regression.
|
| 172 |
+
- Replace the forced-negative ablation prose: p_neg=0.0 does not remove the saturated skip regime in the 500 run.
|
| 173 |
+
- Keep the component ablation; it is strong and clean.
|
| 174 |
+
- Mention prompt tuning collapse explicitly or consider moving prompt tuning to a failure-mode baseline if the huge WER hurts presentation.
|
| 175 |
+
- For multilingual, say the primary aggregate and FLEURS 30 improve, while MLS public7 is near-neutral/slightly worse.
|
| 176 |
+
- For robustness, say ASCEND and most zero-SNR English conditions improve, but MInDS-14 SNR=0 is a small regression.
|
results/paper500_more_testing_stress_20260525/ami_ihm_sdm_snr0_500.json
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "edinburghcstr/ami",
|
| 4 |
+
"dataset_name_input": "edinburghcstr/ami",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"ihm",
|
| 8 |
+
"sdm"
|
| 9 |
+
],
|
| 10 |
+
"streaming": true,
|
| 11 |
+
"text_normalizer": "english",
|
| 12 |
+
"base_model_enabled": true,
|
| 13 |
+
"baseline_checkpoint": "",
|
| 14 |
+
"prompt_tuning_checkpoint": "",
|
| 15 |
+
"lora_r16_checkpoint": "",
|
| 16 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 18 |
+
"rows": [
|
| 19 |
+
{
|
| 20 |
+
"config": "ihm",
|
| 21 |
+
"samples_total": 500,
|
| 22 |
+
"samples_used": 498,
|
| 23 |
+
"text_column": "text",
|
| 24 |
+
"base_model_wer": 0.6516084873374401,
|
| 25 |
+
"base_model_cer": 0.5181941718254568,
|
| 26 |
+
"baseline_ft_wer": null,
|
| 27 |
+
"baseline_ft_cer": null,
|
| 28 |
+
"prompt_tuning_wer": null,
|
| 29 |
+
"prompt_tuning_cer": null,
|
| 30 |
+
"lora_r16_wer": null,
|
| 31 |
+
"lora_r16_cer": null,
|
| 32 |
+
"latent_reasoning_wer": 0.6379192334017796,
|
| 33 |
+
"latent_reasoning_cer": 0.5201570285369168,
|
| 34 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 35 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 37 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 38 |
+
"lora_r16_vs_base_model_wer": null,
|
| 39 |
+
"lora_r16_vs_base_model_cer": null,
|
| 40 |
+
"latent_reasoning_vs_base_model_wer": 0.013689253935660561,
|
| 41 |
+
"latent_reasoning_vs_base_model_cer": -0.001962856711460015,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 43 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 45 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 47 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 48 |
+
},
|
| 49 |
+
{
|
| 50 |
+
"config": "sdm",
|
| 51 |
+
"samples_total": 500,
|
| 52 |
+
"samples_used": 498,
|
| 53 |
+
"text_column": "text",
|
| 54 |
+
"base_model_wer": 0.9040667361835245,
|
| 55 |
+
"base_model_cer": 0.7600583538083538,
|
| 56 |
+
"baseline_ft_wer": null,
|
| 57 |
+
"baseline_ft_cer": null,
|
| 58 |
+
"prompt_tuning_wer": null,
|
| 59 |
+
"prompt_tuning_cer": null,
|
| 60 |
+
"lora_r16_wer": null,
|
| 61 |
+
"lora_r16_cer": null,
|
| 62 |
+
"latent_reasoning_wer": 0.8724365658672228,
|
| 63 |
+
"latent_reasoning_cer": 0.738713144963145,
|
| 64 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 65 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_cer": null,
|
| 70 |
+
"latent_reasoning_vs_base_model_wer": 0.031630170316301665,
|
| 71 |
+
"latent_reasoning_vs_base_model_cer": 0.021345208845208852,
|
| 72 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 73 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 74 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 75 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 76 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 77 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 78 |
+
}
|
| 79 |
+
],
|
| 80 |
+
"summary": {
|
| 81 |
+
"base_model_weighted_wer": 0.7778376117604824,
|
| 82 |
+
"base_model_weighted_cer": 0.6391262628169052,
|
| 83 |
+
"baseline_ft_weighted_wer": null,
|
| 84 |
+
"baseline_ft_weighted_cer": null,
|
| 85 |
+
"prompt_tuning_weighted_wer": null,
|
| 86 |
+
"prompt_tuning_weighted_cer": null,
|
| 87 |
+
"lora_r16_weighted_wer": null,
|
| 88 |
+
"lora_r16_weighted_cer": null,
|
| 89 |
+
"latent_reasoning_weighted_wer": 0.7551778996345011,
|
| 90 |
+
"latent_reasoning_weighted_cer": 0.6294350867500308,
|
| 91 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 92 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 93 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.02265971212598128,
|
| 94 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.009691176066874418,
|
| 95 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 96 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 97 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 98 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 99 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 100 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 101 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 102 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 103 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 104 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 105 |
+
}
|
| 106 |
+
}
|
results/paper500_more_testing_stress_20260525/common_voice_en_snr0_500.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "SpeechTest/common_voice_16_0",
|
| 4 |
+
"dataset_name_input": "SpeechTest/common_voice_16_0",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"en"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "en",
|
| 20 |
+
"samples_total": 500,
|
| 21 |
+
"samples_used": 500,
|
| 22 |
+
"text_column": "sentence",
|
| 23 |
+
"base_model_wer": 0.3161731207289294,
|
| 24 |
+
"base_model_cer": 0.21590186278964107,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.3193621867881549,
|
| 32 |
+
"latent_reasoning_cer": 0.21353930031803725,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": -0.0031890660592255315,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.002362562471603824,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.3161731207289294,
|
| 51 |
+
"base_model_weighted_cer": 0.21590186278964107,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.3193621867881549,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.21353930031803725,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": -0.0031890660592255315,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.002362562471603824,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_more_testing_stress_20260525/extreme_asr_pony_500.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "SpeechTest/extreme_asr_pony",
|
| 4 |
+
"dataset_name_input": "SpeechTest/extreme_asr_pony",
|
| 5 |
+
"split": "train",
|
| 6 |
+
"configs": [
|
| 7 |
+
"default"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "basic",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "default",
|
| 20 |
+
"samples_total": 500,
|
| 21 |
+
"samples_used": 499,
|
| 22 |
+
"text_column": "text",
|
| 23 |
+
"base_model_wer": 0.722823147647377,
|
| 24 |
+
"base_model_cer": 0.41449099372596643,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.6649540292049757,
|
| 32 |
+
"latent_reasoning_cer": 0.3808495806067148,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.05786911844240128,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.03364141311925162,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.722823147647377,
|
| 51 |
+
"base_model_weighted_cer": 0.41449099372596643,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.6649540292049757,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.38084958060671475,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.05786911844240128,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.03364141311925167,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_more_testing_stress_20260525/gigaspeech_test_snr0_500.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "SpeechTest/gigaspeech",
|
| 4 |
+
"dataset_name_input": "SpeechTest/gigaspeech",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"test"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "test",
|
| 20 |
+
"samples_total": 500,
|
| 21 |
+
"samples_used": 391,
|
| 22 |
+
"text_column": "text",
|
| 23 |
+
"base_model_wer": 0.2093853608773272,
|
| 24 |
+
"base_model_cer": 0.12873092389539675,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.20874776842642184,
|
| 32 |
+
"latent_reasoning_cer": 0.12885580837725105,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.000637592450905361,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": -0.00012488448185429823,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.2093853608773272,
|
| 51 |
+
"base_model_weighted_cer": 0.12873092389539675,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.20874776842642184,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.12885580837725105,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.000637592450905361,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": -0.00012488448185429823,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_more_testing_stress_20260525/librispeech_clean_other_snr0_500.json
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "openslr/librispeech_asr",
|
| 4 |
+
"dataset_name_input": "openslr/librispeech_asr",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"clean",
|
| 8 |
+
"other"
|
| 9 |
+
],
|
| 10 |
+
"streaming": true,
|
| 11 |
+
"text_normalizer": "english",
|
| 12 |
+
"base_model_enabled": true,
|
| 13 |
+
"baseline_checkpoint": "",
|
| 14 |
+
"prompt_tuning_checkpoint": "",
|
| 15 |
+
"lora_r16_checkpoint": "",
|
| 16 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 18 |
+
"rows": [
|
| 19 |
+
{
|
| 20 |
+
"config": "clean",
|
| 21 |
+
"samples_total": 500,
|
| 22 |
+
"samples_used": 500,
|
| 23 |
+
"text_column": "text",
|
| 24 |
+
"base_model_wer": 0.10200129115558425,
|
| 25 |
+
"base_model_cer": 0.054846653234833224,
|
| 26 |
+
"baseline_ft_wer": null,
|
| 27 |
+
"baseline_ft_cer": null,
|
| 28 |
+
"prompt_tuning_wer": null,
|
| 29 |
+
"prompt_tuning_cer": null,
|
| 30 |
+
"lora_r16_wer": null,
|
| 31 |
+
"lora_r16_cer": null,
|
| 32 |
+
"latent_reasoning_wer": 0.10255464354883335,
|
| 33 |
+
"latent_reasoning_cer": 0.054777772037678016,
|
| 34 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 35 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 37 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 38 |
+
"lora_r16_vs_base_model_wer": null,
|
| 39 |
+
"lora_r16_vs_base_model_cer": null,
|
| 40 |
+
"latent_reasoning_vs_base_model_wer": -0.0005533523932491069,
|
| 41 |
+
"latent_reasoning_vs_base_model_cer": 6.888119715520863e-05,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 43 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 45 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 47 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 48 |
+
},
|
| 49 |
+
{
|
| 50 |
+
"config": "other",
|
| 51 |
+
"samples_total": 500,
|
| 52 |
+
"samples_used": 500,
|
| 53 |
+
"text_column": "text",
|
| 54 |
+
"base_model_wer": 0.26742456438589035,
|
| 55 |
+
"base_model_cer": 0.16974330287774256,
|
| 56 |
+
"baseline_ft_wer": null,
|
| 57 |
+
"baseline_ft_cer": null,
|
| 58 |
+
"prompt_tuning_wer": null,
|
| 59 |
+
"prompt_tuning_cer": null,
|
| 60 |
+
"lora_r16_wer": null,
|
| 61 |
+
"lora_r16_cer": null,
|
| 62 |
+
"latent_reasoning_wer": 0.26912452188695285,
|
| 63 |
+
"latent_reasoning_cer": 0.16781652060413949,
|
| 64 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 65 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_cer": null,
|
| 70 |
+
"latent_reasoning_vs_base_model_wer": -0.001699957501062499,
|
| 71 |
+
"latent_reasoning_vs_base_model_cer": 0.0019267822736030726,
|
| 72 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 73 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 74 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 75 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 76 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 77 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 78 |
+
}
|
| 79 |
+
],
|
| 80 |
+
"summary": {
|
| 81 |
+
"base_model_weighted_wer": 0.18471292777073728,
|
| 82 |
+
"base_model_weighted_cer": 0.11229497805628791,
|
| 83 |
+
"baseline_ft_weighted_wer": null,
|
| 84 |
+
"baseline_ft_weighted_cer": null,
|
| 85 |
+
"prompt_tuning_weighted_wer": null,
|
| 86 |
+
"prompt_tuning_weighted_cer": null,
|
| 87 |
+
"lora_r16_weighted_wer": null,
|
| 88 |
+
"lora_r16_weighted_cer": null,
|
| 89 |
+
"latent_reasoning_weighted_wer": 0.1858395827178931,
|
| 90 |
+
"latent_reasoning_weighted_cer": 0.11129714632090876,
|
| 91 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 92 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 93 |
+
"latent_reasoning_vs_base_model_weighted_wer": -0.0011266549471558307,
|
| 94 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.000997831735379151,
|
| 95 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 96 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 97 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 98 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 99 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 100 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 101 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 102 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 103 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 104 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 105 |
+
}
|
| 106 |
+
}
|
results/paper500_more_testing_stress_20260525/peoples_speech_test_snr0_500.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "SpeechTest/peoples_speech",
|
| 4 |
+
"dataset_name_input": "SpeechTest/peoples_speech",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"test"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "test",
|
| 20 |
+
"samples_total": 500,
|
| 21 |
+
"samples_used": 500,
|
| 22 |
+
"text_column": "text",
|
| 23 |
+
"base_model_wer": 0.31687113348013,
|
| 24 |
+
"base_model_cer": 0.224544784759499,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.31770997168921045,
|
| 32 |
+
"latent_reasoning_cer": 0.22947058204399537,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": -0.0008388382090804281,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": -0.0049257972844963804,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.31687113348013,
|
| 51 |
+
"base_model_weighted_cer": 0.224544784759499,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.31770997168921045,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.22947058204399537,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": -0.0008388382090804281,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": -0.0049257972844963804,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|
results/paper500_more_testing_stress_20260525/showcase_report.md
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# LR HuggingFace ASR Showcase Report
|
| 2 |
+
|
| 3 |
+
- Generated UTC: 2026-05-25T16:47:09+00:00
|
| 4 |
+
- Output directory: `/user_data/lr_whisper/eval_runs/paper500_more_testing_stress_20260525`
|
| 5 |
+
- Delta WER is `base_model_wer - latent_reasoning_wer`; positive means LR is better.
|
| 6 |
+
|
| 7 |
+
## Best LR Wins
|
| 8 |
+
|
| 9 |
+
| Rank | Dataset | Configs | Condition | N | Base WER | LR WER | Delta WER | Relative | Delta CER | JSON |
|
| 10 |
+
|---:|---|---|---|---:|---:|---:|---:|---:|---:|---|
|
| 11 |
+
| 1 | SpeechTest/extreme_asr_pony | default | unknown | 499 | 0.722823 | 0.664954 | 0.057869 | +8.01% | 0.033641 | `extreme_asr_pony_500.json` |
|
| 12 |
+
| 2 | edinburghcstr/ami | ihm,sdm | unknown | 996 | 0.777838 | 0.755178 | 0.022660 | +2.91% | 0.009691 | `ami_ihm_sdm_snr0_500.json` |
|
| 13 |
+
| 3 | TwinkStart/tedlium | release1 | unknown | 500 | 0.148490 | 0.142912 | 0.005578 | +3.76% | 0.001232 | `tedlium_release1_snr0_500.json` |
|
| 14 |
+
| 4 | SpeechTest/gigaspeech | test | unknown | 391 | 0.209385 | 0.208748 | 0.000638 | +0.30% | -0.000125 | `gigaspeech_test_snr0_500.json` |
|
| 15 |
+
|
| 16 |
+
## All Cases
|
| 17 |
+
|
| 18 |
+
| Dataset | Configs | Condition | N | Base WER | LR WER | Delta WER | Relative | Base CER | LR CER | Delta CER | JSON |
|
| 19 |
+
|---|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---|
|
| 20 |
+
| SpeechTest/common_voice_16_0 | en | unknown | 500 | 0.316173 | 0.319362 | -0.003189 | -1.01% | 0.215902 | 0.213539 | 0.002363 | `common_voice_en_snr0_500.json` |
|
| 21 |
+
| SpeechTest/extreme_asr_pony | default | unknown | 499 | 0.722823 | 0.664954 | 0.057869 | +8.01% | 0.414491 | 0.380850 | 0.033641 | `extreme_asr_pony_500.json` |
|
| 22 |
+
| SpeechTest/gigaspeech | test | unknown | 391 | 0.209385 | 0.208748 | 0.000638 | +0.30% | 0.128731 | 0.128856 | -0.000125 | `gigaspeech_test_snr0_500.json` |
|
| 23 |
+
| SpeechTest/peoples_speech | test | unknown | 500 | 0.316871 | 0.317710 | -0.000839 | -0.26% | 0.224545 | 0.229471 | -0.004926 | `peoples_speech_test_snr0_500.json` |
|
| 24 |
+
| TwinkStart/tedlium | release1 | unknown | 500 | 0.148490 | 0.142912 | 0.005578 | +3.76% | 0.088532 | 0.087300 | 0.001232 | `tedlium_release1_snr0_500.json` |
|
| 25 |
+
| edinburghcstr/ami | ihm,sdm | unknown | 996 | 0.777838 | 0.755178 | 0.022660 | +2.91% | 0.639126 | 0.629435 | 0.009691 | `ami_ihm_sdm_snr0_500.json` |
|
| 26 |
+
| openslr/librispeech_asr | clean,other | unknown | 1000 | 0.184713 | 0.185840 | -0.001127 | -0.61% | 0.112295 | 0.111297 | 0.000998 | `librispeech_clean_other_snr0_500.json` |
|
| 27 |
+
|
| 28 |
+
## Regressions To Check
|
| 29 |
+
|
| 30 |
+
| Dataset | Configs | Condition | N | Delta WER | Relative | JSON |
|
| 31 |
+
|---|---|---|---:|---:|---:|---|
|
| 32 |
+
| SpeechTest/common_voice_16_0 | en | unknown | 500 | -0.003189 | -1.01% | `common_voice_en_snr0_500.json` |
|
| 33 |
+
| openslr/librispeech_asr | clean,other | unknown | 1000 | -0.001127 | -0.61% | `librispeech_clean_other_snr0_500.json` |
|
| 34 |
+
| SpeechTest/peoples_speech | test | unknown | 500 | -0.000839 | -0.26% | `peoples_speech_test_snr0_500.json` |
|
| 35 |
+
|
results/paper500_more_testing_stress_20260525/tedlium_release1_snr0_500.json
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_id": "Qwen/Qwen3-ASR-0.6B",
|
| 3 |
+
"dataset_name": "TwinkStart/tedlium",
|
| 4 |
+
"dataset_name_input": "TwinkStart/tedlium",
|
| 5 |
+
"split": "test",
|
| 6 |
+
"configs": [
|
| 7 |
+
"release1"
|
| 8 |
+
],
|
| 9 |
+
"streaming": true,
|
| 10 |
+
"text_normalizer": "english",
|
| 11 |
+
"base_model_enabled": true,
|
| 12 |
+
"baseline_checkpoint": "",
|
| 13 |
+
"prompt_tuning_checkpoint": "",
|
| 14 |
+
"lora_r16_checkpoint": "",
|
| 15 |
+
"latent_reasoning_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 16 |
+
"latent_checkpoint": "eval_runs/paper500_retrain_20260518/checkpoints/n4/n4_epoch10.pth",
|
| 17 |
+
"rows": [
|
| 18 |
+
{
|
| 19 |
+
"config": "release1",
|
| 20 |
+
"samples_total": 500,
|
| 21 |
+
"samples_used": 500,
|
| 22 |
+
"text_column": "text",
|
| 23 |
+
"base_model_wer": 0.148490094248894,
|
| 24 |
+
"base_model_cer": 0.08853179716516335,
|
| 25 |
+
"baseline_ft_wer": null,
|
| 26 |
+
"baseline_ft_cer": null,
|
| 27 |
+
"prompt_tuning_wer": null,
|
| 28 |
+
"prompt_tuning_cer": null,
|
| 29 |
+
"lora_r16_wer": null,
|
| 30 |
+
"lora_r16_cer": null,
|
| 31 |
+
"latent_reasoning_wer": 0.142912098480477,
|
| 32 |
+
"latent_reasoning_cer": 0.0873000833775487,
|
| 33 |
+
"baseline_ft_vs_base_model_wer": null,
|
| 34 |
+
"baseline_ft_vs_base_model_cer": null,
|
| 35 |
+
"prompt_tuning_vs_base_model_wer": null,
|
| 36 |
+
"prompt_tuning_vs_base_model_cer": null,
|
| 37 |
+
"lora_r16_vs_base_model_wer": null,
|
| 38 |
+
"lora_r16_vs_base_model_cer": null,
|
| 39 |
+
"latent_reasoning_vs_base_model_wer": 0.005577995768416999,
|
| 40 |
+
"latent_reasoning_vs_base_model_cer": 0.0012317137876146494,
|
| 41 |
+
"latent_reasoning_vs_baseline_ft_wer": null,
|
| 42 |
+
"latent_reasoning_vs_baseline_ft_cer": null,
|
| 43 |
+
"latent_reasoning_vs_prompt_tuning_wer": null,
|
| 44 |
+
"latent_reasoning_vs_prompt_tuning_cer": null,
|
| 45 |
+
"latent_reasoning_vs_lora_r16_wer": null,
|
| 46 |
+
"latent_reasoning_vs_lora_r16_cer": null
|
| 47 |
+
}
|
| 48 |
+
],
|
| 49 |
+
"summary": {
|
| 50 |
+
"base_model_weighted_wer": 0.148490094248894,
|
| 51 |
+
"base_model_weighted_cer": 0.08853179716516335,
|
| 52 |
+
"baseline_ft_weighted_wer": null,
|
| 53 |
+
"baseline_ft_weighted_cer": null,
|
| 54 |
+
"prompt_tuning_weighted_wer": null,
|
| 55 |
+
"prompt_tuning_weighted_cer": null,
|
| 56 |
+
"lora_r16_weighted_wer": null,
|
| 57 |
+
"lora_r16_weighted_cer": null,
|
| 58 |
+
"latent_reasoning_weighted_wer": 0.142912098480477,
|
| 59 |
+
"latent_reasoning_weighted_cer": 0.0873000833775487,
|
| 60 |
+
"baseline_ft_vs_base_model_weighted_wer": null,
|
| 61 |
+
"baseline_ft_vs_base_model_weighted_cer": null,
|
| 62 |
+
"latent_reasoning_vs_base_model_weighted_wer": 0.005577995768416999,
|
| 63 |
+
"latent_reasoning_vs_base_model_weighted_cer": 0.0012317137876146494,
|
| 64 |
+
"latent_reasoning_vs_baseline_ft_weighted_wer": null,
|
| 65 |
+
"latent_reasoning_vs_baseline_ft_weighted_cer": null,
|
| 66 |
+
"prompt_tuning_vs_base_model_weighted_wer": null,
|
| 67 |
+
"prompt_tuning_vs_base_model_weighted_cer": null,
|
| 68 |
+
"lora_r16_vs_base_model_weighted_wer": null,
|
| 69 |
+
"lora_r16_vs_base_model_weighted_cer": null,
|
| 70 |
+
"latent_reasoning_vs_prompt_tuning_weighted_wer": null,
|
| 71 |
+
"latent_reasoning_vs_prompt_tuning_weighted_cer": null,
|
| 72 |
+
"latent_reasoning_vs_lora_r16_weighted_wer": null,
|
| 73 |
+
"latent_reasoning_vs_lora_r16_weighted_cer": null
|
| 74 |
+
}
|
| 75 |
+
}
|