Initial release of Notesoop-STT-Ko-v1 (Private)

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,12 @@
+# Python caches
+__pycache__/
+*.py[cod]
+*$py.class
+
+# OS artifacts
+.DS_Store
+Thumbs.db
+
+# Logs & temp files
+*.log
+*.tmp

LICENSE

@@ -0,0 +1,176 @@
+                                 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

README.md

@@ -0,0 +1,248 @@
+---
+language:
+- ko
+license: apache-2.0
+tags:
+- audio
+- speech-recognition
+- automatic-speech-recognition
+- conformer
+- ctc
+- apple-silicon
+- metal
+- on-device
+- korean
+- embeddinggemma-2
+- notesoop
+pipeline_tag: automatic-speech-recognition
+metrics:
+- cer
+- wer
+library_name: pytorch
+---
+
+# Notesoop-STT-Ko-v1
+
+<p align="center">
+  <img src="https://img.shields.io/badge/Language-Korean%20(한국어)-blue" alt="Language" />
+  <img src="https://img.shields.io/badge/Speedup-39.9x%20Real--time-green" alt="Speedup" />
+  <img src="https://img.shields.io/badge/Architecture-Conformer%20CTC%20(21.5M)-orange" alt="Architecture" />
+  <img src="https://img.shields.io/badge/VRAM-1.2GB%20~%201.5GB-purple" alt="VRAM" />
+  <img src="https://img.shields.io/badge/Platform-Apple%20Silicon%20(PyTorch%20MPS)-black" alt="Platform" />
+  <img src="https://img.shields.io/badge/License-Apache%202.0-lightgrey" alt="License" />
+</p>
+
+<p align="center">
+  <a href="https://apps.apple.com/kr/app/notesoop/id6760303593" target="_blank"><img src="https://img.shields.io/badge/App_Store-%EB%85%B8%ED%8A%B8%EC%88%B2_%EB%8B%A4%EC%9A%B4%EB%A1%9C%EB%93%9C-0D96F6?style=for-the-badge&logo=apple&logoColor=white" alt="Download on the App Store" /></a>&nbsp;&nbsp;<a href="https://play.google.com/store/apps/details?id=com.phsoft.notesoop&utm_source=ig&utm_medium=social&utm_content=link_in_bio&fbclid=PAdGRleAS711NwZG9mAmV4dG4DYWVtAjExAHNydGMGYXBwX2lkDzEyNDAyNDU3NDI4NzQxNAABpxFu1PSQrzp-9WgAihbuFrY4gre7MP6MH5Xx4lfKa9FjndnyVoMXb8awipoa_aem_5ccpc4qbd5STxaPUlQif3g" target="_blank"><img src="https://img.shields.io/badge/Google_Play-%EB%85%B8%ED%8A%B8%EC%88%B2_%EB%8B%A4%EC%9A%B4%EB%A1%9C%EB%93%9C-01875F?style=for-the-badge&logo=googleplay&logoColor=white" alt="Get it on Google Play" /></a>
+</p>
+
+**Notesoop-STT-Ko-v1**은 **노트숲(Notesoop)**에서 개발한 Apple Silicon(PyTorch MPS) 가속 한국어 온디바이스 음성인식(ASR) 모델입니다.
+
+동결된(Frozen) `EmbeddingGemma-2` 오디오 인코더의 256차원 Matryoshka 표현(MRL) 위에 21.5M 파라미터 규모의 경량 10-Layer Conformer CTC 리더와 C++ KenLM 3-Gram 디코더를 결합했습니다. 약 1.2~1.5GB 수준의 낮은 VRAM 점유율로 **실시간 대비 약 40배속(`RTF 0.0251`, 41분 분량 음성을 약 63초에 처리)**의 추론 성능을 제공합니다.
+
+> 본 모델(**Notesoop-STT-Ko-v1**)은 **올해 안에 노트숲(Notesoop) 앱에서도 만나보실 수 있습니다.**
+
+---
+
+## Background & Design Goals
+
+노트숲은 **"학습의 비효율을 줄여, 사용자가 도구가 아닌 배움에 집중할 수 있는 경험을 만든다"**는 미션 아래 흩어진 학습 도구와 AI를 자연스럽게 연결하는 기술을 연구하고 있습니다.
+
+실제 학습 환경에서는 수십 분에서 수 시간 분량의 강의 및 학습 음성을 다루게 됩니다. 이때 네트워크 상태에 의존하거나 긴 변환 시간을 기다려야 하고, 무거운 모델 때문에 기기 리소스가 잠식되어 정작 PDF 열람이나 필기 같은 본연의 학습 활동이 방해받는 것은 큰 비효율입니다. **Notesoop-STT-Ko-v1**은 기술적 복잡함은 모델 내부로 숨기고 사용자에게는 끊김 없는 학습 환경을 제공하기 위해 다음 세 가지 목표를 중심으로 설계되었습니다.
+
+* **실질적인 학습 효율 (Practicality & Speed):** 긴 녹음 파일도 대기 시간 없이 즉시 복습과 정리에 활용할 수 있도록, 자기회귀(Autoregressive) 디코딩을 배제한 단일 패스 CTC 구조와 GPU-CPU 비동기 파이프라인을 적용했습니다.
+* **모바일 및 온디바이스 최적화 (Simplicity & Accessibility):** 별도의 서버나 고사양 장비 없이 태블릿·스마트폰 등 제한된 연산 자원을 가진 모바일 디바이스에서도 충분한 추론 속도를 제공할 수 있도록, 핵심 리더 모델을 21.5M 파라미터 규모로 경량화하고 메모리 점유율을 최소화했습니다.
+* **열린 기술 실험 (Open Ecosystem):** 누구나 온디바이스 환경에서 한국어 음성인식 기술을 부담 없이 실험하고 더 나은 학습 경험을 만드는 데 활용할 수 있도록, 모델 가중치와 PyTorch MPS 고속 추론 파이프라인 전체를 Apache-2.0 라이선스로 공개합니다.
+
+---
+
+## Model Specifications
+
+| 항목 | 상세 사양 |
+| :--- | :--- |
+| **Audio Encoder** | `google/embeddinggemma-2` (Frozen 300M Audio Tower, 16kHz Mono, 128-Mel, 40ms/token) |
+| **Feature Projection** | Matryoshka Representation Learning (MRL) 768-dim → 256-dim (L2 Normalized) |
+| **Acoustic Reader** | 10-Layer Conformer (`d_model=384`, `nhead=6`, `ffn_dim=1536`, **21.5M parameters**) |
+| **Output Head** | Cosine Similarity CTC Head (Frozen 25,000-token Spherical Vector Bank) |
+| **Vocabulary** | 25,000 한국어 형태소 + 음절 Fallback (`<blank>` ID = 0) |
+| **LM Decoder** | C++ KenLM 3-Gram Beam Search (`alpha=0.12`, `beta=0.40`) / Greedy CTC 선택 가능 |
+| **Supported Backends** | PyTorch MPS (`bfloat16`), CUDA / CPU |
+| **Training Data** | AI Hub 한국어 강의 음성 데이터 (약 300GB) |
+| **License** | Apache-2.0 |
+
+---
+
+## Architecture
+
+Notesoop-STT-Ko-v1은 비자기회귀(Non-Autoregressive) Conformer CTC 구조와 사전 정렬된 벡터 공간 매칭을 결합하여 파라미터 수와 추론 지연을 동시에 줄였습니다.
+
+```mermaid
+graph TD
+    A["Audio Input (16kHz Mono)"] --> B["Chunking (16.0s window, 1.0s overlap)"]
+    B --> C["EmbeddingGemma-2 Audio Tower (Frozen, 300M)"]
+    C --> D["MRL 256-d Projection & L2 Norm"]
+    D --> E["10-Layer Conformer Reader (21.5M)"]
+    E --> F["Cosine CTC Head (25,000 Vector Bank)"]
+    F --> G["Async C++ KenLM 3-Gram Beam Search"]
+    G --> H["Overlap Deduplication & Final Transcript"]
+```
+
+1. **MRL 256차원 특징 추출 및 구형 벡터 매칭:**
+   `EmbeddingGemma-2` 오디오 인코더의 768차원 출력을 MRL 기반 256차원으로 축소·정규화하여 사용합니다. 출력단에서는 25,000개 한국어 형태소·음절 임베딩으로 구성된 동결 구형 벡터 뱅크(Spherical Vector Bank)와의 코사인 유사도를 통해 CTC 로짓을 계산합니다.
+2. **경량 Conformer Reader (21.5M):**
+   음향 특징을 텍스트 토큰 공간으로 매핑하는 핵심 리더 모델을 2,150만 파라미터 규모로 경량화하여 온디바이스 환경의 메모리 대역폭 병목을 최소화했습니다.
+3. **Apple Silicon MPS 고속 가속:**
+   PyTorch MPS 환경에서 불필요한 FP32/FP16 형변환 오버헤드를 제거한 `bfloat16` 네이티브 텐서 연산 경로를 적용하여, M1/M2/M3/M4 Mac에서 대기 시간 없이 실시간 대비 약 40배속의 빠른 추론 성능을 발휘합니다.
+4. **GPU-CPU 비동기 파이프라이닝:**
+   GPU가 다음 오디오 청크 배치를 연산하는 동안 CPU 스레드 풀에서 C++ KenLM 3-Gram 언어 모델 빔 서치와 청크 경계 중복 제거(Overlap Deduplication)를 병렬로 처리합니다.
+
+---
+
+## Data
+
+**Notesoop-STT-Ko-v1**은 실제 학습 및 교육 현장의 음향·언어적 특성을 반영하기 위해 **AI Hub의 약 300GB 분량 한국어 강의 음성 데이터**로 학습되었습니다.
+
+* **실전 강의 도메인 특화:** 약 300GB 규모의 방대한 실제 강의 음성을 학습에 활용하여, 정제된 낭독체뿐만 아니라 강의 현장의 발화 속도 변화, 설명조 어투, 학술·전공 용어, 장시간 연속 발화 환경에서도 안정적인 인식 성능을 제공합니다.
+* **형태소·음절 하이브리드 토크나이징:** 한국어 교착어 특성에 맞춰 자주 쓰이는 표면 형태소(Surface Morpheme)와 음절 단위 Fallback을 결합한 25,000개 어휘 사전을 구축하여 미등록 어휘(OOV) 손실을 최소화했습니다.
+
+---
+
+## Benchmark
+
+`Notesoop-STT-Ko-v1`의 성능을 **(1) 공인 표준 한국어 학술 벤치마크**와 **(2) 실제 대학 전공 강의 롱폼 음성** 두 가지 환경에서 실측한 결과입니다. (Apple M1 Pro 16GB 기준)
+
+### 1. 표준 한국어 STT 벤치마크 (Zeroth-Korean Test Set, 457개 문장)
+한국어 음성인식 학술 평가의 공식 표준 데이터셋인 `Bingsu/zeroth-korean` 테스트 분할(총 457개 발화 문장, 1시간 11분 27초 분량)에 대한 전수 평가 결과입니다.
+
+| 데이터셋 | 총 오디오 길이 | 전체 전사 시간 ⏱️ | 체감 속도 ⚡ | RTF | 초당 처리량 (Throughput) | 음절 오차율 (CER) 🎯 | 단어 오차율 (WER) |
+| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| **Zeroth-Korean Test (457문장)** | **1시간 11분 (4,287초)** | **2분 30초 (150.7초)** | **28.45배속** | **0.0352** | **3.03 문장/초** | **16.21%** | 45.40% |
+
+* **초당 3문장 이상의 초고속 전사**: 1시간 11분 분량의 방대한 오디오를 단 2분 30초 만에 완주하며, 낭독체 음성 환경에서 **16.21%의 준수한 음절 오차율(CER)**을 달성합니다.
+* **실제 전사 샘플 예시**:
+  * **정답 (REF)**: *"한 잔의 커피가 소비자에게 전해지기 위해선 여러 차례의 공정을 거쳐야 한다"*
+  * **예측 (HYP)**: *"한 잔의 커피가 소비자에게 전해지기 위해선 여러 차례의 공정을 거쳐야 한다"* (CER: **0.00%**)
+  * **정답 (REF)**: *"또한 박 전 대통령 측은 앞으로 검찰 수사와 관련해 수사 내용에 대한 입장뿐 아니라 변호인단의 활동 상황 등도 소상하게 공개한다는 방침이다"*
+  * **예측 (HYP)**: *"또한 박 전 대통령 주변에 앞으로 검찰수사와 관련된 수사 내용에 대한 입장뿐 아니라 변호인단의 활동 상황 등도 소상하게 공개한다는 방식이다"* (CER: **10.39%**)
+
+---
+
+### 2. 실전 전공 강의 녹음 (41분 43초, Field Test)
+대학 강의 환경(배경 잡음, 반향, 말더듬 등이 포함된 실제 41분 43초 녹음 음성)에서 글로벌 표준 ASR 모델들과 비교 측정한 실측 결과입니다.
+
+| Model | Backend / Device | Processing Time | Speedup | RTF | CER (%) 🎯 | WER (%) | Host RAM | GPU Memory (VRAM) 🖥️ |
+| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| **Notesoop-STT-Ko-v1** | **PyTorch MPS (`bf16`)** | **1분 2초 (62.7s)** | **39.9x** | **0.0251** | **22.46** | 39.25 | **1,076 MB** | 1,537 MB |
+| Whisper-medium | MPS | 4분 11초 (251.5s) | 10.0x | 0.1005 | 26.68 | 45.82 | 1,050 MB | 3,467 MB |
+| Whisper-large-v3 | MPS | 11분 45초 (704.9s) | 3.6x | 0.2816 | 25.55 | 42.10 | 1,243 MB | 6,325 MB |
+
+* **`Whisper-large-v3` 대비 11.2배 고속 & 더 우수한 정확도:** Whisper-large-v3(11분 45초, CER 25.55%) 대비 **약 11.2배 빠른 속도(1분 2초)**와 **더 낮은 음절 오차율(CER 22.46%, 3.09%p 향상)**을 동시에 제공합니다.
+* **1/4 수준의 경량 VRAM (1.2~1.5GB):** Whisper-large-v3(6.3GB)의 1/4도 안 되는 가벼운 메모리 점유율 덕분에, 기본형 맥북 에어/프로(8GB/16GB)에서도 발열과 팬 소음 없이 다른 학습 앱과 동시에 쾌적하게 구동됩니다.
+* **비자기회귀(Non-Autoregressive) Conformer의 강점:** 문장을 한 토큰씩 순차 생성하는 Whisper와 달리, Conformer CTC 구조로 오디오 전체 음향 스펙트로그램을 단일 패스($O(T)$)로 병렬 처리하여 장시간 음성에서도 처리 속도가 기하급수적으로 빨라집니다.
+
+---
+
+## Quickstart
+
+### 1. Installation
+
+macOS (Apple Silicon M1 이상 권장) 및 Python 3.10+ 환경을 권장합니다. 오디오 디코딩을 위해 시스템에 `ffmpeg`가 설치되어 있어야 합니다(`brew install ffmpeg`).
+
+```bash
+git clone https://huggingface.co/Phsoft-ai/Notesoop-STT-Ko-v1
+cd Notesoop-STT-Ko-v1
+
+# 의존성 패키지 설치
+pip install -r requirements.txt
+
+# KenLM 3-Gram 디코더 설치
+pip install https://github.com/kpu/kenlm/archive/master.zip
+```
+
+### 2. CLI Usage
+
+`transcribe.py`를 통해 단일 명령어로 오디오 파일을 전사할 수 있습니다. (`.wav`, `.m4a`, `.mp3`, `.flac` 지원)
+
+```bash
+# 기본 실행 (PyTorch MPS 백엔드, batch_size=4)
+python transcribe.py sample_audio.wav
+
+# 장시간 강의 오디오 전사 및 결과 파일 저장
+python transcribe.py my_lecture.m4a --batch-size 4 --chunk-sec 16.0 --output transcript.txt
+
+# KenLM 빔 서치 비활성화 (Greedy CTC 디코딩)
+python transcribe.py sample_audio.wav --no-kenlm
+```
+
+### 3. Python API
+
+```python
+from run_inference import EmbeddingGemmaSTT
+
+stt = EmbeddingGemmaSTT(
+    device="mps",
+    use_kenlm=True,
+    lm_alpha=0.12,
+    lm_beta=0.40,
+)
+
+transcript, rtf = stt.transcribe(
+    "sample_audio.wav",
+    chunk_sec=16.0,
+    overlap_sec=1.0,
+    batch_size=4,
+)
+
+print(f"RTF: {rtf:.4f} ({1/rtf:.1f}x real-time)")
+print(transcript)
+```
+
+---
+
+## Repository Structure
+
+```text
+Notesoop-STT-Ko-v1/
+├── README.md                        # Model Card 및 기술 문서
+├── configs/
+│   └── config.yaml                  # 모델 아키텍처 및 하이퍼파라미터 설정
+├── checkpoints/
+│   └── best_reader_inference.pt     # PyTorch Conformer Reader 가중치 (137MB)
+├── data/
+│   └── vocab/
+│       ├── vocab_25k.json           # 25,000 한국어 형태소/음절 어휘 사전
+│       ├── vector_bank_25k.pt       # MRL 256차원 동결 벡터 뱅크
+│       └── korean_lm.bin            # KenLM 3-Gram 한국어 언어 모델 (82.7MB)
+├── src/
+│   ├── model.py                     # EmbeddingGemmaCTCReader 모델 정의
+│   ├── dataset.py                   # 한국어 형태소 토크나이저 및 정규화
+│   ├── decoder.py                   # KenLM 빔 서치 디코더 래퍼
+│   └── step2_extract_features.py    # EmbeddingGemma-2 오디오 특징 추출기
+├── run_inference.py                 # 추론 파이프라인 및 Python API
+├── transcribe.py                    # CLI 실행 스크립트
+└── sample_audio.wav                 # 테스트용 샘플 오디오
+```
+
+---
+
+## License & Citation
+
+본 모델과 코드는 **[Apache License 2.0](LICENSE)** 하에 공개되어 연구 및 상업적 용도로 자유롭게 활용할 수 있습니다.
+
+```bibtex
+@misc{notesoop2026sttkov1,
+  author = {Notesoop AI Team},
+  title = {Notesoop-STT-Ko-v1: Lightweight Korean On-Device Speech Recognition},
+  year = {2026},
+  publisher = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/Phsoft-ai/Notesoop-STT-Ko-v1}},
+  note = {Apple Silicon MPS Accelerated Korean ASR}
+}
+```
+
+---
+
+<p align="center">
+  <b>Notesoop (노트숲)</b> — 학습의 비효율을 줄여, 도구가 아닌 배움에 집중할 수 있는 경험을 만듭니다.<br>
+  음성 기록과 PDF 필기 연동, 핵심 정리, 복습 문제 생성까지 하나로 연결된 학습 환경을 노트숲 앱에서 만나보세요.
+</p>

checkpoints/best_reader_inference.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:879b41bb654da9384e21ccf04076f065f9bcead96fb104c0835ddd8efb326c8c
+size 137381751

configs/config.yaml

@@ -0,0 +1,84 @@
+# ==============================================================================
+# EmbeddingGemma-2 Korean STT Pipeline Configuration
+# Optimized for Apple Silicon M2 Pro (16GB Unified Memory) & Mobile Deployment
+# ==============================================================================
+
+system:
+  device: "mps"                # Forward & Matmul device (auto falls back to cpu if mps unavailable)
+  ctc_device: "cpu"            # CTC Loss device (MPS CTCLoss has known stability issues)
+  precision: "bfloat16"        # "bfloat16" or "float32" (NEVER float16 due to NaN overflow)
+  seed: 42
+
+model:
+  name: "google/embeddinggemma-2"
+  original_dim: 768            # Full output dimension of EmbeddingGemma-2
+  embedding_dim: 256           # Matryoshka Representation Learning (MRL) sliced dim
+  normalize_embeddings: true   # L2 normalization (F.normalize)
+  
+  reader:
+    input_dim: 256             # Cached audio feature dimension
+    output_dim: 256            # Vector Bank matching dimension
+    d_model: 384               # Expanded Conformer latent dimension
+    nhead: 6                   # 6 attention heads (64 dim per head)
+    num_layers: 10             # 10-layer Conformer (~21.5M params total)
+    dim_feedforward: 1536      # 4x d_model FFN dimension
+    dropout: 0.1
+    spec_augment: true
+    trainable_vector_bank: false   # Strictly FROZEN to preserve EmbeddingGemma-2 Zero-Shot OTA Hot-Swap anchor
+    init_temperature: 0.05
+    max_params: 25000000
+
+audio:
+  sample_rate: 16000           # 16 kHz Mono
+  n_mels: 128                  # Mel filterbanks matching Gemma4AudioFeatureExtractor
+  hop_length: 160              # 10ms hop
+  win_length: 320              # 20ms window
+  audio_ms_per_token: 40       # 40ms per audio token (subsampled 4x from 10ms frame)
+  augmentation:
+    enabled: false             # Pure raw feature extraction (artificial noise disabled)
+    rir_prob: 0.0              # Reverb (RIR) disabled (0.0)
+    noise_prob: 0.0            # No artificial background noise
+    noise_snr_db_low: 15.0     # Low SNR for ambient lecture noise
+    noise_snr_db_high: 30.0
+
+vocab:
+  target_vocab_size: 25000     # Lossless surface morpheme + syllable fallback (25,000 tokens)
+  blank_token: "<blank>"       # Index 0 reserved for CTC blank
+  blank_id: 0
+  unk_token: "<unk>"
+  vocab_txt_path: "data/vocab/vocab_25k.txt"
+  vocab_json_path: "data/vocab/vocab_25k.json"
+  vector_bank_pt_path: "data/vocab/vector_bank_25k.pt"
+  vector_bank_bin_path: "data/vocab/vector_bank_25k.bin"
+
+data:
+  cache_dir: "data/cache_features"
+  hf_dataset_name: "Bingsu/zeroth-korean"
+  max_audio_length_sec: 20.0
+  min_audio_length_sec: 0.5
+
+train:
+  batch_size: 4                # Safe batch size for 16GB Mac mini (prevents MPS peak memory spikes)
+  grad_accum_steps: 4          # Effective batch size = 4 * 4 = 16 (prevents MPS OOM on 16GB Mac)
+  val_batch_size: 4
+  num_workers: 4               # 4 parallel worker processes for asynchronous feature prefetching
+  learning_rate: 2.0e-4
+  weight_decay: 1.0e-2
+  max_grad_norm: 1.0
+  num_epochs: 3                # Phase 1: 3 Epochs frozen anchor pre-alignment
+  scheduler: "cosine"          # "cosine" or "onecycle"
+  warmup_ratio: 0.05
+  min_lr: 1.0e-6
+  eval_interval_epochs: 1
+  save_interval_epochs: 1
+  checkpoint_dir: "checkpoints"
+
+export:
+  export_dir: "exports"
+  coreml:
+    model_name: "EmbeddingGemma2_KoreanSTT"
+    quantization: "int8"       # "int8", "fp16", or "none"
+  onnx:
+    model_name: "embeddinggemma2_korean_stt"
+    opset_version: 18
+    quantization: "int8"       # "int8" dynamic quantization

data/vocab/korean_lm.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26d946900210343069e364edaf15dabdd2f89c4d60ac719a10a572254f91d4ac
+size 86738030

data/vocab/vector_bank.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecf5bdeb93156b5e19b120066a2c14c2ffd487f3d8d59a83099c2695572603d8
+size 12801093

data/vocab/vector_bank_25k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecf5bdeb93156b5e19b120066a2c14c2ffd487f3d8d59a83099c2695572603d8
+size 12801093

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.2.0
+torchaudio>=2.2.0
+transformers>=4.40.0
+safetensors>=0.4.0
+pyyaml>=6.0
+numpy>=1.24.0
+scipy>=1.11.0
+tqdm>=4.66.0
+psutil>=5.9.0
+jiwer>=3.0.0
+# kenlm: pip install https://github.com/kpu/kenlm/archive/master.zip

run_inference.py

@@ -0,0 +1,399 @@
+#!/usr/bin/env python3
+"""
+run_inference.py: Production-ready Korean STT Inference Engine for Notesoop-STT-Ko-v1.
+
+- Google EmbeddingGemma-2 Frozen 300M Audio Tower (MRL 256-dim feature extraction)
+- 10-Layer Conformer Reader (21.5M params) on Apple Silicon MPS / CUDA / CPU
+- 25,000 Morpheme Vocabulary with Cosine Similarity Matching
+- Asynchronous KenLM 3-Gram Language Model Beam Search Rescoring
+"""
+
+from __future__ import annotations
+
+import argparse
+import difflib
+import os
+import re
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Optional, List, Tuple
+from concurrent.futures import ThreadPoolExecutor, Future
+
+_ROOT = Path(__file__).resolve().parent
+_SRC_DIR = _ROOT / "src"
+if str(_SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(_SRC_DIR))
+if str(_ROOT) not in sys.path:
+    sys.path.insert(0, str(_ROOT))
+
+import numpy as np
+import torch
+import yaml
+
+from dataset import KoreanMorphemeTokenizer, normalize_text
+from model import EmbeddingGemmaCTCReader
+from step2_extract_features import GenuineAudioFeatureExtractor
+from decoder import KenLMBeamSearchDecoder
+
+
+def load_audio(
+    file_path: str,
+    target_sr: int = 16000,
+    start_sec: float = 0.0,
+    duration_limit: float = 0.0
+) -> Tuple[torch.Tensor, float]:
+    """
+    Loads any audio container (.wav, .m4a, .mp3, .flac) into 16kHz float32 mono tensor via ffmpeg.
+    Uses ultra-fast seeking (-ss / -t) directly in FFmpeg to eliminate reading massive containers into memory.
+    """
+    p = Path(file_path)
+    if not p.exists():
+        raise FileNotFoundError(f"Audio file not found: {file_path}")
+
+    cmd = ["ffmpeg", "-nostdin", "-threads", "0"]
+    if start_sec > 0:
+        cmd.extend(["-ss", str(start_sec)])
+    if duration_limit > 0:
+        cmd.extend(["-t", str(duration_limit)])
+    cmd.extend([
+        "-i", str(p),
+        "-f", "f32le", "-ac", "1", "-ar", str(target_sr), "-"
+    ])
+    try:
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+        out, _ = proc.communicate()
+        if proc.returncode != 0 or len(out) == 0:
+            raise RuntimeError("FFmpeg decode returned empty data")
+        arr = np.frombuffer(out, dtype=np.float32).copy()
+    except Exception:
+        # Fallback to soundfile if ffmpeg is not installed
+        import soundfile as sf
+        arr, sr = sf.read(str(p), dtype="float32")
+        if arr.ndim > 1:
+            arr = arr.mean(axis=1)
+        if start_sec > 0:
+            arr = arr[int(start_sec * sr):]
+        if duration_limit > 0:
+            arr = arr[:int(duration_limit * sr)]
+        if sr != target_sr:
+            import math
+            gcd = math.gcd(sr, target_sr)
+            up = target_sr // gcd
+            down = sr // gcd
+            import scipy.signal as sps
+            arr = sps.resample_poly(arr, up, down).astype(np.float32)
+
+    wav = torch.from_numpy(arr).unsqueeze(0)
+    duration_sec = wav.size(-1) / float(target_sr)
+    return wav, duration_sec
+
+
+def merge_chunk_transcripts(chunks: List[str]) -> str:
+    """
+    Merges overlapping transcript segments by finding maximum character/word overlap.
+    Prevents stuttering/duplication across sliding audio window boundaries.
+    """
+    if not chunks:
+        return ""
+    if len(chunks) == 1:
+        return chunks[0]
+
+    merged = chunks[0]
+    for nxt in chunks[1:]:
+        if not nxt:
+            continue
+        max_search = min(len(merged), len(nxt), 40)
+        found_overlap = 0
+
+        # Exact substring suffix-prefix match
+        for k in range(max_search, 2, -1):
+            if merged.endswith(nxt[:k]):
+                found_overlap = k
+                break
+
+        # Fuzzy sequence matcher
+        if found_overlap == 0 and len(merged) >= 10 and len(nxt) >= 10:
+            tail = merged[-30:]
+            head = nxt[:30]
+            matcher = difflib.SequenceMatcher(None, tail, head)
+            match = matcher.find_longest_match(0, len(tail), 0, len(head))
+            if match.size >= 4 and match.a + match.size == len(tail) and match.b == 0:
+                found_overlap = match.size
+
+        if found_overlap > 0:
+            merged = merged + nxt[found_overlap:]
+        else:
+            merged = merged + " " + nxt
+
+    return re.sub(r"\s+", " ", merged).strip()
+
+
+class EmbeddingGemmaSTT:
+    """
+    High-performance, standalone Korean STT Pipeline for Notesoop-STT-Ko-v1.
+    """
+    def __init__(
+        self,
+        checkpoint_path: Optional[str] = None,
+        config_path: Optional[str] = None,
+        vocab_path: Optional[str] = None,
+        vector_bank_path: Optional[str] = None,
+        lm_path: Optional[str] = None,
+        device: str = "auto",
+        use_kenlm: bool = True,
+        lm_alpha: float = 0.12,
+        lm_beta: float = 0.40,
+        **kwargs  # Ignore legacy kwargs like engine
+    ):
+        self.root = _ROOT
+        self.lm_alpha = lm_alpha
+        self.lm_beta = lm_beta
+
+        # 0. Determine Device
+        if device == "auto":
+            dev_str = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            dev_str = "mps" if (device == "metal" and torch.backends.mps.is_available()) else device
+
+        self.device = torch.device(dev_str)
+
+        # 1. Resolve paths & Tokenizer
+        cfg_p = Path(config_path) if config_path else self.root / "configs/config.yaml"
+        with open(cfg_p, "r", encoding="utf-8") as f:
+            self.cfg = yaml.safe_load(f)
+
+        v_p = Path(vocab_path) if vocab_path else self.root / "data/vocab/vocab_25k.json"
+        if not v_p.exists():
+            v_p = self.root / "data/vocab/vocab.json"
+        self.tokenizer = KoreanMorphemeTokenizer(str(v_p))
+
+        vb_p = Path(vector_bank_path) if vector_bank_path else self.root / "data/vocab/vector_bank_25k.pt"
+        if not vb_p.exists():
+            vb_p = self.root / "data/vocab/vector_bank.pt"
+        self.vector_bank = torch.load(vb_p, map_location="cpu", weights_only=True)
+
+        # 2. Initialize Conformer Reader
+        if checkpoint_path:
+            ckpt_p = Path(checkpoint_path)
+        else:
+            p_inf = self.root / "checkpoints/best_reader_inference.pt"
+            p_full = self.root / "checkpoints/best_reader.pt"
+            ckpt_p = p_inf if p_inf.exists() else p_full
+
+        if not ckpt_p.exists():
+            raise FileNotFoundError(f"Reader checkpoint not found at: {ckpt_p}")
+
+        ckpt = torch.load(ckpt_p, map_location="cpu", weights_only=False)
+        state_dict = ckpt.get("model_state_dict", ckpt)
+
+        vb_model = state_dict.get("ctc_head.vector_bank", self.vector_bank)
+        d_model = state_dict["reader.input_proj.weight"].shape[0] if "reader.input_proj.weight" in state_dict else 384
+        input_dim = state_dict["reader.input_proj.weight"].shape[1] if "reader.input_proj.weight" in state_dict else 256
+        detected_layers = len([k for k in state_dict.keys() if k.startswith("reader.layers.") and k.endswith(".final_norm.weight")])
+        num_layers = detected_layers if detected_layers > 0 else 10
+        dim_ffn = state_dict["reader.layers.0.ffn1.0.weight"].shape[0] if "reader.layers.0.ffn1.0.weight" in state_dict else 1536
+        nhead = d_model // 64
+        output_dim = state_dict["reader.output_proj.0.weight"].shape[0] if "reader.output_proj.0.weight" in state_dict else 256
+
+        self.reader = EmbeddingGemmaCTCReader(
+            input_dim=input_dim,
+            output_dim=output_dim,
+            d_model=d_model,
+            nhead=nhead,
+            num_layers=num_layers,
+            dim_feedforward=dim_ffn,
+            dropout=0.0,
+            init_temperature=0.05,
+            vector_bank=vb_model,
+            use_spec_augment=False,
+            trainable_vector_bank=False
+        ).to(self.device)
+        self.reader.load_state_dict(state_dict, strict=False)
+        self.reader.eval()
+
+        # 3. Audio Feature Extractor (Google EmbeddingGemma-2 300M Audio Tower)
+        model_name = self.cfg["model"].get("name", "google/embeddinggemma-2")
+        self.extractor = GenuineAudioFeatureExtractor(device=self.device, model_name=model_name)
+
+        # 4. Optional KenLM Decoder
+        self.decoder_kenlm = None
+        if use_kenlm:
+            lm_file = Path(lm_path) if lm_path else self.root / "data/vocab/korean_lm.bin"
+            if lm_file.exists():
+                self.decoder_kenlm = KenLMBeamSearchDecoder(
+                    tokenizer=self.tokenizer,
+                    lm_path=str(lm_file),
+                    beam_size=8,
+                    alpha=lm_alpha,
+                    beta=lm_beta
+                )
+
+    @torch.no_grad()
+    def transcribe(
+        self,
+        audio_path: str,
+        chunk_sec: float = 16.0,
+        overlap_sec: float = 1.0,
+        start_sec: float = 0.0,
+        duration_limit: float = 0.0,
+        batch_size: int = 4
+    ) -> Tuple[str, float]:
+        """
+        Transcribes audio file of arbitrary length into Korean text.
+        Optimized with Batched GPU Forward Passes and Asynchronous KenLM Rescoring:
+        - CPU Mel Preprocessing
+        - GPU 300M Audio Tower Forward (Batched)
+        - GPU Conformer Reader + Top-32 Reduction (Batched)
+        - Asynchronous KenLM 3-Gram Beam Search Rescoring
+        Returns (transcript_text, rtf_speed).
+        """
+        wav, eval_dur = load_audio(
+            audio_path,
+            target_sr=16000,
+            start_sec=start_sec,
+            duration_limit=duration_limit
+        )
+        sr = 16000
+        step_sec = max(chunk_sec - overlap_sec, 0.5) if overlap_sec > 0 else chunk_sec
+        step_samples = int(step_sec * sr)
+        chunk_samples = int(chunk_sec * sr)
+        total_samples = wav.size(-1)
+
+        chunk_ranges: List[Tuple[int, int]] = []
+        c_start = 0
+        while c_start < total_samples:
+            c_end = min(c_start + chunk_samples, total_samples)
+            if (c_end - c_start) / sr >= 0.5:
+                chunk_ranges.append((c_start, c_end))
+            if c_end == total_samples:
+                break
+            c_start += step_samples
+
+        if not chunk_ranges:
+            return "", 0.0
+
+        t0 = time.time()
+
+        # 1. Preprocess Mel-spectrogram on CPU (pad last chunk to chunk_samples for uniform shapes)
+        preprocessed_chunks = []
+        for cs, ce in chunk_ranges:
+            c_wav = wav[:, cs:ce]
+            if c_wav.size(-1) < chunk_samples:
+                c_wav = torch.nn.functional.pad(c_wav, (0, chunk_samples - c_wav.size(-1)))
+            preprocessed_chunks.append(self.extractor.preprocess_audio_cpu(c_wav))
+
+        # 2. Pipelined Batched Inference
+        chunk_texts: List[str] = []
+        kenlm_futures: List[Future] = []
+        use_kenlm_async = (self.decoder_kenlm is not None)
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            for b_idx in range(0, len(preprocessed_chunks), batch_size):
+                b_items = preprocessed_chunks[b_idx : b_idx + batch_size]
+
+                # Step A: Extract Audio Tower Features (GPU)
+                if len(b_items) > 1 and len(set(it[3] for it in b_items)) == 1:
+                    feats = self.extractor.extract_batch_same_bucket(b_items)
+                else:
+                    feats = [self.extractor.extract_from_preprocessed(*it) for it in b_items]
+
+                # Step B: Conformer Reader Forward (Batched GPU)
+                if len(feats) > 1 and len(set(f.size(0) for f in feats)) == 1:
+                    feats_batch = torch.stack(feats, dim=0).to(device=self.device, dtype=torch.float32)
+                    logits_batch = self.reader(feats_batch)
+                else:
+                    logits_batch = [
+                        self.reader(f.unsqueeze(0).to(device=self.device, dtype=torch.float32)).squeeze(0)
+                        for f in feats
+                    ]
+
+                # Step C: Decoding (KenLM Beam Search or Greedy CTC)
+                for b in range(len(feats)):
+                    logits = logits_batch[b]
+                    if use_kenlm_async:
+                        topk_lprobs, topk_ids = torch.topk(logits, k=32, dim=-1)
+                        blank_lprobs = logits[:, 0]
+                        fut = executor.submit(
+                            self.decoder_kenlm.decode_topk,
+                            topk_ids.detach().cpu().to(torch.int32).numpy(),
+                            topk_lprobs.detach().cpu().float().numpy(),
+                            blank_lprobs.detach().cpu().float().numpy(),
+                            alpha=self.lm_alpha,
+                            beta=self.lm_beta
+                        )
+                        kenlm_futures.append(fut)
+                    else:
+                        pred_ids = torch.argmax(logits, dim=-1).cpu().tolist()
+                        txt = self.tokenizer.decode(pred_ids, remove_blank=True, collapse_repeats=True)
+                        chunk_texts.append(txt)
+
+            if use_kenlm_async:
+                chunk_texts = [fut.result() for fut in kenlm_futures]
+
+        t_elapsed = time.time() - t0
+        rtf = t_elapsed / max(eval_dur, 0.001)
+
+        valid_texts = [t.strip() for t in chunk_texts if t.strip()]
+        if overlap_sec > 0:
+            full_text = merge_chunk_transcripts(valid_texts)
+        else:
+            full_text = " ".join(valid_texts)
+
+        return full_text, rtf
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Notesoop-STT-Ko-v1 Standalone Inference")
+    parser.add_argument("--audio", type=str, default="sample_audio.wav", help="Path to audio file (.wav, .m4a, .mp3, etc.)")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Path to checkpoint (.pt)")
+    parser.add_argument("--device", type=str, default="auto", choices=["auto", "mps", "cuda", "cpu"], help="Inference device")
+    parser.add_argument("--kenlm", action="store_true", default=True, help="Enable KenLM 3-gram language model rescoring")
+    parser.add_argument("--no_kenlm", dest="kenlm", action="store_false", help="Disable KenLM (Greedy decoding only)")
+    parser.add_argument("--duration", type=float, default=0.0, help="Duration limit in seconds (0 for full audio)")
+    parser.add_argument("--chunk_sec", type=float, default=16.0, help="Chunk length in seconds")
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size for parallel processing (default: 4)")
+    parser.add_argument("--output", type=str, default=None, help="Optional output .txt file path")
+    args = parser.parse_args()
+
+    print("\n" + "=" * 70)
+    print("🌲 Notesoop-STT-Ko-v1: Korean Speech Recognition Engine")
+    print("=" * 70)
+    print(f"  • Audio File : {args.audio}")
+    print(f"  • Device     : {args.device.upper()}")
+    print(f"  • Batch Size : {args.batch_size}")
+    print(f"  • KenLM LM   : {'Enabled (8-beam rescoring)' if args.kenlm else 'Disabled (Greedy)'}")
+    print("-" * 70)
+
+    stt = EmbeddingGemmaSTT(
+        checkpoint_path=args.checkpoint,
+        device=args.device,
+        use_kenlm=args.kenlm
+    )
+
+    print("\n🚀 Transcribing...")
+    text, rtf = stt.transcribe(
+        args.audio,
+        chunk_sec=args.chunk_sec,
+        duration_limit=args.duration,
+        batch_size=args.batch_size
+    )
+
+    print("\n" + "=" * 70)
+    print("📝 [Transcription Result]:")
+    print("=" * 70)
+    print(text)
+    print("-" * 70)
+    print(f"⏱️  Speed / RTF: {rtf:.4f} ({1.0 / max(rtf, 1e-4):.1f}x Real-time on {stt.device.type.upper()})")
+    print("=" * 70 + "\n")
+
+    if args.output:
+        out_p = Path(args.output)
+        out_p.parent.mkdir(parents=True, exist_ok=True)
+        out_p.write_text(text, encoding="utf-8")
+        print(f"💾 Saved transcription to: {out_p}\n")
+
+
+if __name__ == "__main__":
+    main()

src/dataset.py

@@ -0,0 +1,774 @@
+"""
+src/dataset.py
+
+Dataset Loader and Dynamic Padding Collate Function for Pre-cached Audio Features.
+Handles cached .pt representations, token ID targets, attention masks, and CTC sequence lengths.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+import re
+import unicodedata
+import zipfile
+from pathlib import Path
+from typing import List, Dict, Any, Tuple, Optional
+
+import torch
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+
+# ==============================================================================
+# 1. Korean Morpheme Tokenizer & Normalization
+# ==============================================================================
+
+# SentencePiece word prefix marker (U+2581 ' ')
+WORD_PREFIX = " "
+
+# Korean compatibility consonant (U+3131~U+314E) to standard conjoining jongseong (U+11A8~U+11C2) mapping
+COMPAT_TO_JONGSEONG: Dict[str, str] = {
+    "ㄱ": "\u11a8", "ㄲ": "\u11a9", "ㄳ": "\u11aa", "ㄴ": "\u11ab", "ㄵ": "\u11ac",
+    "ㄶ": "\u11ad", "ㄷ": "\u11ae", "ㄹ": "\u11af", "ㄺ": "\u11b0", "ㄻ": "\u11b1",
+    "ㄼ": "\u11b2", "ㄽ": "\u11b3", "ㄾ": "\u11b4", "ㄿ": "\u11b5", "ㅀ": "\u11b6",
+    "ㅁ": "\u11b7", "ㅂ": "\u11b8", "ㅄ": "\u11b9", "ㅅ": "\u11ba", "ㅆ": "\u11bb",
+    "ㅇ": "\u11bc", "ㅈ": "\u11bd", "ㅊ": "\u11be", "ㅋ": "\u11bf", "ㅌ": "\u11c0",
+    "ㅍ": "\u11c1", "ㅎ": "\u11c2"
+}
+
+_PUNCT_RE = re.compile(r"[?!.,\"~:;()\[\]{}<>@#$%^&*_+=\/\\|`·…“”‘’]")
+_SPACE_RE = re.compile(r"\s+")
+
+def normalize_text(text: str) -> str:
+    """
+    Normalizes spoken Korean text by stripping punctuation marks (?, ., ,, !, etc.),
+    preserving Korean characters, alphanumeric characters, and whitespace.
+    """
+    text = _PUNCT_RE.sub(" ", text)
+    return _SPACE_RE.sub(" ", text).strip()
+
+import sqlite3
+import threading
+
+
+def canonicalize_aihub_key(key: str) -> str:
+    """Normalizes any AI Hub path or key (e.g. 'art/desi/C00890/U00004.wav') to 'C00890_U00004'."""
+    k = key.replace("\\", "/").strip("/")
+    if k.endswith((".wav", ".pcm", ".flac", ".json")):
+        k = k.rsplit(".", 1)[0]
+    if "/" in k:
+        parts = k.split("/")
+        if len(parts) >= 2:
+            return f"{parts[-2]}_{parts[-1]}"
+        return parts[-1]
+    return k
+
+
+class AIHubLabelIndexDB:
+    """
+    Zero-RAM SQLite-backed dictionary interface for AI Hub 1.71M lecture transcriptions.
+    Replaces the 681MB in-memory JSON dict (~2.5-4.5GB RAM) with an on-disk B-Tree index (<5MB RAM).
+    Thread-safe across main thread and PrefetchedStream worker threads.
+    """
+    def __init__(self, db_path: Path):
+        self.db_path = Path(db_path)
+        self._local = threading.local()
+        self._count: Optional[int] = None
+
+    def _get_conn(self) -> sqlite3.Connection:
+        conn = getattr(self._local, "conn", None)
+        if conn is None:
+            conn = sqlite3.connect(f"file:{self.db_path}?mode=ro", uri=True, check_same_thread=False)
+            conn.execute("PRAGMA cache_size = -4000")  # Cap SQLite page cache to ~4MB
+            conn.execute("PRAGMA temp_store = MEMORY")
+            self._local.conn = conn
+        return conn
+
+    def get(self, key: str, default: Optional[str] = None) -> Optional[str]:
+        if not key:
+            return default
+        ckey = canonicalize_aihub_key(key)
+        cur = self._get_conn().execute("SELECT text FROM labels WHERE key = ?", (ckey,))
+        row = cur.fetchone()
+        return row[0] if row is not None else default
+
+    def __getitem__(self, key: str) -> str:
+        val = self.get(key)
+        if val is None:
+            raise KeyError(key)
+        return val
+
+    def __contains__(self, key: str) -> bool:
+        return self.get(key) is not None
+
+    def __len__(self) -> int:
+        if self._count is None:
+            cur = self._get_conn().execute("SELECT COUNT(*) FROM labels")
+            row = cur.fetchone()
+            self._count = int(row[0]) if row else 0
+        return self._count
+
+    def values(self):
+        cur = self._get_conn().execute("SELECT text FROM labels")
+        for (txt,) in cur:
+            yield txt
+
+
+def load_or_build_aihub_label_index(
+    tl_zip_path: str = "/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/02.라벨링데이터/TL.zip",
+    index_cache_path: str = "data/vocab/aihub_label_index.db",
+    max_scan_entries: Optional[int] = None
+) -> Any:
+    """
+    Scans TL.zip directly in memory without disk extraction and stores canonical keys
+    in a lightweight SQLite database (data/vocab/aihub_label_index.db) consuming <5MB RAM.
+    """
+    project_root = Path(__file__).resolve().parent.parent
+    cache_path = Path(index_cache_path)
+    if not cache_path.is_absolute():
+        cache_path = project_root / cache_path
+    db_path = cache_path.with_suffix(".db")
+    legacy_json_path = cache_path.with_suffix(".json")
+
+    if db_path.exists() and db_path.stat().st_size > 1024 * 1024:
+        db_index = AIHubLabelIndexDB(db_path)
+        print(f"[AIHub Index] Opened low-memory SQLite label index ({len(db_index):,} utterances, <5MB RAM) at {db_path.name}")
+        return db_index
+
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_db_path = db_path.with_suffix(".db.tmp")
+    if tmp_db_path.exists():
+        tmp_db_path.unlink()
+
+    conn = sqlite3.connect(str(tmp_db_path))
+    conn.execute("PRAGMA journal_mode = OFF")
+    conn.execute("PRAGMA synchronous = OFF")
+    conn.execute("CREATE TABLE labels (key TEXT PRIMARY KEY, text TEXT) WITHOUT ROWID")
+
+    tl_zip = Path(tl_zip_path)
+    if not tl_zip.exists():
+        conn.close()
+        raise FileNotFoundError(f"AI Hub TL.zip not found at {tl_zip}")
+
+    print(f"[AIHub Index] Building low-memory SQLite index from {tl_zip.name} without disk extraction...")
+    batch_rows: List[Tuple[str, str]] = []
+    inserted = 0
+
+    with zipfile.ZipFile(tl_zip, "r") as zf:
+        infolist = zf.infolist()
+        valid_infos = [
+            info for info in infolist
+            if info.filename.endswith(".json")
+            and not info.filename.startswith("__MACOSX")
+            and ".irx" not in info.filename
+            and ".part" not in info.filename
+            and ".DS_Store" not in info.filename
+        ]
+        if max_scan_entries and len(valid_infos) > max_scan_entries:
+            valid_infos = valid_infos[:max_scan_entries]
+
+        for info in tqdm(valid_infos, desc="Indexing TL.zip -> SQLite"):
+            try:
+                raw_bytes = zf.read(info)
+                jdata = json.loads(raw_bytes.decode("utf-8"))
+                raw_txt = jdata.get("06_transcription", {}).get("1_text", "").strip()
+                if not raw_txt:
+                    continue
+
+                clean_txt = re.sub(r"\([^)]*\)/\(([^)]*)\)", r"\1", raw_txt)
+                clean_txt = normalize_text(clean_txt)
+                if not clean_txt:
+                    continue
+
+                src_path = jdata.get("01_dataset", {}).get("3_src_path", "") or info.filename
+                ckey = canonicalize_aihub_key(src_path)
+                batch_rows.append((ckey, clean_txt))
+
+                if len(batch_rows) >= 20000:
+                    conn.executemany("INSERT OR REPLACE INTO labels (key, text) VALUES (?, ?)", batch_rows)
+                    conn.commit()
+                    inserted += len(batch_rows)
+                    batch_rows.clear()
+            except Exception:
+                continue
+
+    if batch_rows:
+        conn.executemany("INSERT OR REPLACE INTO labels (key, text) VALUES (?, ?)", batch_rows)
+        conn.commit()
+        inserted += len(batch_rows)
+        batch_rows.clear()
+
+    conn.close()
+    tmp_db_path.replace(db_path)
+
+    if legacy_json_path.exists():
+        try:
+            legacy_json_path.unlink()
+            print(f"[AIHub Index] Removed legacy 681MB JSON cache ({legacy_json_path.name}) to free disk & RAM.")
+        except Exception:
+            pass
+
+    db_index = AIHubLabelIndexDB(db_path)
+    print(f"[AIHub Index] SQLite index ready at {db_path.name} ({len(db_index):,} unique utterances).")
+    return db_index
+
+
+def is_consonant_leading(token: str) -> bool:
+    """
+    Checks if a token starts with a standalone consonant / jamo (e.g. ᆫ, ᆯ, ᆷ, ᆸ, ᆺ, or compatibility ㄱ-ㅎ).
+    """
+    if not token:
+        return False
+    c = token[0]
+    # Compatibility consonants: U+3131 ~ U+314E
+    if "\u3131" <= c <= "\u314e":
+        return True
+    # Conjoining jamo: choseong (U+1100~U+115F) or jongseong (U+11A8~U+11FF)
+    if "\u1100" <= c <= "\u11ff":
+        return True
+    return False
+
+def merge_jamo_to_syllable(prev: str, curr: str) -> str:
+    """
+    Composes a preceding morpheme with a leading standalone jamo into a complete syllable:
+      e.g. '쌓이' + 'ᆫ' -> '쌓인'
+           '이' + 'ᆸ니다' -> '입니다'
+           '만들' + '어' + '지' + 'ᆫ' -> '만들', '어', '진'
+           '참여' + '하' + 'ᆯ' -> '참여', '할'
+           '만들' + 'ᆫ다' -> '만든다' (ㄹ-탈락 규칙)
+           '살' + 'ᆸ니다' -> '삽니다' (ㄹ-탈락 규칙)
+           '그렇' + 'ᆫ가' -> '그런가' (ㅎ-탈락 규칙)
+    Uses unicodedata NFD decomposition and NFC re-composition.
+    """
+    if not prev:
+        return curr
+    first_char = curr[0]
+    first_mapped = COMPAT_TO_JONGSEONG.get(first_char, first_char)
+    curr_mapped = first_mapped + curr[1:]
+
+    prev_nfd = list(unicodedata.normalize("NFD", prev))
+    curr_nfd = list(unicodedata.normalize("NFD", curr_mapped))
+
+    # If prev ends in a jongseong and curr starts with a jongseong:
+    # in Korean phonological rules (e.g. ㄹ-탈락, ㅎ-탈락), the stem jongseong is dropped
+    if prev_nfd and ("\u11a8" <= prev_nfd[-1] <= "\u11c2"):
+        if curr_nfd and ("\u11a8" <= curr_nfd[0] <= "\u11c2"):
+            test_merged = unicodedata.normalize("NFC", "".join(prev_nfd + curr_nfd))
+            if any(("\u1100" <= c <= "\u11ff" or "\u3131" <= c <= "\u314e") for c in test_merged):
+                prev_nfd = prev_nfd[:-1]
+
+    merged = unicodedata.normalize("NFC", "".join(prev_nfd + curr_nfd))
+    return merged
+
+def strip_word_prefix(tok: str) -> Tuple[bool, str]:
+    """
+    Checks if a token has a word-boundary prefix (' ', ' ▁', or space) and strips it.
+    Returns (is_word_start, clean_token).
+    """
+    if tok.startswith(WORD_PREFIX) or tok.startswith(" ") or tok.startswith(" ▁") or tok.startswith(" "):
+        clean = tok.lstrip(" \u2581▁")
+        return True, clean
+    return False, tok
+
+def tokenize_korean_syllables(text: str, word_prefix: str = WORD_PREFIX) -> List[str]:
+    clean_text = normalize_text(text)
+    if not clean_text:
+        return []
+    tokens: List[str] = []
+    for word in clean_text.split(" "):
+        if not word:
+            continue
+        for idx, ch in enumerate(word):
+            tokens.append((word_prefix + ch) if idx == 0 else ch)
+    return tokens
+
+
+def tokenize_korean_hybrid(
+    text: str,
+    token_to_id: Dict[str, int],
+    word_prefix: str = WORD_PREFIX,
+    max_tok_len: int = 16
+) -> List[str]:
+    """
+    Greedy Longest-Match (Maximal Matching) Surface Morpheme Tokenizer with Syllable Fallback.
+    Matches the longest registered surface morpheme/word in vocabulary first.
+    Falls back to single-character KS X 1001 syllables when a word/morpheme is unregistered,
+    guaranteeing 0.00% OOV with maximum semantic chunking efficiency.
+    """
+    clean_text = normalize_text(text)
+    if not clean_text:
+        return []
+    tokens: List[str] = []
+    for word in clean_text.split(" "):
+        if not word:
+            continue
+        w_len = len(word)
+        pos = 0
+        first_in_word = True
+        while pos < w_len:
+            matched = False
+            max_end = w_len if w_len - pos <= max_tok_len else pos + max_tok_len
+            for end in range(max_end, pos, -1):
+                sub = word[pos:end]
+                candidate = (word_prefix + sub) if first_in_word else sub
+                if candidate in token_to_id:
+                    tokens.append(candidate)
+                    pos = end
+                    first_in_word = False
+                    matched = True
+                    break
+            if not matched:
+                sub = word[pos:pos + 1]
+                candidate = (word_prefix + sub) if first_in_word else sub
+                tokens.append(candidate)
+                pos += 1
+                first_in_word = False
+    return tokens
+
+
+# Backward compatibility alias
+tokenize_korean_morphemes = tokenize_korean_hybrid
+
+
+class KoreanMorphemeTokenizer:
+    """
+    Surface Morpheme + Syllable Fallback Hybrid Tokenizer using generated vocab.json with word boundary prefix.
+    Utilizes O(L) Greedy Longest-Match dictionary lookup for 0.00% OOV and maximum semantic representation.
+    """
+    def __init__(self, vocab_json_path: str):
+        with open(vocab_json_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        self.token_to_id: Dict[str, int] = data["token_to_id"]
+        # Convert integer keys back if serialized as strings in json
+        self.id_to_token: Dict[int, str] = {int(k): v for k, v in data["id_to_token"].items()}
+        self.blank_id: int = data.get("blank_id", 0)
+        self.blank_token: str = data.get("blank_token", "<blank>")
+        self.unk_id: int = data.get("unk_id", 1)
+        self.unk_token: str = data.get("unk_token", "<unk>")
+        self.word_prefix: str = WORD_PREFIX
+        self._max_tok_len: int = max((len(k) for k in self.token_to_id.keys()), default=16)
+        self._word_id_cache: Dict[str, Tuple[int, ...]] = {}
+
+    def _encode_word(self, word: str) -> Tuple[int, ...]:
+        cached = self._word_id_cache.get(word)
+        if cached is not None:
+            return cached
+        w_len = len(word)
+        pos = 0
+        first_in_word = True
+        ids: List[int] = []
+        t2i = self.token_to_id
+        wp = self.word_prefix
+        max_tl = self._max_tok_len
+        unk = self.unk_id
+
+        while pos < w_len:
+            matched = False
+            max_end = w_len if w_len - pos <= max_tl else pos + max_tl
+            for end in range(max_end, pos, -1):
+                sub = word[pos:end]
+                cand = (wp + sub) if first_in_word else sub
+                tid = t2i.get(cand)
+                if tid is not None:
+                    ids.append(tid)
+                    pos = end
+                    first_in_word = False
+                    matched = True
+                    break
+            if not matched:
+                sub = word[pos:pos + 1]
+                cand = (wp + sub) if first_in_word else sub
+                ids.append(t2i.get(cand, unk))
+                pos += 1
+                first_in_word = False
+
+        res = tuple(ids)
+        if len(self._word_id_cache) < 65536:
+            self._word_id_cache[word] = res
+        return res
+
+    def encode(self, text: str) -> List[int]:
+        """
+        Tokenizes text into surface morpheme tokens (with syllable fallback) and maps them to token IDs.
+        """
+        clean_text = normalize_text(text)
+        if not clean_text:
+            return []
+        out_ids: List[int] = []
+        for word in clean_text.split(" "):
+            if word:
+                out_ids.extend(self._encode_word(word))
+        return out_ids
+
+    def decode(self, token_ids: List[int], remove_blank: bool = True, collapse_repeats: bool = True) -> str:
+        """
+        Decodes a sequence of token IDs into text.
+        For CTC output:
+        - collapse_repeats: collapses consecutive duplicate token IDs
+        - remove_blank: removes blank_id (0)
+        - word boundary reconstruction: reconstructs spaces between words using prefix ' ' or ' ▁'
+        """
+        if collapse_repeats:
+            collapsed = []
+            prev = None
+            for tid in token_ids:
+                if tid != prev:
+                    collapsed.append(tid)
+                    prev = tid
+            token_ids = collapsed
+
+        words = []
+        for tid in token_ids:
+            if remove_blank and tid == self.blank_id:
+                continue
+            token_str = self.id_to_token.get(tid, "")
+            if token_str and token_str != self.blank_token and token_str != self.unk_token:
+                words.append(token_str)
+
+        res = ""
+        for w in words:
+            is_ws, clean_w = strip_word_prefix(w)
+            if is_ws:
+                if res:
+                    res += " " + clean_w
+                else:
+                    res = clean_w
+            else:
+                res += clean_w
+
+        return res.strip()
+
+
+_CONTRACTIONS: Tuple[Tuple[str, str], ...] = (
+    ("하었", "했"), ("하 었", "했"), ("되었", "됐"), ("되 었", "됐"),
+    ("나오었", "나왔"), ("나오 었", "나왔"), ("보았", "봤"), ("보 었", "봤"),
+    ("주었", "줬"), ("주 었", "줬"), ("오었", "왔"), ("오 었", "왔"),
+    ("가었", "갔"), ("가 었", "갔"), ("지었", "졌"), ("지 었", "졌"),
+    ("받아들이어", "받아들여")
+)
+
+
+# ==============================================================================
+# 2. Cached Audio Feature Dataset & Length-Bucketed Batch Sampler
+# ==============================================================================
+
+class _SharedFilePathAccessor:
+    """Zero-RAM sequence wrapper over shared-memory byte & offset tensors."""
+    def __init__(self, base_dir: str, names_buf: torch.Tensor, offsets_buf: torch.Tensor):
+        self.base_dir = base_dir
+        self.names_buf = names_buf
+        self.offsets_buf = offsets_buf
+
+    def __len__(self) -> int:
+        return max(0, int(self.offsets_buf.numel()) - 1)
+
+    def get_name(self, idx: int) -> str:
+        s = int(self.offsets_buf[idx].item())
+        e = int(self.offsets_buf[idx + 1].item())
+        return bytes(self.names_buf[s:e].numpy()).decode("utf-8")
+
+    def __getitem__(self, idx: int) -> str:
+        return os.path.join(self.base_dir, self.get_name(idx))
+
+
+class CachedAudioFeatureDataset(Dataset):
+    """
+    Loads pre-extracted (Time, 256) audio encoder representations and target token IDs.
+    Uses OS shared-memory tensors (.share_memory_()) for file names and sizes so DataLoader
+    worker processes share a single 68MB index with zero Copy-on-Write or pickling bloat.
+    """
+    def __init__(
+        self,
+        cache_dir: str,
+        tokenizer: Optional[KoreanMorphemeTokenizer] = None,
+        max_files: Optional[int] = None
+    ):
+        super().__init__()
+        self.cache_dir = Path(cache_dir)
+        self.cache_dir_str = str(self.cache_dir)
+        self.tokenizer = tokenizer
+
+        index_cache_file = self.cache_dir.parent / "vocab" / "feature_cache_index.pt"
+        loaded_from_cache = False
+
+        if max_files is None and index_cache_file.exists():
+            try:
+                cached_idx = torch.load(index_cache_file, map_location="cpu", weights_only=False)
+                if (
+                    isinstance(cached_idx, dict)
+                    and cached_idx.get("cache_dir") == self.cache_dir_str
+                    and int(cached_idx.get("count", 0)) > 0
+                ):
+                    self._names_buf = cached_idx["names_buf"].share_memory_()
+                    self._offsets_buf = cached_idx["offsets_buf"].share_memory_()
+                    self.file_sizes = cached_idx["sizes_buf"].share_memory_()
+                    loaded_from_cache = True
+            except Exception:
+                loaded_from_cache = False
+
+        if not loaded_from_cache:
+            names_list: List[bytes] = []
+            sizes_list: List[int] = []
+            if self.cache_dir.exists():
+                entries = [e for e in os.scandir(self.cache_dir) if e.name.endswith(".pt")]
+                entries.sort(key=lambda e: e.name)
+                if max_files is not None and max_files > 0:
+                    entries = entries[:max_files]
+                names_list = [e.name.encode("utf-8") for e in entries]
+                sizes_list = [e.stat().st_size for e in entries]
+
+            if names_list:
+                lengths = [len(b) for b in names_list]
+                offsets = [0]
+                acc = 0
+                for ln in lengths:
+                    acc += ln
+                    offsets.append(acc)
+                joined = b"".join(names_list)
+                self._names_buf = torch.frombuffer(bytearray(joined), dtype=torch.uint8).share_memory_()
+                self._offsets_buf = torch.tensor(offsets, dtype=torch.int64).share_memory_()
+                self.file_sizes = torch.tensor(sizes_list, dtype=torch.int32).share_memory_()
+
+                if max_files is None:
+                    try:
+                        index_cache_file.parent.mkdir(parents=True, exist_ok=True)
+                        tmp_idx = index_cache_file.with_suffix(".pt.tmp")
+                        torch.save(
+                            {
+                                "cache_dir": self.cache_dir_str,
+                                "count": len(names_list),
+                                "names_buf": self._names_buf,
+                                "offsets_buf": self._offsets_buf,
+                                "sizes_buf": self.file_sizes,
+                            },
+                            tmp_idx
+                        )
+                        tmp_idx.replace(index_cache_file)
+                    except Exception:
+                        pass
+            else:
+                self._names_buf = torch.empty((0,), dtype=torch.uint8).share_memory_()
+                self._offsets_buf = torch.zeros((1,), dtype=torch.int64).share_memory_()
+                self.file_sizes = torch.empty((0,), dtype=torch.int32).share_memory_()
+                print(f"[Dataset Warning] No cached .pt files found in {self.cache_dir}.")
+
+        self.file_paths = _SharedFilePathAccessor(self.cache_dir_str, self._names_buf, self._offsets_buf)
+
+    def __len__(self) -> int:
+        return int(self.file_sizes.numel())
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        fname = self.file_paths.get_name(idx)
+        file_path = os.path.join(self.cache_dir_str, fname)
+        data = torch.load(file_path, map_location="cpu", weights_only=False)
+        feature = data["feature"]  # Shape: (T, 256)
+        raw_text = data.get("raw_text", "")
+
+        if self.tokenizer is not None and raw_text:
+            tok_ids = self.tokenizer.encode(raw_text)
+            targets = torch.tensor(tok_ids, dtype=torch.long)
+        elif "targets" in data and isinstance(data["targets"], torch.Tensor):
+            targets = data["targets"].to(torch.long)
+        else:
+            targets = torch.tensor([], dtype=torch.long)
+
+        if not isinstance(feature, torch.Tensor):
+            feature = torch.tensor(feature, dtype=torch.bfloat16)
+
+        # Strictly enforce bfloat16 (prevent MPS float16 underflow and float32 NaN instability)
+        if feature.dtype != torch.bfloat16:
+            feature = feature.to(torch.bfloat16)
+
+        return {
+            "feature": feature,
+            "targets": targets,
+            "raw_text": raw_text,
+            "file_name": fname
+        }
+
+
+class LengthBucketBatchSampler(torch.utils.data.Sampler[List[int]]):
+    """
+    Groups samples of similar sequence lengths into batches using cached .pt file sizes
+    as a zero-I/O O(1) proxy for frame length T (~512 bytes per time frame).
+    Eliminates ~40% random-batching padding overhead while preserving stochasticity
+    via mega-batch pool sorting and batch-order shuffling each epoch.
+    """
+    def __init__(
+        self,
+        dataset: Dataset,
+        batch_size: int,
+        shuffle: bool = True,
+        pool_batches: int = 128,
+        seed: int = 42
+    ):
+        super().__init__()
+        self.dataset = dataset
+        self.batch_size = max(1, batch_size)
+        self.shuffle = shuffle
+        self.pool_size = self.batch_size * max(1, pool_batches)
+        self.seed = seed
+        self.epoch = 0
+        self.start_batch = 0
+
+        # Resolve per-item file size proxy as compact int32 tensor
+        if hasattr(dataset, "file_sizes"):
+            fs = dataset.file_sizes
+            self.lengths_t = fs if isinstance(fs, torch.Tensor) else torch.tensor(fs, dtype=torch.int32)
+        elif hasattr(dataset, "dataset") and hasattr(dataset, "indices") and hasattr(dataset.dataset, "file_sizes"):
+            base_sizes = dataset.dataset.file_sizes
+            base_t = base_sizes if isinstance(base_sizes, torch.Tensor) else torch.tensor(base_sizes, dtype=torch.int32)
+            idx_t = dataset.indices if isinstance(dataset.indices, torch.Tensor) else torch.tensor(dataset.indices, dtype=torch.long)
+            self.lengths_t = base_t[idx_t]
+        else:
+            self.lengths_t = torch.zeros((len(dataset),), dtype=torch.int32)
+
+    def set_epoch(self, epoch: int, start_batch: int = 0) -> None:
+        """Sets current epoch (0-indexed) and optional starting batch offset for zero-I/O resume."""
+        self.epoch = max(0, int(epoch))
+        self.start_batch = max(0, int(start_batch))
+
+    def __len__(self) -> int:
+        n = int(self.lengths_t.numel())
+        total = (n + self.batch_size - 1) // self.batch_size
+        return max(0, total - self.start_batch)
+
+    def __iter__(self):
+        n = int(self.lengths_t.numel())
+        start_b = self.start_batch
+        self.start_batch = 0  # Reset after consuming once
+
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            self.epoch += 1
+            perm = torch.randperm(n, generator=g)
+
+            # Vectorized mega-pool sorting in C++ (torch.argsort) instead of Python list sorting
+            num_full_pools = n // self.pool_size
+            if num_full_pools > 0:
+                full_len = num_full_pools * self.pool_size
+                perm_main = perm[:full_len].view(num_full_pools, self.pool_size)
+                pool_lengths = self.lengths_t[perm_main]
+                sort_idx = torch.argsort(pool_lengths, dim=1, stable=True)
+                sorted_main = torch.gather(perm_main, 1, sort_idx).reshape(-1)
+                if full_len < n:
+                    rem = perm[full_len:]
+                    rem_sorted = rem[torch.argsort(self.lengths_t[rem], stable=True)]
+                    sorted_indices = torch.cat([sorted_main, rem_sorted], dim=0)
+                else:
+                    sorted_indices = sorted_main
+            else:
+                sorted_indices = perm[torch.argsort(self.lengths_t[perm], stable=True)]
+
+            all_indices = sorted_indices.tolist()
+            batches = [
+                all_indices[b_start : b_start + self.batch_size]
+                for b_start in range(0, n, self.batch_size)
+            ]
+
+            b_perm = torch.randperm(len(batches), generator=g).tolist()
+            for bi in b_perm[start_b:]:
+                yield batches[bi]
+        else:
+            sorted_indices = torch.argsort(self.lengths_t, stable=True).tolist()
+            all_batches = [
+                sorted_indices[b_start : b_start + self.batch_size]
+                for b_start in range(0, n, self.batch_size)
+            ]
+            for b in all_batches[start_b:]:
+                yield b
+
+
+# ==============================================================================
+# 3. Dynamic Padding Collate Function for CTC
+# ==============================================================================
+
+class PadCollate:
+    """
+    Collates a list of variable-length audio features and targets:
+    - Caps outlier audio lengths to max_frames (default 512 = ~20.5s) to prevent MPS OOM spikes
+    - Pads features to maximum time length in the batch -> (Batch, T_max, 256)
+    - Generates key_padding_mask -> (Batch, T_max) for Conformer attention
+    - Concatenates flat_targets for CTC Loss
+    - Pre-computes active_ids (int32) and compact_targets (int64) in background DataLoader workers
+      for FusedMetalCosineCTCLoss
+    """
+    def __init__(
+        self,
+        pad_feature_value: float = 0.0,
+        pad_target_value: int = 0,
+        pad_multiple: int = 16,
+        max_frames: int = 512
+    ):
+        self.pad_feature_val = pad_feature_value
+        self.pad_target_val = pad_target_value
+        self.pad_multiple = pad_multiple
+        self.max_frames = max_frames
+
+    def __call__(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
+        # Filter out corrupted, empty, or CTC-invalid (T < L) entries
+        batch = [
+            item for item in batch
+            if item["feature"].size(0) > 0 and 0 < item["targets"].size(0) <= item["feature"].size(0)
+        ]
+        if not batch:
+            return {}
+
+        # Cap outlier sequence lengths while guaranteeing CTC constraint (T >= target_len)
+        if self.max_frames and self.max_frames > 0:
+            for item in batch:
+                t_len = item["feature"].size(0)
+                if t_len > self.max_frames:
+                    safe_cap = max(self.max_frames, int(item["targets"].size(0)) + 16)
+                    if t_len > safe_cap:
+                        item["feature"] = item["feature"][:safe_cap]
+
+        batch_size = len(batch)
+        feat_dim = batch[0]["feature"].size(-1)
+        dtype = batch[0]["feature"].dtype
+
+        # Extract sequence lengths
+        input_lengths = torch.tensor([item["feature"].size(0) for item in batch], dtype=torch.long)
+        target_lengths = torch.tensor([item["targets"].size(0) for item in batch], dtype=torch.long)
+
+        actual_max_t = int(input_lengths.max().item())
+        max_t = actual_max_t
+        if self.pad_multiple > 1:
+            max_t = ((max_t + self.pad_multiple - 1) // self.pad_multiple) * self.pad_multiple
+
+        # Pad features: (B, max_t, 256)
+        padded_features = torch.full((batch_size, max_t, feat_dim), self.pad_feature_val, dtype=dtype)
+        # Key padding mask: True for padded positions (B, max_t)
+        key_padding_mask = torch.zeros((batch_size, max_t), dtype=torch.bool)
+
+        for i, item in enumerate(batch):
+            t_len = item["feature"].size(0)
+            padded_features[i, :t_len] = item["feature"]
+            if t_len < max_t:
+                key_padding_mask[i, t_len:] = True
+
+        # Flattened targets for standard PyTorch CTCLoss compatibility: (sum(target_lengths),)
+        flat_targets = torch.cat([item["targets"] for item in batch], dim=0)
+
+        # Pre-compute unique active target IDs and compact targets in DataLoader worker thread
+        uniq_t, inv_idx = torch.unique(flat_targets, sorted=True, return_inverse=True)
+        active_ids = torch.cat([torch.zeros(1, dtype=torch.int32), uniq_t.to(torch.int32)])
+        compact_targets = inv_idx + 1
+
+        raw_texts = [item.get("raw_text", "") for item in batch]
+
+        return {
+            "features": padded_features,          # (B, T_max, 256)
+            "input_lengths": input_lengths,        # (B,)
+            "max_valid_len": actual_max_t,         # int (unbucketed max T in batch)
+            "key_padding_mask": key_padding_mask,  # (B, T_max)
+            "flat_targets": flat_targets,          # (sum(L),)
+            "active_ids": active_ids,              # (U,) int32
+            "compact_targets": compact_targets,    # (sum(L),) int64
+            "target_lengths": target_lengths,      # (B,)
+            "raw_texts": raw_texts
+        }

src/decoder.py

@@ -0,0 +1,336 @@
+"""
+src/decoder.py
+
+Production On-Device CTC Beam Search Decoder with KenLM Language Model Rescoring.
+Integrates Korean Morpheme Tokenizer, Trie N-Gram language model, and contraction rules.
+"""
+
+from __future__ import annotations
+
+import math
+import ctypes
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional, Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+try:
+    import kenlm
+except ImportError:
+    kenlm = None
+
+from dataset import KoreanMorphemeTokenizer, normalize_text
+
+# CTC Beam Search Engine with KenLM Language Model
+_FAST_BEAM_LIB = None
+
+
+def _log_add(a: float, b: float) -> float:
+    """Numerically stable log-sum-exp for two scalars: log(exp(a) + exp(b))."""
+    if a <= -1e20:
+        return b
+    if b <= -1e20:
+        return a
+    diff = a - b
+    if diff > 50.0:
+        return a
+    if diff < -50.0:
+        return b
+    return a + math.log1p(math.exp(-diff)) if diff >= 0 else b + math.log1p(math.exp(diff))
+
+
+class KenLMBeamSearchDecoder:
+    """
+    On-device CTC Beam Search Decoder equipped with KenLM Language Model:
+    - Maintains top-K acoustic beams across temporal frames
+    - Rescores word sequences using KenLM 3-gram language model
+    - Integrates pure syllable decoding and batchim/contraction rules
+    - Runs in pure CPU C++ bindings in ~5-15ms per utterance
+    """
+    def __init__(
+        self,
+        tokenizer: KoreanMorphemeTokenizer,
+        lm_path: Optional[str] = "data/vocab/korean_lm.bin",
+        beam_size: int = 8,
+        alpha: float = 0.5,      # Language model weight
+        beta: float = 1.0,       # Word insertion bonus
+        blank_id: int = 0
+    ):
+        self.tokenizer = tokenizer
+        self.beam_size = max(1, beam_size)
+        self.alpha = alpha
+        self.beta = beta
+        self.blank_id = blank_id
+
+        self.lm = None
+        if lm_path and kenlm is not None:
+            p = Path(lm_path)
+            if not p.is_absolute():
+                p = Path(__file__).resolve().parent.parent / p
+            if p.exists():
+                self.lm = kenlm.LanguageModel(str(p))
+                print(f"[Decoder] Loaded KenLM language model ({p.stat().st_size / (1024*1024):.1f} MB) from {p.name}")
+            else:
+                print(f"[Decoder Warning] KenLM file not found at {p}. Falling back to acoustic beam search.")
+
+    def decode_greedy(self, logits: torch.Tensor) -> str:
+        """Standard Argmax Greedy Decoding (baseline)."""
+        preds = torch.argmax(logits, dim=-1).cpu().tolist()
+        return self.tokenizer.decode(preds, remove_blank=True, collapse_repeats=True)
+
+    def decode_beam(
+        self,
+        logits: torch.Tensor,
+        beam_size: Optional[int] = None,
+        alpha: Optional[float] = None,
+        beta: Optional[float] = None
+    ) -> str:
+        """
+        Executes Time-Synchronous CTC Beam Search with KenLM Rescoring.
+        Args:
+            logits: Tensor of shape (Time, Vocab) or (1, Time, Vocab)
+        Returns:
+            Best decoded Korean text
+        """
+        if logits.dim() == 3:
+            logits = logits.squeeze(0)
+
+        k_beam = beam_size or self.beam_size
+        a = alpha if alpha is not None else self.alpha
+        b = beta if beta is not None else self.beta
+
+        # Convert logits to log-probabilities (Time, Vocab)
+        log_probs = F.log_softmax(logits, dim=-1).cpu()
+        T, V = log_probs.shape
+
+        # Beam state: (tuple_of_non_blank_tokens, last_token) -> (p_blank, p_non_blank)
+        # Empty hypothesis at t=0
+        beams: Dict[Tuple[int, ...], Tuple[float, float, int]] = {
+            (): (0.0, -float("inf"), self.blank_id)
+        }
+
+        # Consider top-N most likely tokens per frame to keep beam search blazing fast (<10ms)
+        top_k_candidates = min(40, V)
+
+        for t in range(T):
+            frame_probs = log_probs[t]
+            top_vals, top_indices = torch.topk(frame_probs, top_k_candidates)
+            top_vals = top_vals.tolist()
+            top_indices = top_indices.tolist()
+
+            next_beams: Dict[Tuple[int, ...], Tuple[float, float, int]] = {}
+
+            for prefix, (p_b, p_nb, last_tok) in beams.items():
+                p_total = math.log(math.exp(p_b) + (math.exp(p_nb) if p_nb > -float("inf") else 0.0) + 1e-30)
+
+                # 1. Transition to Blank
+                prob_blank = frame_probs[self.blank_id].item()
+                nb_p_b, nb_p_nb, _ = next_beams.get(prefix, (-float("inf"), -float("inf"), self.blank_id))
+                new_p_b = math.log(math.exp(nb_p_b) + math.exp(p_total + prob_blank) + 1e-30)
+                next_beams[prefix] = (new_p_b, nb_p_nb, self.blank_id)
+
+                # 2. Transition to Non-Blank tokens
+                for p_val, tok_id in zip(top_vals, top_indices):
+                    if tok_id == self.blank_id:
+                        continue
+
+                    prob_tok = p_val
+                    if tok_id == last_tok:
+                        # Same token repeated immediately without blank: only extends if coming from blank
+                        new_prefix = prefix
+                        nb_p_b, nb_p_nb, _ = next_beams.get(new_prefix, (-float("inf"), -float("inf"), tok_id))
+                        add_p = math.log(math.exp(nb_p_nb) + math.exp(p_nb + prob_tok) + 1e-30)
+                        next_beams[new_prefix] = (nb_p_b, add_p, tok_id)
+
+                        # Coming from blank creates new token
+                        extended_prefix = prefix + (tok_id,)
+                        enb_p_b, enb_p_nb, _ = next_beams.get(extended_prefix, (-float("inf"), -float("inf"), tok_id))
+                        new_enb_p = math.log(math.exp(enb_p_nb) + math.exp(p_b + prob_tok) + 1e-30)
+                        next_beams[extended_prefix] = (enb_p_b, new_enb_p, tok_id)
+                    else:
+                        new_prefix = prefix + (tok_id,)
+                        nb_p_b, nb_p_nb, _ = next_beams.get(new_prefix, (-float("inf"), -float("inf"), tok_id))
+                        new_p_nb = math.log(math.exp(nb_p_nb) + math.exp(p_total + prob_tok) + 1e-30)
+                        next_beams[new_prefix] = (nb_p_b, new_p_nb, tok_id)
+
+            # Prune beams to top-k
+            def _score_beam(item):
+                pref, (pb, pnb, _) = item
+                acoustic_score = math.log(
+                    (math.exp(pb) if pb > -float("inf") else 0.0) +
+                    (math.exp(pnb) if pnb > -float("inf") else 0.0) +
+                    1e-30
+                )
+                return acoustic_score
+
+            sorted_items = sorted(next_beams.items(), key=_score_beam, reverse=True)[: k_beam * 3]
+            beams = {pref: val for pref, val in sorted_items}
+
+        # Final Rescoring with KenLM
+        best_hypothesis = ""
+        best_total_score = -float("inf")
+
+        for prefix, (p_b, p_nb, _) in beams.items():
+            acoustic_score = _log_add(p_b, p_nb)
+
+            decoded_text = self.tokenizer.decode(list(prefix), remove_blank=True, collapse_repeats=False)
+            if not decoded_text:
+                continue
+
+            lm_score = 0.0
+            num_words = len(decoded_text.split())
+            if self.lm is not None:
+                # KenLM log10 score
+                lm_score = self.lm.score(decoded_text)
+
+            total_score = acoustic_score + (a * lm_score) + (b * num_words)
+            if total_score > best_total_score:
+                best_total_score = total_score
+                best_hypothesis = decoded_text
+
+        if not best_hypothesis:
+            # Fallback to greedy if all beams empty
+            best_hypothesis = self.decode_greedy(logits)
+
+        return best_hypothesis
+
+    def decode_topk(
+        self,
+        topk_ids: np.ndarray,
+        topk_logprobs: np.ndarray,
+        blank_logprobs: np.ndarray,
+        beam_size: Optional[int] = None,
+        alpha: Optional[float] = None,
+        beta: Optional[float] = None
+    ) -> str:
+        """
+        Ultra-fast Time-Synchronous CTC Beam Search:
+        - Uses Native C++ Rolling-Hash Engine (<2ms) if available
+        - Direct KenLM rescoring on candidate beams
+        """
+        k_beam = beam_size or self.beam_size
+        a = alpha if alpha is not None else self.alpha
+        b = beta if beta is not None else self.beta
+
+        T, K = topk_ids.shape
+
+        # -------------------------------------------------------------
+        # C++ Ultra-Fast Fast-Path (< 2ms for 375 frames)
+        # -------------------------------------------------------------
+        if _FAST_BEAM_LIB is not None:
+            max_beams = max(16, k_beam * 2)
+            max_seq_len = 512
+            out_tokens = np.zeros((max_beams, max_seq_len), dtype=np.int32)
+            out_lens = np.zeros(max_beams, dtype=np.int32)
+            out_scores = np.zeros(max_beams, dtype=np.float32)
+
+            # Ensure contiguous memory
+            t_ids_c = np.ascontiguousarray(topk_ids, dtype=np.int32)
+            t_lprob_c = np.ascontiguousarray(topk_logprobs, dtype=np.float32)
+            blank_lp_c = np.ascontiguousarray(blank_logprobs, dtype=np.float32)
+
+            n_hyp = _FAST_BEAM_LIB.ctc_beam_search_fast_c(
+                t_ids_c.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
+                t_lprob_c.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
+                blank_lp_c.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
+                T, K, k_beam, self.blank_id,
+                out_tokens.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
+                out_lens.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
+                out_scores.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
+                max_beams, max_seq_len
+            )
+
+            best_hypothesis = ""
+            best_total_score = -float("inf")
+
+            for i in range(n_hyp):
+                seq_l = out_lens[i]
+                if seq_l == 0:
+                    continue
+                toks = out_tokens[i, :seq_l].tolist()
+                text = self.tokenizer.decode(toks, remove_blank=True, collapse_repeats=True)
+                lm_score = self.lm.score(text) if self.lm else 0.0
+                words = len(text.split())
+                score = out_scores[i] + (a * lm_score) + (b * words)
+                if score > best_total_score:
+                    best_total_score = score
+                    best_hypothesis = text
+
+            if best_hypothesis:
+                return best_hypothesis
+
+        # -------------------------------------------------------------
+        # Fallback Python Path
+        # -------------------------------------------------------------
+        beams: Dict[Tuple[int, ...], Tuple[float, float, int]] = {
+            (): (0.0, -float("inf"), self.blank_id)
+        }
+
+        for t in range(T):
+            p_blank = float(blank_logprobs[t])
+            t_ids = topk_ids[t]
+            t_lprobs = topk_logprobs[t]
+
+            next_beams: Dict[Tuple[int, ...], Tuple[float, float, int]] = {}
+
+            for prefix, (p_b, p_nb, last_tok) in beams.items():
+                p_total = _log_add(p_b, p_nb)
+
+                # 1. Transition to Blank
+                nb_p_b, nb_p_nb, _ = next_beams.get(prefix, (-float("inf"), -float("inf"), self.blank_id))
+                new_p_b = _log_add(nb_p_b, p_total + p_blank)
+                next_beams[prefix] = (new_p_b, nb_p_nb, self.blank_id)
+
+                # 2. Transition to Top-K Non-Blank tokens
+                for k_idx in range(K):
+                    tok_id = int(t_ids[k_idx])
+                    if tok_id == self.blank_id:
+                        continue
+                    prob_tok = float(t_lprobs[k_idx])
+
+                    if tok_id == last_tok:
+                        new_prefix = prefix
+                        nb_p_b, nb_p_nb, _ = next_beams.get(new_prefix, (-float("inf"), -float("inf"), tok_id))
+                        add_p = _log_add(nb_p_nb, p_nb + prob_tok)
+                        next_beams[new_prefix] = (nb_p_b, add_p, tok_id)
+
+                        extended_prefix = prefix + (tok_id,)
+                        enb_p_b, enb_p_nb, _ = next_beams.get(extended_prefix, (-float("inf"), -float("inf"), tok_id))
+                        new_enb_p = _log_add(enb_p_nb, p_b + prob_tok)
+                        next_beams[extended_prefix] = (enb_p_b, new_enb_p, tok_id)
+                    else:
+                        new_prefix = prefix + (tok_id,)
+                        nb_p_b, nb_p_nb, _ = next_beams.get(new_prefix, (-float("inf"), -float("inf"), tok_id))
+                        new_p_nb = _log_add(nb_p_nb, p_total + prob_tok)
+                        next_beams[new_prefix] = (nb_p_b, new_p_nb, tok_id)
+
+            def _score_beam(item):
+                pref, (pb, pnb, _) = item
+                return _log_add(pb, pnb)
+
+            sorted_items = sorted(next_beams.items(), key=_score_beam, reverse=True)[: k_beam * 3]
+            beams = {pref: val for pref, val in sorted_items}
+
+        best_hypothesis = ""
+        best_total_score = -float("inf")
+
+        for prefix, (p_b, p_nb, _) in beams.items():
+            acoustic_score = _log_add(p_b, p_nb)
+
+            decoded_text = self.tokenizer.decode(list(prefix), remove_blank=True, collapse_repeats=False)
+            if not decoded_text:
+                continue
+
+            lm_score = 0.0
+            num_words = len(decoded_text.split())
+            if self.lm is not None:
+                lm_score = self.lm.score(decoded_text)
+
+            total_score = acoustic_score + (a * lm_score) + (b * num_words)
+            if total_score > best_total_score:
+                best_total_score = total_score
+                best_hypothesis = decoded_text
+
+        return best_hypothesis

src/model.py

@@ -0,0 +1,783 @@
+"""
+src/model.py
+
+Non-Autoregressive On-Device CTC Reader Architecture for EmbeddingGemma-2.
+Maps frozen 300M Audio Encoder representations (MRL 256-dim) to the shared multimodal
+vector space, matching fixed text embeddings via cosine similarity projection.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from typing import Optional, Tuple, Dict, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Standard PyTorch Execution Mode
+MetalFusedConvMidFunction = None
+MetalFusedQKVLinearFunction = None
+MetalFusedResL2NormFunction = None
+
+
+# ==============================================================================
+# 1. Conformer / Transformer Building Blocks (Mobile-Friendly Standard Ops)
+# ==============================================================================
+
+class DepthwiseConv1d(nn.Module):
+    """
+    Depthwise separable 1D convolution block with GLU and BatchNorm,
+    providing temporal context modeling without heavy parameter overhead.
+    Designed with standard PyTorch ops for seamless CoreML / ONNX export.
+    """
+    def __init__(self, channels: int, kernel_size: int = 15, dropout: float = 0.1):
+        super().__init__()
+        assert kernel_size % 2 == 1, "Kernel size must be odd for same-padding"
+        self.padding = kernel_size // 2
+
+        self.pointwise_in = nn.Conv1d(channels, channels * 2, kernel_size=1)
+        self.act_glu = nn.GLU(dim=1)
+        self.depthwise = nn.Conv1d(
+            channels, channels,
+            kernel_size=kernel_size,
+            padding=self.padding,
+            groups=channels,
+            bias=False
+        )
+        self.bn = nn.BatchNorm1d(channels)
+        self.act_swish = nn.SiLU()
+        self.pointwise_out = nn.Conv1d(channels, channels, kernel_size=1)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Fast-path on Apple Silicon MPS with bfloat16: F.linear pointwise + Fused Metal GLU + DWConv15 + BN + SiLU
+        if (
+            x.device.type == "mps"
+            and x.dtype == torch.bfloat16
+            and self.depthwise.kernel_size[0] == 15
+            and MetalFusedConvMidFunction is not None
+            and not torch.jit.is_tracing()
+            and not torch.onnx.is_in_onnx_export()
+        ):
+            res = x
+            u = F.linear(x, self.pointwise_in.weight.squeeze(-1), self.pointwise_in.bias)
+            y = MetalFusedConvMidFunction.apply(
+                u,
+                self.depthwise.weight,
+                self.bn.weight,
+                self.bn.bias,
+                self.bn.running_mean,
+                self.bn.running_var,
+                self.training
+            )
+            y = F.linear(y, self.pointwise_out.weight.squeeze(-1), self.pointwise_out.bias)
+            return res + self.dropout(y)
+
+        # Standard mobile-exportable path: (Batch, Time, Channels) -> transpose to (Batch, Channels, Time)
+        res = x
+        x = x.transpose(1, 2)
+        x = self.pointwise_in(x)
+        x = self.act_glu(x)
+        x = self.depthwise(x)
+        x = self.bn(x)
+        x = self.act_swish(x)
+        x = self.pointwise_out(x)
+        x = self.dropout(x)
+        x = x.transpose(1, 2)
+        return res + x
+
+
+class ExplicitMultiHeadAttention(nn.Module):
+    """
+    Explicit Multi-Head Attention implemented using standard tensor operations
+    (MatMul, Scale, Softmax, MatMul) without proprietary PyTorch C++ kernels.
+    Guarantees 100% compatibility with ONNX export (opset 17) and CoreML conversion.
+    """
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        attn_bias: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        b, t_q, _ = q.size()
+        _, t_k, _ = k.size()
+
+        # Fused QKV self-attention fast-path on MPS bfloat16
+        if (
+            q is k
+            and k is v
+            and q.device.type == "mps"
+            and not torch.jit.is_tracing()
+            and not torch.onnx.is_in_onnx_export()
+        ):
+            qkv_w_view = getattr(self, "_qkv_w_view", None)
+            if qkv_w_view is not None and MetalFusedQKVLinearFunction is not None:
+                if torch.is_grad_enabled() and q.requires_grad:
+                    qkv_lin = MetalFusedQKVLinearFunction.apply(
+                        q, qkv_w_view, self._qkv_b_view, self._qkv_w_grad_view, self._qkv_b_grad_view
+                    )
+                else:
+                    qkv_lin = F.linear(q, qkv_w_view, self._qkv_b_view)
+            else:
+                w_qkv = torch.cat([self.q_proj.weight, self.k_proj.weight, self.v_proj.weight], dim=0)
+                b_qkv = torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias], dim=0)
+                qkv_lin = F.linear(q, w_qkv, b_qkv)
+            qkv = qkv_lin.view(b, t_q, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+            q_proj, k_proj, v_proj = qkv[0], qkv[1], qkv[2]
+
+            scores = torch.matmul(q_proj, k_proj.transpose(-2, -1)) * self.scale
+            if attn_bias is not None:
+                scores = scores + attn_bias
+            elif key_padding_mask is not None:
+                scores = scores.masked_fill(key_padding_mask[:, None, None, :], -10000.0)
+
+            probs = F.softmax(scores.float(), dim=-1).to(dtype=q.dtype)
+            if self.training and self.dropout.p > 0.0:
+                probs = F.dropout(probs, p=self.dropout.p, training=True)
+            context = torch.matmul(probs, v_proj).transpose(1, 2).reshape(b, t_q, self.embed_dim)
+            return self.out_proj(context)
+
+        q_proj = self.q_proj(q).view(b, t_q, self.num_heads, self.head_dim).transpose(1, 2)
+        k_proj = self.k_proj(k).view(b, t_k, self.num_heads, self.head_dim).transpose(1, 2)
+        v_proj = self.v_proj(v).view(b, t_k, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Scaled dot-product attention using MPS-native fused SDPA kernel
+        attn_mask = None
+        if attn_bias is not None:
+            attn_mask = attn_bias
+        elif key_padding_mask is not None:
+            # key_padding_mask: True for pad positions (b, t_k)
+            # In SDPA boolean mask: True means attend (valid), False means mask out (pad)
+            attn_mask = (~key_padding_mask).unsqueeze(1).unsqueeze(2)
+
+        dropout_p = self.dropout.p if self.training else 0.0
+        context = F.scaled_dot_product_attention(
+            q_proj, k_proj, v_proj,
+            attn_mask=attn_mask,
+            dropout_p=dropout_p
+        )  # (b, heads, t_q, head_dim)
+        context = context.transpose(1, 2).contiguous().view(b, t_q, self.embed_dim)
+        return self.out_proj(context)
+
+
+class ConformerBlock(nn.Module):
+    """
+    A lightweight Conformer block combining:
+    - Pre-norm Multi-Head Self-Attention (MHSA)
+    - Depthwise Convolution for local speech features
+    - Feed-Forward Network (FFN) with half-step Macaron-style scaling
+    """
+    def __init__(
+        self,
+        d_model: int = 256,
+        nhead: int = 4,
+        dim_feedforward: int = 1024,
+        dropout: float = 0.1,
+        conv_kernel_size: int = 15
+    ):
+        super().__init__()
+        # 1. First FFN (Macaron style)
+        self.ffn1_norm = nn.LayerNorm(d_model)
+        self.ffn1 = nn.Sequential(
+            nn.Linear(d_model, dim_feedforward),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim_feedforward, d_model),
+            nn.Dropout(dropout)
+        )
+
+        # 2. Multi-Head Self-Attention (Explicit ops for mobile deployment)
+        self.attn_norm = nn.LayerNorm(d_model)
+        self.self_attn = ExplicitMultiHeadAttention(
+            embed_dim=d_model,
+            num_heads=nhead,
+            dropout=dropout
+        )
+        self.attn_dropout = nn.Dropout(dropout)
+
+        # 3. Convolution Module
+        self.conv_norm = nn.LayerNorm(d_model)
+        self.conv_module = DepthwiseConv1d(
+            channels=d_model,
+            kernel_size=conv_kernel_size,
+            dropout=dropout
+        )
+
+        # 4. Second FFN
+        self.ffn2_norm = nn.LayerNorm(d_model)
+        self.ffn2 = nn.Sequential(
+            nn.Linear(d_model, dim_feedforward),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim_feedforward, d_model),
+            nn.Dropout(dropout)
+        )
+
+        self.final_norm = nn.LayerNorm(d_model)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        attn_bias: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        # FFN 1 with 0.5 residual scale
+        x = x + 0.5 * self.ffn1(self.ffn1_norm(x))
+
+        # Multi-Head Attention
+        normed = self.attn_norm(x)
+        attn_out = self.self_attn(
+            normed, normed, normed,
+            key_padding_mask=key_padding_mask,
+            attn_bias=attn_bias
+        )
+        x = x + self.attn_dropout(attn_out)
+
+        # Convolution Module
+        x = x + self.conv_module(self.conv_norm(x))
+
+        # FFN 2 with 0.5 residual scale
+        x = x + 0.5 * self.ffn2(self.ffn2_norm(x))
+
+        return self.final_norm(x)
+
+
+# ==============================================================================
+# 2. Audio-to-Vector Reader Module
+# ==============================================================================
+
+class AudioToVectorReader(nn.Module):
+    """
+    Reader v2 Conformer Architecture (~21.5M parameters).
+    Input: Audio Encoder features of shape (Batch, Time, input_dim=256)
+    Internal: Conformer latent space (d_model=384, nhead=6, num_layers=10, dim_feedforward=1536)
+    Output: Normalized vectors of shape (Batch, Time, output_dim=256) aligned with EmbeddingGemma-2 text space.
+    """
+    def __init__(
+        self,
+        input_dim: int = 256,
+        output_dim: int = 256,
+        d_model: int = 384,
+        nhead: int = 6,
+        num_layers: int = 10,
+        dim_feedforward: int = 1536,
+        dropout: float = 0.1,
+        conv_kernel_size: int = 15
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.d_model = d_model
+
+        self.input_proj = nn.Linear(input_dim, d_model)
+        self.pos_emb = PositionalEncoding(d_model, max_len=8192)
+
+        self.layers = nn.ModuleList([
+            ConformerBlock(
+                d_model=d_model,
+                nhead=nhead,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                conv_kernel_size=conv_kernel_size
+            )
+            for _ in range(num_layers)
+        ])
+
+        self.output_proj = nn.Sequential(
+            nn.Linear(d_model, output_dim),
+            nn.LayerNorm(output_dim)
+        )
+        # Small std=0.02 normal initialization to enable smooth early gradient flow
+        nn.init.normal_(self.output_proj[0].weight, std=0.02)
+        nn.init.zeros_(self.output_proj[0].bias)
+
+        # Learnable global residual scale to balance acoustic adaptation and pretrained alignment
+        self.res_scale = nn.Parameter(torch.tensor(0.5, dtype=torch.float32))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape (Batch, Time, input_dim)
+            key_padding_mask: Optional bool tensor (Batch, Time) where True denotes padded positions
+        Returns:
+            L2-normalized vectors (Batch, Time, output_dim)
+        """
+        res = x
+        h = self.input_proj(x)
+        h = self.pos_emb(h)
+
+        # Precompute additive attention mask once across all Conformer layers
+        attn_bias: Optional[torch.Tensor] = None
+        if key_padding_mask is not None:
+            attn_bias = key_padding_mask[:, None, None, :].to(dtype=h.dtype) * -10000.0
+
+        for layer in self.layers:
+            h = layer(h, attn_bias=attn_bias)
+
+        out = self.output_proj(h)
+
+        # Fused Metal Residual Scale + Mask + FP32-Accumulated L2 Normalization fast-path
+        if (
+            out.device.type == "mps"
+            and out.dtype == torch.bfloat16
+            and res.dtype == torch.bfloat16
+            and res.shape[-1] == out.shape[-1]
+            and out.shape[-1] == 256
+            and key_padding_mask is not None
+            and MetalFusedResL2NormFunction is not None
+            and not torch.jit.is_tracing()
+            and not torch.onnx.is_in_onnx_export()
+        ):
+            return MetalFusedResL2NormFunction.apply(out, res, self.res_scale, key_padding_mask)
+
+        # Global residual connection with learnable scale parameter
+        if res.shape[-1] == out.shape[-1]:
+            combined = self.res_scale * res + out
+        else:
+            combined = out
+
+        # Zero out padding positions before normalization if key_padding_mask provided
+        if key_padding_mask is not None:
+            combined = combined.masked_fill(key_padding_mask.unsqueeze(-1), 0.0)
+
+        # Use FP32 accumulation + safe eps=1e-5 to guarantee zero underflow/NaN in bfloat16
+        normed = F.normalize(combined.float(), p=2, dim=-1, eps=1e-5).to(dtype=x.dtype)
+        if key_padding_mask is not None:
+            normed = normed.masked_fill(key_padding_mask.unsqueeze(-1), 0.0)
+        return normed
+
+
+class PositionalEncoding(nn.Module):
+    """Sinusoidal Positional Encoding compatible with dynamic sequence lengths."""
+    def __init__(self, d_model: int, max_len: int = 8192):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe.unsqueeze(0), persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        seq_len = x.size(1)
+        return x + self.pe[:, :seq_len].to(dtype=x.dtype, device=x.device)
+
+
+# ==============================================================================
+# 3. Cosine Similarity Matching & Blank Logit Projection Head
+# ==============================================================================
+
+class CosineCTCHead(nn.Module):
+    """
+    Unified Cosine Similarity Matching & Blank Head:
+    - Learns a unit vector in the spherical space for <blank>: self.blank_vec
+    - Combines [norm_blank, vector_bank] -> all_vectors: (V + 1, embedding_dim)
+    - Computes single matrix multiplication: (embeddings @ all_vectors.T) * logit_scale
+    - Applies learnable scalar bias offset to <blank> (index 0) and per-token vocab_bias (index 1..V)
+    - Supports trainable vector bank mode for end-to-end representation adaptation into acoustic space
+    """
+    def __init__(
+        self,
+        embedding_dim: int = 256,
+        init_temperature: float = 0.05,
+        vector_bank: Optional[torch.Tensor] = None,
+        trainable_vector_bank: bool = True
+    ):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.trainable_vector_bank = trainable_vector_bank
+        # logit_scale initialized to ln(1 / init_temperature), safely bounded in [-5.0, 5.2983]
+        init_logit_scale = math.log(1.0 / max(init_temperature, 1e-4))
+        init_logit_scale = max(min(init_logit_scale, 5.2983), -5.0)
+        self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale, dtype=torch.float32))
+
+        # Learnable parameter vector in R^embedding_dim for <blank> token (Index 0)
+        self.blank_vec = nn.Parameter(torch.randn(1, embedding_dim))
+        nn.init.normal_(self.blank_vec, std=0.02)
+
+        # Learnable scalar bias for <blank> token (Index 0)
+        self.blank_bias = nn.Parameter(torch.zeros(1, dtype=torch.float32))
+
+        # Vector bank buffer or parameter (V, embedding_dim) & token prior bias
+        if vector_bank is not None:
+            vb_f32 = vector_bank.to(torch.float32)
+            norm_vb = F.normalize(vb_f32, p=2, dim=-1)
+            if trainable_vector_bank:
+                self.vector_bank = nn.Parameter(norm_vb)
+            else:
+                self.register_buffer("vector_bank", norm_vb, persistent=False)
+            self.vocab_bias = nn.Parameter(torch.zeros(norm_vb.size(0), dtype=torch.float32))
+        else:
+            self.register_buffer("vector_bank", None, persistent=False)
+            self.vocab_bias = nn.Parameter(torch.zeros(0, dtype=torch.float32))
+
+    def set_vector_bank(self, vector_bank: torch.Tensor):
+        """Sets or updates the target vector bank (V, embedding_dim)."""
+        vb_f32 = vector_bank.to(torch.float32)
+        norm_vb = F.normalize(vb_f32, p=2, dim=-1)
+        if self.trainable_vector_bank:
+            self.vector_bank = nn.Parameter(norm_vb)
+        else:
+            self.register_buffer("vector_bank", norm_vb, persistent=False)
+        self.vocab_bias = nn.Parameter(torch.zeros(norm_vb.size(0), dtype=torch.float32))
+
+    def forward(
+        self,
+        embeddings: torch.Tensor,
+        external_vector_bank: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            embeddings: Normalized Reader output (Batch, Time, embedding_dim)
+            external_vector_bank: Optional runtime vector bank (V, embedding_dim)
+        Returns:
+            Logits of shape (Batch, Time, V + 1), where index 0 is <blank>
+        """
+        vb = external_vector_bank if external_vector_bank is not None else self.vector_bank
+        if vb is None:
+            raise ValueError("Vector bank must be provided either in constructor, set_vector_bank(), or as an argument.")
+
+        # On MPS, execute the large (B*T, 256) x (256, V+1) projection strictly in bfloat16
+        # (with FP32 L2-normalization and clamped scale) to prevent MPS float32/float16 NaN bugs.
+        compute_dtype = torch.bfloat16 if embeddings.device.type == "mps" else torch.float32
+
+        emb_f32 = embeddings.to(dtype=torch.float32)
+        norm_blank = F.normalize(self.blank_vec.to(device=embeddings.device, dtype=torch.float32), p=2, dim=-1, eps=1e-5)
+
+        if external_vector_bank is not None or self.trainable_vector_bank:
+            vb_matched = vb.to(device=embeddings.device, dtype=torch.float32)
+            norm_vb = F.normalize(vb_matched, p=2, dim=-1, eps=1e-5)
+        else:
+            # Frozen vector_bank is already float32 and L2-normalized in __init__ / set_vector_bank
+            norm_vb = vb if (vb.device == embeddings.device and vb.dtype == torch.float32) else vb.to(device=embeddings.device, dtype=torch.float32)
+
+        all_vectors = torch.cat([norm_blank, norm_vb], dim=0).to(dtype=compute_dtype)  # (V + 1, embedding_dim)
+
+        # Safely clamp logit_scale parameter directly BEFORE exp to prevent math overflow
+        # ln(200) = ~5.2983. Clamping between -5.0 and 5.2983 bounds scale in [0.0067, 200.0]
+        clamped_scale_log = torch.clamp(self.logit_scale, min=-5.0, max=5.2983)
+        scale = torch.exp(clamped_scale_log)
+        scaled_emb = (emb_f32 * scale).to(dtype=compute_dtype)
+
+        # Fused linear projection + per-token prior bias [blank_bias, vocab_bias]
+        blank_b = self.blank_bias.to(device=embeddings.device, dtype=compute_dtype)
+        if self.vocab_bias is not None and self.vocab_bias.size(0) == norm_vb.size(0):
+            vocab_b = self.vocab_bias.to(device=embeddings.device, dtype=compute_dtype)
+        else:
+            vocab_b = torch.zeros(norm_vb.size(0), device=embeddings.device, dtype=compute_dtype)
+        bias = torch.cat([blank_b, vocab_b], dim=0)
+
+        return F.linear(scaled_emb, all_vectors, bias)
+
+
+# ==============================================================================
+# 4. SpecAugment for 256-dim Conformer Features
+# ==============================================================================
+
+class SpecAugment1D(nn.Module):
+    """
+    Time & Feature Masking for 256-dim feature sequences:
+    - Time Masking: 2 masks, random length between 5 and 15 frames
+    - Feature Masking: 2 masks, random length between 16 and 32 dimensions
+    Applied strictly during model.train() mode.
+    """
+    def __init__(
+        self,
+        time_mask_min: int = 5,
+        time_mask_max: int = 15,
+        num_time_masks: int = 2,
+        freq_mask_min: int = 16,
+        freq_mask_max: int = 32,
+        num_freq_masks: int = 2
+    ):
+        super().__init__()
+        self.time_mask_min = time_mask_min
+        self.time_mask_max = time_mask_max
+        self.num_time_masks = num_time_masks
+        self.freq_mask_min = freq_mask_min
+        self.freq_mask_max = freq_mask_max
+        self.num_freq_masks = num_freq_masks
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training:
+            return x
+
+        B, T, D = x.shape
+        x_aug = x.clone()
+
+        # Time masking (2 times, 5~15 frames)
+        for _ in range(self.num_time_masks):
+            t_len = random.randint(self.time_mask_min, min(self.time_mask_max, max(T - 1, self.time_mask_min)))
+            t_start = random.randint(0, max(0, T - t_len))
+            x_aug[:, t_start : t_start + t_len, :] = 0.0
+
+        # Feature masking (2 times, 16~32 dims)
+        for _ in range(self.num_freq_masks):
+            f_len = random.randint(self.freq_mask_min, min(self.freq_mask_max, max(D - 1, self.freq_mask_min)))
+            f_start = random.randint(0, max(0, D - f_len))
+            x_aug[:, :, f_start : f_start + f_len] = 0.0
+
+        return x_aug
+
+
+# ==============================================================================
+# 5. Complete Trainable Reader System
+# ==============================================================================
+
+class EmbeddingGemmaCTCReader(nn.Module):
+    """
+    Combined Module for training:
+    AudioToVectorReader + CosineCTCHead with SpecAugment and Trainable Vector Bank options.
+    """
+    def __init__(
+        self,
+        input_dim: int = 256,
+        output_dim: int = 256,
+        d_model: int = 384,
+        nhead: int = 6,
+        num_layers: int = 10,
+        dim_feedforward: int = 1536,
+        dropout: float = 0.1,
+        conv_kernel_size: int = 15,
+        init_temperature: float = 0.05,
+        vector_bank: Optional[torch.Tensor] = None,
+        use_spec_augment: bool = True,
+        trainable_vector_bank: bool = True
+    ):
+        super().__init__()
+        self.reader = AudioToVectorReader(
+            input_dim=input_dim,
+            output_dim=output_dim,
+            d_model=d_model,
+            nhead=nhead,
+            num_layers=num_layers,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            conv_kernel_size=conv_kernel_size
+        )
+        self.ctc_head = CosineCTCHead(
+            embedding_dim=output_dim,
+            init_temperature=init_temperature,
+            vector_bank=vector_bank,
+            trainable_vector_bank=trainable_vector_bank
+        )
+        self.spec_aug = SpecAugment1D(
+            time_mask_min=5,
+            time_mask_max=15,
+            num_time_masks=2,
+            freq_mask_min=16,
+            freq_mask_max=32,
+            num_freq_masks=2
+        ) if use_spec_augment else None
+
+    def enable_mps_bfloat16_fastpath(self) -> None:
+        """
+        Configures Reader Conformer layers in native bfloat16 while preserving float32
+        BatchNorm1d parameters/stats, res_scale, and CosineCTCHead parameters.
+        Eliminates 200+ per-step autocast weight-casting kernels on Apple Silicon MPS.
+        """
+        self.reader.to(dtype=torch.bfloat16)
+        for layer in self.reader.layers:
+            layer.conv_module.bn.to(dtype=torch.float32)
+        self.reader.res_scale.data = self.reader.res_scale.data.to(dtype=torch.float32)
+        self.ctc_head.to(dtype=torch.float32)
+
+    def forward_reader(
+        self,
+        features: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Runs SpecAugment (if training) + AudioToVectorReader without CosineCTCHead."""
+        if self.training and self.spec_aug is not None:
+            features = self.spec_aug(features)
+        return self.reader(features, key_padding_mask=key_padding_mask)
+
+    def forward(
+        self,
+        features: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        external_vector_bank: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            features: Cached audio features (Batch, Time, 256)
+            key_padding_mask: Padding mask (Batch, Time)
+            external_vector_bank: Optional custom vector bank (V, 256)
+        Returns:
+            Logits: (Batch, Time, V + 1)
+        """
+        if self.training and self.spec_aug is not None:
+            features = self.spec_aug(features)
+
+        embeddings = self.reader(features, key_padding_mask=key_padding_mask)
+        logits = self.ctc_head(embeddings, external_vector_bank=external_vector_bank)
+        return logits
+
+    def extract_reader_embeddings(
+        self,
+        features: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Outputs normalized embeddings and blank logit separately
+        (Used for Mobile OTA / Plug-and-Play Dictionary mode).
+        """
+        embeddings = self.reader(features, key_padding_mask=key_padding_mask)
+        norm_blank = F.normalize(self.ctc_head.blank_vec.to(device=embeddings.device), p=2, dim=-1)
+        scale = torch.clamp(self.ctc_head.logit_scale.exp(), max=200.0)
+        blank_logit = torch.matmul(embeddings, norm_blank.t()) * scale + self.ctc_head.blank_bias
+        return embeddings, blank_logit
+
+    def count_parameters(self) -> Dict[str, int]:
+        """Returns train parameter counts."""
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        reader_params = sum(p.numel() for p in self.reader.parameters())
+        head_params = sum(p.numel() for p in self.ctc_head.parameters())
+        return {
+            "total_params": total,
+            "trainable_params": trainable,
+            "reader_params": reader_params,
+            "ctc_head_params": head_params
+        }
+
+
+# ==============================================================================
+# 5. Native Fallback Audio Encoder (Matches EmbeddingGemma-2 Audio Front-End)
+# ==============================================================================
+
+class StandardTransformerBlock(nn.Module):
+    """Clean, standard Transformer block without proprietary fastpath C++ kernels."""
+    def __init__(self, d_model: int = 1024, nhead: int = 8, dim_feedforward: int = 2048, dropout: float = 0.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model)
+        self.self_attn = ExplicitMultiHeadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, dim_feedforward),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim_feedforward, d_model),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        nx = self.norm1(x)
+        attn_out = self.self_attn(nx, nx, nx)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class RealEmbeddingGemma2AudioEncoder(nn.Module):
+    """
+    Official 305M Audio Encoder from google/embeddinggemma-2:
+    - audio_tower: 304.8M Gemma 4 Audio Conformer Encoder
+    - embed_audio: 0.8M Multimodal Projector
+    - Matryoshka Slicing to 256 dimensions + L2 Normalization
+    Strictly uses real pretrained weights; raises RuntimeError on failure.
+    """
+    def __init__(self, audio_tower: nn.Module, embed_audio: nn.Module):
+        super().__init__()
+        self.audio_tower = audio_tower
+        self.embed_audio = embed_audio
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name: str = "google/embeddinggemma-2",
+        dtype: torch.dtype = torch.bfloat16
+    ) -> "RealEmbeddingGemma2AudioEncoder":
+        """Loads and freezes the genuine 300M Audio Encoder from Hugging Face checkpoint."""
+        try:
+            from transformers import AutoModel
+            full_model = AutoModel.from_pretrained(model_name, dtype=dtype)
+            audio_tower = full_model.audio_tower
+            embed_audio = full_model.embed_audio
+
+            # Freeze all parameters
+            for p in audio_tower.parameters():
+                p.requires_grad = False
+            for p in embed_audio.parameters():
+                p.requires_grad = False
+
+            encoder = cls(audio_tower=audio_tower, embed_audio=embed_audio)
+            encoder.eval()
+            return encoder
+        except Exception as e:
+            raise RuntimeError(
+                f"[FATAL ERROR] Failed to load genuine '{model_name}' Audio Encoder: {e}\n"
+                f"Make sure transformers is up-to-date: pip install git+https://github.com/huggingface/transformers.git"
+            )
+
+    def forward(self, input_features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input_features: (Batch, Time_Frames, 128) 128-bin log-mel filterbanks
+        Returns:
+            (Batch, Time_Tokens, 256) MRL-sliced, L2-normalized representations
+        """
+        tower_out = self.audio_tower(input_features, attention_mask=None)
+        audio_hs = tower_out.last_hidden_state  # (B, T, 1536)
+        audio_proj = self.embed_audio(audio_hs)  # (B, T, 512)
+        mrl_256 = audio_proj[:, :, :256]
+        return F.normalize(mrl_256.float(), p=2, dim=-1)
+
+
+# Quick Self-Test when executed directly
+if __name__ == "__main__":
+    print("[Testing model.py - Reader v2 Conformer]")
+    batch_size = 2
+    time_steps = 150
+    input_dim = 256
+    output_dim = 256
+    vocab_size = 5999  # Sliced vocab size (5,999 non-blank tokens -> 6,000 total with <blank>)
+
+    dummy_features = torch.randn(batch_size, time_steps, input_dim)
+    dummy_vector_bank = torch.randn(vocab_size, output_dim)
+
+    model = EmbeddingGemmaCTCReader(
+        input_dim=input_dim,
+        output_dim=output_dim,
+        d_model=384,
+        nhead=6,
+        num_layers=10,
+        dim_feedforward=1536,
+        conv_kernel_size=15,
+        init_temperature=0.05,
+        vector_bank=dummy_vector_bank,
+        trainable_vector_bank=True
+    )
+
+    stats = model.count_parameters()
+    print("Parameter Stats:")
+    for k, v in stats.items():
+        print(f"  {k}: {v:,} ({v / 1e6:.2f}M)")
+
+    logits = model(dummy_features)
+    print(f"Input features shape: {dummy_features.shape}")
+    print(f"Logits output shape:  {logits.shape} (Expected: ({batch_size}, {time_steps}, {vocab_size + 1}))")
+    assert logits.shape == (batch_size, time_steps, vocab_size + 1)
+    print("Model test PASSED successfully!")

src/step2_extract_features.py

@@ -0,0 +1,1476 @@
+"""
+src/step2_extract_features.py
+
+Step 2: Pre-extract and Cache Audio Features from Genuine Pretrained 300M Audio Encoder
+Loads the official google/embeddinggemma-2 audio_tower (304.8M) + embed_audio (0.8M)
++ language_model (271.0M) in torch.bfloat16 and freezes all parameters.
+Includes Lecture Room Augmentation (Ambient Noise, p=0.4).
+Supports three streaming data sources:
+- Mode A: Hugging Face Streaming (Bingsu/zeroth-korean)
+- Mode B: AI Hub ZIP In-Memory Reader (KsponSpeech / Lecture speech)
+- Mode C: AI Hub Relay Reader (100GB split zip volumes TS.z01, TS.z02, ... or folders)
+Slices 768d output to 256d (MRL) and caches bfloat16 .pt files in data/cache_features/.
+
+Optimizations (100% Bit-Exact & Training-Compatible):
+1. 3-Stage Asynchronous Pipeline:
+   - Stage 1 (CPU Producer Thread): Archive streaming + resampling + noise augmentation + Mel-spectrogram extraction & bucket padding + early CTC length validation (T >= L).
+   - Stage 2 (Main MPS GPU Thread): Dedicated 100% to GPU forward pass without CPU/disk stalls or per-sample empty_cache() thrashing.
+   - Stage 3 (Async Disk Writer Thread): Non-blocking atomic .pt.tmp -> .pt disk writes.
+2. Bit-Exact Model & Kernel Acceleration:
+   - Bypasses redundant AutoTokenizer string construction/parsing per audio sample and extracts only valid Mel frames via zero-copy strided STFT.
+   - Caches 5D boolean Audio Attention Masks, 4D LM Full/Sliding Masks, RoPE (cos, sin), and prompt token embeddings per (n_audio, n_audio_pad) bucket.
+   - Fuses Q/K/V projections across all 12 Audio Attention layers and all 24 LM Attention layers into single contiguous GEMMs, and fuses Gate/Up projections across all 24 LM MLP layers.
+   - Replaces Tensor-bound clamps in Gemma4ClippableLinear with direct scalar-bound clamps + F.linear, and dispatches native F.scaled_dot_product_attention directly.
+   - Supports dynamic same-bucket GPU batching (--batch_size) to amortize 577M model weight reads across multiple samples sharing the same 64-token bucket.
+   - Precomputes frozen Gemma4AudioAttention relative key projections and softplus scales across all 12 audio layers.
+   - Pre-caches float32 RMSNorm weights with native torch.rsqrt across all 302 RMSNorm modules and eliminates redundant 1e10 activation clamps in eval mode.
+   - Pre-slices LM final embedding_projection weights to 256 dims (MRL) before D2H copy.
+3. O(1) In-Memory Set Resume & Multi-Volume Split ZIP Scanner:
+   - Eliminates hundreds of thousands of disk stat() calls and SQLite lookups when resuming.
+   - Continuously scans across split ZIP boundaries (e.g. TS.z02) where leading bytes contain continuation data from the previous volume.
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import gc
+import io
+import json
+import math
+import os
+import queue
+import random
+import struct
+import sys
+import threading
+import time
+import zipfile
+import zlib
+from collections import defaultdict
+from pathlib import Path
+from typing import Generator, Tuple, Dict, Any, Optional, Iterator, Set, List
+
+# Auto-detect and switch to project virtualenv if running outside of it
+_project_root = Path(__file__).resolve().parent.parent
+_venv_python = _project_root / "venv" / "bin" / "python"
+if _venv_python.exists():
+    try:
+        if Path(sys.executable).resolve() != _venv_python.resolve():
+            print(f"🔄 [Auto-Env] Switching Python interpreter:\n    System:  {sys.executable}\n    Virtual: {_venv_python}\n", flush=True)
+            os.execv(str(_venv_python), [str(_venv_python)] + sys.argv)
+    except Exception as _e:
+        print(f"⚠️ [Auto-Env] Warning: Could not auto-switch to venv: {_e}", flush=True)
+
+_src_dir = str(Path(__file__).resolve().parent)
+_proj_dir = str(Path(__file__).resolve().parent.parent)
+if _src_dir not in sys.path:
+    sys.path.insert(0, _src_dir)
+if _proj_dir not in sys.path:
+    sys.path.insert(0, _proj_dir)
+
+# Strictly cap macOS MPS Unified Memory cache to prevent swap explosion
+os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.7")
+os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.5")
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import torchaudio
+import torchaudio.functional as AF
+import torchaudio.transforms as T
+import yaml
+from tqdm import tqdm
+
+from dataset import normalize_text
+
+
+# ==============================================================================
+# 1. Lecture Room Augmenter (Reverb & Ambient Noise)
+# ==============================================================================
+
+class LectureRoomAugmenter:
+    """
+    Simulates lecture hall ambient acoustics:
+    - Probabilistic ambient HVAC / audience background noise with controllable SNR
+    - (Reverb / RIR removed per configuration)
+    """
+    def __init__(
+        self,
+        sample_rate: int = 16000,
+        rir_prob: float = 0.0,
+        noise_prob: float = 0.4,
+        snr_db_low: float = 15.0,
+        snr_db_high: float = 30.0
+    ):
+        self.sr = sample_rate
+        self.rir_prob = 0.0
+        self.noise_prob = noise_prob
+        self.snr_db_low = snr_db_low
+        self.snr_db_high = snr_db_high
+
+    def add_background_noise(self, waveform: torch.Tensor, snr_db: float) -> torch.Tensor:
+        """Adds ambient white/pink noise matching lecture hall background."""
+        signal_power = torch.mean(waveform ** 2) + 1e-12
+        snr_linear = 10.0 ** (snr_db / 10.0)
+        noise_power = signal_power / snr_linear
+
+        noise = torch.randn_like(waveform)
+        noise = 0.7 * noise + 0.3 * torch.roll(noise, 1, dims=-1)
+        current_noise_power = torch.mean(noise ** 2) + 1e-12
+        scaled_noise = noise * torch.sqrt(noise_power / current_noise_power)
+        return waveform + scaled_noise
+
+    def __call__(self, waveform: torch.Tensor) -> torch.Tensor:
+        if waveform.dim() == 1:
+            waveform = waveform.unsqueeze(0)
+
+        if random.random() < self.noise_prob:
+            snr = random.uniform(self.snr_db_low, self.snr_db_high)
+            waveform = self.add_background_noise(waveform, snr_db=snr)
+
+        max_val = torch.max(torch.abs(waveform))
+        if max_val > 0.95:
+            waveform = waveform / (max_val / 0.95)
+
+        return waveform
+
+
+# ==============================================================================
+# 2. Mode A: Hugging Face Streaming Loader (Zeroth-Korean)
+# ==============================================================================
+
+def stream_zeroth_korean(
+    dataset_name: str = "Bingsu/zeroth-korean",
+    split: str = "train",
+    max_samples: Optional[int] = None
+) -> Generator[Tuple[torch.Tensor, int, str], None, None]:
+    """
+    Streams audio directly from Hugging Face Datasets without downloading whole dataset to disk.
+    Yields (waveform_tensor, sample_rate, transcription_text).
+    Uses Audio(decode=False) + soundfile to decode raw bytes cleanly without torchcodec dependency.
+    """
+    print(f"[Step 2] Initializing Hugging Face Streaming on '{dataset_name}' (split='{split}')...")
+    from datasets import load_dataset, Audio
+
+    ds = load_dataset(dataset_name, split=split, streaming=True)
+    # Disable torchcodec auto-decoder by receiving raw audio bytes
+    ds = ds.cast_column("audio", Audio(decode=False))
+    count = 0
+
+    for sample in ds:
+        audio_info = sample.get("audio", {})
+        raw_bytes = audio_info.get("bytes")
+        text = sample.get("text", "").strip()
+
+        if not text:
+            continue
+
+        if raw_bytes is not None:
+            try:
+                data, sr = sf.read(io.BytesIO(raw_bytes), dtype="float32")
+                if data.ndim > 1:
+                    data = data.mean(axis=1)
+                wav = torch.from_numpy(data).unsqueeze(0)
+            except Exception:
+                continue
+        elif "array" in audio_info:
+            wav = torch.from_numpy(np.array(audio_info["array"], dtype=np.float32)).unsqueeze(0)
+            sr = audio_info.get("sampling_rate", 16000)
+        else:
+            continue
+
+        yield wav, sr, text
+
+        count += 1
+        if max_samples is not None and count >= max_samples:
+            break
+
+
+class PrefetchedStream:
+    """
+    Asynchronous background prefetcher that pipelines I/O and CPU preprocessing
+    with MPS GPU inference, eliminating CPU and disk/network stalls.
+    """
+    def __init__(self, generator: Iterator, max_queue_size: int = 16):
+        self.generator = generator
+        self.q: queue.Queue = queue.Queue(maxsize=max_queue_size)
+        self.stop_event = threading.Event()
+        self.worker = threading.Thread(target=self._produce, daemon=True)
+        self.worker.start()
+
+    def _produce(self):
+        try:
+            for item in self.generator:
+                if self.stop_event.is_set():
+                    break
+                self.q.put(item)
+        except Exception as e:
+            print(f"[Step 2 Prefetch Warning] Streaming producer exception: {e}")
+        finally:
+            self.q.put(None)
+
+    def __iter__(self):
+        while True:
+            item = self.q.get()
+            if item is None:
+                break
+            yield item
+
+    def close(self):
+        self.stop_event.set()
+
+
+class AsyncDiskSaver:
+    """
+    Asynchronous background disk writer that performs atomic .pt.tmp -> .pt saves
+    and optional source file cleanup off the main MPS GPU thread.
+    """
+    def __init__(self, max_queue_size: int = 32, delete_source_after_extract: bool = False):
+        self.q: queue.Queue = queue.Queue(maxsize=max_queue_size)
+        self.delete_source = delete_source_after_extract
+        self.error: Optional[Exception] = None
+        self.worker = threading.Thread(target=self._consume, daemon=True)
+        self.worker.start()
+
+    def _consume(self):
+        while True:
+            item = self.q.get()
+            if item is None:
+                self.q.task_done()
+                break
+            payload, cache_file, src_path = item
+            try:
+                tmp_file = cache_file.with_suffix(".pt.tmp")
+                torch.save(payload, tmp_file)
+                os.replace(tmp_file, cache_file)
+                if self.delete_source and src_path is not None:
+                    try:
+                        src_path.unlink()
+                    except Exception as de:
+                        print(f"[Step 2 Warning] Could not delete {src_path}: {de}")
+            except Exception as e:
+                self.error = e
+                print(f"[Step 2 Disk Write Error] Failed to save {cache_file}: {e}")
+            finally:
+                self.q.task_done()
+
+    def save(self, payload: Dict[str, Any], cache_file: Path, src_path: Optional[Path] = None):
+        if self.error is not None:
+            raise RuntimeError(f"AsyncDiskSaver encountered fatal write error: {self.error}")
+        self.q.put((payload, cache_file, src_path))
+
+    def close(self):
+        self.q.put(None)
+        self.q.join()
+        self.worker.join(timeout=30.0)
+
+
+# ==============================================================================
+# 3. Mode B & C: AI Hub ZIP & Split-Volume Relay Readers
+# ==============================================================================
+
+class AIHubZipReader:
+    """
+    In-memory reader for large AI Hub speech archives (KsponSpeech, Lecture speech).
+    Directly extracts and decodes .wav / .pcm audio and corresponding .txt transcripts
+    from memory buffer without extracting archives to disk.
+    """
+    def __init__(self, zip_path: str):
+        self.zip_path = Path(zip_path)
+        if not self.zip_path.exists():
+            raise FileNotFoundError(f"AI Hub ZIP archive not found at {self.zip_path}")
+
+        print(f"[Step 2] Indexing AI Hub ZIP archive: {self.zip_path.name}...")
+        self.zf = zipfile.ZipFile(self.zip_path, "r")
+        self.file_list = self.zf.namelist()
+
+        self.pairs: Dict[str, Dict[str, str]] = {}
+        for fname in self.file_list:
+            if fname.endswith("/") or fname.startswith("__MACOSX"):
+                continue
+            base, ext = os.path.splitext(fname)
+            ext = ext.lower()
+            if ext in [".wav", ".pcm", ".flac"]:
+                self.pairs.setdefault(base, {})["audio"] = fname
+            elif ext in [".txt", ".trn"]:
+                self.pairs.setdefault(base, {})["text"] = fname
+
+        self.valid_keys = [k for k, v in self.pairs.items() if "audio" in v and "text" in v]
+        print(f"[Step 2] Discovered {len(self.valid_keys):,} valid audio-text pairs in ZIP.")
+
+    def __len__(self) -> int:
+        return len(self.valid_keys)
+
+    def stream_samples(self, max_samples: Optional[int] = None) -> Generator[Tuple[torch.Tensor, int, str], None, None]:
+        keys = self.valid_keys
+        if max_samples is not None:
+            keys = keys[:max_samples]
+
+        for k in keys:
+            audio_fname = self.pairs[k]["audio"]
+            text_fname = self.pairs[k]["text"]
+
+            try:
+                raw_txt = self.zf.read(text_fname)
+                try:
+                    text = raw_txt.decode("utf-8").strip()
+                except UnicodeDecodeError:
+                    text = raw_txt.decode("cp949", errors="ignore").strip()
+            except Exception:
+                continue
+
+            try:
+                raw_audio = self.zf.read(audio_fname)
+                if audio_fname.lower().endswith(".pcm"):
+                    pcm_arr = np.frombuffer(raw_audio, dtype=np.int16).astype(np.float32) / 32768.0
+                    wav = torch.from_numpy(pcm_arr).unsqueeze(0)
+                    sr = 16000
+                else:
+                    audio_buf = io.BytesIO(raw_audio)
+                    data, sr = sf.read(audio_buf, dtype="float32")
+                    if data.ndim > 1:
+                        data = data.mean(axis=1)
+                    wav = torch.from_numpy(data).unsqueeze(0)
+            except Exception:
+                continue
+
+            yield wav, sr, text
+
+
+class AIHubLectureRelayReader:
+    """
+    Reader for AI Hub Korean College Lecture dataset in relay extraction mode.
+    Matches audio files (.wav in folder or in split zip / zip) with labels from aihub_label_index.db / TL.zip.
+    - Uses O(1) in-memory set lookup (existing_files) to skip already cached features with zero disk I/O.
+    - Continuously scans across split ZIP volume boundaries (e.g. TS.z02) where leading bytes
+      contain compressed data from the previous volume.
+    - Strictly skips junk files (*.irx*, *.part, .DS_Store, __MACOSX).
+    - Yields (waveform, sr, text, short_key, optional_source_path).
+    """
+    def __init__(
+        self,
+        source_dir: str,
+        labels_source: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        max_samples: Optional[int] = None,
+        existing_files: Optional[Set[str]] = None
+    ):
+        self.source_dir = Path(source_dir)
+        self.max_samples = max_samples
+        self.cache_dir = Path(cache_dir) if cache_dir else None
+        self.existing_files = existing_files
+
+        # Auto-resolve parent Training folder (contains 01.원천데이터 and 02.라벨링데이터)
+        if self.source_dir.is_dir():
+            candidate_raw = self.source_dir / "01.원천데이터"
+            candidate_lbl = self.source_dir / "02.라벨링데이터" / "TL.zip"
+            if not labels_source and candidate_lbl.exists():
+                labels_source = str(candidate_lbl)
+            if candidate_raw.is_dir():
+                print(f"[Step 2] Auto-detected AI Hub raw audio folder: {candidate_raw}")
+                self.source_dir = candidate_raw
+
+        # Load label index
+        if not labels_source:
+            default_labels_zip = Path("/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/02.라벨링데이터/TL.zip")
+            labels_source = str(default_labels_zip)
+
+        print(f"[Step 2] Loading AI Hub label index from {labels_source}...")
+        from dataset import load_or_build_aihub_label_index
+        self.labels_map = load_or_build_aihub_label_index(tl_zip_path=labels_source)
+        print(f"[Step 2] Label index ready with {len(self.labels_map):,} entries.")
+
+    @staticmethod
+    def _seek_next_local_header(f: io.BufferedReader, initial_bytes: bytes = b"") -> bool:
+        """
+        Scans forward in 1MB chunks until a valid ZIP local file header (PK\\x03\\x04) is found,
+        positioned right at the start of the 30-byte header.
+        Returns True if found, False on EOF or Central Directory (PK\\x01\\x02).
+        """
+        buf = initial_bytes
+        chunk_size = 1024 * 1024  # 1MB scan chunks
+        while True:
+            pos = 0
+            while True:
+                idx = buf.find(b"PK\x03\x04", pos)
+                if idx == -1:
+                    break
+                # Ensure we have at least 30 bytes to validate header fields
+                if len(buf) - idx < 30:
+                    extra = f.read(30)
+                    buf += extra
+                    if len(buf) - idx < 30:
+                        return False
+                hdr_cand = buf[idx : idx + 30]
+                _, ver, flag, method, _, _, _, _, _, fname_len, extra_len = struct.unpack("<4sHHHHHIIIHH", hdr_cand)
+                if method in (0, 8) and 0 < fname_len <= 512 and extra_len <= 4096:
+                    # Valid local file header found; seek exactly to its start
+                    f.seek(f.tell() - len(buf) + idx)
+                    return True
+                pos = idx + 4
+
+            # Check if we reached Central Directory at the end of the final volume
+            if b"PK\x01\x02" in buf:
+                cd_idx = buf.find(b"PK\x01\x02")
+                if cd_idx != -1:
+                    hdr_check = buf[cd_idx : cd_idx + 4]
+                    if hdr_check == b"PK\x01\x02":
+                        return False
+
+            chunk = f.read(chunk_size)
+            if not chunk:
+                return False
+            buf = buf[-30:] + chunk if len(buf) >= 30 else buf + chunk
+
+    def stream_samples(self) -> Generator[Tuple[torch.Tensor, int, str, str, Optional[Path]], None, None]:
+        if not self.source_dir.exists():
+            print(f"[Step 2 Warning] Source dir {self.source_dir} does not exist.")
+            return
+
+        import re
+        def _is_zip_archive(name: str) -> bool:
+            nl = name.lower()
+            return nl.endswith(".zip") or bool(re.match(r"^.+\.z\d{2,3}$", nl))
+
+        # 1. Discover split zip volumes / zip files (*.z01, *.z02, *.zip, etc.), strictly skipping .irx, .part, .DS_Store
+        archive_candidates = []
+        if self.source_dir.is_file():
+            archive_candidates.append(self.source_dir)
+        elif self.source_dir.is_dir():
+            for f in sorted(self.source_dir.iterdir()):
+                fname = f.name
+                if fname.startswith(".") or ".irx" in fname or ".part" in fname:
+                    continue
+                if _is_zip_archive(fname):
+                    archive_candidates.append(f)
+
+        count = 0
+
+        # Stream from archives (split zip / zip) without disk extraction
+        for arch in archive_candidates:
+            print(f"[Step 2] Streaming audio from archive: {arch.name}...")
+            try:
+                with open(arch, "rb", buffering=4 * 1024 * 1024) as f:
+                    sig = f.read(4)
+                    if sig != b"PK\x07\x08":
+                        f.seek(0)
+
+                    while True:
+                        hdr = f.read(30)
+                        if len(hdr) < 30:
+                            break
+                        if not hdr.startswith(b"PK\x03\x04"):
+                            if not self._seek_next_local_header(f, hdr):
+                                break
+                            continue
+
+                        sig, ver, flag, method, t_val, d_val, crc, comp_size, uncomp_size, fname_len, extra_len = struct.unpack("<4sHHHHHIIIHH", hdr)
+                        if method not in (0, 8) or fname_len == 0 or fname_len > 512 or extra_len > 4096:
+                            if not self._seek_next_local_header(f, hdr[4:]):
+                                break
+                            continue
+
+                        fname_bytes = f.read(fname_len)
+                        if len(fname_bytes) < fname_len:
+                            break
+                        if extra_len > 0:
+                            f.seek(extra_len, os.SEEK_CUR)
+                        fname = fname_bytes.decode("utf-8", errors="replace")
+
+                        if not fname.lower().endswith((".wav", ".pcm", ".flac")) or fname.startswith("__MACOSX"):
+                            if comp_size > 0:
+                                f.seek(comp_size, os.SEEK_CUR)
+                            continue
+
+                        # Extract relative key and short key
+                        rel_key = fname.rsplit(".", 1)[0]
+                        parts = rel_key.split("/")
+                        short_key = f"{parts[-2]}_{parts[-1]}" if len(parts) >= 2 else parts[-1]
+                        target_pt_name = f"feat_aihub_{short_key}.pt"
+
+                        # Fast O(1) in-memory Resume check: skip reading audio & SQLite lookup if already cached!
+                        if self.existing_files is not None:
+                            if target_pt_name in self.existing_files:
+                                f.seek(comp_size, os.SEEK_CUR)
+                                continue
+                        elif self.cache_dir and (self.cache_dir / target_pt_name).exists():
+                            f.seek(comp_size, os.SEEK_CUR)
+                            continue
+
+                        # Silent skip if label not in index (short_key is already the canonical key)
+                        text = self.labels_map.get(short_key)
+                        if not text:
+                            f.seek(comp_size, os.SEEK_CUR)
+                            continue
+
+                        try:
+                            comp_data = f.read(comp_size)
+                            if len(comp_data) < comp_size:
+                                # Truncated entry at the end of a split volume
+                                break
+                            decomp = zlib.decompress(comp_data, -15) if method == 8 else comp_data
+                            data, sr = sf.read(io.BytesIO(decomp), dtype="float32")
+                            if data.ndim > 1:
+                                data = data.mean(axis=1)
+                            wav = torch.from_numpy(data).unsqueeze(0)
+                        except Exception:
+                            # Silent skip on decode errors (e.g., split boundary)
+                            continue
+
+                        yield wav, sr, text, short_key, None
+                        count += 1
+                        if self.max_samples and count >= self.max_samples:
+                            return
+            except Exception as e:
+                print(f"[Step 2 Warning] Error reading archive {arch}: {e}. Proceeding...")
+
+        # 2. Discover uncompressed standalone .wav files in directory
+        if self.source_dir.is_dir():
+            import glob
+            wav_patterns = [str(self.source_dir / "**/*.wav"), str(self.source_dir / "**/*.pcm")]
+            wav_files = []
+            for pat in wav_patterns:
+                for fpath in glob.glob(pat, recursive=True):
+                    if ".irx" not in fpath and ".part" not in fpath and ".DS_Store" not in fpath and "__MACOSX" not in fpath:
+                        wav_files.append(Path(fpath))
+
+            wav_files.sort()
+            for wf in wav_files:
+                try:
+                    rel_key = str(wf.relative_to(self.source_dir)).rsplit(".", 1)[0]
+                except Exception:
+                    rel_key = wf.stem
+                parts = rel_key.split("/")
+                short_key = f"{parts[-2]}_{parts[-1]}" if len(parts) >= 2 else parts[-1]
+                target_pt_name = f"feat_aihub_{short_key}.pt"
+
+                # Fast O(1) Resume check
+                if self.existing_files is not None:
+                    if target_pt_name in self.existing_files:
+                        continue
+                elif self.cache_dir and (self.cache_dir / target_pt_name).exists():
+                    continue
+
+                # Silent skip if label not found
+                text = self.labels_map.get(short_key)
+                if not text:
+                    continue
+
+                try:
+                    data, sr = sf.read(str(wf), dtype="float32")
+                    if data.ndim > 1:
+                        data = data.mean(axis=1)
+                    wav = torch.from_numpy(data).unsqueeze(0)
+                except Exception:
+                    continue
+
+                yield wav, sr, text, short_key, wf
+                count += 1
+                if self.max_samples and count >= self.max_samples:
+                    return
+
+
+# ==============================================================================
+# 4. Genuine 300M Audio Encoder Feature Extractor (Bit-Exact Accelerated)
+# ==============================================================================
+
+class GenuineAudioFeatureExtractor:
+    """
+    Extracts features using the official google/embeddinggemma-2 multimodal pipeline:
+    - Audio processed by audio_tower (304.8M) + embed_audio (0.8M)
+    - Passed through language_model (271.0M) into the shared 768-dim multimodal embedding space
+    - Audio token representations sliced to 256-dim (MRL) and L2 normalized
+    This guarantees audio and text representations share the exact same geometric vector space,
+    producing bit-exact identical outputs while accelerating execution via precomputation and caching.
+    """
+    def __init__(self, device: torch.device, model_name: str = "google/embeddinggemma-2"):
+        self.device = device
+        self.target_sr = 16000
+        self.audio_bucket = 64
+
+        print(f"[Step 2] Loading genuine '{model_name}' Multimodal Model on {device} (bfloat16)...")
+        try:
+            from transformers import AutoModel, AutoProcessor
+            self.processor = AutoProcessor.from_pretrained(model_name)
+            self.feature_extractor = self.processor.feature_extractor
+            self.model = AutoModel.from_pretrained(model_name, dtype=torch.bfloat16)
+
+            if not hasattr(self.model, "audio_tower") or not hasattr(self.model, "embed_audio") or not hasattr(self.model, "language_model"):
+                raise AttributeError(f"Model missing required modules. Available: {[n for n, _ in self.model.named_children()]}")
+
+            # Delete vision_tower to save unified memory
+            if hasattr(self.model, "vision_tower"):
+                del self.model.vision_tower
+            gc.collect()
+
+            if hasattr(self.model, "config"):
+                self.model.config.use_cache = False
+            if hasattr(self.model, "language_model") and hasattr(self.model.language_model, "config"):
+                self.model.language_model.config.use_cache = False
+
+            self.model = self.model.to(device)
+            self.model.eval()
+
+            #Strictly freeze all parameters
+            for p in self.model.parameters():
+                p.requires_grad = False
+
+            total_p = sum(p.numel() for p in self.model.parameters())
+            self._apply_bit_exact_optimizations()
+            print(f"[Step 2] Successfully loaded, FROZE, and optimized genuine Multimodal Model ({total_p:,} params, {total_p/1e6:.1f}M)!")
+
+            if torch.backends.mps.is_available():
+                torch.mps.empty_cache()
+
+        except Exception as e:
+            raise RuntimeError(
+                f"\n[FATAL ERROR] Failed to load genuine '{model_name}' Multimodal Model: {e}\n"
+                f"Underlying error: {e}\n"
+                f"Ensure you are running with updated transformers: pip install git+https://github.com/huggingface/transformers.git"
+            )
+
+    def _apply_bit_exact_optimizations(self) -> None:
+        """
+        Applies mathematically and numerically bit-exact optimizations to the frozen
+        Audio Tower and Language Model layers:
+        1. Pre-caches float32 RMSNorm weights and uses native torch.rsqrt across all 302 RMSNorm modules.
+        2. Optimizes Gemma4ClippableLinear with pre-extracted scalar bounds and direct F.linear.
+        3. Fuses Q/K/V weights and precomputes relative positional key projections + softplus scales across all 12 Audio Attention layers.
+        4. Removes redundant 1e10 activation clamps in eval mode (values are already bounded by ClippableLinear).
+        5. Fuses Q/K/V weights + direct SDPA dispatch in all 24 LM Attention layers, and fuses Gate/Up weights in all 24 LM MLP layers.
+        6. Precomputes prompt token embeddings ([BOS, BOA], [EOA, EOS], [PAD]) and MRL 256d projection slice.
+        """
+        from transformers.models.gemma4.modeling_gemma4 import (
+            Gemma4RMSNorm,
+            Gemma4ClippableLinear,
+            Gemma4AudioAttention,
+            Gemma4AudioLayer,
+            Gemma4AudioFeedForward,
+            Gemma4AudioLightConv1d,
+            create_bidirectional_mask as g4_create_mask,
+            sliding_window_mask_function,
+        )
+        from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import (
+            EmbeddingGemma2RMSNorm,
+            EmbeddingGemma2Attention,
+            EmbeddingGemma2MLP,
+            EmbeddingGemma2TextPLEBlock,
+            EmbeddingGemma2EncoderLayer,
+            apply_rotary_pos_emb,
+            create_bidirectional_mask as eg2_create_mask,
+            create_bidirectional_sliding_window_mask,
+        )
+
+        self._g4_create_mask = g4_create_mask
+        self._sliding_window_mask_fn = sliding_window_mask_function
+        self._eg2_create_mask = eg2_create_mask
+        self._eg2_create_sliding_mask = create_bidirectional_sliding_window_mask
+
+        # 1. Pre-cache float32 weights in all RMSNorm layers & replace torch.pow(-0.5) with torch.rsqrt
+        for m in self.model.modules():
+            if isinstance(m, (Gemma4RMSNorm, EmbeddingGemma2RMSNorm)):
+                if getattr(m, "with_scale", True) and hasattr(m, "weight"):
+                    m.register_buffer("_weight_f32", m.weight.detach().float(), persistent=False)
+
+        def _exact_fast_rmsnorm(norm_self, hidden_states: torch.Tensor) -> torch.Tensor:
+            x = hidden_states.float()
+            normed = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + norm_self.eps)
+            if norm_self.with_scale:
+                normed = normed * norm_self._weight_f32
+            return normed.type_as(hidden_states)
+
+        Gemma4RMSNorm.forward = _exact_fast_rmsnorm
+        EmbeddingGemma2RMSNorm.forward = _exact_fast_rmsnorm
+
+        # 2. Optimize Gemma4ClippableLinear with Python float scalar bounds & direct F.linear
+        at = self.model.audio_tower
+        lm = self.model.language_model
+        for m in at.modules():
+            if isinstance(m, Gemma4ClippableLinear):
+                m._w = m.linear.weight
+                if m.use_clipped_linears:
+                    m._imin = float(m.input_min.item())
+                    m._imax = float(m.input_max.item())
+                    m._omin = float(m.output_min.item())
+                    m._omax = float(m.output_max.item())
+
+        def _fast_clippable_linear(c_self, hidden_states: torch.Tensor) -> torch.Tensor:
+            if c_self.use_clipped_linears:
+                hidden_states = torch.clamp(hidden_states, c_self._imin, c_self._imax)
+                return torch.clamp(F.linear(hidden_states, c_self._w), c_self._omin, c_self._omax)
+            return F.linear(hidden_states, c_self._w)
+
+        Gemma4ClippableLinear.forward = _fast_clippable_linear
+
+        # 3. Precompute constant relative position embeddings & fuse Q/K/V projections in Audio Tower
+        with torch.no_grad():
+            self._pos_emb_const = at.rel_pos_enc(
+                torch.empty(1, 1, at.config.hidden_size, dtype=torch.bfloat16, device=self.device)
+            )
+            for layer in at.layers:
+                attn = layer.self_attn
+                attn.register_buffer("_softplus_scale", F.softplus(attn.per_dim_scale), persistent=False)
+                rel_k = (
+                    attn.relative_k_proj(self._pos_emb_const)
+                    .view(-1, attn.num_heads, attn.head_dim)
+                    .float()
+                    .permute(1, 2, 0)
+                )
+                attn.register_buffer("_rel_k_perm", rel_k, persistent=False)
+
+                q_w = attn.q_proj.linear.weight
+                k_w = attn.k_proj.linear.weight
+                v_w = attn.v_proj.linear.weight
+                attn.register_buffer("_qkv_w", torch.cat([q_w, k_w, v_w], dim=0).contiguous(), persistent=False)
+                attn._qkv_dim = q_w.shape[0]
+                attn._qkv_imin = float(attn.q_proj.input_min.item())
+                attn._qkv_imax = float(attn.q_proj.input_max.item())
+                attn._q_omin = float(attn.q_proj.output_min.item())
+                attn._q_omax = float(attn.q_proj.output_max.item())
+                attn._k_omin = float(attn.k_proj.output_min.item())
+                attn._k_omax = float(attn.k_proj.output_max.item())
+                attn._v_omin = float(attn.v_proj.output_min.item())
+                attn._v_omax = float(attn.v_proj.output_max.item())
+                attn._invalid_val = float(attn.config.attention_invalid_logits_value)
+                del attn.q_proj.linear.weight, attn.k_proj.linear.weight, attn.v_proj.linear.weight
+
+        def _fast_audio_attn_forward(attn_self, hidden_states: torch.Tensor, position_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
+            batch_size, seq_length, _ = hidden_states.shape
+            hidden_shape = (batch_size, seq_length, attn_self.num_heads, attn_self.head_dim)
+            D = attn_self._qkv_dim
+
+            x_c = torch.clamp(hidden_states, attn_self._qkv_imin, attn_self._qkv_imax)
+            qkv = F.linear(x_c, attn_self._qkv_w)
+            query_states = torch.clamp(qkv[..., :D], attn_self._q_omin, attn_self._q_omax).float().view(hidden_shape)
+            key_states = torch.clamp(qkv[..., D:2*D], attn_self._k_omin, attn_self._k_omax).float().view(hidden_shape)
+            value_states = torch.clamp(qkv[..., 2*D:], attn_self._v_omin, attn_self._v_omax).float().view(hidden_shape)
+
+            query_states = query_states * attn_self.q_scale * attn_self._softplus_scale
+            key_states = key_states * attn_self.k_scale
+
+            query_states = attn_self._convert_to_block(query_states)
+            key_states = attn_self._extract_block_context(key_states)
+            value_states = attn_self._extract_block_context(value_states)
+            num_blocks = query_states.shape[1]
+
+            queries = query_states.permute(0, 3, 1, 2, 4)
+            matrix_ac = queries @ key_states.permute(0, 3, 1, 4, 2)
+
+            queries_flat = queries.reshape(batch_size, attn_self.num_heads, -1, attn_self.head_dim)
+            matrix_bd = (queries_flat @ attn_self._rel_k_perm).reshape(
+                batch_size, attn_self.num_heads, num_blocks, attn_self.chunk_size, -1
+            )
+            matrix_bd = attn_self._rel_shift(matrix_bd)
+
+            attn_weights = (matrix_ac + matrix_bd) / attn_self.softcap
+            attn_weights = torch.tanh(attn_weights) * attn_self.softcap
+
+            if attention_mask is not None:
+                attn_weights = attn_weights.masked_fill(attention_mask, attn_self._invalid_val)
+
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
+            attn_output = attn_weights @ value_states.permute(0, 3, 1, 2, 4)
+            attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, num_blocks * attn_self.chunk_size, -1)
+            attn_output = attn_output[:, :seq_length].contiguous()
+            return attn_self.post(attn_output.to(hidden_states.dtype)), None
+
+        Gemma4AudioAttention.forward = _fast_audio_attn_forward
+
+        # 4. Eliminate redundant 1e10 clamp kernels in Audio FFN, LightConv1d, and AudioLayer
+        def _fast_ffn_forward(ffn_self, hidden_states: torch.Tensor) -> torch.Tensor:
+            residual = hidden_states
+            hidden_states = ffn_self.pre_layer_norm(hidden_states)
+            hidden_states = ffn_self.ffw_layer_2(ffn_self.act_fn(ffn_self.ffw_layer_1(hidden_states)))
+            hidden_states = ffn_self.post_layer_norm(hidden_states) * ffn_self.post_layer_scale
+            return hidden_states + residual
+
+        def _fast_lconv_forward(conv_self, hidden_states: torch.Tensor) -> torch.Tensor:
+            residual = hidden_states
+            hidden_states = F.glu(conv_self.linear_start(conv_self.pre_layer_norm(hidden_states)), dim=-1)
+            hidden_states = conv_self.depthwise_conv1d(hidden_states.transpose(1, 2)).transpose(1, 2)
+            hidden_states = conv_self.linear_end(conv_self.act_fn(conv_self.conv_norm(hidden_states)))
+            return hidden_states + residual
+
+        def _fast_layer_forward(layer_self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor], position_embeddings: torch.Tensor, **kwargs) -> torch.Tensor:
+            hidden_states = layer_self.feed_forward1(hidden_states)
+            residual = hidden_states
+            h_attn, _ = layer_self.self_attn(
+                hidden_states=layer_self.norm_pre_attn(hidden_states),
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+            )
+            hidden_states = layer_self.norm_post_attn(h_attn) + residual
+            hidden_states = layer_self.lconv1d(hidden_states)
+            hidden_states = layer_self.feed_forward2(hidden_states)
+            return layer_self.norm_out(hidden_states)
+
+        Gemma4AudioFeedForward.forward = _fast_ffn_forward
+        Gemma4AudioLightConv1d.forward = _fast_lconv_forward
+        Gemma4AudioLayer.forward = _fast_layer_forward
+
+        # 5. Fuse Q/K/V and Gate/Up projections across all 24 Language Model layers + direct SDPA
+        with torch.no_grad():
+            for layer in lm.layers:
+                attn = layer.self_attn
+                q_w = attn.q_proj.weight
+                k_w = attn.k_proj.weight
+                v_w = attn.v_proj.weight
+                attn.register_buffer("_qkv_w", torch.cat([q_w, k_w, v_w], dim=0).contiguous(), persistent=False)
+                attn._q_dim = q_w.shape[0]
+                attn._kv_dim = k_w.shape[0]
+                attn._o_w = attn.o_proj.weight
+                del attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight
+
+                mlp = layer.mlp
+                g_w = mlp.gate_proj.weight
+                u_w = mlp.up_proj.weight
+                mlp.register_buffer("_gu_w", torch.cat([g_w, u_w], dim=0).contiguous(), persistent=False)
+                mlp._inter_dim = mlp.intermediate_size
+                mlp._down_w = mlp.down_proj.weight
+                del mlp.gate_proj.weight, mlp.up_proj.weight
+
+                ple_b = layer.ple_block
+                ple_b._gate_w = ple_b.per_layer_input_gate.weight
+                ple_b._proj_w = ple_b.per_layer_projection.weight
+                layer._scalar_val = float(layer.layer_scalar.item())
+
+        def _fused_lm_mlp_forward(mlp_self, x: torch.Tensor) -> torch.Tensor:
+            gu = F.linear(x, mlp_self._gu_w)
+            I = mlp_self._inter_dim
+            return F.linear(mlp_self.act_fn(gu[..., :I]) * gu[..., I:], mlp_self._down_w)
+
+        def _fused_lm_attn_forward(attn_self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs):
+            input_shape = hidden_states.shape[:-1]
+            hidden_shape = (*input_shape, -1, attn_self.head_dim)
+            Dq = attn_self._q_dim
+            Dkv = attn_self._kv_dim
+
+            qkv = F.linear(hidden_states, attn_self._qkv_w)
+            query_states = attn_self.q_norm(qkv[..., :Dq].view(hidden_shape).transpose(1, 2))
+            key_states = attn_self.k_norm(qkv[..., Dq:Dq+Dkv].view(hidden_shape).transpose(1, 2))
+            value_states = attn_self.v_norm(qkv[..., Dq+Dkv:].view(hidden_shape).transpose(1, 2))
+
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                scale=attn_self.scaling,
+                is_causal=False,
+                enable_gqa=attn_self.num_key_value_groups > 1,
+            )
+            attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+            return F.linear(attn_output, attn_self._o_w), None
+
+        def _fused_ple_block_forward(ple_self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor) -> torch.Tensor:
+            h = ple_self.act_fn(F.linear(hidden_states, ple_self._gate_w)) * per_layer_input
+            h = ple_self.post_per_layer_input_norm(F.linear(h, ple_self._proj_w))
+            return hidden_states + h
+
+        def _fused_lm_layer_forward(layer_self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs) -> torch.Tensor:
+            h_attn, _ = layer_self.self_attn(
+                hidden_states=layer_self.input_layernorm(hidden_states),
+                attention_mask=attention_mask,
+                position_embeddings=position_embeddings,
+            )
+            hidden_states = hidden_states + layer_self.post_attention_layernorm(h_attn)
+            h_mlp = layer_self.mlp(layer_self.pre_feedforward_layernorm(hidden_states))
+            hidden_states = hidden_states + layer_self.post_feedforward_layernorm(h_mlp)
+            hidden_states = layer_self.ple_block(hidden_states, per_layer_input)
+            if layer_self._scalar_val != 1.0:
+                hidden_states = hidden_states * layer_self.layer_scalar
+            return hidden_states
+
+        EmbeddingGemma2MLP.forward = _fused_lm_mlp_forward
+        EmbeddingGemma2Attention.forward = _fused_lm_attn_forward
+        EmbeddingGemma2TextPLEBlock.forward = _fused_ple_block_forward
+        EmbeddingGemma2EncoderLayer.forward = _fused_lm_layer_forward
+
+        # 6. Precompute prompt token embeddings & MRL 256d projection matrix slice
+        self.audio_token_id = self.model.config.audio_token_id
+        self.pad_id = getattr(self.model.config.text_config, "pad_token_id", 0) or 0
+
+        with torch.no_grad():
+            sample_prompt = self.processor(
+                text="<|audio|>",
+                audio=np.zeros(1600, dtype=np.float32),
+                sampling_rate=self.target_sr,
+                return_tensors="pt"
+            )["input_ids"][0]
+            prefix_ids = sample_prompt[:2].unsqueeze(0).to(self.device)
+            suffix_ids = sample_prompt[-2:].unsqueeze(0).to(self.device)
+            pad_ids = torch.tensor([[self.pad_id]], dtype=torch.long, device=self.device)
+
+            embed_fn = self.model.get_input_embeddings()
+            self._prefix_emb = embed_fn(prefix_ids).detach()  # (1, 2, 512)
+            self._suffix_emb = embed_fn(suffix_ids).detach()  # (1, 2, 512)
+            self._pad_emb = embed_fn(pad_ids).detach()        # (1, 1, 512)
+            self._proj_w_256 = lm.embedding_projection.weight[:256, :].detach().contiguous()
+
+        # Pre-cache STFT / Mel parameters for zero-copy CPU Mel extraction
+        fe = self.feature_extractor
+        self._mel_pad_left = fe.frame_length // 2
+        self._mel_frame_length = fe.frame_length
+        self._mel_hop_length = fe.hop_length
+        self._mel_fft_length = fe.fft_length
+        self._mel_window = fe.window
+        self._mel_filters = fe.mel_filters
+        self._mel_floor = fe.mel_floor
+
+        # Caches for attention masks, base embeddings, and RoPE (cos, sin)
+        self._audio_mask_cache: Dict[Tuple[int, int], Optional[torch.Tensor]] = {}
+        self._batched_win_cache: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}
+        self._lm_entry_cache: Dict[Tuple[int, int], Tuple[Dict[str, Optional[torch.Tensor]], torch.Tensor]] = {}
+        self._rope_cache: Dict[int, Dict[str, Tuple[torch.Tensor, torch.Tensor]]] = {}
+        gc.collect()
+
+    def resample_if_needed(self, wav: torch.Tensor, orig_sr: int) -> torch.Tensor:
+        if orig_sr != self.target_sr:
+            if not hasattr(self, "_resamplers"):
+                self._resamplers: Dict[int, T.Resample] = {}
+            resampler = self._resamplers.get(orig_sr)
+            if resampler is None:
+                resampler = T.Resample(orig_sr, self.target_sr)
+                self._resamplers[orig_sr] = resampler
+            wav = resampler(wav)
+        return wav
+
+    def preprocess_audio_cpu(self, wav: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, int, int]:
+        """
+        CPU-only preprocessing stage (safe to run in background producer thread):
+        Computes Mel-spectrogram features and masks matching Gemma4AudioFeatureExtractor bit-for-bit,
+        bypassing redundant string construction, tokenizer calls, and padded STFT frames,
+        and pads to 64-token buckets.
+        Returns (f_in_bf16, fm_in_bool, n_audio, n_audio_pad).
+        """
+        if isinstance(wav, torch.Tensor):
+            wav_np = wav.detach().cpu().numpy().reshape(-1)
+        else:
+            wav_np = np.asarray(wav, dtype=np.float32).reshape(-1)
+
+        if wav_np.dtype != np.float32:
+            wav_np = wav_np.astype(np.float32)
+
+        raw_len = min(wav_np.shape[0], 480_000)
+        hop_length = self._mel_hop_length
+        pad_left = self._mel_pad_left
+        frame_length = self._mel_frame_length
+
+        valid_frames = max(0, (raw_len - 160 + hop_length - 1) // hop_length)
+        n_audio = (valid_frames + 3) // 4
+        n_audio_pad = ((n_audio + self.audio_bucket - 1) // self.audio_bucket) * self.audio_bucket
+        mel_target = n_audio_pad * 4
+        copy_len = min(valid_frames, mel_target)
+
+        f_in = torch.zeros((1, mel_target, 128), dtype=torch.bfloat16)
+        fm_in = torch.zeros((1, mel_target), dtype=torch.bool)
+
+        if copy_len > 0:
+            needed_len = pad_left + (copy_len - 1) * hop_length + frame_length
+            buf = np.zeros(needed_len, dtype=np.float32)
+            buf[pad_left : pad_left + raw_len] = wav_np[:raw_len]
+            frames = np.lib.stride_tricks.as_strided(
+                buf,
+                shape=(copy_len, frame_length),
+                strides=(buf.strides[0] * hop_length, buf.strides[0]),
+            )
+            stft = np.fft.rfft(frames * self._mel_window, n=self._mel_fft_length, axis=-1)
+            mel_spec = np.matmul(np.abs(stft), self._mel_filters)
+            log_mel = np.log(mel_spec + self._mel_floor).astype(np.float32)
+            f_in[0, :copy_len] = torch.from_numpy(log_mel).to(torch.bfloat16)
+            fm_in[0, :copy_len] = True
+
+        return f_in, fm_in, n_audio, n_audio_pad
+
+    def _get_batched_audio_mask(self, n_t: torch.Tensor, n_audio_pad: int) -> torch.Tensor:
+        """
+        Builds the 5D boolean blocked local-attention mask (B, 1, num_blocks, chunk_size, context_size)
+        directly on device for a same-bucket batch.
+        """
+        at = self.model.audio_tower
+        entry = self._batched_win_cache.get(n_audio_pad)
+        if entry is None:
+            C = at.config.attention_chunk_size
+            L = at.config.attention_context_left - 1
+            R = at.config.attention_context_right
+            CTX = C + L + R
+            NB = (n_audio_pad + C - 1) // C
+            b = torch.arange(NB, device=self.device).view(NB, 1, 1)
+            c = torch.arange(C, device=self.device).view(1, C, 1)
+            j = torch.arange(CTX, device=self.device).view(1, 1, CTX)
+            q = b * C + c
+            k = b * C + j - L
+            dist = q - k
+            win = ((dist >= 0) & (dist < L)) | ((dist < 0) & (-dist < R))
+            base_valid = win & (q < n_audio_pad) & (k >= 0) & (k < n_audio_pad)
+            entry = (base_valid.unsqueeze(0), k.unsqueeze(0))
+            self._batched_win_cache[n_audio_pad] = entry
+        base_valid, k_idx = entry
+        valid = base_valid & (k_idx < n_t.view(-1, 1, 1, 1))
+        return (~valid).unsqueeze(1)
+
+    def extract_from_preprocessed(
+        self,
+        f_in: torch.Tensor,
+        fm_in: torch.Tensor,
+        n_audio: int,
+        n_audio_pad: int
+    ) -> torch.Tensor:
+        """
+        Executes MPS GPU forward pass on pre-padded Mel features using cached masks, RoPE,
+        and prompt embeddings, returning the bit-exact (n_audio, 256) bfloat16 tensor on CPU.
+        """
+        f_mps = f_in.to(self.device, dtype=torch.bfloat16)
+        fm_mps = fm_in.to(self.device)
+        tok_target = n_audio_pad + 4
+        am_key = (n_audio, n_audio_pad)
+
+        at = self.model.audio_tower
+        lm = self.model.language_model
+
+        with torch.no_grad():
+            # 1. Audio Tower Forward
+            hidden_states, output_mask = at.subsample_conv_projection(f_mps, fm_mps)
+            if am_key in self._audio_mask_cache:
+                audio_attn_mask = self._audio_mask_cache[am_key]
+            else:
+                m4d = self._g4_create_mask(
+                    config=at.config,
+                    inputs_embeds=hidden_states,
+                    attention_mask=output_mask,
+                    and_mask_function=self._sliding_window_mask_fn(
+                        (at.config.attention_context_left - 1, at.config.attention_context_right)
+                    ),
+                )
+                audio_attn_mask = (at._convert_4d_mask_to_blocked_5d(m4d) != 0) if m4d is not None else None
+                self._audio_mask_cache[am_key] = audio_attn_mask
+
+            for encoder_layer in at.layers:
+                hidden_states = encoder_layer(
+                    hidden_states,
+                    attention_mask=audio_attn_mask,
+                    position_embeddings=self._pos_emb_const,
+                )
+            audio_proj = self.model.embed_audio(inputs_embeds=at.output_proj(hidden_states))
+
+            # 2. Language Model Forward with cached masks, base embeddings, and RoPE
+            lm_entry = self._lm_entry_cache.get(am_key)
+            if lm_entry is None:
+                am_mps = torch.zeros((1, tok_target), dtype=torch.long, device=self.device)
+                am_mps[:, : n_audio + 4] = 1
+                dummy_emb = torch.empty((1, tok_target, lm.config.hidden_size), dtype=torch.bfloat16, device=self.device)
+                mk = {"config": lm.config, "inputs_embeds": dummy_emb, "attention_mask": am_mps}
+                am_map = {
+                    "full_attention": self._eg2_create_mask(**mk),
+                    "sliding_attention": self._eg2_create_sliding_mask(**mk),
+                }
+                base_embeds = self._pad_emb.expand(1, tok_target, -1).clone()
+                base_embeds[:, :2, :] = self._prefix_emb
+                base_embeds[:, 2 + n_audio : 4 + n_audio, :] = self._suffix_emb
+                lm_entry = (am_map, base_embeds)
+                self._lm_entry_cache[am_key] = lm_entry
+
+            am_map, base_embeds = lm_entry
+            inputs_embeds = base_embeds.clone()
+            inputs_embeds[:, 2 : 2 + n_audio, :] = audio_proj[:, :n_audio, :]
+
+            pos_emb = self._rope_cache.get(tok_target)
+            if pos_emb is None:
+                pos_ids = torch.arange(tok_target, device=self.device).unsqueeze(0)
+                pos_emb = {lt: lm.rotary_emb(inputs_embeds, pos_ids, lt) for lt in lm.unique_layer_types}
+                self._rope_cache[tok_target] = pos_emb
+
+            per_layer_inputs = lm.ple(inputs_embeds)
+            h = inputs_embeds
+            for i, encoder_layer in enumerate(lm.layers):
+                lt = lm.config.layer_types[i]
+                h = encoder_layer(
+                    h,
+                    per_layer_inputs[:, :, i, :],
+                    attention_mask=am_map[lt],
+                    position_embeddings=pos_emb[lt],
+                )
+
+            h_normed = lm.norm(h)
+            mrl_256 = F.linear(h_normed[:, 2 : 2 + n_audio, :], self._proj_w_256).cpu()[0]
+
+        norm_256 = F.normalize(mrl_256.float(), p=2, dim=-1).to(torch.bfloat16)
+        return norm_256
+
+    def extract_batch_same_bucket(
+        self,
+        batch_items: List[Tuple[torch.Tensor, torch.Tensor, int, int]]
+    ) -> List[torch.Tensor]:
+        """
+        Executes a batched MPS GPU forward pass for multiple preprocessed audio samples that share
+        the exact same `n_audio_pad` bucket length (zero cross-bucket padding distortion).
+        Falls back to the bit-exact single-sample path when len(batch_items) == 1.
+        """
+        B = len(batch_items)
+        if B == 1:
+            return [self.extract_from_preprocessed(*batch_items[0])]
+
+        n_audio_pad = batch_items[0][3]
+        n_list = [it[2] for it in batch_items]
+        tok_target = n_audio_pad + 4
+
+        f_mps = torch.cat([it[0] for it in batch_items], dim=0).to(self.device, dtype=torch.bfloat16)
+        fm_mps = torch.cat([it[1] for it in batch_items], dim=0).to(self.device)
+        n_t = torch.tensor(n_list, device=self.device)
+
+        at = self.model.audio_tower
+        lm = self.model.language_model
+
+        with torch.no_grad():
+            hidden_states, _ = at.subsample_conv_projection(f_mps, fm_mps)
+            audio_attn_mask = self._get_batched_audio_mask(n_t, n_audio_pad)
+
+            for encoder_layer in at.layers:
+                hidden_states = encoder_layer(
+                    hidden_states,
+                    attention_mask=audio_attn_mask,
+                    position_embeddings=self._pos_emb_const,
+                )
+            audio_proj = self.model.embed_audio(inputs_embeds=at.output_proj(hidden_states))
+
+            t = torch.arange(tok_target, device=self.device).view(1, tok_target, 1)
+            n3 = n_t.view(B, 1, 1)
+            inputs_embeds = self._pad_emb.expand(B, tok_target, -1)
+            shifted = F.pad(audio_proj, (0, 0, 2, 2))
+            inputs_embeds = torch.where((t >= 2) & (t < 2 + n3), shifted, inputs_embeds)
+            inputs_embeds = torch.where(t == 0, self._prefix_emb[:, 0:1], inputs_embeds)
+            inputs_embeds = torch.where(t == 1, self._prefix_emb[:, 1:2], inputs_embeds)
+            inputs_embeds = torch.where(t == 2 + n3, self._suffix_emb[:, 0:1], inputs_embeds)
+            inputs_embeds = torch.where(t == 3 + n3, self._suffix_emb[:, 1:2], inputs_embeds)
+
+            am_mps = (torch.arange(tok_target, device=self.device).unsqueeze(0) < (n_t + 4).unsqueeze(1)).long()
+            mk = {"config": lm.config, "inputs_embeds": inputs_embeds, "attention_mask": am_mps}
+            am_map = {
+                "full_attention": self._eg2_create_mask(**mk),
+                "sliding_attention": self._eg2_create_sliding_mask(**mk),
+            }
+
+            pos_emb = self._rope_cache.get(tok_target)
+            if pos_emb is None:
+                pos_ids = torch.arange(tok_target, device=self.device).unsqueeze(0)
+                pos_emb = {lt: lm.rotary_emb(inputs_embeds[0:1], pos_ids, lt) for lt in lm.unique_layer_types}
+                self._rope_cache[tok_target] = pos_emb
+
+            per_layer_inputs = lm.ple(inputs_embeds)
+            h = inputs_embeds
+            for i, encoder_layer in enumerate(lm.layers):
+                lt = lm.config.layer_types[i]
+                h = encoder_layer(
+                    h,
+                    per_layer_inputs[:, :, i, :],
+                    attention_mask=am_map[lt],
+                    position_embeddings=pos_emb[lt],
+                )
+
+            mrl_256 = F.linear(lm.norm(h)[:, 2 : 2 + n_audio_pad, :], self._proj_w_256).cpu()
+
+        return [
+            F.normalize(mrl_256[i, :n].float(), p=2, dim=-1).to(torch.bfloat16)
+            for i, n in enumerate(n_list)
+        ]
+
+    def extract_features(self, wav: torch.Tensor) -> torch.Tensor:
+        """
+        Extracts genuine (Time, 256) bfloat16 feature tensor from raw waveform in the shared multimodal space.
+        Maintains full backward compatibility with standalone callers.
+        """
+        f_in, fm_in, n_audio, n_audio_pad = self.preprocess_audio_cpu(wav)
+        return self.extract_from_preprocessed(f_in, fm_in, n_audio, n_audio_pad)
+
+
+def format_time_hms(seconds: Optional[float]) -> str:
+    """초 단위를 'X시간 Y분 Z초', 'Y분 Z초', 'Z초' 형태로 사람이 읽기 쉽게 변환합니다."""
+    if seconds is None or seconds < 0 or not math.isfinite(seconds):
+        return "계산 중..."
+    s = int(round(seconds))
+    h = s // 3600
+    m = (s % 3600) // 60
+    sec = s % 60
+    if h > 0:
+        return f"{h}시간 {m:02d}분 {sec:02d}초"
+    elif m > 0:
+        return f"{m}분 {sec:02d}초"
+    else:
+        return f"{sec}초"
+
+
+# ==============================================================================
+# 5. Main Execution Entry
+# ==============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="Step 2: Pre-extract and Cache Audio Features from Real 300M Audio Encoder")
+    parser.add_argument("--config", type=str, default="configs/config.yaml", help="Path to config file")
+    parser.add_argument("--mode", type=str, choices=["streaming", "zip", "aihub_relay"], default="streaming",
+                        help="Data source mode: 'streaming' (Zeroth-Korean), 'zip' (AI Hub single zip), or 'aihub_relay' (100GB batch relay)")
+    parser.add_argument("--source_dir", type=str, default=None, help="Path to source directory or split zip volumes (e.g. TS.z01)")
+    parser.add_argument("--zip_path", type=str, default=None, help="Path to AI Hub ZIP archive (required if mode=zip)")
+    parser.add_argument("--audio_dir", type=str, default=None, help="Path to directory containing batch .wav audio files")
+    parser.add_argument("--aihub_labels_zip", type=str, default=None, help="Path to AI Hub labeling TL.zip archive")
+    parser.add_argument("--aihub_labels_dir", type=str, default=None, help="Path to AI Hub labeling directory containing json files")
+    parser.add_argument("--delete_source_after_extract", action="store_true", help="Delete source .wav file after feature extraction to free disk space")
+    parser.add_argument("--clean_cache", action="store_true", help="Clean existing cached features in data/cache_features before starting")
+    parser.add_argument("--max_samples", type=int, default=None, help="Number of audio samples to extract and cache (default: all available in archive)")
+    parser.add_argument("--split", type=str, default="train", help="Dataset split for streaming mode")
+    parser.add_argument("--no_augment", action="store_true", help="Disable lecture room acoustic augmentation")
+    parser.add_argument("--batch_size", type=int, default=8, help="Max same-bucket GPU batch size (default: 8; set 1 for strict single-sample mode)")
+    args = parser.parse_args()
+
+    config_path = Path(args.config)
+    if not config_path.is_absolute():
+        config_path = Path(__file__).resolve().parent.parent / config_path
+
+    with open(config_path, "r", encoding="utf-8") as f:
+        cfg = yaml.safe_load(f)
+
+    base_dir = Path(__file__).resolve().parent.parent
+    cache_dir = base_dir / cfg["data"].get("cache_dir", "data/cache_features")
+
+    # Clean previous cached features only if explicitly requested
+    if args.clean_cache and cache_dir.exists():
+        print(f"[Step 2] Cleaning previous cached features in {cache_dir}...")
+        for old_f in cache_dir.glob("*.pt"):
+            try:
+                old_f.unlink()
+            except Exception:
+                pass
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # Build O(1) in-memory set of existing .pt filenames in a single scandir pass (no extra stat() calls)
+    with os.scandir(cache_dir) as it:
+        existing_pt_files: Set[str] = {entry.name for entry in it if entry.name.endswith(".pt")}
+    start_index = len(existing_pt_files)
+    if start_index > 0:
+        print(f"[Step 2] Incremental Relay Mode: {start_index:,} existing features indexed in RAM (O(1) resume lookup).")
+
+    # Initialize Augmenter (Reverb disabled; ambient noise only)
+    aug_cfg = cfg["audio"].get("augmentation", {})
+    do_augment = aug_cfg.get("enabled", True) and not args.no_augment
+    augmenter = LectureRoomAugmenter(
+        sample_rate=16000,
+        rir_prob=0.0,
+        noise_prob=aug_cfg.get("noise_prob", 0.4) if do_augment else 0.0,
+        snr_db_low=aug_cfg.get("noise_snr_db_low", 15.0),
+        snr_db_high=aug_cfg.get("noise_snr_db_high", 30.0)
+    )
+
+    # Initialize Genuine Feature Extractor (300M Audio Encoder + 271M LM)
+    device_name = cfg["system"].get("device", "mps")
+    device = torch.device(device_name if (device_name == "mps" and torch.backends.mps.is_available()) else "cpu")
+    extractor = GenuineAudioFeatureExtractor(device=device, model_name=cfg["model"].get("name", "google/embeddinggemma-2"))
+
+    # Determine Data Stream
+    if args.mode == "aihub_relay":
+        audio_src = args.source_dir or args.audio_dir or args.zip_path
+        if not audio_src:
+            default_audio_src = Path("/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/01.원천데이터")
+            if default_audio_src.exists():
+                audio_src = str(default_audio_src)
+            else:
+                raise ValueError("--source_dir or --audio_dir required when --mode aihub_relay is specified.")
+        labels_src = args.aihub_labels_zip or args.aihub_labels_dir
+        if not labels_src:
+            default_labels_zip = Path("/Users/yeongchan/Downloads/137.한국어 대학 강의 데이터/01-1.정식개방데이터/Training/02.라벨링데이터/TL.zip")
+            if default_labels_zip.exists():
+                labels_src = str(default_labels_zip)
+            else:
+                raise ValueError("AI Hub labels source (TL.zip or folder) must be provided via --aihub_labels_zip or --aihub_labels_dir.")
+        relay_reader = AIHubLectureRelayReader(
+            source_dir=audio_src,
+            labels_source=labels_src,
+            cache_dir=str(cache_dir),
+            max_samples=args.max_samples,
+            existing_files=existing_pt_files
+        )
+        raw_stream = relay_reader.stream_samples()
+    elif args.mode == "zip":
+        if not args.zip_path:
+            raise ValueError("--zip_path is required when --mode zip is specified.")
+        zip_reader = AIHubZipReader(args.zip_path)
+        raw_stream = ((w, s, t, None, None) for w, s, t in zip_reader.stream_samples(max_samples=args.max_samples))
+    else:
+        dataset_name = cfg["data"].get("hf_dataset_name", "Bingsu/zeroth-korean")
+        raw_stream = ((w, s, t, None, None) for w, s, t in stream_zeroth_korean(dataset_name=dataset_name, split=args.split, max_samples=args.max_samples))
+
+    # Stage 1 Generator: Perform all CPU-bound work (resampling, noise augmentation,
+    # Mel-spectrogram extraction, bucket padding, and early CTC length validation) in the background thread!
+    def _cpu_preprocess_generator():
+        seq_counter = start_index
+        for wav, sr, text, short_key, src_path in raw_stream:
+            clean_text = normalize_text(text)
+            if not clean_text:
+                continue
+
+            num_syllables = len(clean_text.replace(" ", ""))
+            if num_syllables == 0:
+                continue
+
+            if short_key:
+                target_fname = f"feat_aihub_{short_key}.pt"
+            else:
+                while f"feat_{seq_counter:06d}.pt" in existing_pt_files:
+                    seq_counter += 1
+                target_fname = f"feat_{seq_counter:06d}.pt"
+
+            if target_fname in existing_pt_files:
+                continue
+
+            # 1. Resample to 16,000Hz mono
+            wav = extractor.resample_if_needed(wav, orig_sr=sr)
+            if wav.dim() > 1 and wav.size(0) > 1:
+                wav = wav.mean(dim=0, keepdim=True)
+
+            # 2. Apply Lecture Room Augmentation (Ambient Noise only)
+            if do_augment:
+                wav = augmenter(wav)
+
+            # 3. Extract Mel-spectrogram & bucket padding on CPU
+            f_in, fm_in, n_audio, n_audio_pad = extractor.preprocess_audio_cpu(wav)
+
+            # 4. Early CTC validity check BEFORE running expensive GPU inference (T >= L)
+            if n_audio < num_syllables:
+                continue
+
+            cache_file = cache_dir / target_fname
+            if not short_key:
+                seq_counter += 1
+
+            yield f_in, fm_in, n_audio, n_audio_pad, num_syllables, clean_text, cache_file, target_fname, src_path
+
+    # Wrap CPU preprocessing generator in asynchronous prefetcher (Stage 1) and disk saver (Stage 3)
+    preprocessed_stream = PrefetchedStream(_cpu_preprocess_generator(), max_queue_size=32)
+    disk_saver = AsyncDiskSaver(max_queue_size=64, delete_source_after_extract=args.delete_source_after_extract)
+
+    saved_count = start_index
+    new_saved = 0
+    max_bs = max(1, int(args.batch_size))
+    token_budget = 1536  # Memory-safe token budget for 16GB Unified Memory
+    bucket_queues: Dict[int, List[Tuple[Any, ...]]] = defaultdict(list)
+    buffered_count = 0
+    max_buffered = max_bs * 6
+
+    # Determine total expected extraction count for progress bar & remaining time calculation
+    target_new: Optional[int] = None
+    if args.max_samples and args.max_samples > 0:
+        target_new = args.max_samples
+    elif args.mode == "aihub_relay" and "relay_reader" in locals() and hasattr(relay_reader, "labels_map") and relay_reader.labels_map:
+        target_new = max(0, len(relay_reader.labels_map) - start_index)
+    elif args.mode == "zip" and "zip_reader" in locals() and hasattr(zip_reader, "valid_keys") and zip_reader.valid_keys:
+        target_new = max(0, len(zip_reader.valid_keys) - start_index)
+
+    pbar = tqdm(
+        total=target_new,
+        desc="Extracting Genuine Features",
+        dynamic_ncols=True
+    )
+
+    start_time = time.time()
+
+    def _flush_bucket(n_audio_pad_key: int) -> bool:
+        nonlocal saved_count, new_saved, buffered_count
+        items = bucket_queues.get(n_audio_pad_key)
+        if not items:
+            return False
+
+        bs_cap = max(1, min(max_bs, token_budget // max(n_audio_pad_key, 64)))
+        batch = items[:bs_cap]
+        del items[:bs_cap]
+        buffered_count -= len(batch)
+
+        # Stage 2: Dedicated MPS GPU Forward Pass (Same-Bucket Batched or Bit-Exact Single)
+        pre_inputs = [(it[0], it[1], it[2], it[3]) for it in batch]
+        feats_list = extractor.extract_batch_same_bucket(pre_inputs)
+
+        for it, features_256 in zip(batch, feats_list):
+            _, _, _, _, num_syllables, clean_text, cache_file, target_fname, src_path = it
+            if features_256.size(0) < num_syllables:
+                continue
+
+            # Stage 3: Enqueue non-blocking atomic disk write
+            existing_pt_files.add(target_fname)
+            disk_saver.save(
+                {
+                    "feature": features_256,
+                    "raw_text": clean_text
+                },
+                cache_file,
+                src_path
+            )
+
+            saved_count += 1
+            new_saved += 1
+            pbar.update(1)
+
+            if new_saved % 250 == 0:
+                gc.collect()
+                if torch.backends.mps.is_available():
+                    torch.mps.empty_cache()
+
+            if new_saved % 500 == 0:
+                import shutil
+                free_gb = shutil.disk_usage(cache_dir).free / (1024 ** 3)
+                if free_gb < 5.0:
+                    print(f"\n⚠️ [Step 2 Disk Safety Alert] Remaining disk space is {free_gb:.2f} GB (< 5GB safety threshold). Gracefully concluding extraction.")
+                    return True
+
+            if args.max_samples and args.max_samples > 0 and new_saved >= args.max_samples:
+                return True
+
+        return False
+
+    try:
+        stop_early = False
+        for item in preprocessed_stream:
+            n_audio_pad = item[3]
+            if max_bs <= 1:
+                bucket_queues[n_audio_pad].append(item)
+                buffered_count += 1
+                if _flush_bucket(n_audio_pad):
+                    stop_early = True
+                    break
+            else:
+                bucket_queues[n_audio_pad].append(item)
+                buffered_count += 1
+                bs_cap = max(1, min(max_bs, token_budget // max(n_audio_pad, 64)))
+                if len(bucket_queues[n_audio_pad]) >= bs_cap:
+                    if _flush_bucket(n_audio_pad):
+                        stop_early = True
+                        break
+                elif buffered_count >= max_buffered:
+                    # Flush the fullest bucket to bound latency and RAM
+                    oldest_bucket = max(bucket_queues.keys(), key=lambda k: len(bucket_queues[k]))
+                    if _flush_bucket(oldest_bucket):
+                        stop_early = True
+                        break
+
+        if not stop_early:
+            for k in sorted(list(bucket_queues.keys())):
+                while bucket_queues[k]:
+                    if _flush_bucket(k):
+                        break
+    finally:
+        preprocessed_stream.close()
+        disk_saver.close()
+        pbar.close()
+
+    total_elapsed = time.time() - start_time
+    avg_speed = new_saved / max(total_elapsed, 1e-5)
+    print(f"\n[Step 2 Completed Successfully!]")
+    print(f"  Newly extracted features: {new_saved:,} (Total in cache: {saved_count:,})")
+    print(f"  Total time elapsed: {format_time_hms(total_elapsed)} (Average speed: {avg_speed:.2f} samples/s)")
+    print(f"  Target dimension: 256d (bfloat16), Sample rate: 16kHz")
+
+
+if __name__ == "__main__":
+    main()

transcribe.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""
+Notesoop-STT-Ko-v1: Standalone CLI Inference Script
+
+Usage:
+    python transcribe.py audio.mp3
+    python transcribe.py audio.wav --batch-size 4 --chunk-sec 16.0
+    python transcribe.py audio.m4a --output result.txt
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+# Add current directory to path
+_ROOT = Path(__file__).resolve().parent
+if str(_ROOT) not in sys.path:
+    sys.path.insert(0, str(_ROOT))
+
+from run_inference import EmbeddingGemmaSTT
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Notesoop-STT-Ko-v1: Ultra-Fast Korean On-Device Speech Recognition"
+    )
+    parser.add_argument("audio", type=str, help="Path to audio file (.wav, .m4a, .mp3, .flac, etc.)")
+    parser.add_argument("--batch-size", type=int, default=4, help="Batch size for parallel chunk processing (default: 4)")
+    parser.add_argument("--chunk-sec", type=float, default=16.0, help="Chunk length in seconds (default: 16.0s)")
+    parser.add_argument("--overlap-sec", type=float, default=1.0, help="Overlap length in seconds (default: 1.0s)")
+    parser.add_argument("--output", type=str, default=None, help="Optional text file path to save transcription")
+    parser.add_argument("--no-kenlm", action="store_true", help="Disable KenLM 3-Gram beam search and use Greedy CTC")
+    args = parser.parse_args()
+
+    audio_path = Path(args.audio)
+    if not audio_path.exists():
+        print(f"❌ Error: Audio file not found at '{audio_path}'", file=sys.stderr)
+        sys.exit(1)
+
+    print("=" * 70)
+    print("🌲 Notesoop-STT-Ko-v1: Korean On-Device Speech Recognition")
+    print(f"  • Audio File : {audio_path.name}")
+    print(f"  • Batch Size : {args.batch_size}")
+    print(f"  • Chunking   : {args.chunk_sec}s chunks (overlap {args.overlap_sec}s)")
+    print(f"  • KenLM      : {'Enabled (alpha=0.12, beta=0.40)' if not args.no_kenlm else 'Disabled (Greedy)'}")
+    print("=" * 70)
+
+    # Initialize Notesoop STT Engine
+    stt = EmbeddingGemmaSTT(
+        device="auto",
+        use_kenlm=(not args.no_kenlm),
+        lm_alpha=0.12,
+        lm_beta=0.40
+    )
+
+    print("\n🎙️ Transcribing audio...")
+    t0 = time.perf_counter()
+    transcript, rtf = stt.transcribe(
+        str(audio_path),
+        chunk_sec=args.chunk_sec,
+        overlap_sec=args.overlap_sec,
+        batch_size=args.batch_size
+    )
+    elapsed = time.perf_counter() - t0
+
+    print("\n" + "=" * 70)
+    print(f"✨ Transcription Completed in {elapsed:.2f}s (RTF: {rtf:.4f}, Speedup: {1/max(rtf, 1e-4):.1f}x)")
+    print("=" * 70)
+    print(transcript)
+    print("=" * 70)
+
+    if args.output:
+        out_p = Path(args.output)
+        with open(out_p, "w", encoding="utf-8") as f:
+            f.write(transcript)
+        print(f"💾 Transcript saved to: {out_p}")
+
+
+if __name__ == "__main__":
+    main()