Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* 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
+audio_encoder_int4.onnx.data filter=lfs diff=lfs merge=lfs -text
+decoder_prefill_int4.onnx.data filter=lfs diff=lfs merge=lfs -text
+decoder_step_int4.onnx.data filter=lfs diff=lfs merge=lfs -text

audio_encoder_int4.onnx

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

audio_encoder_int4.onnx.data

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

audio_encoder_int8_static.onnx

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

decoder_prefill_int4.onnx

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

decoder_prefill_int4.onnx.data

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

decoder_step_int4.onnx

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

decoder_step_int4.onnx.data

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

inference.py

@@ -0,0 +1,150 @@
+"""Standalone Mega-ASR inference using the INT4 ONNX models.
+
+This script demonstrates how to run Mega-ASR end-to-end from raw audio
+using only the ONNX models in this repo plus the qwen-asr tokenizer.
+
+Requirements:
+    pip install onnxruntime numpy soundfile transformers qwen-asr
+
+Usage:
+    python inference.py --audio examples/noise.wav
+"""
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+import soundfile as sf
+
+
+HERE = Path(__file__).parent
+
+
+def _ort(path, providers=("CPUExecutionProvider",)):
+    so = ort.SessionOptions()
+    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    return ort.InferenceSession(str(path), so, providers=list(providers))
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--audio", required=True, type=Path)
+    ap.add_argument("--max-new-tokens", type=int, default=80)
+    ap.add_argument("--providers", default="cuda,cpu",
+                    help="Comma-separated ORT providers, e.g. cuda,cpu")
+    args = ap.parse_args()
+
+    providers = []
+    for p in args.providers.split(","):
+        p = p.strip().lower()
+        if p == "cuda": providers.append("CUDAExecutionProvider")
+        elif p == "cpu": providers.append("CPUExecutionProvider")
+        elif p == "webgpu": providers.append("WebGpuExecutionProvider")
+
+    # 1) Tokenizer + processor (needs HF transformers + qwen-asr installed)
+    from transformers import AutoTokenizer
+    from qwen_asr.core.transformers_backend.processing_qwen3_asr import Qwen3ASRProcessor
+    tokenizer = AutoTokenizer.from_pretrained(str(HERE))
+    processor = Qwen3ASRProcessor.from_pretrained(str(HERE))
+
+    # 2) Read audio + build mel features
+    audio, sr = sf.read(str(args.audio))
+    if sr != 16000:
+        import librosa
+        audio = librosa.resample(audio.astype(np.float32), orig_sr=sr, target_sr=16000)
+
+    # The processor expects a prompt with the audio placeholder. Build it as
+    # the Qwen3-ASR chat template would.
+    prompt = (
+        "<|im_start|>system\\nYou are a speech recognition model.<|im_end|>\\n"
+        "<|im_start|>user\\nDetect the language and recognize the speech: "
+        "<|audio_start|><|audio_pad|><|audio_end|><|im_end|>\\n"
+        "<|im_start|>assistant\\n"
+    )
+    inputs = processor(text=[prompt], audio=[audio], return_tensors="np", padding=True)
+    text_ids = inputs["input_ids"]                 # (1, L)
+    attention_mask = inputs["attention_mask"]      # (1, L)
+    input_features = inputs["input_features"]      # (1, 128, T_mel)
+
+    # 3) Audio encoder
+    T_mel = input_features.shape[-1]
+    if T_mel > 3000:
+        input_features = input_features[..., :3000]
+        T_mel = 3000
+    mel_padded = np.pad(input_features.astype(np.float32),
+                         ((0, 0), (0, 0), (0, 3000 - T_mel)))
+    enc = _ort(HERE / "onnx" / "audio_encoder_int4.onnx", providers)
+    audio_embeds = enc.run(["audio_embeds"], {"mel": mel_padded})[0]   # (1, 390, 2048)
+    real_chunks = (T_mel + 99) // 100
+    last_chunk_mel = T_mel - (real_chunks - 1) * 100
+    real_audio_frames = (real_chunks - 1) * 13 + (last_chunk_mel + 7) // 8
+    audio_embeds = audio_embeds[:, :real_audio_frames]                 # (1, N, 2048)
+
+    # 4) Embed text tokens + scatter audio into placeholder positions
+    # The embedding table is published as embed_tokens.npy in the repo
+    # (extracted from the original Qwen3-ASR weights for portability).
+    embed_path = HERE / "onnx" / "embed_tokens.npy"
+    if embed_path.exists():
+        embed_w = np.load(embed_path).astype(np.float16)
+    else:
+        raise FileNotFoundError(
+            "embed_tokens.npy missing — re-extract from the source PT model."
+        )
+    inputs_embeds = embed_w[text_ids[0]][None]                          # (1, L, hidden)
+    audio_token_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
+    placeholder_mask = (text_ids == audio_token_id)                     # (1, L)
+    # Replace placeholder embeddings with audio_embeds[0] in order
+    placeholder_idx = np.where(placeholder_mask[0])[0]
+    n_replace = min(len(placeholder_idx), audio_embeds.shape[1])
+    inputs_embeds[0, placeholder_idx[:n_replace]] = audio_embeds[0, :n_replace]
+
+    # 5) Position ids (1D)
+    pos_ids = np.arange(text_ids.shape[1])[None].astype(np.int64)
+
+    # 6) Prefill ONNX
+    prefill = _ort(HERE / "onnx" / "decoder_prefill_int4.onnx", providers)
+    feeds = {
+        "inputs_embeds": inputs_embeds.astype(np.float16),
+        "attention_mask": attention_mask.astype(np.int64),
+        "position_ids": pos_ids,
+    }
+    out = prefill.run(None, feeds)
+    logits, *past_kvs = out
+    eos = tokenizer.eos_token_id
+
+    # 7) Greedy decode with step ONNX
+    step = _ort(HERE / "onnx" / "decoder_step_int4.onnx", providers)
+    num_layers = 28
+    cur_len = inputs_embeds.shape[1]
+    gen = []
+    nid = int(np.argmax(logits[0, -1, :]))
+    gen.append(nid)
+
+    for _ in range(args.max_new_tokens - 1):
+        if nid == eos: break
+        new_embed = embed_w[nid][None, None].astype(np.float16)
+        new_attn = np.ones((1, cur_len + 1), dtype=np.int64)
+        new_pos = np.array([[cur_len]], dtype=np.int64)
+        f = {"inputs_embeds": new_embed,
+             "attention_mask": new_attn,
+             "position_ids": new_pos}
+        for i in range(num_layers):
+            f[f"past.{i}.key"] = past_kvs[2 * i]
+            f[f"past.{i}.value"] = past_kvs[2 * i + 1]
+        out = step.run(None, f)
+        logits, *past_kvs = out
+        nid = int(np.argmax(logits[0, -1, :]))
+        gen.append(nid)
+        cur_len += 1
+
+    text = tokenizer.decode(gen, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    print(f"TRANSCRIPTION: {text}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

preprocessor_config.json

@@ -0,0 +1,14 @@
+{
+  "chunk_length": 30,
+  "dither": 0.0,
+  "feature_extractor_type": "WhisperFeatureExtractor",
+  "feature_size": 128,
+  "hop_length": 160,
+  "n_fft": 400,
+  "n_samples": 480000,
+  "nb_max_frames": 3000,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "processor_class": "Qwen3ASRProcessor",
+  "return_attention_mask": true
+}

tokenizer_config.json

@@ -0,0 +1,549 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151669": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151671": {
+      "content": "<tts_pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151672": {
+      "content": "<tts_text_bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151673": {
+      "content": "<tts_text_eod>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151674": {
+      "content": "<tts_text_bos_single>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151675": {
+      "content": "<non_speech>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151676": {
+      "content": "<|audio_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151677": {
+      "content": "<blank1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151678": {
+      "content": "<blank2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151679": {
+      "content": "<blank3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151680": {
+      "content": "<blank4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151681": {
+      "content": "<blank5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151682": {
+      "content": "<blank6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151683": {
+      "content": "<blank7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151684": {
+      "content": "<blank8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151685": {
+      "content": "<blank9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151686": {
+      "content": "<blank10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151687": {
+      "content": "<blank11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151688": {
+      "content": "<blank12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151689": {
+      "content": "<blank13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151690": {
+      "content": "<blank14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151691": {
+      "content": "<blank15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151692": {
+      "content": "<blank16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151693": {
+      "content": "<blank17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151694": {
+      "content": "<blank18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151695": {
+      "content": "<blank19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151696": {
+      "content": "<blank20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151697": {
+      "content": "<blank21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151698": {
+      "content": "<blank22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151699": {
+      "content": "<blank23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151700": {
+      "content": "<blank24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151701": {
+      "content": "<blank25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151702": {
+      "content": "<blank26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151703": {
+      "content": "<blank27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151704": {
+      "content": "<asr_text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<tts_pad>",
+    "<tts_text_bos>",
+    "<tts_text_bos_single>",
+    "<|audio_pad|>"
+  ],
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "image_token": "<|image_pad|>",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "Qwen3ASRProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}