# Audex-2B audio understanding in vLLM 0.20 (Nemotron Nano Omni-style)

Production-style integration: **audio -> text** through vLLM's own multimodal
pipeline. A registered model wraps the vLLM-native NemotronDense (2B dense) LLM
and adds the NV-Whisper audio encoder + Audex projector; vLLM owns placeholder
merging and serving (offline `LLM.generate` and the OpenAI `audio_url` server).

This is the 2B-dense counterpart to `Audex-30B-A3B/.../audioqa_scripts`. The only
differences from the 30B package: the backbone is the plain attention
`NemotronDenseForCausalLM` (no Mamba/hybrid cache) and the LLM hidden size is
2048 (the Audex projector output matches `config.hidden_size`). Audio
preprocessing, 30s chunking, `<so_embedding>` expansion, caps, and the vLLM
processor logic are identical.

Architecture name (from `checkpoint_folder_full/config.json`):
`NemotronDenseAudexForConditionalGeneration` (`model_type=nemotron_dense_audex`).

## Layout
```
inference_scripts_vllm/audioqa_scripts/
  README.md
  pyproject.toml                       # installs the model as a vLLM plugin
  run_audioqa_vllm.py                  # offline LLM.generate runner
  serve_audioqa_vllm.sh                # OpenAI-compatible server
  client_audioqa.py                    # OpenAI audio_url client (one request)
  audex_2b_vllm/                       # import package (distinct from 30B audex_30b_a3b_vllm)
    modeling_audex_vllm.py   # NemotronDenseAudex model (SupportsMultiModal + SupportsPP)
    processing_audex_vllm.py # processor: 30s chunking, <so_embedding> expansion, caps
    audio_features.py        # waveform -> NV-Whisper features
    audio_encoder.py         # Audex projector + Qwen2AudioEncoder factory
    plugin.py / register.py  # registration in every TP worker
```
The import package is `audex_2b_vllm` (the 30B package is `audex_30b_a3b_vllm`), so
both audioqa plugins can be co-installed without one shadowing the other.

## Install (once)
The model must be registered in every tensor-parallel worker, which vLLM does via
the `vllm.general_plugins` entry points. Install the dense backbone plugin and
this folder as editable plugins:
```bash
pip install -e ../../nemotron_dense_vllm_plugin --no-deps --no-build-isolation
pip install -e . --no-deps --no-build-isolation
```

## Run offline
```bash
python run_audioqa_vllm.py \
  --model-path "$(cd ../.. && pwd)/checkpoint_folder_full" \
  --input-json ./inputs.json \
  --output-jsonl ./audioqa_outputs/results.jsonl \
  --tensor-parallel-size 1
```
`inputs.json`: `[{"id", "sound": "/abs/path.wav", "conversations":[{"from":"human","value":"<sound>\nDescribe this audio."}]}]`.

## Serve + query (OpenAI audio_url)
```bash
# Safe defaults: HOST=127.0.0.1, local audio restricted to the Audex-2B release root.
bash serve_audioqa_vllm.sh "$(cd ../.. && pwd)/checkpoint_folder_full" 8000
python client_audioqa.py --audio /path/to/audio.wav --prompt "Describe this audio."
```
To expose externally or widen file access (advanced):
```bash
HOST=0.0.0.0 ALLOWED_MEDIA_PATH=/data bash serve_audioqa_vllm.sh ... 8000
# ALLOWED_MEDIA_PATH= (empty) disables local-file audio entirely.
```

## Benchmark recipes
Both `run_audioqa_vllm.py` and `client_audioqa.py` take `--recipe`, which sets
reproduction-safe sampling. No audio benchmark uses thinking mode; **default is
`audio-understanding`** (non-thinking).

| recipe | thinking | temperature | top_p | top_k |
| --- | --- | --- | --- | --- |
| `audio-understanding` (default) | off | 0.7 | 0.9 | 0 |
| `speech-recognition-translation` | off | 0.0 | 1.0 | 0 |
| `custom` | on | 0.7 | 0.9 | 0 |

- `top_k=0` means "disabled" in vLLM (consider all tokens).
- **Audio understanding is non-thinking**: `audio-understanding`
  (`temperature=0.7, top_p=0.9, top_k=0`) covers audio understanding/reasoning.
- **Greedy translation**: `speech-recognition-translation` is true greedy for
  ASR/AST — `temperature=0.0` triggers vLLM's greedy path (it normalizes
  `top_p=1.0, top_k=0`).
- **Precedence**: recipe defaults < explicit CLI override (`--reasoning/--no-reasoning`,
  `--temperature`, `--top-p`, `--top-k`). `custom` is the manual escape hatch
  (thinking on by default; override as needed).

## Key facts / gotchas
- **Backbone**: 2B dense (`NemotronDenseForCausalLM`) — plain RMSNorm / relu^2 /
  GQA, no Mamba; the wrapper is `SupportsMultiModal + SupportsPP` (not hybrid).
- **Hidden size**: 2048. The Audex projector's `fc2` maps the 1280-d encoder
  features to `config.hidden_size` (2048), not 2688.
- **Long audio**: non-overlapping 30s windows, padded tail, `N = num_clips*750`
  placeholders. Caps (fail loud): `MAX_AUDIO_SECONDS=900`, `MAX_AUDIO_CLIPS=30`,
  `MAX_AUDIO_TOKENS=22500` (in `processing_audex_vllm.py`). The offline runner and
  serve script default to `--max-model-len 32768` so the full 22500-embedding cap
  fits in the context window.
- **Placeholder contract**: `<so_embedding>` -> `<so_start>` + N*`<so_embedding>`
  + `<so_end>`; a placeholder/token count mismatch fails loud.
- **No audio-token leakage** (offline *and* served): generation is masked to text
  ids — `allowed_token_ids = range(131072)` minus the sound placeholder ids
  (`<so_embedding>`/`<so_start>`/`<so_end>` = 29/30/31). The offline runner passes
  this to `SamplingParams`; the client passes it via `extra_body`. All audio
  codec/gen tokens are id >= 131072. Both paths also scan the output text for
  `<audiocodec_`/`<speechcodec_`/`<audiogen_`/`<speechgen_`/`<so_*>` leakage.
- **Self-contained preflight**: the offline runner and serve script fail early
  with a clear message if `model.safetensors.index.json` references shards that
  are missing/unresolvable.
- **Reasoning / prompt format**: the `audio-understanding` recipe is
  non-thinking; the offline runner and client default to it. These scripts use
  the Audex audio-understanding evaluation prompt format — the non-thinking
  generation prompt uses the `<think></think>` assistant prefix, consistently
  offline (`run_audioqa_vllm.py`) and served (`checkpoint_folder_full/chat_template.jinja`,
  `enable_thinking=False`). Use `--recipe custom` for a thinking-capable manual setup.

## Integration notes (why the model code looks the way it does)
1. Out-of-tree registration must reach TP workers -> done via the plugin entry
   points (registering only in the main process raises "unsupported arch" in workers).
2. vLLM streams *all* `.safetensors` in the model dir; `load_weights` splits the
   stream: `model.*`/`lm_head.*` -> dense language model (its AutoWeightsLoader
   fuses q/k/v into qkv_proj), `audio_encoder.*`/`audio_projector.*` -> audio.
3. The dense backbone is registered by the separate `nemotron-dense-vllm` plugin;
   `register_audex` also registers it defensively so the wrapped `language_model`
   arch always resolves.