Complete OmniVoice source mirror

runtime/omnivoice-python/omnivoice-source/.gitignore

@@ -0,0 +1,30 @@
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+*.egg
+dist/
+build/
+.venv/
+.env
+.DS_Store
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+*.so
+/.cache*
+/exp*/
+/.tmp/
+/results/
+/data/
+/download
+/local/
+/run*
+example.py
+results/
+examples/data*
+examples/download*
+examples/exp*
+.claude/
+*.wav
+*.jsonl

runtime/omnivoice-python/omnivoice-source/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Xiaomi Corp.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

runtime/omnivoice-python/omnivoice-source/docs/OmniVoice.ipynb

@@ -0,0 +1,144 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "# OmniVoice Quick Start\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/k2-fsa/OmniVoice/blob/master/docs/OmniVoice.ipynb)\n\nThis notebook demonstrates the basic usage of [OmniVoice](https://github.com/k2-fsa/OmniVoice), a massively multilingual zero-shot TTS model supporting 600+ languages.\n\n**Contents:**\n1. Installation\n2. Option A — Gradio Demo (interactive web UI, no code needed)\n3. Option B — Python API\n   - 3.1 Load Model\n   - 3.2 Voice Cloning\n   - 3.3 Voice Design\n   - 3.4 Auto Voice"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Installation\n",
+    "\n",
+    "Colab already provides a compatible PyTorch + CUDA environment, so we only need to install OmniVoice."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install omnivoice"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 2. Option A — Gradio Demo\n\nLaunch an interactive web UI with a public Gradio link. The `--share` flag creates a temporary public URL so you can access the demo from any browser.\n\n> **If you prefer to use the Python API directly, skip to Option B below.**"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!omnivoice-demo --share"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 3. Option B — Python API\n\n### 3.1 Load Model"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "from omnivoice import OmniVoice\nimport soundfile as sf\nimport torch\nfrom IPython.display import Audio, display\n\nmodel = OmniVoice.from_pretrained(\n    \"k2-fsa/OmniVoice\",\n    device_map=\"cuda:0\",\n    dtype=torch.float16,\n    load_asr=True,\n)"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "### 3.2 Voice Cloning\n\nClone a voice from a short (3-10s) reference audio clip. Upload your own `ref.wav` or use any audio file.\n\n`ref_text` is optional — if omitted, the model uses Whisper ASR to auto-transcribe it."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from google.colab import files\n",
+    "\n",
+    "print(\"Upload a reference audio file (wav/mp3/flac):\")\n",
+    "uploaded = files.upload()\n",
+    "ref_audio_path = list(uploaded.keys())[0]\n",
+    "print(f\"Uploaded: {ref_audio_path}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio = model.generate(\n",
+    "    text=\"Hello, this is a test of zero-shot voice cloning.\",\n",
+    "    ref_audio=ref_audio_path,\n",
+    "    # ref_text=\"Transcription of the reference audio.\",  # optional\n",
+    ")\n",
+    "\n",
+    "sf.write(\"clone_out.wav\", audio[0], 24000)\n",
+    "display(Audio(audio[0], rate=24000))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "### 3.3 Voice Design\n\nDescribe the desired voice with speaker attributes — no reference audio needed.\n\nSupported attributes: gender, age, pitch, style (whisper), English accent, Chinese dialect. See [docs/voice-design.md](https://github.com/k2-fsa/OmniVoice/blob/master/docs/voice-design.md) for the full list."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio = model.generate(\n",
+    "    text=\"Hello, this is a test of zero-shot voice design.\",\n",
+    "    instruct=\"female, low pitch, british accent\",\n",
+    ")\n",
+    "\n",
+    "sf.write(\"design_out.wav\", audio[0], 24000)\n",
+    "display(Audio(audio[0], rate=24000))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "### 3.4 Auto Voice\n\nLet the model choose a voice automatically — no reference audio or instruct needed."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio = model.generate(\n",
+    "    text=\"This is a sentence generated with automatic voice selection.\",\n",
+    ")\n",
+    "\n",
+    "sf.write(\"auto_out.wav\", audio[0], 24000)\n",
+    "display(Audio(audio[0], rate=24000))"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

runtime/omnivoice-python/omnivoice-source/docs/lang_id_name_map.tsv

@@ -0,0 +1,647 @@
+language_id	language_name	iso_639_3_id	train_data_duration
+aae	Arbëreshë Albanian	aae	6.11
+aal	Afade	aal	10.19
+aao	Algerian Saharan Arabic	aao	2.02
+ab	Abkhazian	abk	57.27
+abb	Bankon	abb	11.2
+abn	Abua	abn	10.27
+abr	Abron	abr	9.22
+abs	Ambonese Malay	abs	10.03
+abv	Baharna Arabic	abv	10.41
+acm	Mesopotamian Arabic	acm	3.78
+acw	Hijazi Arabic	acw	22.32
+acx	Omani Arabic	acx	22.03
+adf	Dhofari Arabic	adf	0.31
+adx	Amdo Tibetan	adx	56.94
+ady	Adyghe	ady	32.6
+aeb	Tunisian Arabic	aeb	21.63
+aec	Saidi Arabic	aec	9.28
+af	Afrikaans	afr	4.4
+afb	Gulf Arabic	afb	98.55
+afo	Eloyi	afo	11.21
+ahl	Igo	ahl	9.22
+ahs	Ashe	ahs	10.62
+ajg	Aja (Benin)	ajg	5.63
+aju	Judeo-Moroccan Arabic	aju	7.21
+ala	Alago	ala	11.04
+aln	Gheg Albanian	aln	3.92
+alo	Larike-Wakasihu	alo	9.97
+am	Amharic	amh	12.83
+amu	Guerrero Amuzgo	amu	10.1
+an	Aragonese	arg	16.4
+anc	Ngas	anc	10.14
+ank	Goemai	ank	10.0
+anp	Angika	anp	10.65
+anw	Anaang	anw	9.65
+aom	Ömie	aom	8.19
+apc	Levantine Arabic	apc	15.65
+apd	Sudanese Arabic	apd	9.93
+arb	Standard Arabic	arb	1483.53
+arq	Algerian Arabic	arq	9.64
+ars	Najdi Arabic	ars	203.54
+ary	Moroccan Arabic	ary	104.67
+arz	Egyptian Arabic	arz	23.23
+as	Assamese	asm	270.85
+ast	Asturian	ast	8.48
+avl	Eastern Egyptian Bedawi Arabic	avl	1.86
+awo	Awak	awo	10.22
+ayl	Libyan Arabic	ayl	20.13
+ayp	North Mesopotamian Arabic	ayp	10.92
+az	Azerbaijani	aze	9.84
+ba	Bashkir	bak	249.1
+bag	Tuki	bag	10.97
+bas	Basa (Cameroon)	bas	10.66
+bax	Bamun	bax	10.24
+bba	Baatonum	bba	10.53
+bbj	Ghomálá'	bbj	7.32
+bbl	Bats	bbl	11.22
+bbu	Kulung (Nigeria)	bbu	10.39
+bce	Bamenyam	bce	9.9
+bci	Baoulé	bci	10.21
+bcs	Kohumono	bcs	10.45
+bcy	Bacama	bcy	9.94
+bda	Bayot	bda	9.47
+bde	Bade	bde	9.89
+bdm	Buduma	bdm	10.17
+be	Belarusian	bel	1809.43
+beb	Bebele	beb	7.52
+bew	Betawi	bew	11.15
+bfd	Bafut	bfd	9.03
+bft	Balti	bft	16.28
+bg	Bulgarian	bul	2190.76
+bgp	Eastern Balochi	bgp	10.98
+bhb	Bhili	bhb	9.98
+bhh	Bukharic	bhh	11.38
+bho	Bhojpuri	bho	10.05
+bhp	Bima	bhp	10.67
+bhr	Bara Malagasy	bhr	12.14
+bjj	Kanauji	bjj	11.01
+bjk	Barok	bjk	10.16
+bjn	Banjar	bjn	11.68
+bjt	Balanta-Ganja	bjt	9.41
+bkh	Bakoko	bkh	6.0
+bkm	Kom (Cameroon)	bkm	10.76
+bky	Bokyi	bky	9.85
+bmm	Northern Betsimisaraka Malagasy	bmm	19.12
+bmq	Bomu	bmq	10.68
+bn	Bengali	ben	271.76
+bnm	Batanga	bnm	15.01
+bnn	Bunun	bnn	9.26
+bns	Bundeli	bns	10.88
+bo	Tibetan	bod	82.27
+bou	Bondei	bou	9.98
+bqg	Bago-Kusuntu	bqg	8.86
+br	Breton	bre	25.48
+bra	Braj	bra	10.68
+brh	Brahui	brh	19.89
+bri	Mokpwe	bri	7.53
+brx	Bodo	brx	231.57
+bs	Bosnian	bos	690.73
+bsh	Kati	bsh	8.77
+bsj	Bangwinji	bsj	10.0
+bsk	Burushaski	bsk	9.14
+btm	Batak Mandailing	btm	11.09
+btv	Bateri	btv	9.8
+bug	Buginese	bug	11.09
+bum	Bulu (Cameroon)	bum	9.06
+buo	Terei	buo	9.48
+bux	Boghom	bux	10.48
+bwr	Bura-Pabir	bwr	10.4
+bxf	Bilur	bxf	10.84
+byc	Ubaghara	byc	11.11
+bys	Burak	bys	9.92
+byv	Medumba	byv	10.95
+byx	Qaqet	byx	9.79
+bzc	Southern Betsimisaraka Malagasy	bzc	17.45
+bzw	Basa (Nigeria)	bzw	10.27
+ca	Catalan	cat	3358.6
+ccg	Samba Daka	ccg	10.11
+ceb	Cebuano	ceb	12.17
+cen	Cen	cen	9.85
+cfa	Dijim-Bwilim	cfa	10.32
+cgg	Chiga	cgg	10.84
+chq	Quiotepec Chinantec	chq	9.76
+cjk	Chokwe	cjk	11.01
+ckb	Central Kurdish	ckb	137.52
+ckl	Cibak	ckl	10.91
+ckr	Kairak	ckr	10.51
+cky	Cakfem-Mushere	cky	8.96
+cnh	Hakha Chin	cnh	2.24
+cpy	South Ucayali Ashéninka	cpy	9.15
+cs	Czech	ces	148.13
+cte	Tepinapa Chinantec	cte	9.54
+ctl	Tlacoatzintepec Chinantec	ctl	10.04
+cut	Teutila Cuicatec	cut	8.04
+cux	Tepeuxila Cuicatec	cux	7.83
+cv	Chuvash	chv	23.96
+cy	Welsh	cym	131.21
+da	Danish	dan	1665.98
+dag	Dagbani	dag	10.14
+dar	Dargwa	dar	1.22
+dav	Taita	dav	9.12
+dbd	Dadiya	dbd	9.61
+dcc	Deccan	dcc	10.38
+de	German	deu	21927.13
+deg	Degema	deg	11.07
+dgh	Dghwede	dgh	9.95
+dgo	Dogri	dgo	117.04
+dje	Zarma	dje	10.72
+dmk	Domaaki	dmk	6.38
+dml	Dameli	dml	9.18
+dru	Rukai	dru	9.26
+dty	Dotyali	dty	10.85
+dua	Duala	dua	12.13
+dv	Dhivehi	div	38.61
+dyu	Dyula	dyu	0.34
+dzg	Dazaga	dzg	9.96
+ebr	Ebrié	ebr	1.5
+ebu	Embu	ebu	9.81
+ego	Eggon	ego	9.95
+eiv	Askopan	eiv	10.44
+eko	Koti	eko	8.15
+ekr	Yace	ekr	10.76
+el	Greek	ell	2412.54
+elm	Eleme	elm	11.27
+en	English	eng	206061.1
+eo	Esperanto	epo	1396.64
+es	Spanish	spa	27559.74
+esu	Central Yupik	esu	2.18
+et	Estonian	est	960.37
+eto	Eton (Cameroon)	eto	7.43
+ets	Yekhee	ets	10.11
+etu	Ejagham	etu	10.3
+eu	Basque	eus	479.86
+ewo	Ewondo	ewo	12.71
+ext	Extremaduran	ext	13.59
+eyo	Keiyo	eyo	9.24
+fa	Persian	fas	366.07
+fan	Fang (Equatorial Guinea)	fan	3.51
+fat	Fanti	fat	11.38
+ff	Fulah	ful	13.84
+ffm	Maasina Fulfulde	ffm	10.46
+fi	Finnish	fin	468.62
+fia	Nobiin	fia	9.96
+fil	Filipino	fil	7.71
+fip	Fipa	fip	10.55
+fkk	Kirya-Konzəl	fkk	9.98
+fmp	Fe'fe'	fmp	9.86
+fr	French	fra	23675.32
+fub	Adamawa Fulfulde	fub	13.12
+fuc	Pulaar	fuc	14.77
+fue	Borgu Fulfulde	fue	20.1
+fuf	Pular	fuf	13.77
+fuh	Western Niger Fulfulde	fuh	9.69
+fui	Bagirmi Fulfulde	fui	15.04
+fuq	Central-Eastern Niger Fulfulde	fuq	9.28
+fuv	Nigerian Fulfulde	fuv	9.97
+fy	Western Frisian	fry	70.41
+ga	Irish	gle	21.4
+gbm	Garhwali	gbm	19.14
+gbr	Gbagyi	gbr	12.12
+gby	Gbari	gby	12.59
+gcc	Mali	gcc	9.87
+gdf	Guduf-Gava	gdf	12.21
+gej	Gen	gej	5.39
+ges	Geser-Gorom	ges	10.08
+ggg	Gurgula	ggg	7.12
+gid	Gidar	gid	10.06
+gig	Goaria	gig	9.41
+giz	South Giziga	giz	10.03
+gjk	Kachi Koli	gjk	20.83
+gju	Gujari	gju	8.66
+gl	Galician	glg	208.81
+glw	Glavda	glw	10.51
+gn	Guarani	grn	4.06
+gol	Gola	gol	9.26
+gom	Goan Konkani	gom	9.82
+gsl	Gusilay	gsl	10.0
+gu	Gujarati	guj	91.18
+gui	Eastern Bolivian Guaraní	gui	22.72
+gur	Farefare	gur	9.24
+guz	Gusii	guz	9.5
+gv	Manx	glv	10.07
+gwc	Gawri	gwc	10.83
+gwe	Gweno	gwe	8.87
+gwt	Gawar-Bati	gwt	12.16
+gya	Northwest Gbaya	gya	8.45
+gyz	Geji	gyz	10.49
+ha	Hausa	hau	17.75
+hah	Hahon	hah	9.64
+hao	Hakö	hao	8.56
+haw	Hawaiian	haw	11.79
+haz	Hazaragi	haz	9.69
+hbb	Huba	hbb	10.7
+he	Hebrew	heb	13.4
+hem	Hemba	hem	9.53
+hi	Hindi	hin	117.17
+hia	Lamang	hia	11.07
+hkk	Hunjara-Kaina Ke	hkk	8.69
+hla	Halia	hla	9.86
+hno	Northern Hindko	hno	20.04
+hoj	Hadothi	hoj	10.08
+hr	Croatian	hrv	2795.31
+hsb	Upper Sorbian	hsb	2.71
+ht	Haitian	hat	0.04
+hu	Hungarian	hun	255.83
+hue	San Francisco Del Mar Huave	hue	9.45
+hul	Hula	hul	10.33
+hux	Nüpode Huitoto	hux	9.04
+hwo	Hwana	hwo	11.23
+hy	Armenian	hye	42.15
+hz	Herero	her	9.59
+ia	Interlingua (International Auxiliary Language Association)	ina	13.48
+ibb	Ibibio	ibb	7.38
+id	Indonesian	ind	6327.87
+ida	Idakho-Isukha-Tiriki	ida	9.31
+idu	Idoma	idu	11.16
+ig	Igbo	ibo	13.69
+ijc	Izon	ijc	9.95
+ijn	Kalabari	ijn	11.04
+ik	Inupiaq	ipk	2.11
+ikw	Ikwere	ikw	10.0
+is	Icelandic	isl	647.29
+ish	Esan	ish	10.05
+iso	Isoko	iso	10.33
+it	Italian	ita	9402.46
+its	Isekiri	its	11.85
+itw	Ito	itw	9.19
+itz	Itzá	itz	7.08
+ja	Japanese	jpn	36914.4
+jal	Yalahatan	jal	11.18
+jax	Jambi Malay	jax	10.29
+jgo	Ngomba	jgo	10.15
+jmx	Western Juxtlahuaca Mixtec	jmx	10.01
+jns	Jaunsari	jns	11.25
+jqr	Jaqaru	jqr	9.32
+juk	Wapan	juk	10.22
+juo	Jiba	juo	10.43
+jv	Javanese	jav	11.19
+ka	Georgian	kat	156.96
+kab	Kabyle	kab	529.52
+kai	Karekare	kai	10.52
+kaj	Jju	kaj	10.16
+kam	Kamba	kam	14.72
+kbd	Kabardian	kbd	108.35
+kbl	Kanembu	kbl	10.19
+kbt	Abadi	kbt	9.73
+kcq	Kamo	kcq	10.49
+kdh	Tem	kdh	4.07
+kea	Kabuverdianu	kea	10.51
+keu	Akebu	keu	9.1
+kfe	Kota (India)	kfe	10.25
+kfk	Kinnauri	kfk	10.32
+kfp	Korwa	kfp	11.87
+khg	Khams Tibetan	khg	6.38
+khw	Khowar	khw	15.55
+kj	Kuanyama	kua	9.88
+kjc	Coastal Konjo	kjc	10.18
+kjk	Highland Konjo	kjk	10.21
+kk	Kazakh	kaz	1537.29
+kln	Kalenjin	kln	40.42
+kls	Kalasha	kls	9.11
+km	Khmer	khm	7.1
+kmr	Northern Kurdish	kmr	69.59
+kmy	Koma	kmy	10.28
+kn	Kannada	kan	128.06
+kna	Dera (Nigeria)	kna	11.91
+knn	Konkani	knn	112.83
+ko	Korean	kor	8609.28
+kol	Kol (Papua New Guinea)	kol	9.95
+koo	Konzo	koo	13.23
+kpo	Ikposo	kpo	7.83
+kqo	Eastern Krahn	kqo	9.28
+ks	Kashmiri	kas	110.42
+ksd	Kuanua	ksd	9.91
+ksf	Bafia	ksf	16.43
+kto	Kuot	kto	9.77
+kuh	Kushi	kuh	10.35
+kvx	Parkari Koli	kvx	11.04
+kw	Cornish	cor	12.15
+kwm	Kwambi	kwm	9.9
+kxp	Wadiyara Koli	kxp	20.0
+ky	Kirghiz	kir	46.63
+kyx	Rapoisi	kyx	9.17
+lag	Rangi	lag	9.47
+lb	Luxembourgish	ltz	8.46
+lcm	Tungag	lcm	9.77
+ldb	Dũya	ldb	11.31
+lg	Ganda	lug	447.82
+lij	Ligurian	lij	15.97
+lir	Liberian English	lir	10.26
+lkb	Kabras	lkb	9.99
+lla	Lala-Roba	lla	10.38
+ln	Lingala	lin	17.99
+lnu	Longuda	lnu	10.46
+lo	Lao	lao	7.63
+loa	Loloda	loa	9.31
+lrk	Loarki	lrk	10.5
+lss	Lasi	lss	6.53
+lt	Lithuanian	lit	2629.45
+ltg	Latgalian	ltg	27.23
+lto	Tsotso	lto	9.77
+lua	Luba-Lulua	lua	8.47
+luo	Luo	luo	36.17
+lus	Lushai	lus	20.24
+lv	Latvian	lav	1441.58
+lwg	Wanga	lwg	9.36
+mab	Yutanduchi Mixtec	mab	9.26
+maf	Mafa	maf	9.97
+mai	Maithili	mai	131.37
+mau	Huautla Mazatec	mau	6.39
+max	North Moluccan Malay	max	9.43
+mbo	Mbo (Cameroon)	mbo	9.51
+mcf	Matsés	mcf	9.61
+mcn	Masana	mcn	10.09
+mcx	Mpiemo	mcx	9.88
+mdd	Mbum	mdd	9.82
+mde	Maba (Chad)	mde	9.5
+mdf	Moksha	mdf	0.47
+mek	Mekeo	mek	9.18
+mer	Meru	mer	9.89
+meu	Motu	meu	9.88
+mfm	Marghi South	mfm	10.05
+mfn	Cross River Mbembe	mfn	10.03
+mfo	Mbe	mfo	10.24
+mfv	Mandjak	mfv	9.55
+mgg	Mpumpong	mgg	4.94
+mgi	Lijili	mgi	10.89
+mhk	Mungaka	mhk	7.53
+mhr	Eastern Mari	mhr	272.31
+mi	Maori	mri	18.02
+mig	San Miguel El Grande Mixtec	mig	9.66
+miu	Cacaloxtepec Mixtec	miu	9.18
+mk	Macedonian	mkd	27.21
+mkf	Miya	mkf	10.16
+mki	Dhatki	mki	8.83
+ml	Malayalam	mal	166.57
+mlq	Western Maninkakan	mlq	9.83
+mn	Mongolian	mon	269.08
+mne	Naba	mne	10.37
+mni	Manipuri	mni	44.46
+mqy	Manggarai	mqy	10.5
+mr	Marathi	mar	156.71
+mrj	Western Mari	mrj	32.26
+mrr	Maria (India)	mrr	11.0
+mrt	Marghi Central	mrt	10.36
+ms	Malay	msa	9.57
+mse	Musey	mse	7.21
+msh	Masikoro Malagasy	msh	14.16
+msw	Mansoanka	msw	9.32
+mt	Maltese	mlt	630.29
+mtr	Mewari	mtr	10.58
+mtu	Tututepec Mixtec	mtu	10.13
+mtx	Tidaá Mixtec	mtx	9.09
+mua	Mundang	mua	9.2
+mug	Musgu	mug	4.74
+mui	Musi	mui	10.52
+mve	Marwari (Pakistan)	mve	9.96
+mvy	Indus Kohistani	mvy	21.64
+mxs	Huitepec Mixtec	mxs	9.64
+mxu	Mada (Cameroon)	mxu	12.0
+mxy	Southeastern Nochixtlán Mixtec	mxy	9.48
+my	Burmese	mya	12.14
+myv	Erzya	myv	3.1
+mzl	Mazatlán Mixe	mzl	10.05
+nal	Nalik	nal	10.33
+nan	Min Nan Chinese	nan	17.55
+nap	Neapolitan	nap	9.97
+nb	Norwegian Bokmål	nob	12.7
+nbh	Ngamo	nbh	10.04
+ncf	Notsi	ncf	9.84
+nco	Sibe	nco	9.96
+ncx	Central Puebla Nahuatl	ncx	9.86
+ndi	Samba Leko	ndi	11.27
+ng	Ndonga	ndo	9.08
+ngi	Ngizim	ngi	10.06
+nhg	Tetelcingo Nahuatl	nhg	8.92
+nhi	Zacatlán-Ahuacatlán-Tepetzintla Nahuatl	nhi	0.05
+nhn	Central Nahuatl	nhn	9.51
+nhq	Huaxcaleca Nahuatl	nhq	5.07
+nja	Nzanyi	nja	10.02
+nl	Dutch	nld	2264.13
+nla	Ngombale	nla	8.79
+nlv	Orizaba Nahuatl	nlv	11.42
+nmg	Kwasio	nmg	10.39
+nmz	Nawdm	nmz	6.3
+nn	Norwegian Nynorsk	nno	1.54
+nnh	Ngiemboon	nnh	16.15
+no	Norwegian	nor	3849.8
+noe	Nimadi	noe	11.12
+npi	Nepali	npi	171.5
+nso	Pedi	nso	12.64
+ny	Chichewa	nya	10.8
+nyu	Nyungwe	nyu	8.98
+oc	Occitan	oci	16.8
+odk	Od	odk	20.26
+odu	Odual	odu	10.57
+ogo	Khana	ogo	10.51
+om	Oromo	orm	6.6
+orc	Orma	orc	22.01
+oru	Ormuri	oru	16.74
+ory	Odia	ory	144.81
+os	Iron Ossetic	oss	1.38
+pa	Panjabi	pan	147.37
+pbs	Central Pame	pbs	9.69
+pbt	Southern Pashto	pbt	11.6
+pbu	Northern Pashto	pbu	11.03
+pcm	Nigerian Pidgin	pcm	11.04
+pex	Petats	pex	10.2
+phl	Phalura	phl	20.69
+phr	Pahari-Potwari	phr	24.03
+pip	Pero	pip	9.85
+piy	Piya-Kwonci	piy	10.38
+pko	Pökoot	pko	10.4
+pl	Polish	pol	911.68
+plk	Kohistani Shina	plk	12.75
+plt	Plateau Malagasy	plt	19.39
+pmq	Northern Pame	pmq	10.24
+pms	Piemontese	pms	16.01
+pmy	Papuan Malay	pmy	10.17
+pnb	Western Panjabi	pnb	10.0
+poc	Poqomam	poc	9.63
+poe	San Juan Atzingo Popoloca	poe	10.01
+pow	San Felipe Otlaltepec Popoloca	pow	8.84
+prq	Ashéninka Perené	prq	7.16
+ps	Pushto	pus	88.62
+pst	Central Pashto	pst	11.4
+pt	Portuguese	por	16855.05
+pua	Western Highland Purepecha	pua	10.17
+pwn	Paiwan	pwn	13.76
+qug	Chimborazo Highland Quichua	qug	10.12
+qum	Sipacapense	qum	9.37
+qup	Southern Pastaza Quechua	qup	11.13
+qur	Yanahuanca Pasco Quechua	qur	9.95
+qus	Santiago del Estero Quichua	qus	9.55
+quv	Sacapulteco	quv	8.9
+qux	Yauyos Quechua	qux	9.35
+quy	Ayacucho Quechua	quy	0.05
+qva	Ambo-Pasco Quechua	qva	9.59
+qvi	Imbabura Highland Quichua	qvi	11.0
+qvj	Loja Highland Quichua	qvj	10.59
+qvl	Cajatambo North Lima Quechua	qvl	9.95
+qwa	Corongo Ancash Quechua	qwa	9.72
+qws	Sihuas Ancash Quechua	qws	10.18
+qxa	Chiquián Ancash Quechua	qxa	9.99
+qxp	Puno Quechua	qxp	9.81
+qxt	Santa Ana de Tusi Pasco Quechua	qxt	10.05
+qxu	Arequipa-La Unión Quechua	qxu	10.12
+qxw	Jauja Wanca Quechua	qxw	11.42
+rag	Logooli	rag	9.39
+rm	Romansh	roh	9.21
+ro	Romanian	ron	70.23
+rob	Tae'	rob	9.02
+rof	Rombo	rof	18.9
+roo	Rotokas	roo	9.07
+rth	Ratahan	rth	9.34
+ru	Russian	rus	20338.5
+rup	Macedo-Romanian	rup	0.02
+rw	Kinyarwanda	kin	2021.66
+sa	Sanskrit	san	84.44
+sah	Yakut	sah	16.08
+sat	Santali	sat	98.37
+sau	Saleman	sau	10.53
+say	Saya	say	10.02
+sbn	Sindhi Bhil	sbn	10.53
+sc	Sardinian	srd	2.77
+scl	Shina	scl	9.84
+scn	Sicilian	scn	13.35
+sd	Sindhi	snd	46.27
+sei	Seri	sei	9.81
+shu	Chadian Arabic	shu	2.29
+si	Sinhala	sin	11.98
+sip	Sikkimese	sip	10.07
+siw	Siwai	siw	10.47
+sjr	Siar-Lak	sjr	9.87
+sk	Slovak	slk	2478.46
+skg	Sakalava Malagasy	skg	9.02
+skr	Saraiki	skr	4.13
+sl	Slovenian	slv	1172.61
+sn	Shona	sna	9.96
+snc	Sinaugoro	snc	10.38
+snk	Soninke	snk	10.04
+so	Somali	som	13.22
+sol	Solos	sol	9.95
+sps	Saposa	sps	9.81
+sq	Albanian	sqi	8.59
+sr	Serbian	srp	1855.33
+src	Logudorese Sardinian	src	10.67
+sro	Campidanese Sardinian	sro	10.16
+ssi	Sansi	ssi	10.47
+ste	Liana-Seti	ste	10.43
+sua	Sulka	sua	10.12
+sv	Swedish	swe	2453.14
+sva	Svan	sva	15.11
+sw	Swahili	swa	418.41
+szy	Sakizaya	szy	11.47
+ta	Tamil	tam	423.09
+tan	Tangale	tan	10.14
+tar	Central Tarahumara	tar	9.73
+tay	Atayal	tay	7.02
+tbf	Mandara	tbf	10.01
+tcf	Malinaltepec Me'phaa	tcf	9.04
+tcy	Tulu	tcy	11.72
+tdn	Tondano	tdn	9.14
+tdx	Tandroy-Mahafaly Malagasy	tdx	3.81
+te	Telugu	tel	230.21
+tg	Tajik	tgk	9.23
+tgc	Tigak	tgc	9.71
+th	Thai	tha	10499.77
+the	Chitwania Tharu	the	10.06
+thq	Kochila Tharu	thq	10.28
+thr	Rana Tharu	thr	9.99
+thv	Tahaggart Tamahaq	thv	4.25
+ti	Tigrinya	tir	0.08
+tig	Tigre	tig	7.49
+tio	Teop	tio	9.85
+tk	Turkmen	tuk	2.86
+tkg	Tesaka Malagasy	tkg	17.86
+tkt	Kathoriya Tharu	tkt	10.64
+tli	Tlingit	tli	0.41
+tlp	Filomena Mata-Coahuitlán Totonac	tlp	11.35
+tn	Tswana	tsn	4.24
+tok	Toki Pona	tok	13.51
+tpl	Tlacoapa Me'phaa	tpl	9.28
+tpz	Tinputz	tpz	9.33
+tqp	Tomoip	tqp	10.1
+tr	Turkish	tur	125.36
+trp	Kok Borok	trp	10.74
+trq	San Martín Itunyoso Triqui	trq	8.29
+trv	Sediq	trv	7.77
+trw	Torwali	trw	14.98
+tt	Tatar	tat	30.03
+ttj	Tooro	ttj	10.31
+ttr	Tera	ttr	9.89
+ttu	Torau	ttu	9.87
+tui	Tupuri	tui	9.26
+tul	Tula	tul	9.79
+tuq	Tedaga	tuq	10.0
+tuv	Turkana	tuv	10.17
+tuy	Tugen	tuy	8.79
+tvo	Tidore	tvo	10.31
+tvu	Tunen	tvu	9.85
+tw	Twi	twi	0.25
+twu	Termanu	twu	11.45
+txs	Tonsea	txs	9.32
+txy	Tanosy Malagasy	txy	12.07
+udl	Wuzlam	udl	9.23
+ug	Uighur	uig	428.77
+uk	Ukrainian	ukr	1851.97
+uki	Kui (India)	uki	10.77
+umb	Umbundu	umb	10.59
+ur	Urdu	urd	211.27
+ush	Ushojo	ush	6.36
+uz	Uzbek	uzb	115.28
+uzn	Northern Uzbek	uzn	15.23
+vai	Vai	vai	8.76
+var	Huarijio	var	9.28
+ver	Mom Jango	ver	10.93
+vi	Vietnamese	vie	8481.98
+vmc	Juxtlahuaca Mixtec	vmc	9.43
+vmj	Ixtayutla Mixtec	vmj	10.17
+vmm	Mitlatongo Mixtec	vmm	9.95
+vmp	Soyaltepec Mazatec	vmp	10.17
+vmz	Mazatlán Mazatec	vmz	9.82
+vot	Votic	vot	0.1
+vro	Võro	vro	15.66
+wbl	Wakhi	wbl	11.67
+wci	Waci Gbe	wci	8.02
+weo	Wemale	weo	9.09
+wes	Cameroon Pidgin	wes	10.06
+wja	Waja	wja	10.22
+wji	Warji	wji	11.39
+wo	Wolof	wol	8.71
+wof	Gambian Wolof	wof	9.46
+xh	Xhosa	xho	13.35
+xhe	Khetrani	xhe	9.4
+xka	Kalkoti	xka	8.0
+xmf	Mingrelian	xmf	11.47
+xmv	Antankarana Malagasy	xmv	17.9
+xmw	Tsimihety Malagasy	xmw	11.53
+xpe	Liberia Kpelle	xpe	9.5
+xti	Sinicahua Mixtec	xti	9.5
+xtu	Cuyamecalco Mixtec	xtu	9.4
+yaq	Yaqui	yaq	9.93
+yav	Yangben	yav	8.7
+yay	Agwagwune	yay	8.26
+ydd	Eastern Yiddish	ydd	18.43
+ydg	Yidgha	ydg	9.89
+yer	Tarok	yer	10.08
+yes	Nyankpa	yes	10.26
+yi	Yiddish	yid	1.81
+yo	Yoruba	yor	15.66
+yue	Cantonese	yue	13302.38
+zga	Kinga	zga	9.5
+zgh	Standard Moroccan Tamazight	zgh	1.19
+zh	Chinese	cmn	111343.3
+zoc	Copainalá Zoque	zoc	10.07
+zoh	Chimalapa Zoque	zoh	9.35
+zor	Rayón Zoque	zor	9.04
+zpv	Chichicapan Zapotec	zpv	9.85
+zpy	Mazaltepec Zapotec	zpy	9.47
+ztg	Xanaguía Zapotec	ztg	9.86
+ztn	Santa Catarina Albarradas Zapotec	ztn	10.02
+ztp	Loxicha Zapotec	ztp	9.62
+zts	Tilquiapan Zapotec	zts	9.33
+ztu	Güilá Zapotec	ztu	9.17
+zu	Zulu	zul	14.83
+zza	Zaza	zza	1.52

runtime/omnivoice-python/omnivoice-source/examples/run_emilia.sh

@@ -0,0 +1,115 @@
+#!/bin/bash
+
+# This script demonstrates how to run the full training pipeline on the Emilia dataset.
+
+set -euo pipefail
+
+stage=0
+stop_stage=2
+
+# ====== Modify as needed ======
+# GPUs to use
+GPU_IDS="0,1,2,3,4,5,6,7"
+NUM_GPUS=8
+
+# Download directory for raw Emilia data
+dl_dir="download"
+
+# Directory containing JSONL manifests for train/dev splits
+# Stage 0 will check for the presence of the following files:
+#   data/emilia/manifests/emilia_en_train.jsonl
+#   data/emilia/manifests/emilia_en_dev.jsonl
+#   data/emilia/manifests/emilia_zh_train.jsonl
+#   data/emilia/manifests/emilia_zh_dev.jsonl
+MANIFEST_DIR="data/emilia/manifests"
+
+# Directory to write tokenized WebDataset shards
+TOKEN_DIR="data/emilia/tokens"
+
+# Audio tokenizer model (HuggingFace repo or local path)
+TOKENIZER_PATH="eustlb/higgs-audio-v2-tokenizer"
+
+# Training config file
+TRAIN_CONFIG="config/train_config_emilia.json"
+
+# Data config file
+data_config="config/data_config_emilia.json"
+
+# Output directory for checkpoints
+OUTPUT_DIR="exp/omnivoice_emilia"
+# =================================
+
+export PYTHONPATH="$(cd "$(dirname "$0")/.." && pwd):${PYTHONPATH:-}"
+
+
+# Stage 0: Download data
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+    echo "Stage 0: Download data"
+
+    # You should manually download the Emilia dataset from
+    # https://openxlab.org.cn/datasets/Amphion/Emilia
+    # or https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07
+    # and place it in the download directory.
+    # Your download directory should at least contain the following structure:
+    #
+    #    download/Amphion___Emilia
+    #    ├── raw
+    #    │   ├── EN
+    #    │   └── ZH
+
+    if [ ! -d "$dl_dir"/Amphion___Emilia/raw ]; then
+        echo "Please refer https://openxlab.org.cn/datasets/Amphion/Emilia to download the dataset."
+        exit 1
+    fi
+
+    # We require JSONL manifests for the training and dev splits. You can
+    # either generate them yourself using the raw data and the provided
+    # metadata, or download our processed JSONL manifests from HuggingFace.
+    # https://huggingface.co/datasets/zhu-han/Emilia-Manifests
+    #
+    # Place them as data/emilia/manifests/{emilia_en_train,emilia_en_dev,emilia_zh_train,emilia_zh_dev}.jsonl
+
+    for split in emilia_en_dev emilia_zh_dev emilia_en_train emilia_zh_train; do
+        if [ ! -f "${MANIFEST_DIR}/${split}.jsonl" ]; then
+            echo "Please download the manifest for ${split} and place it in ${MANIFEST_DIR}/${split}.jsonl"
+            exit 1
+        fi
+    done
+
+    echo "  Done. All manifests and data are in place."
+fi
+
+
+# Stage 1: Tokenize splits into directories matching data_config_emilia.json
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+    echo "Stage 1: Tokenizing audio"
+
+    for split in emilia_en_dev emilia_zh_dev emilia_en_train emilia_zh_train; do
+        echo "  Tokenizing ${split} from ${MANIFEST_DIR}/${split}.jsonl"
+
+        CUDA_VISIBLE_DEVICES=${GPU_IDS} \
+            python -m omnivoice.scripts.extract_audio_tokens \
+            --input_jsonl "${MANIFEST_DIR}/${split}.jsonl" \
+            --tar_output_pattern "${TOKEN_DIR}/${split}/audios/shard-%06d.tar" \
+            --jsonl_output_pattern "${TOKEN_DIR}/${split}/txts/shard-%06d.jsonl" \
+            --tokenizer_path "${TOKENIZER_PATH}" \
+            --nj_per_gpu 3 \
+            --shuffle True
+
+        echo "  Done. Tokens written to ${TOKEN_DIR}/${split}"
+    done
+fi
+
+
+# Stage 2: Train
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
+    echo "Stage 2: Training"
+
+    accelerate launch \
+        --gpu_ids "${GPU_IDS}" \
+        --num_processes ${NUM_GPUS} \
+        -m omnivoice.cli.train \
+        --train_config ${TRAIN_CONFIG} \
+        --data_config ${data_config} \
+        --output_dir ${OUTPUT_DIR}
+fi

runtime/omnivoice-python/omnivoice-source/examples/run_eval.sh

@@ -0,0 +1,283 @@
+#!/bin/bash
+
+# Evaluate OmniVoice models on TTS benchmarks.
+
+# Stage 1: Download the test sets and evaluation models.
+# Stage 2: LibriSpeech-PC
+# Stage 3: seedtts_en
+# Stage 4: seedtts_zh
+# Stage 5: fleurs
+# Stage 6: minimax
+
+set -euo pipefail
+
+# Specify the stages to run by setting the `stage` and `stop_stage` variables. 
+stage=1
+stop_stage=6
+
+# Available GPUs for evaluation. Adjust this according to your setup.
+export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
+
+# Specify the checkpoint to evaluate.
+CHECKPOINT=k2-fsa/OmniVoice
+emilia_checkpoint=false
+
+# CHECKPOINT=k2-fsa/OmniVoice
+# emilia_checkpoint=true
+
+# For the OmniVoice-Emilia checkpoint, we set denoise to False and lang_id to None
+#, as the model is trained without prompt denoising or language id.
+
+if [ "${emilia_checkpoint}" = true ]; then
+    infer_options="--preprocess_prompt False \
+        --postprocess_output False \
+        --batch_duration 600 \
+        --denoise False \
+        --lang_id None \
+        --audio_chunk_threshold 1000"
+else
+    infer_options="--preprocess_prompt False \
+        --postprocess_output False \
+        --batch_duration 600 \
+        --audio_chunk_threshold 1000"
+fi
+
+export PYTHONPATH="$(cd "$(dirname "$0")/.." && pwd):${PYTHONPATH:-}"
+
+download_dir="download"
+TTS_EVAL_MODEL_DIR="${download_dir}/tts_eval_models/"
+TTS_EVAL_DATA_DIR="${download_dir}/tts_eval_datasets/"
+
+# Map test_name to its test.jsonl path.
+get_test_list() {
+    case "$1" in
+        librispeech_pc) echo "${TTS_EVAL_DATA_DIR}/librispeech_pc_test_clean.jsonl" ;;
+        seedtts_en)     echo "${TTS_EVAL_DATA_DIR}/seedtts_test_en.jsonl" ;;
+        seedtts_zh)     echo "${TTS_EVAL_DATA_DIR}/seedtts_test_zh.jsonl" ;;
+        minimax)        echo "${TTS_EVAL_DATA_DIR}/minimax_multilingual_24.jsonl" ;;
+        fleurs)         echo "${TTS_EVAL_DATA_DIR}/fleurs_multilingual_102.jsonl" ;;
+        *)              echo ""; return 1 ;;
+    esac
+}
+
+# ============================================================
+# Stage 1: Prepare the test sets and evaluation models
+# ============================================================
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+    echo "Stage 1: Download test sets and evaluation models"
+
+    hf_repo=k2-fsa/TTS_eval_datasets
+    mkdir -p ${TTS_EVAL_DATA_DIR}/
+    for file in \
+        librispeech_pc_test_clean.jsonl \
+        librispeech_pc_test_clean_transcript.jsonl \
+        seedtts_test_en.jsonl \
+        seedtts_test_zh.jsonl \
+        minimax_multilingual_24.jsonl \
+        fleurs_multilingual_102.jsonl; do
+        echo "Downloading ${file}..."
+        huggingface-cli download \
+                --repo-type dataset \
+                --local-dir ${TTS_EVAL_DATA_DIR}/ \
+                ${hf_repo} \
+                ${file}
+    done
+
+    for file in \
+        librispeech_pc_testset.tar.gz \
+        seedtts_testset.tar.gz \
+        minimax_multilingual_24.tar.gz \
+        fleurs_multilingual_102.tar.gz; do
+        echo "Downloading ${file}..."
+        huggingface-cli download \
+                --repo-type dataset \
+                --local-dir ${TTS_EVAL_DATA_DIR}/ \
+                ${hf_repo} \
+                ${file}
+
+        echo "Extracting ${file}..."
+        tar -xzf ${TTS_EVAL_DATA_DIR}/${file} -C ${TTS_EVAL_DATA_DIR}/
+    done
+
+    echo "Download all evaluation models"
+    hf_repo=k2-fsa/TTS_eval_models
+    mkdir -p ${TTS_EVAL_MODEL_DIR}
+    huggingface-cli download \
+        --local-dir ${TTS_EVAL_MODEL_DIR} \
+        ${hf_repo}
+fi
+
+# ============================================================
+# Stage 2: Evaluation on LibriSpeech-PC
+# ============================================================
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+    echo "Stage 2: Evaluation on LibriSpeech-PC"
+    wav_path="results/librispeech_pc"
+    test_jsonl="$(get_test_list librispeech_pc)"
+    transcript_jsonl="${TTS_EVAL_DATA_DIR}/librispeech_pc_test_clean_transcript.jsonl"
+
+    python -m omnivoice.cli.infer_batch \
+        --model "${CHECKPOINT}" \
+        --test_list "${test_jsonl}" \
+        --res_dir "${wav_path}" ${infer_options}
+
+    python -m omnivoice.eval.speaker_similarity.sim \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.sim.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+
+    python -m omnivoice.eval.wer.hubert \
+        --wav-path "${wav_path}" \
+        --test-list "${transcript_jsonl}" \
+        --decode-path "${wav_path}.wer.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+
+    python -m omnivoice.eval.mos.utmos \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.mos.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+fi
+
+
+# ============================================================
+# Stage 3: Evaluation on Seed-TTS en
+# ============================================================
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+    echo "Stage 3: Evaluation on Seed-TTS en"
+    wav_path="results/seedtts_en"
+    test_jsonl="$(get_test_list seedtts_en)"
+
+    python -m omnivoice.cli.infer_batch \
+        --model "${CHECKPOINT}" \
+        --test_list "${test_jsonl}" \
+        --res_dir "${wav_path}"  ${infer_options}
+
+
+    python -m omnivoice.eval.speaker_similarity.sim \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.sim.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+
+    python -m omnivoice.eval.wer.seedtts \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.wer.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}" \
+        --lang en
+
+    python -m omnivoice.eval.mos.utmos \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.mos.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+fi
+
+
+# ============================================================
+# Stage 4: Evaluation on Seed-TTS zh
+# ============================================================
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+    echo "Stage 4: Evaluation on Seed-TTS zh"
+    wav_path="results/seedtts_zh"
+    test_jsonl="$(get_test_list seedtts_zh)"
+
+    python -m omnivoice.cli.infer_batch \
+        --model "${CHECKPOINT}" \
+        --test_list "${test_jsonl}" \
+        --res_dir "${wav_path}"  ${infer_options}
+
+
+    python -m omnivoice.eval.speaker_similarity.sim \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.sim.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+
+    python -m omnivoice.eval.wer.seedtts \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.wer.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}" \
+        --lang zh
+
+    python -m omnivoice.eval.mos.utmos \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.mos.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+fi
+
+
+
+# ============================================================
+# Stage 5: Evaluation on MiniMax multilingual
+# ============================================================
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+    echo "Stage 5: Evaluation on MiniMax multilingual"
+    wav_path="results/minimax"
+    test_jsonl="$(get_test_list minimax)"
+
+    python -m omnivoice.cli.infer_batch \
+        --model "${CHECKPOINT}" \
+        --test_list "${test_jsonl}" \
+        --res_dir "${wav_path}"  ${infer_options}
+
+    python -m omnivoice.eval.speaker_similarity.sim \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.sim.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+
+    python -m omnivoice.eval.wer.minimax \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.wer.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+fi
+
+
+# ============================================================
+# Stage 6: Evaluation on FLEURS multilingual
+# ============================================================
+
+if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
+    echo "Stage 6: Evaluation on FLEURS multilingual"
+    wav_path="results/fleurs"
+    test_jsonl="$(get_test_list fleurs)"
+
+    python -m omnivoice.cli.infer_batch \
+        --model "${CHECKPOINT}" \
+        --test_list "${test_jsonl}" \
+        --res_dir "${wav_path}"  ${infer_options}
+
+
+    python -m omnivoice.eval.speaker_similarity.sim \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.sim.log" \
+        --model-dir "${TTS_EVAL_MODEL_DIR}"
+
+    # Evaluation on FLEURS requires omnilingual-asr, which has dependencies that
+    # conflict with other packages (at least the transformers package) in our project.
+
+    # To evaluate on FLEURS, we suggest users to set up a separate virtual
+    # environment to install omnilingual-asr. Install instructions can be found in
+    # https://github.com/facebookresearch/omnilingual-asr
+
+    python ${PWD}/../omnivoice/eval/wer/fleurs.py \
+        --wav-path "${wav_path}" \
+        --test-list "${test_jsonl}" \
+        --decode-path "${wav_path}.wer.log" \
+        --model-card omniASR_LLM_Unlimited_7B_v2 \
+        --chunk-size 100 \
+        --batch-size 50
+fi

runtime/omnivoice-python/omnivoice-source/examples/run_finetune.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+
+# This script demonstrates how to fine-tune OmniVoice from a JSONL manifest.
+
+set -euo pipefail
+
+stage=0
+stop_stage=1
+
+# ====== Modify as needed ======
+# GPUs to use
+GPU_IDS="0,1"
+NUM_GPUS=2
+
+# Path to your input JSONL file
+# (each line: {"id": ..., "audio_path": ..., "text": ..., "language_id": ...})
+TRAIN_JSONL="data/my_data_train.jsonl"
+
+# Path to your dev JSONL file. Set to empty string to skip dev set.
+DEV_JSONL="data/my_data_dev.jsonl"
+
+# Directory to write tokenized WebDataset shards
+TOKEN_DIR="data/finetune/tokens"
+
+# Audio tokenizer model (HuggingFace repo or local path)
+TOKENIZER_PATH="eustlb/higgs-audio-v2-tokenizer"
+
+# Training config file
+# If you encounter issues with flex_attention on your GPU, use the SDPA config instead:
+# TRAIN_CONFIG="config/train_config_finetune_sdpa.json"
+TRAIN_CONFIG="config/train_config_finetune.json"
+
+# Data config file
+data_config="config/data_config_finetune.json"
+
+# Output directory for fine-tuned checkpoints
+OUTPUT_DIR="exp/omnivoice_finetune"
+# =================================
+
+export PYTHONPATH="$(cd "$(dirname "$0")/.." && pwd):${PYTHONPATH:-}"
+
+
+# Stage 0: Tokenize audio into WebDataset shards
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+    echo "Stage 0: Tokenizing audio"
+
+    for split_jsonl_path in ${TRAIN_JSONL} ${DEV_JSONL}; do
+        if [ -z "${split_jsonl_path}" ]; then
+            continue
+        fi
+
+        if [ "${split_jsonl_path}" = "${TRAIN_JSONL}" ]; then
+            split="train"
+        else
+            split="dev"
+        fi
+
+        echo "  Tokenizing ${split} from ${split_jsonl_path}"
+
+        CUDA_VISIBLE_DEVICES=${GPU_IDS} \
+            python -m omnivoice.scripts.extract_audio_tokens \
+            --input_jsonl "${split_jsonl_path}" \
+            --tar_output_pattern "${TOKEN_DIR}/${split}/audios/shard-%06d.tar" \
+            --jsonl_output_pattern "${TOKEN_DIR}/${split}/txts/shard-%06d.jsonl" \
+            --tokenizer_path "${TOKENIZER_PATH}" \
+            --nj_per_gpu 3 \
+            --shuffle True
+
+        echo "  Done. Manifest written to ${TOKEN_DIR}/${split}/data.lst"
+    done
+fi
+
+
+# Stage 1: Fine-tune
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+    echo "Stage 1: Fine-tuning"
+
+    accelerate launch \
+        --gpu_ids "${GPU_IDS}" \
+        --num_processes ${NUM_GPUS} \
+        -m omnivoice.cli.train \
+        --train_config ${TRAIN_CONFIG} \
+        --data_config ${data_config} \
+        --output_dir ${OUTPUT_DIR}
+fi

runtime/omnivoice-python/omnivoice-source/omnivoice/eval/wer/punctuations.lst

@@ -0,0 +1,188 @@
+	7355	INVALID UNICODE	0x81
+	5265	INVALID UNICODE	0x90
+	75	INVALID UNICODE	0x8
+	31	INVALID UNICODE	0x8d
+”	3	INVALID UNICODE	0x94
+	2	INVALID UNICODE	0x8f
+	2	INVALID UNICODE	0x1a
+	1	INVALID UNICODE	0x9d
+“	1	INVALID UNICODE	0x93
+’	1	INVALID UNICODE	0x92
+	8647	INVALID UNICODE	0xe295
+	6650	INVALID UNICODE	0xf21d
+	6234	INVALID UNICODE	0xf62d
+	4815	INVALID UNICODE	0xf173
+	4789	INVALID UNICODE	0xe514
+	4409	INVALID UNICODE	0xe293
+	3881	INVALID UNICODE	0xf523
+	3788	INVALID UNICODE	0xe233
+	2448	INVALID UNICODE	0xf50f
+	2177	INVALID UNICODE	0xe232
+	1955	INVALID UNICODE	0xea7b
+	1926	INVALID UNICODE	0xf172
+	973	INVALID UNICODE	0xe290
+	972	INVALID UNICODE	0xf519
+	661	INVALID UNICODE	0xe292
+	591	INVALID UNICODE	0xe328
+	509	INVALID UNICODE	0xe2fa
+	458	INVALID UNICODE	0xe234
+	446	INVALID UNICODE	0xe043
+	419	INVALID UNICODE	0xe040
+	399	INVALID UNICODE	0xe2fb
+	387	INVALID UNICODE	0xe32b
+	381	INVALID UNICODE	0xe236
+	374	INVALID UNICODE	0xf511
+	314	INVALID UNICODE	0xe517
+	296	INVALID UNICODE	0xe2fe
+	293	INVALID UNICODE	0xe492
+	291	INVALID UNICODE	0xf52d
+	289	INVALID UNICODE	0xe2fc
+	195	INVALID UNICODE	0xf521
+	190	INVALID UNICODE	0xe516
+	182	INVALID UNICODE	0xe041
+	178	INVALID UNICODE	0xf529
+	113	INVALID UNICODE	0xe2f9
+	87	INVALID UNICODE	0xe2d9
+	78	INVALID UNICODE	0xe32a
+	76	INVALID UNICODE	0xe291
+	74	INVALID UNICODE	0xe296
+	66	INVALID UNICODE	0xe518
+	52	INVALID UNICODE	0xe32c
+	46	INVALID UNICODE	0xe2db
+	41	INVALID UNICODE	0xe231
+	34	INVALID UNICODE	0xf522
+	33	INVALID UNICODE	0xf518
+	32	INVALID UNICODE	0xf513
+	27	INVALID UNICODE	0xe32d
+	25	INVALID UNICODE	0xe32e
+	23	INVALID UNICODE	0xe06b
+	15	INVALID UNICODE	0xea01
+	12	INVALID UNICODE	0xe294
+	11	INVALID UNICODE	0xe203
+	8	INVALID UNICODE	0xf218
+	7	INVALID UNICODE	0xe070
+	7	INVALID UNICODE	0xe013
+	5	INVALID UNICODE	0xe2de
+	4	INVALID UNICODE	0xe493
+	3	INVALID UNICODE	0xf7e8
+	3	INVALID UNICODE	0xf7d0
+	3	INVALID UNICODE	0xe313
+	2	INVALID UNICODE	0xe329
+	2	INVALID UNICODE	0xe06d
+	2	INVALID UNICODE	0xe003
+	1	INVALID UNICODE	0xf50e
+	1	INVALID UNICODE	0xf171
+	1	INVALID UNICODE	0xe01d
+	71	NOMINAL DIGIT SHAPES	0x206f
+⁠	3	WORD JOINER	0x2060
+―	126545	HORIZONTAL BAR	0x2015
+־	1028	HEBREW PUNCTUATION MAQAF	0x5be
+)	98429	RIGHT PARENTHESIS	0x29
+]	27108	RIGHT SQUARE BRACKET	0x5d
+⌋	1567	RIGHT FLOOR	0x230b
+〕	97	RIGHT TORTOISE SHELL BRACKET	0x3015
+】	36	RIGHT BLACK LENTICULAR BRACKET	0x3011
+﴾	14	ORNATE LEFT PARENTHESIS	0xfd3e
+&	170517	AMPERSAND	0x26
+།	106330	TIBETAN MARK SHAD	0xf0d
+።	90203	ETHIOPIC FULL STOP	0x1362
+፥	60484	ETHIOPIC COLON	0x1365
+༌	60464	TIBETAN MARK DELIMITER TSHEG BSTAR	0xf0c
+။	51567	MYANMAR SIGN SECTION	0x104b
+/	46929	SOLIDUS	0x2f
+၊	38042	MYANMAR SIGN LITTLE SECTION	0x104a
+·	37985	MIDDLE DOT	0xb7
+‸	36310	CARET	0x2038
+*	34793	ASTERISK	0x2a
+۔	32432	ARABIC FULL STOP	0x6d4
+፤	31906	ETHIOPIC SEMICOLON	0x1364
+၏	21519	MYANMAR SYMBOL GENITIVE	0x104f
+។	20834	KHMER SIGN KHAN	0x17d4
+꓾	15773	LISU PUNCTUATION COMMA	0xa4fe
+᙮	13473	CANADIAN SYLLABICS FULL STOP	0x166e
+꤯	12892	KAYAH LI SIGN SHYA	0xa92f
+⵰	11478	TIFINAGH SEPARATOR MARK	0x2d70
+꓿	11118	LISU PUNCTUATION FULL STOP	0xa4ff
+॥	10763	DEVANAGARI DOUBLE DANDA	0x965
+؞	10403	ARABIC TRIPLE DOT PUNCTUATION MARK	0x61e
+၍	8936	MYANMAR SYMBOL COMPLETED	0x104d
+·	8431	GREEK ANO TELEIA	0x387
+†	7477	DAGGER	0x2020
+၌	6632	MYANMAR SYMBOL LOCATIVE	0x104c
+፣	5719	ETHIOPIC COMMA	0x1363
+៖	5528	KHMER SIGN CAMNUC PII KUUH	0x17d6
+꤮	4791	KAYAH LI SIGN CWI	0xa92e
+※	3439	REFERENCE MARK	0x203b
+፦	2727	ETHIOPIC PREFACE COLON	0x1366
+•	1749	BULLET	0x2022
+¶	1507	PILCROW SIGN	0xb6
+၎	1386	MYANMAR SYMBOL AFOREMENTIONED	0x104e
+﹖	1224	SMALL QUESTION MARK	0xfe56
+;	975	GREEK QUESTION MARK	0x37e
+…	827	HORIZONTAL ELLIPSIS	0x2026
+%	617	PERCENT SIGN	0x25
+・	468	KATAKANA MIDDLE DOT	0x30fb
+༎	306	TIBETAN MARK NYIS SHAD	0xf0e
+‡	140	DOUBLE DAGGER	0x2021
+#	137	NUMBER SIGN	0x23
+@	125	COMMERCIAL AT	0x40
+፡	121	ETHIOPIC WORDSPACE	0x1361
+៚	55	KHMER SIGN KOOMUUT	0x17da
+៕	49	KHMER SIGN BARIYOOSAN	0x17d5
+﹐	10	SMALL COMMA	0xfe50
+༅	6	TIBETAN MARK CLOSING YIG MGO SGAB MA	0xf05
+༄	6	TIBETAN MARK INITIAL YIG MGO MDUN MA	0xf04
+．	2	FULLWIDTH FULL STOP	0xff0e
+﹗	2	SMALL EXCLAMATION MARK	0xfe57
+﹕	2	SMALL COLON	0xfe55
+‰	2	PER MILLE SIGN	0x2030
+･	1	HALFWIDTH KATAKANA MIDDLE DOT	0xff65
+(	98504	LEFT PARENTHESIS	0x28
+[	27245	LEFT SQUARE BRACKET	0x5b
+⌊	1567	LEFT FLOOR	0x230a
+〔	95	LEFT TORTOISE SHELL BRACKET	0x3014
+【	36	LEFT BLACK LENTICULAR BRACKET	0x3010
+﴿	14	ORNATE RIGHT PARENTHESIS	0xfd3f
+_	4851	LOW LINE	0x5f
+$	72	DOLLAR SIGN	0x24
+€	14	EURO SIGN	0x20ac
+£	2	POUND SIGN	0xa3
+~	27462	TILDE	0x7e
+=	11450	EQUALS SIGN	0x3d
+|	8430	VERTICAL LINE	0x7c
+−	3971	MINUS SIGN	0x2212
+≫	1904	MUCH GREATER-THAN	0x226b
+≪	1903	MUCH LESS-THAN	0x226a
++	1450	PLUS SIGN	0x2b
+＜	345	FULLWIDTH LESS-THAN SIGN	0xff1c
+＞	344	FULLWIDTH GREATER-THAN SIGN	0xff1e
+¬	5	NOT SIGN	0xac
+×	4	MULTIPLICATION SIGN	0xd7
+→	2	RIGHTWARDS ARROW	0x2192
+᙭	537	CANADIAN SYLLABICS CHI SIGN	0x166d
+°	499	DEGREE SIGN	0xb0
+႟	421	MYANMAR SYMBOL SHAN EXCLAMATION	0x109f
+�	192	REPLACEMENT CHARACTER	0xfffd
+⌟	54	BOTTOM RIGHT CORNER	0x231f
+⌞	54	BOTTOM LEFT CORNER	0x231e
+©	2	COPYRIGHT SIGN	0xa9
+ 	40	NARROW NO-BREAK SPACE	0x202f
+ 	1	SIX-PER-EM SPACE	0x2006
+˜	40261	SMALL TILDE	0x2dc
+^	6469	CIRCUMFLEX ACCENT	0x5e
+¯	20	MACRON	0xaf
+ˇ	191442	CARON	0x2c7
+ⁿ	38144	SUPERSCRIPT LATIN SMALL LETTER N	0x207f
+ـ	9440	ARABIC TATWEEL	0x640
+ๆ	6766	THAI CHARACTER MAIYAMOK	0xe46
+ៗ	3310	KHMER SIGN LEK TOO	0x17d7
+々	678	IDEOGRAPHIC ITERATION MARK	0x3005
+ໆ	430	LAO KO LA	0xec6
+ー	319	KATAKANA-HIRAGANA PROLONGED SOUND MARK	0x30fc
+ⁱ	137	SUPERSCRIPT LATIN SMALL LETTER I	0x2071
+৷	11056	BENGALI CURRENCY NUMERATOR FOUR	0x9f7
+⅓	26	VULGAR FRACTION ONE THIRD	0x2153
+½	26	VULGAR FRACTION ONE HALF	0xbd
+¼	4	VULGAR FRACTION ONE QUARTER	0xbc
+⅟	1	FRACTION NUMERATOR ONE	0x215f
+⁄	57	FRACTION SLASH	0x2044

runtime/omnivoice-python/omnivoice-source/omnivoice/scripts/extract_audio_tokens.py

@@ -0,0 +1,625 @@
+#!/usr/bin/env python3
+# Copyright    2026  Xiaomi Corp.        (authors:  Han Zhu)
+#
+# See ../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Extract audio tokens from audio data and pack them into WebDataset shards.
+
+Supports two input modes:
+
+1. WebDataset manifest (data.lst):
+    python extract_audio_tokens.py \
+        --input_manifest data.lst \
+        --tar_output_pattern output/audios/shard-%06d.tar \
+        --jsonl_output_pattern output/txts/shard-%06d.jsonl
+
+2. Raw JSONL (each line: {"id": "...", "audio_path": "...", "text": "...", ...}):
+    python extract_audio_tokens.py \
+        --input_jsonl data.jsonl \
+        --tar_output_pattern output/audios/shard-%06d.tar \
+        --jsonl_output_pattern output/txts/shard-%06d.jsonl
+
+Output structure:
+    output_dir/
+    ├── audios/           # WebDataset tar shards (.npy audio tokens + .json metadata)
+    │   ├── shard_000000.tar
+    │   └── ...
+    ├── txts/             # Per-shard JSONL metadata
+    │   ├── shard_000000.jsonl
+    │   └── ...
+    ├── data.lst          # Manifest: <tar_path> <jsonl_path> <sample_count> <total_duration>
+    └── errors.jsonl      # Failed samples with error details
+"""
+
+import argparse
+import io
+import json
+import logging
+import multiprocessing as mp
+import os
+import warnings
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import webdataset as wds
+from torch.utils.data import DataLoader, IterableDataset
+from tqdm.auto import tqdm
+from transformers import AutoFeatureExtractor, HiggsAudioV2TokenizerModel
+
+from omnivoice.data.dataset import JsonlDatasetReader, WebDatasetReader
+from omnivoice.utils.common import str2bool
+
+warnings.filterwarnings(
+    "ignore", category=FutureWarning, module="torch.nn.utils.weight_norm"
+)
+
+HIGGS_INPUT_SAMPLE_RATE = 24_000
+
+
+# Global variables: Store tokenizer and device for each worker process
+worker_tokenizer = None
+worker_feature_extractor = None
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--input_manifest",
+        default=None,
+        help="Path to input dataset manifest (data.lst).",
+    )
+    parser.add_argument(
+        "--input_jsonl",
+        default=None,
+        help="Path to raw JSONL file (alternative to --input_manifest).",
+    )
+    parser.add_argument(
+        "--tar_output_pattern",
+        required=True,
+        help="Tar shard pattern passed to WebDataset",
+    )
+    parser.add_argument(
+        "--jsonl_output_pattern",
+        required=True,
+        help="Jsonl shard pattern passed to WebDataset",
+    )
+    parser.add_argument(
+        "--samples_per_shard",
+        type=int,
+        default=1000,
+        help="Maximum records per shard",
+    )
+    parser.add_argument(
+        "--min_num_shards",
+        type=int,
+        default=32,
+        help="Minimum number of output shards (use to ensure "
+        "shard count >= num_gpu * num_workers)",
+    )
+    parser.add_argument(
+        "--tokenizer_path",
+        type=str,
+        default="eustlb/higgs-audio-v2-tokenizer",
+        help="Path to audio tokenizer.",
+    )
+    parser.add_argument(
+        "--skip_errors", action="store_true", help="Skip items that fail to process"
+    )
+    parser.add_argument(
+        "--min_length",
+        type=float,
+        default=0.0,
+        help="Minimum audio duration in seconds (e.g. 2.0)",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=float,
+        default=float("inf"),
+        help="Maximum audio duration in seconds (e.g. 15.0)",
+    )
+    parser.add_argument(
+        "--num_machines",
+        type=int,
+        default=1,
+        help="Total number of machines for distributed runs",
+    )
+    parser.add_argument(
+        "--machine_index",
+        type=int,
+        default=0,
+        help="Zero-based machine index when distributing across multiple "
+        "machines (e.g. 0, 1, ... num_machines-1)",
+    )
+    parser.add_argument(
+        "--nj_per_gpu",
+        type=int,
+        default=3,
+        help="Number of worker processes to spawn per GPU.",
+    )
+    parser.add_argument(
+        "--loader_workers",
+        type=int,
+        default=24,
+        help="Number of DataLoader workers for streaming IterableDataset.",
+    )
+    parser.add_argument(
+        "--shuffle",
+        type=str2bool,
+        default=True,
+        help="Shuffle data by default.",
+    )
+    parser.add_argument(
+        "--shuffle-seed",
+        type=int,
+        default=42,
+        help="Random seed for shuffle (default: 42).",
+    )
+    return parser
+
+
+def count_lines(path):
+    with open(path, "rb") as f:
+        return sum(buf.count(b"\n") for buf in iter(lambda: f.read(1 << 20), b""))
+
+
+def serialise_numpy(key: str, tokens: np.ndarray) -> dict:
+    buffer = io.BytesIO()
+    np.save(buffer, tokens)
+    return {"__key__": key, "npy": buffer.getvalue()}
+
+
+def process_init(rank_queue, tokenizer_path):
+    """
+    Initialization function for each worker process.
+    Assigns a specific GPU to the process and loads the tokenizer.
+    """
+    global worker_tokenizer, worker_feature_extractor
+
+    # Configure worker process logging
+    formatter = (
+        "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d]"
+        " [Worker %(process)d] %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO, force=True)
+
+    # Get assigned GPU rank
+    rank = rank_queue.get()
+    # Determine device
+    if rank != -1 and torch.cuda.is_available():
+        worker_device = torch.device(f"cuda:{rank}")
+    else:
+        worker_device = torch.device("cpu")
+
+    logging.debug(f"Worker process initialized with device: {worker_device}")
+    # Load tokenizer onto the specified device
+    worker_feature_extractor = AutoFeatureExtractor.from_pretrained(tokenizer_path)
+    worker_tokenizer = HiggsAudioV2TokenizerModel.from_pretrained(
+        tokenizer_path, device_map=worker_device
+    )
+    logging.debug(f"Tokenizer loaded successfully on device {worker_device}")
+
+
+def process_single_sample(sample: dict[str, Any]) -> dict[str, Any]:
+    """
+    Single-sample processing function executed in worker processes.
+    Skips invalid samples during streaming processing.
+    """
+    try:
+        audio_tensor = sample.get("audio", None)  # shape (1, T)
+        if audio_tensor is None:
+            raise ValueError("Sample missing 'audio' field")
+
+        with torch.inference_mode():
+            key = sample["label"]["id"]
+            inputs = worker_feature_extractor(
+                raw_audio=audio_tensor.squeeze(0).numpy(),
+                sampling_rate=HIGGS_INPUT_SAMPLE_RATE,
+                return_tensors="pt",
+            ).to(worker_tokenizer.device)
+            audio_tokens = worker_tokenizer.encode(
+                inputs["input_values"],
+            ).audio_codes.squeeze(0)
+
+            assert len(audio_tokens.shape) == 2
+            assert audio_tokens.size(0) == 8
+
+            num_tokens = audio_tokens.size(1)
+            metadata = sample["label"]
+            metadata["num_tokens"] = num_tokens
+
+            # Convert to numpy format for subsequent serialization (int16 to save space)
+            audio_tokens_np = audio_tokens.to(torch.int16).cpu().numpy()
+
+            return {
+                "status": "success",
+                "key": key,
+                "audio_tokens": audio_tokens_np,
+                "metadata": metadata,
+                "error_msg": None,
+            }
+    except Exception as e:
+        sample_id = sample.get("label", {}).get("id", "unknown")
+        logging.error(f"Failed to process sample {sample_id}: {e}")
+        return {
+            "status": "error",
+            "key": sample_id,
+            "audio_tokens": None,
+            "metadata": None,
+            "error_msg": str(e),
+        }
+
+
+def _normalise_value(value: Any) -> Any:
+    """Convert tensors and NumPy scalars to serialisable Python objects."""
+    if isinstance(value, torch.Tensor):
+        if value.ndim == 0:
+            return value.item()
+        return value.cpu().tolist()
+    if isinstance(value, np.generic):
+        return value.item()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    return value
+
+
+def _encode_metadata(metadata: dict[str, Any]) -> bytes:
+    cleaned: dict[str, Any] = {}
+    for key, value in metadata.items():
+        if value is None:
+            continue
+        cleaned[key] = _normalise_value(value)
+    return json.dumps(cleaned, ensure_ascii=False).encode("utf-8")
+
+
+class StreamingLengthFilteredDataset(IterableDataset):
+    def __init__(
+        self,
+        base_iterable,
+        min_len: float,
+        max_len: float,
+        sr: int,
+    ):
+        self.base_iterable = base_iterable
+        self.min_len = min_len
+        self.max_len = max_len
+        self.sr = sr
+        self.filtered_count = 0
+
+    def __iter__(self):
+        """Stream samples one by one and filter on the fly."""
+        for sample in self.base_iterable:
+            try:
+                duration = sample["audio"].size(-1) / self.sr
+                if self.min_len <= duration <= self.max_len:
+                    yield sample
+                else:
+                    self.filtered_count += 1
+                    logging.warning(
+                        f"Filtered sample (duration out of range): "
+                        f"{sample['label']['id']} ({duration:.2f}s)"
+                    )
+            except Exception as e:
+                logging.warning(f"Skipped invalid sample during streaming: {e}")
+                continue
+
+
+def main() -> None:
+    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+    logging.basicConfig(format=formatter, level=logging.INFO, force=True)
+    parser = build_parser()
+    args = parser.parse_args()
+    mp.set_start_method("spawn", force=True)
+
+    # Validate input arguments
+    assert bool(args.input_manifest) != bool(
+        args.input_jsonl
+    ), "Exactly one of --input_manifest or --input_jsonl must be provided."
+
+    if args.num_machines > 1:
+        assert (
+            0 <= args.machine_index < args.num_machines
+        ), f"machine_index {args.machine_index} must be in [0, {args.num_machines})"
+
+    # Build base dataset and count total samples based on input mode
+    if args.input_jsonl:
+        logging.info(f"Input mode: raw JSONL ({args.input_jsonl})")
+        total_samples = count_lines(args.input_jsonl)
+        base_dataset = JsonlDatasetReader(
+            args.input_jsonl,
+            sample_rate=HIGGS_INPUT_SAMPLE_RATE,
+            shuffle=args.shuffle,
+            shuffle_seed=args.shuffle_seed,
+        )
+        loader_workers = args.loader_workers
+    else:
+        logging.info(f"Input mode: WebDataset manifest ({args.input_manifest})")
+        manifest_num_lines = count_lines(args.input_manifest)
+        loader_workers = min(args.loader_workers, manifest_num_lines)
+        total_samples = 0
+        manifests = []
+        with open(args.input_manifest, "r", encoding="utf-8") as f:
+            for line_id, line in tqdm(
+                enumerate(f),
+                total=manifest_num_lines,
+                desc="Calculating dataset length",
+            ):
+                items = line.strip().split(" ")
+                tar_path, jsonl_path, num_items, duration = (
+                    items[0],
+                    items[1],
+                    int(items[2]),
+                    float(items[3]),
+                )
+                assert os.path.exists(tar_path), f"File {tar_path} does not exist."
+                assert os.path.exists(jsonl_path), f"File {jsonl_path} does not exist."
+                assert jsonl_path.endswith(
+                    ".jsonl"
+                ), f"File {jsonl_path} is not a .jsonl file."
+                if (
+                    args.num_machines > 1
+                    and line_id % args.num_machines != args.machine_index
+                ):
+                    continue
+                total_samples += num_items
+                manifests.append((tar_path, jsonl_path, num_items, duration))
+        logging.info(
+            f"Total shards: {manifest_num_lines}, "
+            f"Shards for current index: {len(manifests)}"
+        )
+        base_dataset = WebDatasetReader(
+            manifests=manifests,
+            sample_rate=HIGGS_INPUT_SAMPLE_RATE,
+            evaluation=True,
+        )
+
+    # Adjust samples_per_shard if min_num_shards would be violated
+    samples_per_shard = args.samples_per_shard
+    if total_samples > 0:
+        estimated_shards = max(
+            1, (total_samples + samples_per_shard - 1) // samples_per_shard
+        )
+        if estimated_shards < args.min_num_shards:
+            samples_per_shard = max(1, total_samples // args.min_num_shards)
+            logging.info(
+                f"Adjusted samples_per_shard from {args.samples_per_shard} to "
+                f"{samples_per_shard} to meet min_num_shards={args.min_num_shards} "
+                f"(total_samples={total_samples})"
+            )
+
+    # Apply length filter and create DataLoader
+    filtered_dataset = StreamingLengthFilteredDataset(
+        base_iterable=base_dataset,
+        min_len=args.min_length,
+        max_len=args.max_length,
+        sr=HIGGS_INPUT_SAMPLE_RATE,
+    )
+    dataloader = DataLoader(
+        dataset=filtered_dataset,
+        batch_size=None,
+        num_workers=loader_workers,
+        persistent_workers=loader_workers > 0,
+        pin_memory=False,
+    )
+
+    # Configure multi-GPU multi-process setup
+    num_devices = torch.cuda.device_count()
+    if num_devices == 0:
+        logging.warning("No GPUs detected - using CPU for processing")
+        num_processes = args.nj_per_gpu
+    else:
+        num_processes = num_devices * args.nj_per_gpu
+    logging.info(
+        f"GPU count: {num_devices}, Processes per GPU: {args.nj_per_gpu}, "
+        f"Total processes: {num_processes}"
+    )
+
+    # Shared GPU rank queue for process assignment
+    manager = mp.Manager()
+    rank_queue = manager.Queue()
+    for rank in list(range(num_devices)) * args.nj_per_gpu:
+        rank_queue.put(rank)
+    if num_devices == 0:
+        for _ in range(num_processes):
+            rank_queue.put(-1)
+
+    # Prepare output paths
+    tar_output_pattern = str(Path(args.tar_output_pattern).expanduser())
+    jsonl_output_pattern = str(Path(args.jsonl_output_pattern).expanduser())
+    Path(tar_output_pattern).parent.mkdir(parents=True, exist_ok=True)
+    Path(jsonl_output_pattern).parent.mkdir(parents=True, exist_ok=True)
+
+    # Determine output directory from tar_output_pattern
+    output_dir = Path(tar_output_pattern).parent.parent
+    error_log_path = str(output_dir / "errors.jsonl")
+    manifest_path = str(output_dir / "data.lst")
+
+    # Setup error logger (writes to errors.jsonl)
+    error_logger = logging.getLogger("error_log")
+    error_logger.setLevel(logging.ERROR)
+    error_logger.handlers.clear()
+    error_fh = logging.FileHandler(error_log_path, mode="w", encoding="utf-8")
+    error_fh.setFormatter(logging.Formatter("%(message)s"))
+    error_logger.addHandler(error_fh)
+
+    # Progress and error tracking
+    processed_count = 0
+    error_count = 0
+    write_error_count = 0
+    failed_ids = []
+    shard_idx = 0
+    shard_sample_count = 0
+    shard_duration = 0.0
+    shard_manifest = {}  # shard_idx -> (tar_path, jsonl_path, count, duration)
+
+    tar_writer = None
+    jsonl_file = None
+
+    def open_new_shard():
+        nonlocal tar_writer, jsonl_file, shard_idx, shard_sample_count, shard_duration
+        if tar_writer is not None:
+            tar_writer.close()
+        if jsonl_file is not None:
+            jsonl_file.close()
+        # Record manifest for the previous shard
+        if shard_idx > 0 and shard_sample_count > 0:
+            prev_idx = shard_idx - 1
+            shard_manifest[prev_idx] = (
+                os.path.abspath(tar_output_pattern % prev_idx),
+                os.path.abspath(jsonl_output_pattern % prev_idx),
+                shard_sample_count,
+                shard_duration,
+            )
+        tar_fname = tar_output_pattern % shard_idx
+        jsonl_fname = jsonl_output_pattern % shard_idx
+        tar_writer = wds.TarWriter(tar_fname)
+        jsonl_file = open(jsonl_fname, "w", encoding="utf-8")
+        shard_idx += 1
+        shard_sample_count = 0
+        shard_duration = 0.0
+
+    def write_sample(key, audio_tokens_np, metadata):
+        nonlocal shard_sample_count, write_error_count, shard_duration
+        assert tar_writer is not None and jsonl_file is not None
+        try:
+            token_record = serialise_numpy(key, audio_tokens_np)
+            json_record = _encode_metadata(metadata)
+            tar_writer.write(token_record)
+            jsonl_file.write(json_record.decode("utf-8") + "\n")
+            shard_sample_count += 1
+            shard_duration += metadata.get("audio_duration", 0.0)
+        except Exception as exc:
+            write_error_count += 1
+            failed_ids.append(key)
+            error_logger.error(
+                json.dumps({"id": key, "reason": str(exc)}, ensure_ascii=False)
+            )
+            logging.error(f"Write failed for sample {key}: {exc}")
+
+    def handle_result(result):
+        nonlocal processed_count, error_count
+        if result["status"] == "success":
+            # Rotate shard if needed
+            if tar_writer is None or shard_sample_count >= samples_per_shard:
+                open_new_shard()
+            write_sample(result["key"], result["audio_tokens"], result["metadata"])
+            processed_count += 1
+        else:
+            error_count += 1
+            failed_ids.append(result["key"])
+            error_logger.error(
+                json.dumps(
+                    {"id": result["key"], "reason": result["error_msg"]},
+                    ensure_ascii=False,
+                )
+            )
+            if not args.skip_errors:
+                raise RuntimeError(
+                    f"Sample {result['key']} processing failed due "
+                    f"to {result['error_msg']} - terminating"
+                )
+            logging.warning(
+                f"Skipping failed sample {result['key']}: {result['error_msg']}"
+            )
+
+    main_progress = tqdm(total=total_samples, desc="Extracting Audio Tokens")
+
+    try:
+        with ProcessPoolExecutor(
+            max_workers=num_processes,
+            initializer=process_init,
+            initargs=(rank_queue, args.tokenizer_path),
+        ) as executor:
+            logging.info(f"Submitting tasks... ({num_processes} workers)")
+            futures = set()
+            max_pending = num_processes * 10
+
+            def drain_completed():
+                """Wait for at least one future to complete, process all done."""
+                nonlocal futures
+                done, _ = wait(futures, return_when=FIRST_COMPLETED)
+                for f in done:
+                    futures.discard(f)
+                    result = f.result()
+                    main_progress.update(1)
+                    handle_result(result)
+                    main_progress.set_postfix(
+                        Samples=processed_count,
+                        Errors=error_count,
+                    )
+
+            # Stream samples from DataLoader
+            for sample in dataloader:
+                if len(futures) >= max_pending:
+                    drain_completed()
+
+                future = executor.submit(process_single_sample, sample)
+                futures.add(future)
+
+            # Process remaining futures
+            logging.info("Processing remaining pending samples...")
+            while futures:
+                drain_completed()
+
+    except Exception:
+        logging.error("Critical error during processing", exc_info=True)
+        raise
+    finally:
+        main_progress.close()
+        if tar_writer is not None:
+            tar_writer.close()
+        if jsonl_file is not None:
+            jsonl_file.close()
+        # Record the last shard in the manifest
+        if shard_idx > 0 and shard_sample_count > 0:
+            last_idx = shard_idx - 1
+            shard_manifest[last_idx] = (
+                os.path.abspath(tar_output_pattern % last_idx),
+                os.path.abspath(jsonl_output_pattern % last_idx),
+                shard_sample_count,
+                shard_duration,
+            )
+
+    # Write manifest file (data.lst)
+    with open(manifest_path, "w", encoding="utf-8") as mf:
+        for idx in sorted(shard_manifest.keys()):
+            tar_path, jsonl_path, count, duration = shard_manifest[idx]
+            mf.write(f"{tar_path} {jsonl_path} {count} {duration:.3f}\n")
+
+    # Output final statistics
+    total_failed = error_count + write_error_count
+    filtered_and_skipped = total_samples - processed_count - total_failed
+    logging.info(
+        f"Processing Complete - Successful: {processed_count}, Failed: {total_failed}, "
+        f"Filtered/Skipped: {filtered_and_skipped}, Shards written: {shard_idx}"
+    )
+    logging.info(f"Manifest written to: {manifest_path} ({len(shard_manifest)} shards)")
+    if total_failed > 0:
+        logging.info(f"Error details: {error_log_path}")
+    if failed_ids and args.skip_errors:
+        logging.warning(
+            f"Failed sample IDs (count: {len(failed_ids)}): {failed_ids[:100]}..."
+        )
+    if write_error_count > 0 and not args.skip_errors:
+        raise RuntimeError(
+            f"{write_error_count} samples failed to write - check logs for details"
+        )
+
+
+if __name__ == "__main__":
+    main()

runtime/omnivoice-python/omnivoice-source/omnivoice/scripts/extract_audio_tokens_add_noise.py

@@ -0,0 +1,819 @@
+#!/usr/bin/env python3
+# Copyright    2026  Xiaomi Corp.        (authors:  Han Zhu)
+#
+# See ../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Extract audio tokens from audio data and pack them into WebDataset shards.
+
+Extends ``extract_audio_tokens.py`` with optional noise and reverberation
+augmentation on the prompt (reference) portion of the audio. Requires a
+noise manifest and/or RIR manifest.
+
+Supports two input modes:
+
+1. WebDataset manifest (data.lst):
+    python extract_audio_tokens_add_noise.py \\
+        --input_manifest data.lst \\
+        --noise_manifest noise.lst \\
+        --tar_output_pattern output/audios/shard-%06d.tar \\
+        --jsonl_output_pattern output/txts/shard-%06d.jsonl
+
+2. Raw JSONL (each line: {"id": "...", "audio_path": "...", "text": "...", ...}):
+    python extract_audio_tokens_add_noise.py \\
+        --input_jsonl data.jsonl \\
+        --noise_manifest noise.lst \\
+        --tar_output_pattern output/audios/shard-%06d.tar \\
+        --jsonl_output_pattern output/txts/shard-%06d.jsonl
+
+Output structure:
+    output_dir/
+    ├── audios/           # WebDataset tar shards (.npy audio tokens + .json metadata)
+    │   ├── shard_000000.tar
+    │   └── ...
+    ├── txts/             # Per-shard JSONL metadata
+    │   ├── shard_000000.jsonl
+    │   └── ...
+    ├── data.lst          # Manifest: <tar_path> <jsonl_path> <sample_count> <total_duration>
+    └── errors.jsonl      # Failed samples with error details
+"""
+
+import argparse
+import io
+import json
+import logging
+import math
+import multiprocessing as mp
+import os
+import random
+import warnings
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import webdataset as wds
+from torch.utils.data import DataLoader, IterableDataset
+from tqdm.auto import tqdm
+from transformers import AutoFeatureExtractor, HiggsAudioV2TokenizerModel
+
+from omnivoice.data.dataset import JsonlDatasetReader, WebDatasetReader
+from omnivoice.utils.audio import load_audio_bytes
+from omnivoice.utils.common import str2bool
+
+warnings.filterwarnings(
+    "ignore", category=FutureWarning, module="torch.nn.utils.weight_norm"
+)
+
+HIGGS_INPUT_SAMPLE_RATE = 24_000
+
+# Global variables: Store tokenizer and device for each worker process
+worker_tokenizer = None
+worker_feature_extractor = None
+worker_noise_sampler = None
+worker_rir_sampler = None
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--input_manifest",
+        default=None,
+        help="Path to input dataset manifest (data.lst).",
+    )
+    parser.add_argument(
+        "--input_jsonl",
+        default=None,
+        help="Path to raw JSONL file (alternative to --input_manifest).",
+    )
+    parser.add_argument(
+        "--tar_output_pattern",
+        required=True,
+        help="Tar shard pattern passed to WebDataset",
+    )
+    parser.add_argument(
+        "--jsonl_output_pattern",
+        required=True,
+        help="Jsonl shard pattern passed to WebDataset",
+    )
+    parser.add_argument(
+        "--samples_per_shard",
+        type=int,
+        default=1000,
+        help="Maximum records per shard",
+    )
+    parser.add_argument(
+        "--min_num_shards",
+        type=int,
+        default=32,
+        help="Minimum number of output shards (use to ensure "
+        "shard count >= num_gpu * num_workers)",
+    )
+    parser.add_argument(
+        "--tokenizer_path",
+        type=str,
+        default="eustlb/higgs-audio-v2-tokenizer",
+        help="Path to audio tokenizer.",
+    )
+    parser.add_argument(
+        "--skip_errors", action="store_true", help="Skip items that fail to process"
+    )
+    parser.add_argument(
+        "--min_length",
+        type=float,
+        default=0.0,
+        help="Minimum audio duration in seconds (e.g. 2.0)",
+    )
+    parser.add_argument(
+        "--max_length",
+        type=float,
+        default=float("inf"),
+        help="Maximum audio duration in seconds (e.g. 15.0)",
+    )
+    parser.add_argument(
+        "--num_machines",
+        type=int,
+        default=1,
+        help="Total number of machines for distributed runs",
+    )
+    parser.add_argument(
+        "--machine_index",
+        type=int,
+        default=0,
+        help="Zero-based machine index when distributing across multiple "
+        "machines (e.g. 0, 1, ... num_machines-1)",
+    )
+    parser.add_argument(
+        "--nj_per_gpu",
+        type=int,
+        default=3,
+        help="Number of worker processes to spawn per GPU.",
+    )
+    parser.add_argument(
+        "--loader_workers",
+        type=int,
+        default=24,
+        help="Number of DataLoader workers for streaming IterableDataset.",
+    )
+    parser.add_argument(
+        "--shuffle",
+        type=str2bool,
+        default=True,
+        help="Shuffle data by default.",
+    )
+    parser.add_argument(
+        "--shuffle-seed",
+        type=int,
+        default=42,
+        help="Random seed for shuffle (default: 42).",
+    )
+    parser.add_argument(
+        "--noise_manifest",
+        default=None,
+        help="Path to noise manifest (list of tar files). Enables prompt noise augmentation.",
+    )
+    parser.add_argument(
+        "--rir_manifest",
+        default=None,
+        help="Path to RIR manifest (list of tar files). Enables prompt reverb augmentation.",
+    )
+    return parser
+
+
+def count_lines(path):
+    with open(path, "rb") as f:
+        return sum(buf.count(b"\n") for buf in iter(lambda: f.read(1 << 20), b""))
+
+
+def serialise_numpy(key: str, tokens: np.ndarray) -> dict:
+    buffer = io.BytesIO()
+    np.save(buffer, tokens)
+    return {"__key__": key, "npy": buffer.getvalue()}
+
+
+def _load_aug_audio(data, sample_rate=24000):
+    """Simple audio loader for augmentation files."""
+    return torch.from_numpy(load_audio_bytes(data, sample_rate))
+
+
+class SimpleWorkerSampler:
+    """A lightweight infinite sampler for noise/RIR within a worker process."""
+
+    def __init__(self, tar_paths, sample_rate=24000):
+        self.dataset = (
+            wds.WebDataset(
+                tar_paths, shardshuffle=True, nodesplitter=None, workersplitter=None
+            )
+            .decode()
+            .map(lambda s: self._decode(s, sample_rate))
+            .select(lambda x: x is not None)
+            .shuffle(100)
+            .repeat()
+        )
+        self.iterator = iter(self.dataset)
+
+    def _decode(self, sample, sample_rate):
+        for ext in ["wav", "flac", "mp3"]:
+            if ext in sample:
+                return _load_aug_audio(sample[ext], sample_rate)
+        return None
+
+    def sample_segment(self, target_len, allow_repeat=True):
+        """Get a random segment of noise matching the target length."""
+        try:
+            audio = next(self.iterator)
+        except StopIteration:
+            self.iterator = iter(self.dataset)
+            audio = next(self.iterator)
+
+        cur_len = audio.size(-1)
+        if cur_len < target_len and allow_repeat:
+            if cur_len > 0:
+                num_repeats = math.ceil(target_len / cur_len)
+                audio = audio.repeat(1, num_repeats)
+            else:
+                audio = F.pad(audio, (0, target_len), mode="constant")
+            cur_len = audio.size(-1)
+
+        if cur_len > target_len:
+            start = random.randint(0, cur_len - target_len)
+            audio = audio[..., start : start + target_len]
+
+        return audio
+
+
+def _convolve1d(signal: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
+    m = signal.size(-1)
+    n = kernel.size(-1)
+    padded_size = m + n - 1
+    f_signal = torch.fft.rfft(signal, n=padded_size)
+    f_kernel = torch.fft.rfft(kernel, n=padded_size)
+    f_result = f_signal * f_kernel
+    result = torch.fft.irfft(f_result, n=padded_size)
+    return result[:padded_size]
+
+
+def _apply_rir(audio, rir, mix_ratio=0.5):
+    rir_scaling_factor = 0.5**15
+    N_in = audio.shape[-1]
+    rir_d = rir[0, :] * rir_scaling_factor
+    aug_d = _convolve1d(audio[0], rir_d)
+    shift_index = torch.argmax(torch.abs(rir_d))
+    end_index = shift_index + N_in
+    if end_index > aug_d.shape[0]:
+        augmented = F.pad(aug_d[shift_index:], (0, end_index - aug_d.shape[0]))
+    else:
+        augmented = aug_d[shift_index:end_index]
+    power_before = torch.sum(audio[0] ** 2)
+    power_after = torch.sum(augmented**2)
+    if power_after > 0:
+        augmented *= torch.sqrt(power_before / power_after)
+    mixed = (1 - mix_ratio) * audio[0] + mix_ratio * augmented
+    return mixed.unsqueeze(0)
+
+
+def process_init(rank_queue, tokenizer_path, noise_manifest=None, rir_manifest=None):
+    """
+    Initialization function for each worker process.
+    Assigns a specific GPU to the process and loads the tokenizer.
+    """
+    global worker_tokenizer, worker_feature_extractor, worker_noise_sampler, worker_rir_sampler
+
+    # Configure worker process logging
+    formatter = (
+        "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d]"
+        " [Worker %(process)d] %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO, force=True)
+
+    # Get assigned GPU rank
+    rank = rank_queue.get()
+    # Determine device
+    if rank != -1 and torch.cuda.is_available():
+        worker_device = torch.device(f"cuda:{rank}")
+    else:
+        worker_device = torch.device("cpu")
+
+    logging.debug(f"Worker process initialized with device: {worker_device}")
+    # Load tokenizer onto the specified device
+    worker_feature_extractor = AutoFeatureExtractor.from_pretrained(tokenizer_path)
+    worker_tokenizer = HiggsAudioV2TokenizerModel.from_pretrained(
+        tokenizer_path, device_map=worker_device
+    )
+    logging.debug(f"Tokenizer loaded successfully on device {worker_device}")
+
+    # Initialize augmentation samplers (optional)
+    if noise_manifest:
+        try:
+            with open(noise_manifest, "r") as f:
+                tars = [l.strip().split()[0] for l in f if l.strip()]
+            worker_noise_sampler = SimpleWorkerSampler(
+                tars, sample_rate=HIGGS_INPUT_SAMPLE_RATE
+            )
+            logging.debug("Noise sampler initialized.")
+        except Exception as e:
+            logging.warning(f"Failed to load noise manifest: {e}")
+
+    if rir_manifest:
+        try:
+            with open(rir_manifest, "r") as f:
+                tars = [l.strip().split()[0] for l in f if l.strip()]
+            worker_rir_sampler = SimpleWorkerSampler(
+                tars, sample_rate=HIGGS_INPUT_SAMPLE_RATE
+            )
+            logging.debug("RIR sampler initialized.")
+        except Exception as e:
+            logging.warning(f"Failed to load RIR manifest: {e}")
+
+
+def _augment_prompt(audio_tensor: torch.Tensor) -> tuple[torch.Tensor, int]:
+    """Apply noise/reverb augmentation to the front portion of audio.
+
+    Returns the augmented audio and the sample index where clean audio starts.
+    """
+    # Pre-normalization
+    max_val = audio_tensor.abs().max() + 1e-7
+    audio_tensor = (audio_tensor / max_val) * 0.6
+
+    total_len = audio_tensor.size(-1)
+    ratio = random.uniform(0.1, 0.3)
+    split_idx = int(total_len * ratio)
+    front_part = audio_tensor[:, :split_idx].clone()
+
+    # Apply noise
+    if worker_noise_sampler is not None:
+        noise = worker_noise_sampler.sample_segment(split_idx)
+        snr_db = random.uniform(5, 15)
+        sig_rms = front_part.norm(p=2) / (split_idx**0.5)
+        noise_rms = noise.norm(p=2) / (split_idx**0.5)
+        if noise_rms > 1e-9:
+            snr = 10 ** (snr_db / 20)
+            scale = sig_rms / (snr * noise_rms + 1e-8)
+            front_part = front_part + noise * scale
+
+    # Apply RIR (30% probability)
+    if worker_rir_sampler is not None and random.random() < 0.3:
+        rir = worker_rir_sampler.sample_segment(split_idx, allow_repeat=False)
+        reverb_amt = random.uniform(0.3, 1.0)
+        try:
+            front_part = _apply_rir(front_part, rir, reverb_amt)
+        except Exception as e:
+            logging.warning(f"RIR failed: {e}")
+
+    # Merge back
+    if front_part.device != audio_tensor.device:
+        front_part = front_part.to(audio_tensor.device)
+    audio_tensor[:, :split_idx] = front_part
+
+    # Post-normalization
+    max_val = audio_tensor.abs().max() + 1e-7
+    audio_tensor = (audio_tensor / max_val) * 0.9
+
+    return audio_tensor, split_idx
+
+
+def process_single_sample(sample: dict[str, Any]) -> dict[str, Any]:
+    """
+    Single-sample processing function executed in worker processes.
+    Skips invalid samples during streaming processing.
+    """
+    try:
+        audio_tensor = sample.get("audio", None)  # shape (1, T)
+        if audio_tensor is None:
+            raise ValueError("Sample missing 'audio' field")
+
+        # Apply prompt augmentation if noise/rir samplers are available
+        enable_aug = worker_noise_sampler is not None or worker_rir_sampler is not None
+        clean_sample_idx = 0
+        if enable_aug:
+            audio_tensor, clean_sample_idx = _augment_prompt(audio_tensor)
+
+        with torch.inference_mode():
+            key = sample["label"]["id"]
+
+            inputs = worker_feature_extractor(
+                raw_audio=audio_tensor.squeeze(0).numpy(),
+                sampling_rate=HIGGS_INPUT_SAMPLE_RATE,
+                return_tensors="pt",
+            ).to(worker_tokenizer.device)
+            audio_tokens = worker_tokenizer.encode(
+                inputs["input_values"],
+            ).audio_codes.squeeze(0)
+
+            assert len(audio_tokens.shape) == 2
+            assert audio_tokens.size(0) == 8
+
+            num_tokens = audio_tokens.size(1)
+            metadata = sample["label"]
+            metadata["num_tokens"] = num_tokens
+
+            if enable_aug:
+                clean_token_idx = math.ceil(
+                    clean_sample_idx / worker_tokenizer.config.hop_length
+                )
+                metadata["clean_start_token_idx"] = clean_token_idx
+
+            # Convert to numpy format for subsequent serialization (int16 to save space)
+            audio_tokens_np = audio_tokens.to(torch.int16).cpu().numpy()
+
+            return {
+                "status": "success",
+                "key": key,
+                "audio_tokens": audio_tokens_np,
+                "metadata": metadata,
+                "error_msg": None,
+            }
+    except Exception as e:
+        sample_id = sample.get("label", {}).get("id", "unknown")
+        logging.error(f"Failed to process sample {sample_id}: {e}")
+        return {
+            "status": "error",
+            "key": sample_id,
+            "audio_tokens": None,
+            "metadata": None,
+            "error_msg": str(e),
+        }
+
+
+def _normalise_value(value: Any) -> Any:
+    """Convert tensors and NumPy scalars to serialisable Python objects."""
+    if isinstance(value, torch.Tensor):
+        if value.ndim == 0:
+            return value.item()
+        return value.cpu().tolist()
+    if isinstance(value, np.generic):
+        return value.item()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    return value
+
+
+def _encode_metadata(metadata: dict[str, Any]) -> bytes:
+    cleaned: dict[str, Any] = {}
+    for key, value in metadata.items():
+        if value is None:
+            continue
+        cleaned[key] = _normalise_value(value)
+    return json.dumps(cleaned, ensure_ascii=False).encode("utf-8")
+
+
+class StreamingLengthFilteredDataset(IterableDataset):
+    def __init__(
+        self,
+        base_iterable,
+        min_len: float,
+        max_len: float,
+        sr: int,
+    ):
+        self.base_iterable = base_iterable
+        self.min_len = min_len
+        self.max_len = max_len
+        self.sr = sr
+        self.filtered_count = 0
+
+    def __iter__(self):
+        """Stream samples one by one and filter on the fly."""
+        for sample in self.base_iterable:
+            try:
+                duration = sample["audio"].size(-1) / self.sr
+                if self.min_len <= duration <= self.max_len:
+                    yield sample
+                else:
+                    self.filtered_count += 1
+                    logging.warning(
+                        f"Filtered sample (duration out of range): "
+                        f"{sample['label']['id']} ({duration:.2f}s)"
+                    )
+            except Exception as e:
+                logging.warning(f"Skipped invalid sample during streaming: {e}")
+                continue
+
+
+def main() -> None:
+    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+    logging.basicConfig(format=formatter, level=logging.INFO, force=True)
+    parser = build_parser()
+    args = parser.parse_args()
+    mp.set_start_method("spawn", force=True)
+
+    # Validate input arguments
+    assert bool(args.input_manifest) != bool(
+        args.input_jsonl
+    ), "Exactly one of --input_manifest or --input_jsonl must be provided."
+
+    if args.num_machines > 1:
+        assert (
+            0 <= args.machine_index < args.num_machines
+        ), f"machine_index {args.machine_index} must be in [0, {args.num_machines})"
+
+    # Build base dataset and count total samples based on input mode
+    if args.input_jsonl:
+        logging.info(f"Input mode: raw JSONL ({args.input_jsonl})")
+        total_samples = count_lines(args.input_jsonl)
+        base_dataset = JsonlDatasetReader(
+            args.input_jsonl,
+            sample_rate=HIGGS_INPUT_SAMPLE_RATE,
+            shuffle=args.shuffle,
+            shuffle_seed=args.shuffle_seed,
+        )
+        loader_workers = args.loader_workers
+    else:
+        logging.info(f"Input mode: WebDataset manifest ({args.input_manifest})")
+        manifest_num_lines = count_lines(args.input_manifest)
+        loader_workers = min(args.loader_workers, manifest_num_lines)
+        total_samples = 0
+        manifests = []
+        with open(args.input_manifest, "r", encoding="utf-8") as f:
+            for line_id, line in tqdm(
+                enumerate(f),
+                total=manifest_num_lines,
+                desc="Calculating dataset length",
+            ):
+                items = line.strip().split(" ")
+                tar_path, jsonl_path, num_items, duration = (
+                    items[0],
+                    items[1],
+                    int(items[2]),
+                    float(items[3]),
+                )
+                assert os.path.exists(tar_path), f"File {tar_path} does not exist."
+                assert os.path.exists(jsonl_path), f"File {jsonl_path} does not exist."
+                assert jsonl_path.endswith(
+                    ".jsonl"
+                ), f"File {jsonl_path} is not a .jsonl file."
+                if (
+                    args.num_machines > 1
+                    and line_id % args.num_machines != args.machine_index
+                ):
+                    continue
+                total_samples += num_items
+                manifests.append((tar_path, jsonl_path, num_items, duration))
+        logging.info(
+            f"Total shards: {manifest_num_lines}, "
+            f"Shards for current index: {len(manifests)}"
+        )
+        base_dataset = WebDatasetReader(
+            manifests=manifests,
+            sample_rate=HIGGS_INPUT_SAMPLE_RATE,
+            evaluation=True,
+        )
+
+    # Apply length filter and create DataLoader
+    filtered_dataset = StreamingLengthFilteredDataset(
+        base_iterable=base_dataset,
+        min_len=args.min_length,
+        max_len=args.max_length,
+        sr=HIGGS_INPUT_SAMPLE_RATE,
+    )
+    dataloader = DataLoader(
+        dataset=filtered_dataset,
+        batch_size=None,
+        num_workers=loader_workers,
+        persistent_workers=loader_workers > 0,
+        pin_memory=False,
+    )
+
+    # Adjust samples_per_shard if min_num_shards would be violated
+    samples_per_shard = args.samples_per_shard
+    if total_samples > 0:
+        estimated_shards = max(
+            1, (total_samples + samples_per_shard - 1) // samples_per_shard
+        )
+        if estimated_shards < args.min_num_shards:
+            samples_per_shard = max(1, total_samples // args.min_num_shards)
+            logging.info(
+                f"Adjusted samples_per_shard from {args.samples_per_shard} to "
+                f"{samples_per_shard} to meet min_num_shards={args.min_num_shards} "
+                f"(total_samples={total_samples})"
+            )
+
+    # Configure multi-GPU multi-process setup
+    num_devices = torch.cuda.device_count()
+    if num_devices == 0:
+        logging.warning("No GPUs detected - using CPU for processing")
+        num_processes = args.nj_per_gpu
+    else:
+        num_processes = num_devices * args.nj_per_gpu
+    logging.info(
+        f"GPU count: {num_devices}, Processes per GPU: {args.nj_per_gpu}, "
+        f"Total processes: {num_processes}"
+    )
+    if args.noise_manifest or args.rir_manifest:
+        logging.info(
+            f"Prompt augmentation enabled - "
+            f"noise: {args.noise_manifest or 'off'}, rir: {args.rir_manifest or 'off'}"
+        )
+
+    # Shared GPU rank queue for process assignment
+    manager = mp.Manager()
+    rank_queue = manager.Queue()
+    for rank in list(range(num_devices)) * args.nj_per_gpu:
+        rank_queue.put(rank)
+    if num_devices == 0:
+        for _ in range(num_processes):
+            rank_queue.put(-1)
+
+    # Prepare output paths
+    tar_output_pattern = str(Path(args.tar_output_pattern).expanduser())
+    jsonl_output_pattern = str(Path(args.jsonl_output_pattern).expanduser())
+    Path(tar_output_pattern).parent.mkdir(parents=True, exist_ok=True)
+    Path(jsonl_output_pattern).parent.mkdir(parents=True, exist_ok=True)
+
+    # Determine output directory from tar_output_pattern
+    output_dir = Path(tar_output_pattern).parent.parent
+    error_log_path = str(output_dir / "errors.jsonl")
+    manifest_path = str(output_dir / "data.lst")
+
+    # Setup error logger (writes to errors.jsonl)
+    error_logger = logging.getLogger("error_log")
+    error_logger.setLevel(logging.ERROR)
+    error_logger.handlers.clear()
+    error_fh = logging.FileHandler(error_log_path, mode="w", encoding="utf-8")
+    error_fh.setFormatter(logging.Formatter("%(message)s"))
+    error_logger.addHandler(error_fh)
+
+    # Progress and error tracking
+    processed_count = 0
+    error_count = 0
+    write_error_count = 0
+    failed_ids = []
+    shard_idx = 0
+    shard_sample_count = 0
+    shard_duration = 0.0
+    shard_manifest = {}  # shard_idx -> (tar_path, jsonl_path, count, duration)
+
+    tar_writer = None
+    jsonl_file = None
+
+    def open_new_shard():
+        nonlocal tar_writer, jsonl_file, shard_idx, shard_sample_count, shard_duration
+        if tar_writer is not None:
+            tar_writer.close()
+        if jsonl_file is not None:
+            jsonl_file.close()
+        # Record manifest for the previous shard
+        if shard_idx > 0 and shard_sample_count > 0:
+            prev_idx = shard_idx - 1
+            shard_manifest[prev_idx] = (
+                os.path.abspath(tar_output_pattern % prev_idx),
+                os.path.abspath(jsonl_output_pattern % prev_idx),
+                shard_sample_count,
+                shard_duration,
+            )
+        tar_fname = tar_output_pattern % shard_idx
+        jsonl_fname = jsonl_output_pattern % shard_idx
+        tar_writer = wds.TarWriter(tar_fname)
+        jsonl_file = open(jsonl_fname, "w", encoding="utf-8")
+        shard_idx += 1
+        shard_sample_count = 0
+        shard_duration = 0.0
+
+    def write_sample(key, audio_tokens_np, metadata):
+        nonlocal shard_sample_count, write_error_count, shard_duration
+        assert tar_writer is not None and jsonl_file is not None
+        try:
+            token_record = serialise_numpy(key, audio_tokens_np)
+            json_record = _encode_metadata(metadata)
+            tar_writer.write(token_record)
+            jsonl_file.write(json_record.decode("utf-8") + "\n")
+            shard_sample_count += 1
+            shard_duration += metadata.get("audio_duration", 0.0)
+        except Exception as exc:
+            write_error_count += 1
+            failed_ids.append(key)
+            error_logger.error(
+                json.dumps({"id": key, "reason": str(exc)}, ensure_ascii=False)
+            )
+            logging.error(f"Write failed for sample {key}: {exc}")
+
+    def handle_result(result):
+        nonlocal processed_count, error_count
+        if result["status"] == "success":
+            # Rotate shard if needed
+            if tar_writer is None or shard_sample_count >= samples_per_shard:
+                open_new_shard()
+            write_sample(result["key"], result["audio_tokens"], result["metadata"])
+            processed_count += 1
+        else:
+            error_count += 1
+            failed_ids.append(result["key"])
+            error_logger.error(
+                json.dumps(
+                    {"id": result["key"], "reason": result["error_msg"]},
+                    ensure_ascii=False,
+                )
+            )
+            if not args.skip_errors:
+                raise RuntimeError(
+                    f"Sample {result['key']} processing failed due "
+                    f"to {result['error_msg']} - terminating"
+                )
+            logging.warning(
+                f"Skipping failed sample {result['key']}: {result['error_msg']}"
+            )
+
+    main_progress = tqdm(total=total_samples, desc="Extracting Audio Tokens")
+
+    try:
+        with ProcessPoolExecutor(
+            max_workers=num_processes,
+            initializer=process_init,
+            initargs=(
+                rank_queue,
+                args.tokenizer_path,
+                args.noise_manifest,
+                args.rir_manifest,
+            ),
+        ) as executor:
+            logging.info(f"Submitting tasks... ({num_processes} workers)")
+            futures = set()
+            max_pending = num_processes * 10
+
+            def drain_completed():
+                """Wait for at least one future to complete, process all done."""
+                nonlocal futures
+                done, _ = wait(futures, return_when=FIRST_COMPLETED)
+                for f in done:
+                    futures.discard(f)
+                    result = f.result()
+                    main_progress.update(1)
+                    handle_result(result)
+                    main_progress.set_postfix(
+                        Samples=processed_count,
+                        Errors=error_count,
+                    )
+
+            # Stream samples from DataLoader
+            for sample in dataloader:
+                if len(futures) >= max_pending:
+                    drain_completed()
+
+                future = executor.submit(process_single_sample, sample)
+                futures.add(future)
+
+            # Process remaining futures
+            logging.info("Processing remaining pending samples...")
+            while futures:
+                drain_completed()
+
+    except Exception:
+        logging.error("Critical error during processing", exc_info=True)
+        raise
+    finally:
+        main_progress.close()
+        if tar_writer is not None:
+            tar_writer.close()
+        if jsonl_file is not None:
+            jsonl_file.close()
+        # Record the last shard in the manifest
+        if shard_idx > 0 and shard_sample_count > 0:
+            last_idx = shard_idx - 1
+            shard_manifest[last_idx] = (
+                os.path.abspath(tar_output_pattern % last_idx),
+                os.path.abspath(jsonl_output_pattern % last_idx),
+                shard_sample_count,
+                shard_duration,
+            )
+
+    # Write manifest file (data.lst)
+    with open(manifest_path, "w", encoding="utf-8") as mf:
+        for idx in sorted(shard_manifest.keys()):
+            tar_path, jsonl_path, count, duration = shard_manifest[idx]
+            mf.write(f"{tar_path} {jsonl_path} {count} {duration:.3f}\n")
+
+    # Output final statistics
+    total_failed = error_count + write_error_count
+    filtered_and_skipped = total_samples - processed_count - total_failed
+    logging.info(
+        f"Processing Complete - Successful: {processed_count}, Failed: {total_failed}, "
+        f"Filtered/Skipped: {filtered_and_skipped}, Shards written: {shard_idx}"
+    )
+    logging.info(f"Manifest written to: {manifest_path} ({len(shard_manifest)} shards)")
+    if total_failed > 0:
+        logging.info(f"Error details: {error_log_path}")
+    if failed_ids and args.skip_errors:
+        logging.warning(
+            f"Failed sample IDs (count: {len(failed_ids)}): {failed_ids[:100]}..."
+        )
+    if write_error_count > 0 and not args.skip_errors:
+        raise RuntimeError(
+            f"{write_error_count} samples failed to write - check logs for details"
+        )
+
+
+if __name__ == "__main__":
+    main()