Add README

README.md

@@ -0,0 +1,110 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+- ja
+- ko
+- multilingual
+library_name: mlx
+tags:
+- mlx
+- automatic-speech-recognition
+- asr
+- speech-recognition
+- robust-asr
+- quantized
+- int4
+- 4bit
+- dwq
+- distillation-aware-quantization
+- on-device
+- apple-silicon
+- qwen3
+- qwen3-asr
+- mega-asr
+pipeline_tag: automatic-speech-recognition
+base_model: zhifeixie/Mega-ASR
+base_model_relation: quantized
+---
+
+# Mega-ASR — MLX 4-bit DWQ
+
+MLX 4-bit (DWQ-distilled) deployment of [zhifeixie/Mega-ASR](https://huggingface.co/zhifeixie/Mega-ASR),
+the 1.7B-parameter robust multilingual ASR foundation model built on Qwen3-ASR-1.7B.
+
+Quantized via [`mlx_lm.quant.dwq`](https://github.com/ml-explore/mlx-lm) (Distillation-aware
+Weight Quantization) — the 4-bit weights are fine-tuned to match the bf16 teacher's
+logits on a calibration corpus, recovering most of the post-training drift.
+
+## What's in this repo
+
+| File | Size | Role |
+| --- | ---: | --- |
+| `mlx/llm-dwq4/model.safetensors` | **923 MB** | Qwen3 1.7B LLM, MLX 4-bit DWQ, group_size=64 (4.5 bits/weight) |
+| `mlx/llm-dwq4/{config,tokenizer*}.json` | — | mlx-lm load metadata + LLM tokenizer |
+| `onnx/audio_encoder_fp32.onnx` | **1.27 GB** | 24-layer Whisper-style audio encoder (ONNX fp32, run via onnxruntime). MLX port is planned. |
+| `tokenizer/*` | — | Original Qwen3-ASR tokenizer (with audio special tokens). Required because the LLM-only tokenizer drops `<\|audio_pad\|>`, `<\|audio_start\|>`, etc. |
+| `examples/*.wav` | ~3 MB | 8 noisy benchmark clips from Voices-in-the-Wild-Bench |
+| `inference.py` | — | End-to-end ASR pipeline: ONNX encoder + MLX LLM |
+
+## Quality (bench)
+
+8-clip [Voices-in-the-Wild-Bench](https://github.com/xzf-thu/Voices-in-the-Wild-Bench)
+agreement (1 − WER), prompt forced to `language English`:
+
+| Variant | Encoder | Decoder | Agreement | Total size |
+| --- | --- | --- | ---: | ---: |
+| PT bf16 (original) | fp16 | fp16 | 95.1% | 7.5 GB |
+| ONNX recommended (GPTQ) | INT8 ONNX | INT4 GPTQ | 92.7% | 2.3 GB |
+| **MLX 4-bit DWQ (this repo)** | **fp32 ONNX** | **MLX 4-bit DWQ** | **89.9%** | **~2.2 GB** |
+| MLX 4-bit (no DWQ) | fp32 ONNX | MLX 4-bit | 89.1% | ~2.2 GB |
+
+DWQ recovers +0.8% over straight mlx-lm 4-bit quantization. Audio encoder is
+still ONNX fp32 because the `qwen3_asr_audio_encoder` architecture (custom
+chunked Whisper-style with Conv2d front-end) isn't natively supported by
+mlx-lm or mlx-whisper yet — porting it to pure MLX is on the roadmap.
+
+## Inference
+
+```bash
+pip install mlx mlx-lm onnxruntime soundfile transformers librosa numpy
+git clone https://huggingface.co/Reza2kn/mega-asr-mlx
+cd mega-asr-mlx
+python inference.py --encoder-path onnx/audio_encoder_fp32.onnx \
+                    --mlx-llm-path mlx/llm-dwq4 \
+                    --examples-dir examples
+```
+
+The pipeline runs:
+1. mel features (Whisper preprocessor) → ONNX audio encoder (onnxruntime CPU) → audio embeddings (1, F, 2048)
+2. tokenize the Qwen3-ASR chat prompt with `audio_pad_id=151676`, expand the single pad placeholder to F copies
+3. embed all tokens via `model.model.embed_tokens` (MLX), scatter audio embeddings at the audio_pad positions
+4. greedy decode via MLX with `input_embeddings`
+
+## Conversion details
+
+- LLM extracted from `zhifeixie/Mega-ASR/Qwen3-ASR-1.7B/` by stripping the
+  `thinker.model.` prefix from layer weights and dropping the tied `lm_head`
+  (relies on `tie_word_embeddings=True`).
+- `mlx_lm.convert --hf-path <extracted> -q --q-bits 4 --q-group-size 64` →
+  base MLX 4-bit (4.501 bits/weight).
+- `mlx_lm.quant.dwq --model <bf16-teacher> --bits 4 --group-size 64
+  --num-samples 64 --max-seq-length 256 --learning-rate 1e-6` →
+  DWQ-distilled weights. 64 steps of distillation on tulu-3-sft-mixture
+  reduced validation KL loss from ~0.18 to ~0.14.
+- Audio encoder ONNX is reused unchanged from
+  [Reza2kn/mega-asr-onnx](https://huggingface.co/Reza2kn/mega-asr-onnx).
+
+## Performance
+
+| Hardware | Cold load | Warm (3-4s audio) |
+| --- | ---: | ---: |
+| M-series Mac (MLX) | ~3 s | ~1.5 s (LLM @ ~60 tps) |
+
+## Credits
+
+- Original model: [zhifeixie/Mega-ASR](https://huggingface.co/zhifeixie/Mega-ASR) (1.7B params, Apache-2.0)
+- MLX port + DWQ: this repo
+- Benchmark: [Voices-in-the-Wild-Bench](https://github.com/xzf-thu/Voices-in-the-Wild-Bench)
+- DWQ tool: [`mlx_lm.quant.dwq`](https://github.com/ml-explore/mlx-lm) (Apple Inc.)