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.

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

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:

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:

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:

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:

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.