Jais-ASR / vllm_plugin /README.md
EvanD's picture amine-khelif's picture guokan-shang's picture
Squash history
1323995
|
Raw History Blame Contribute Delete
5.57 kB
# Jais-ASR vLLM plugin
The `jais-asr-vllm` package registers the Jais-ASR model, audio processing, the
`jais_asr_tag` reasoning parser, and language controls. Version 2.0.0 changes
`language` to logits pinning and adds `language_prefill` for prompt prefill.
No legacy language option or parser/package aliases are provided.
## Install and serve
Use the same Python environment for vLLM and the plugin. The tested runtime is
vLLM 0.23.0 with Transformers 5.5.4 and PyTorch 2.11.0+cu130. Remove the previous
plugin installation before installing this renamed distribution, then restart
vLLM so the new package entry points are loaded.
```bash
hf download IFM/Jais-ASR --include 'vllm_plugin/*' --local-dir ./jais-asr
pip install ./jais-asr/vllm_plugin
vllm serve IFM/Jais-ASR \
--served-model-name Jais-ASR \
--trust-remote-code --reasoning-parser jais_asr_tag \
--dtype bfloat16 --max-model-len 8192 \
--max-num-batched-tokens 2048 --max-num-seqs 4 \
--limit-mm-per-prompt '{"audio":2}' --enforce-eager
```
The Hub repository ID is `IFM/Jais-ASR`. The public serving ID is
**`Jais-ASR`**: use it in API requests and expect it in `/v1/models` and chat responses.
Authenticate with Hugging Face before downloading/loading if the repository is private.
## Language: pin by default
```bash
curl http://localhost:8000/v1/audio/transcriptions \
-F model=Jais-ASR -F file=@speech.wav -F language=ar
```
`language` constrains the generated tag while retaining the pin's strict
speech-versus-abstention score comparison. `ar` allows all supported Arabic tags;
`en`, `fr`, `EG`, `msa` and other concrete tags pin a single choice. This is a tag
constraint, not a translation operation or a guarantee about every transcript word.
For explicit prompt prefill:
```bash
curl http://localhost:8000/v1/audio/transcriptions \
-F model=Jais-ASR -F file=@speech.wav -F language_prefill=en
```
`language_prefill` must resolve to one concrete tag. It bypasses the normal tag
selection step. Both controls together, unknown languages, a family used as prefill,
and legacy option names are rejected. Omitting both retains auto-detection.
Chat requests use the same top-level JSON fields:
```python
client.chat.completions.create(
model="Jais-ASR", messages=messages,
extra_body={"language": "ar"},
)
```
Use `extra_body={"language_prefill": "en"}` for prefill. Existing chat-template
kwargs can carry the same names; the API moves `language` into sampling controls
before rendering. Thinking and pinning can be combined; thinking and prefill cannot.
Native vLLM beam search does not support this logits pin and is rejected when a pin
is requested. Ordinary sampling/greedy decoding supports it.
Offline vLLM uses the same key in sampling options:
```python
params = SamplingParams(temperature=0, extra_args={"language": "ar"})
outputs = llm.generate(prompts, params)
```
For offline prefill, render the prompt with `language_prefill` first. Do not put
`language_prefill` in `SamplingParams.extra_args`, because sampling cannot rewrite
a prompt. The plugin does not contain weights; load the Jais-ASR checkpoint separately.
## Integration and validation
The API bridge is scoped to configurations with `model_type="jais_asr"`. It exposes
language fields in the native request schemas, forwards pins into the per-request
sampling parameters, and forwards prefill tags into prompt construction. It does not
change ordinary language handling for other model types.
The model adapter preserves the checkpoint's weight paths and K2 grouped RMSNorm.
Transcription responses expose cleaned text; chat responses use `message.content`
for text and `message.reasoning` for decision metadata. Streaming uses delta fields.
The encoder produces about 13 tokens/second. With `max_num_batched_tokens=2048`, a single
request can carry up to 157 seconds of audio (the encoder budget). The transcription
endpoint, however, splits longer audio into 30-second windows by default, because the
model drops text on long segments: on Earnings22-Cleaned-AA, 157 s windows gave 44% WER
against 13% with 30 s windows. Set `JAIS_ASR_MAX_CLIP_S` before starting vLLM to change
the window length (it is capped by the encoder budget). Chat requests are not chunked, so
split long recordings yourself. Other vLLM versions and compiled/parallel execution
must be validated separately.
## Batch reproducibility
The earlier comparison on 12 fixtures found 5/72 default-mode singleton-versus-batch
transcription changes and 0/72 with `VLLM_BATCH_INVARIANT=1`. Set the variable before
starting vLLM. The tested settings used eager execution and prefix caching disabled:
```bash
VLLM_BATCH_INVARIANT=1 vllm serve IFM/Jais-ASR \
--served-model-name Jais-ASR --trust-remote-code --reasoning-parser jais_asr_tag \
--dtype bfloat16 --max-model-len 8192 --max-num-batched-tokens 2048 \
--max-num-seqs 8 --no-enable-prefix-caching --enforce-eager
```
Batch invariance may reduce throughput and can change the baseline result relative
to default kernels. Previous warmed fixed-token tests measured roughly 30–38% lower
throughput on one RTX 6000 Ada GPU; this is workload-specific. See the model README
for the full reproducibility and speed limitations.
Version 2.0.0 validation included four tests using native vLLM request/schema types
and 26 live API cases. These checked the rendered prompt and sampler controls,
public model identity, pin/prefill semantics, invalid/removed parameters, concurrent
pins, both conversations in a batch, and prefill metadata in streamed responses.