Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ 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
+examples/mlk_speech.wav filter=lfs diff=lfs merge=lfs -text
+examples/sample_speech.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -6,7 +6,27 @@ colorTo: indigo
 sdk: gradio
 sdk_version: 6.20.0
 app_file: app.py
-short_description: Probe
+short_description: Audio understanding, speech recognition & translation
 python_version: "3.10"
 startup_duration_timeout: 1h
 ---
+
+# Nemotron-Labs-Audex-30B-A3B
+
+A Gradio demo for [`nvidia/Nemotron-Labs-Audex-30B-A3B`](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B),
+a unified audio-text MoE (30B total / 3B active) built on Nemotron-Cascade-2.
+This demo exposes its **audio understanding, speech recognition, and speech
+translation** capabilities via the official Hugging Face inference path.
+
+## Correct numerics
+
+The model is a Nemotron-H hybrid (Mamba2 + attention + MLP). The Mamba layers
+require the compiled CUDA fast path from `mamba-ssm` (`selective_scan_cuda`) and
+`causal-conv1d`; the pure-torch fallback produces degenerate output. This Space
+builds both extensions from source against the runtime's exact toolchain
+(torch 2.11.0+cu130, CUDA 13.0, Python 3.10, sm_120) on first boot and caches
+the resulting wheels back into the repo (`wheels/`) for fast subsequent boots.
+
+Runs on ZeroGPU `xlarge` (full 96 GB card) since the bf16 weights are ~65 GB.
+
+*Licensed for non-commercial use only (NVIDIA One-Way Noncommercial License).*

app.py

@@ -1,42 +1,355 @@
-import os, sys, subprocess, platform
+"""
+Nemotron-Labs-Audex-30B-A3B — Audio Understanding / Speech Recognition & Translation demo.
 
-# Probe at import so it lands in run logs regardless of UI/SSR state
-def _probe():
-    lines = []
-    lines.append(f"python: {sys.version}")
-    lines.append(f"platform: {platform.platform()}")
+This unified audio-text MoE (30B total, 3B active) is a Nemotron-H hybrid
+(Mamba2 + attention + MLP). Correct numerics REQUIRE the compiled CUDA fast
+path from `mamba-ssm` (selective_scan_cuda) and `causal-conv1d`; the pure-torch
+fallback produces degenerate/repeated tokens. We therefore build both extensions
+from source against the Space's exact torch (2.11.0+cu130) / CUDA 13.0 / py3.10
+/ sm_120 toolchain, cache the resulting wheels back into this Space repo, and
+reuse them on subsequent boots.
+"""
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import sys
+import glob
+import time
+import subprocess
+from pathlib import Path
+
+# Import spaces FIRST (before torch / any CUDA-touching import) so its
+# torch.cuda.* monkey-patch is installed. The kernel build below runs in
+# subprocesses, so it does not initialize CUDA in this process.
+import spaces  # noqa: E402
+
+APP_DIR = Path(__file__).resolve().parent
+WHEELS_DIR = APP_DIR / "wheels"
+REPO_ID = "hugging-apps/nvidia-nemotron-labs-audex-30b-a3b"
+
+# Pinned versions recommended by the model card.
+CAUSAL_CONV1D_SPEC = "causal-conv1d==1.6.2.post1"
+MAMBA_SSM_SPEC = "mamba-ssm==2.3.2.post1"
+# sm_120 (RTX PRO 6000 Blackwell). Build env has no GPU, so arch must be explicit.
+TORCH_ARCH = "12.0"
+
+
+def _run(cmd, env=None, timeout=3000):
+    print(f"[build] $ {' '.join(cmd)}", flush=True)
+    p = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=timeout)
+    if p.stdout:
+        print(p.stdout[-4000:], flush=True)
+    if p.returncode != 0:
+        print(p.stderr[-8000:], flush=True)
+        raise RuntimeError(f"command failed ({p.returncode}): {' '.join(cmd)}")
+    else:
+        # surface tail of stderr (nvcc warnings etc.) but not as failure
+        if p.stderr:
+            print(p.stderr[-2000:], flush=True)
+    return p
+
+
+def _installed(mod):
     try:
-        import torch
-        lines.append(f"torch: {torch.__version__}")
-        lines.append(f"torch.version.cuda: {torch.version.cuda}")
-    except Exception as e:
-        lines.append(f"torch import error: {e!r}")
-    for cmd in (["nvcc", "--version"], ["which", "nvcc"], ["ls", "/usr/local/cuda/bin"]):
+        __import__(mod)
+        return True
+    except Exception:
+        return False
+
+
+def _pip(*args):
+    _run([sys.executable, "-m", "pip"] + list(args))
+
+
+def _build_env():
+    env = dict(os.environ)
+    env["MAMBA_FORCE_BUILD"] = "TRUE"
+    env["CAUSAL_CONV1D_FORCE_BUILD"] = "TRUE"
+    env["TORCH_CUDA_ARCH_LIST"] = TORCH_ARCH
+    env["MAX_JOBS"] = env.get("MAX_JOBS", "4")
+    # Ensure nvcc is discoverable.
+    cuda_home = env.get("CUDA_HOME") or "/cuda-image/usr/local/cuda-13.0"
+    env["CUDA_HOME"] = cuda_home
+    env["PATH"] = f"{cuda_home}/bin:" + env.get("PATH", "")
+    return env
+
+
+def ensure_kernels():
+    """Install compiled causal-conv1d + mamba-ssm. Use cached wheels if present,
+    otherwise build from source and cache the wheels back into the repo."""
+    if _installed("causal_conv1d") and _installed("selective_scan_cuda"):
+        print("[build] kernels already importable", flush=True)
+        return
+
+    WHEELS_DIR.mkdir(exist_ok=True)
+    cached = sorted(glob.glob(str(WHEELS_DIR / "*.whl")))
+    if cached:
+        print(f"[build] installing cached wheels: {cached}", flush=True)
         try:
-            out = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
-            lines.append(f"$ {' '.join(cmd)} (rc={out.returncode})\n{out.stdout}{out.stderr}")
+            _pip("install", "--no-deps", "--no-build-isolation", *cached)
+            if _installed("causal_conv1d") and _installed("selective_scan_cuda"):
+                print("[build] cached wheels installed OK", flush=True)
+                return
+            print("[build] cached wheels imported incompletely; rebuilding", flush=True)
         except Exception as e:
-            lines.append(f"{' '.join(cmd)} error: {e!r}")
-    # look for nvcc from pip nvidia-cuda-nvcc
+            print(f"[build] cached wheel install failed ({e!r}); rebuilding", flush=True)
+
+    env = _build_env()
+    out_dir = APP_DIR / "_wheelout"
+    out_dir.mkdir(exist_ok=True)
+
+    # Build wheels (no-build-isolation => uses the preinstalled torch).
+    t0 = time.time()
+    print("[build] building causal-conv1d + mamba-ssm from source (this can take ~20 min)", flush=True)
+    _pip("wheel", "--no-build-isolation", "--no-deps", "-w", str(out_dir),
+         CAUSAL_CONV1D_SPEC)
+    # mamba-ssm needs causal-conv1d importable during its own build; install it first.
+    built = sorted(glob.glob(str(out_dir / "causal_conv1d*.whl")))
+    _pip("install", "--no-deps", *built)
+    _pip("wheel", "--no-build-isolation", "--no-deps", "-w", str(out_dir),
+         MAMBA_SSM_SPEC)
+    print(f"[build] source build finished in {time.time()-t0:.0f}s", flush=True)
+
+    all_wheels = sorted(glob.glob(str(out_dir / "*.whl")))
+    _pip("install", "--no-deps", *all_wheels)
+
+    if not (_installed("causal_conv1d") and _installed("selective_scan_cuda")):
+        raise RuntimeError("kernel build completed but imports still fail")
+    print("[build] kernels built + installed OK", flush=True)
+
+    # Cache wheels back into the repo for fast subsequent boots.
     try:
-        import glob
-        cands = glob.glob("/usr/local/lib/python3.10/site-packages/nvidia/**/nvcc", recursive=True)
-        cands += glob.glob("/usr/local/lib/python3.10/site-packages/nvidia/cuda_nvcc/bin/*", recursive=True)
-        lines.append("pip nvcc candidates: " + str(cands))
+        for w in all_wheels:
+            dst = WHEELS_DIR / Path(w).name
+            if not dst.exists():
+                import shutil
+                shutil.copy(w, dst)
+        from huggingface_hub import HfApi
+        tok = os.environ.get("HF_TOKEN")
+        if tok:
+            HfApi(token=tok).upload_folder(
+                folder_path=str(WHEELS_DIR),
+                path_in_repo="wheels",
+                repo_id=REPO_ID,
+                repo_type="space",
+                commit_message="cache compiled mamba-ssm + causal-conv1d wheels",
+            )
+            print("[build] cached wheels uploaded to repo", flush=True)
+        else:
+            print("[build] no HF_TOKEN; skipping wheel cache upload", flush=True)
     except Exception as e:
-        lines.append(f"glob err {e!r}")
-    lines.append("CUDA_HOME=" + str(os.environ.get("CUDA_HOME")))
-    return "\n".join(lines)
+        print(f"[build] wheel cache upload skipped ({e!r})", flush=True)
 
-PROBE = _probe()
-print("=========== PROBE START ===========")
-print(PROBE)
-print("=========== PROBE END ===========", flush=True)
 
+ensure_kernels()
+
+# ---- Now safe to bring in torch / model ----
+import torch
 import gradio as gr
+from huggingface_hub import snapshot_download
+from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer
+
+from audio_utils import (
+    IM_END_TOKEN,
+    build_attention_mask,
+    build_prompt_template,
+    expand_sound_placeholder,
+    extract_whisper_features,
+    load_audio,
+    resolve_audio_preprocessor_path,
+    split_thinking,
+)
+
+MODEL_ID = "nvidia/Nemotron-Labs-Audex-30B-A3B"
+SUBFOLDER = "checkpoint_folder_full"
+SAMPLE_RATE = 16000
+
+print("[load] downloading checkpoint…", flush=True)
+LOCAL_DIR = snapshot_download(
+    MODEL_ID,
+    allow_patterns=[f"{SUBFOLDER}/*"],
+    token=os.environ.get("HF_TOKEN"),
+)
+MODEL_PATH = os.path.join(LOCAL_DIR, SUBFOLDER)
+print(f"[load] checkpoint at {MODEL_PATH}", flush=True)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True)
+feature_extractor = AutoFeatureExtractor.from_pretrained(
+    resolve_audio_preprocessor_path(MODEL_PATH, config)
+)
+print("[load] tokenizer/config/feature_extractor loaded", flush=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+).eval()
+model = model.to("cuda")
+print("[load] model loaded and moved to cuda (packed by ZeroGPU)", flush=True)
+
+# Verify fast path is active (import-time flag inside the remote modeling module).
+try:
+    import transformers_modules  # noqa
+except Exception:
+    pass
+
+
+def _find_fast_path_flag():
+    for name, mod in list(sys.modules.items()):
+        if name.endswith("modeling_nemotron_h") and hasattr(mod, "is_fast_path_available"):
+            return bool(mod.is_fast_path_available)
+    return None
+
+
+TASK_PROMPTS = {
+    "Describe the audio": "Describe the audio in detail.",
+    "Transcribe (ASR)": "Transcribe the speech in the input audio.",
+    "Translate speech to English": "Translate the speech in the input audio into English.",
+    "Answer a question about the audio": "What is being said in this audio, and what is the tone?",
+}
+
+
+def _estimate(audio, prompt, reasoning, max_new_tokens, *a, **k):
+    base = 90
+    return int(min(240, base + (int(max_new_tokens) / 1024.0) * 60))
+
+
+@spaces.GPU(duration=_estimate, size="xlarge")
+def transcribe(audio, prompt, reasoning, max_new_tokens,
+               temperature=0.7, top_p=0.9, top_k=0):
+    """Run Nemotron-Labs-Audex audio understanding / ASR / translation.
+
+    Args:
+        audio: path to an input audio file (wav/mp3/flac), 16 kHz mono is ideal.
+        prompt: natural-language instruction about the audio.
+        reasoning: enable the model's <think> reasoning mode.
+        max_new_tokens: maximum number of tokens to generate.
+        temperature: sampling temperature (>0).
+        top_p: nucleus sampling threshold in [0,1].
+        top_k: top-k sampling (0 disables).
+    Returns:
+        (thinking, answer) — the reasoning trace (if any) and the final answer.
+    """
+    if audio is None:
+        return "", "Please provide an audio input."
+    if not prompt or not prompt.strip():
+        prompt = "Describe the audio in detail."
+
+    ff = _find_fast_path_flag()
+    print(f"[gpu] is_fast_path_available={ff}", flush=True)
+
+    wav, sr = load_audio(audio, target_sr=SAMPLE_RATE)
+    input_features = extract_whisper_features(
+        feature_extractor, wav, sample_rate=sr,
+        clip_duration=float(getattr(config, "sound_clip_duration", 30.0)),
+    )
+    num_embeddings = input_features.shape[0] * int(getattr(config, "sound_embedding_size", 750))
+    formatted = build_prompt_template(prompt.strip(), reasoning=bool(reasoning),
+                                      prompt_repitition="none")
+    expanded = expand_sound_placeholder(formatted, num_embeddings)
+    tok = tokenizer(expanded, return_tensors="pt", add_special_tokens=False)
+    input_ids = tok.input_ids.to("cuda")
+    attention_mask = (tok.attention_mask if "attention_mask" in tok
+                      else build_attention_mask(input_ids)).to("cuda")
+    input_features = input_features.to("cuda")
+
+    eos_token_id = tokenizer.convert_tokens_to_ids(IM_END_TOKEN)
+    if eos_token_id is None or eos_token_id == tokenizer.unk_token_id:
+        eos_token_id = getattr(config, "eos_token_id", None)
+
+    temperature = max(float(temperature), 1e-4)
+    top_p = float(top_p)
+    top_k = int(top_k)
+    do_sample = (temperature != 1.0) or (0.0 < top_p < 1.0) or (top_k > 0)
+    gen_kwargs = dict(
+        do_sample=do_sample,
+        eos_token_id=eos_token_id,
+        pad_token_id=tokenizer.pad_token_id or getattr(config, "pad_token_id", 0),
+        max_new_tokens=int(max_new_tokens),
+    )
+    if do_sample:
+        gen_kwargs["temperature"] = temperature
+        if top_p > 0.0:
+            gen_kwargs["top_p"] = top_p
+        if top_k > 0:
+            gen_kwargs["top_k"] = top_k
+
+    with torch.inference_mode():
+        out = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            input_features=input_features,
+            **gen_kwargs,
+        )
+    new_tokens = out[0, input_ids.shape[-1]:]
+    response = tokenizer.decode(new_tokens, skip_special_tokens=False)
+    response = response.split(IM_END_TOKEN, 1)[0].strip()
+    thinking, prediction = split_thinking(response)
+    return thinking, prediction
+
+
+def _ui_run(audio, task, custom_prompt, reasoning, max_new_tokens, temperature, top_p):
+    prompt = custom_prompt.strip() if custom_prompt and custom_prompt.strip() else TASK_PROMPTS.get(task, "")
+    thinking, answer = transcribe(audio, prompt, reasoning, max_new_tokens,
+                                  temperature=temperature, top_p=top_p, top_k=0)
+    return answer, thinking
+
+
+THEME = gr.themes.Citrus()
+
+with gr.Blocks(theme=THEME, title="Nemotron-Labs-Audex-30B-A3B") as demo:
+    gr.Markdown(
+        "# 🎧 Nemotron-Labs-Audex-30B-A3B\n"
+        "Unified audio-text MoE (30B total / 3B active) for **audio understanding, "
+        "speech recognition, and speech translation**. Runs the official Hugging Face "
+        "inference path with compiled `mamba-ssm` + `causal-conv1d` CUDA kernels for "
+        "correct numerics.\n\n"
+        "*Non-commercial use only (NVIDIA One-Way Noncommercial License).*"
+    )
+    with gr.Row():
+        with gr.Column():
+            audio_in = gr.Audio(type="filepath", label="Input audio", sources=["upload", "microphone"])
+            task = gr.Radio(
+                choices=list(TASK_PROMPTS.keys()),
+                value="Transcribe (ASR)",
+                label="Task",
+            )
+            custom_prompt = gr.Textbox(
+                label="Custom instruction (optional — overrides Task)",
+                placeholder="e.g. Summarize what the speaker is saying.",
+                lines=2,
+            )
+            run_btn = gr.Button("Run", variant="primary")
+            with gr.Accordion("Advanced options", open=False):
+                reasoning = gr.Checkbox(value=False, label="Enable reasoning (<think>) mode")
+                max_new_tokens = gr.Slider(16, 1024, value=256, step=16, label="Max new tokens")
+                temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature")
+                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+        with gr.Column():
+            answer_out = gr.Textbox(label="Answer", lines=10)
+            thinking_out = gr.Textbox(label="Reasoning trace (if enabled)", lines=6)
+
+    run_btn.click(
+        _ui_run,
+        inputs=[audio_in, task, custom_prompt, reasoning, max_new_tokens, temperature, top_p],
+        outputs=[answer_out, thinking_out],
+    )
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Probe")
-    gr.Textbox(value=PROBE, label="env", lines=30)
+    gr.Examples(
+        examples=[
+            ["examples/mlk_speech.wav", "Transcribe (ASR)", "", False, 256, 0.7, 0.9],
+            ["examples/sample_speech.wav", "Transcribe (ASR)", "", False, 256, 0.7, 0.9],
+            ["examples/mlk_speech.wav", "Describe the audio", "", False, 256, 0.7, 0.9],
+        ],
+        inputs=[audio_in, task, custom_prompt, reasoning, max_new_tokens, temperature, top_p],
+        outputs=[answer_out, thinking_out],
+        fn=_ui_run,
+        cache_examples=False,
+        run_on_click=True,
+    )
 
-demo.launch(ssr_mode=False)
+if __name__ == "__main__":
+    demo.queue(max_size=8).launch(ssr_mode=False, mcp_server=True)

audio_utils.py

@@ -0,0 +1,176 @@
+# coding=utf-8
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import json
+import math
+import os
+from pathlib import Path
+from typing import Iterable, Optional
+
+import numpy as np
+import torch
+
+
+SOUND_PLACEHOLDER = "<sound>"
+SOUND_TOKEN = "<so_embedding>"
+SOUND_START_TOKEN = "<so_start>"
+SOUND_END_TOKEN = "<so_end>"
+IM_END_TOKEN = "<|im_end|>"
+DEFAULT_SYSTEM_PROMPT = (
+    "<|im_start|>system\n"
+    "You are a helpful and harmless assistant.\n\n"
+    "You are not allowed to use any tools."
+    "<|im_end|>\n"
+)
+
+
+def strip_hf_prefix(path: str) -> str:
+    """Convert Megatron-style hf:// paths into local filesystem paths."""
+    return path[len("hf://") :] if path.startswith("hf://") else path
+
+
+def load_audio(audio_path: str, target_sr: int = 16000) -> tuple[np.ndarray, int]:
+    import librosa
+
+    audio_data, sr = librosa.load(audio_path, sr=target_sr, mono=True)
+    return normalize_audio(audio_data), sr
+
+
+def normalize_audio(audio: np.ndarray) -> np.ndarray:
+    """Return mono float32 audio in [-1, 1], matching the Megatron eval path."""
+    audio = np.asarray(audio)
+    if audio.ndim == 2:
+        if audio.shape[1] <= 2:
+            audio = audio.mean(axis=1)
+        elif audio.shape[0] <= 2:
+            audio = audio.mean(axis=0)
+        else:
+            raise ValueError(f"Unsupported audio shape: {audio.shape}")
+
+    if audio.dtype == np.int16:
+        audio = audio.astype(np.float32) / 32768.0
+    elif audio.dtype != np.float32:
+        audio = audio.astype(np.float32)
+
+    max_abs = float(np.abs(audio).max()) if audio.size else 0.0
+    if max_abs > 1.0:
+        audio = audio / max_abs
+    return audio.astype(np.float32, copy=False)
+
+
+def split_audio_into_clips(
+    audio: np.ndarray,
+    sample_rate: int = 16000,
+    clip_duration: float = 30.0,
+) -> list[np.ndarray]:
+    """Split audio into fixed 30s clips; keep a padded final clip for Whisper."""
+    audio = normalize_audio(audio)
+    clip_samples = int(round(sample_rate * clip_duration))
+    if clip_samples <= 0:
+        raise ValueError(f"Invalid clip_samples: {clip_samples}")
+    if audio.size == 0:
+        audio = np.zeros(1, dtype=np.float32)
+
+    num_clips = max(1, math.ceil(audio.shape[0] / clip_samples))
+    clips: list[np.ndarray] = []
+    for idx in range(num_clips):
+        start = idx * clip_samples
+        clip = audio[start : start + clip_samples]
+        if clip.shape[0] < clip_samples:
+            clip = np.pad(clip, (0, clip_samples - clip.shape[0]))
+        clips.append(clip.astype(np.float32, copy=False))
+    return clips
+
+
+def extract_whisper_features(
+    feature_extractor,
+    audio: np.ndarray,
+    sample_rate: int = 16000,
+    clip_duration: float = 30.0,
+) -> torch.Tensor:
+    """Return NV-Whisper input features shaped (num_clips, 128, 3000)."""
+    clips = split_audio_into_clips(audio, sample_rate=sample_rate, clip_duration=clip_duration)
+    features = feature_extractor(
+        clips,
+        sampling_rate=sample_rate,
+        return_tensors="pt",
+        padding="max_length",
+        return_attention_mask=False,
+    )
+    input_features = features.input_features
+    if input_features.ndim != 3:
+        raise ValueError(f"Expected 3D Whisper features, got {tuple(input_features.shape)}")
+    return input_features
+
+
+def parse_conversation(conversation: list[dict]) -> tuple[str, str]:
+    human_prompt = ""
+    gt_answer = ""
+    for turn in conversation:
+        if turn["from"] == "human":
+            human_prompt = turn["value"].replace("<sound>\n", "").replace("<sound>", "").strip()
+        elif turn["from"] == "gpt":
+            gt_answer = turn["value"]
+    return human_prompt, gt_answer
+
+
+def build_prompt_template(
+    prompt: str,
+    reasoning: bool = False,
+    prompt_repitition: str = "none",
+) -> str:
+    if prompt_repitition not in {"none", "repetition"}:
+        raise ValueError(f"Unknown prompt repetition mode: {prompt_repitition}")
+    if prompt_repitition == "repetition":
+        prompt = f"{prompt}\n{prompt}"
+
+    if reasoning:
+        return f"<|im_start|>user\n<sound>\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n"
+    return f"<|im_start|>user\n<sound>\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def expand_sound_placeholder(prompt: str, num_embeddings: int) -> str:
+    if prompt.count(SOUND_PLACEHOLDER) != 1:
+        raise ValueError(f"Expected exactly one {SOUND_PLACEHOLDER}, found {prompt.count(SOUND_PLACEHOLDER)}")
+    replacement = SOUND_START_TOKEN + (SOUND_TOKEN * num_embeddings) + SOUND_END_TOKEN
+    return prompt.replace(SOUND_PLACEHOLDER, replacement)
+
+
+def build_attention_mask(input_ids: torch.Tensor) -> torch.Tensor:
+    return torch.ones_like(input_ids, dtype=torch.long)
+
+
+def split_thinking(response: str) -> tuple[str, str]:
+    if "</think>" not in response:
+        return "", response.strip()
+    thinking = response.rsplit("</think>", 1)[0].strip() + "</think>"
+    prediction = response.rsplit("</think>", 1)[1].strip()
+    return thinking, prediction
+
+
+def save_results_jsonl(results: Iterable[dict], output_path: str) -> None:
+    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+    with open(output_path, "w", encoding="utf-8") as f:
+        for result in results:
+            f.write(json.dumps(result, ensure_ascii=False) + "\n")
+
+
+def resolve_audio_preprocessor_path(model_path: str, config) -> str:
+    path = getattr(config, "audio_preprocessor_path", None) or "audio_preprocessor"
+    candidate = Path(path)
+    if not candidate.is_absolute():
+        candidate = Path(model_path) / candidate
+    return str(candidate)

examples/mlk_speech.wav

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

examples/sample_speech.wav

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

requirements.txt

@@ -1 +1,12 @@
-transformers
+transformers>=4.53.3
+accelerate
+sentencepiece
+librosa
+soundfile
+numpy
+einops
+# build-time deps for compiling mamba-ssm + causal-conv1d from source (no-build-isolation)
+ninja
+packaging
+setuptools
+wheel