Asr switch

docs/AUDIO.md

@@ -0,0 +1,124 @@
+# Audio Input (Alpha)
+
+Granite Switch can accept **audio input** through a single vLLM model load — no
+separate speech server, no change to how developers deploy or call the model.
+
+This is an **alpha**: a speech-to-text *cascade*. Audio is transcribed to text by
+a small ASR model and the transcript is fed to the LLM as ordinary tokens. It is
+intentionally simple and requires no training. The "proper" upgrade (feeding a
+trained projection of a speech encoder's embeddings straight into the LLM) reuses
+the same hooks — see [Design](#design) below.
+
+## Building an audio-enabled checkpoint
+
+Add `--enable-audio` when composing:
+
+```bash
+python -m granite_switch.composer.compose_granite_switch \
+  --base-model ibm-granite/granite-4.0-micro \
+  --built-in-adapters core \
+  --enable-audio \
+  --output ./granite-switch-audio
+```
+
+This adds the `<|audio|>` marker token to the tokenizer and writes the audio
+settings into `config.json` so the checkpoint is self-describing:
+
+```json
+{ "asr_enabled": true, "asr_model_id": null, "asr_device": "cpu" }
+```
+
+- `asr_model_id` — HF id of the speech-to-text model (default: a small built-in
+  `distil-whisper/distil-small.en`). Override with `--asr-model <hf-id>`, e.g.
+  `openai/whisper-small` for multilingual.
+- `asr_device` — `cpu` (default) keeps vLLM's GPU KV-cache budget clean; set
+  `--asr-device cuda:0` to run transcription on GPU (watch GPU memory).
+
+Audio capability is **gated per checkpoint** by `asr_enabled`: a checkpoint built
+without `--enable-audio` reports no audio modality and never loads the ASR model.
+
+## Calling it
+
+### Python (offline)
+
+```python
+from granite_switch.vllm import register; register()
+from vllm import LLM, SamplingParams
+import soundfile as sf
+
+llm = LLM(model="./granite-switch-audio")          # one model load
+audio, sr = sf.read("question.wav")                # numpy array + sample rate
+
+out = llm.generate({
+    "prompt": "Transcript of the audio: <|audio|>\nAnswer:",
+    "multi_modal_data": {"audio": [(audio, sr)]},
+}, SamplingParams(max_tokens=128))
+print(out[0].outputs[0].text)
+```
+
+The `<|audio|>` marker is where the transcript is spliced in.
+
+### OpenAI-compatible server / chat API
+
+```bash
+vllm serve ./granite-switch-audio --port 8000
+```
+```python
+from openai import OpenAI
+client = OpenAI(base_url="http://localhost:8000/v1", api_key="x")
+resp = client.chat.completions.create(
+    model="granite-switch-audio",
+    messages=[{"role": "user", "content": [
+        {"type": "text", "text": "Answer the question in the audio."},
+        {"type": "input_audio", "input_audio": {"data": "<base64-wav>", "format": "wav"}},
+    ]}],
+)
+print(resp.choices[0].message.content)
+```
+
+The chat template emits the `<|audio|>` marker for audio content parts
+(`audio` / `input_audio` / `audio_url`), so the processor splices the transcript
+in automatically — callers send standard chat messages, no manual marker needed.
+
+## Design
+
+Per request, before the scheduler allocates KV cache:
+
+1. vLLM's multimodal pipeline hands the audio to our processor
+   (`granite_switch.vllm.audio`).
+2. The processor runs ASR → transcript → token ids.
+3. A `PromptReplacement` swaps the `<|audio|>` marker for those transcript token
+   ids. The scheduler then sizes KV for the **real** length — the audio "window"
+   is variable and decided at runtime, not reserved in advance.
+4. The model's `embed_multimodal` supplies embeddings for those positions. In the
+   alpha that is simply the transcript's own token embeddings (identical to
+   embedding them as text). **This is the seam the future encoder reuses:** swap
+   `embed_multimodal` to return `projection(speech_encoder(audio))` and the rest
+   of the machinery is unchanged.
+
+The decoder, switch, and LoRA paths are untouched — they only ever see text
+tokens.
+
+## Limitations (alpha)
+
+- **Cascade, not end-to-end.** Prosody/emotion/uncertainty are lost; ASR errors
+  propagate to the LLM. Two models run sequentially (ASR then LLM).
+- **English by default** (`distil-whisper/distil-small.en`). Use `--asr-model`
+  for multilingual.
+- One audio clip per request.
+
+## Audio + adapters
+
+Audio requests route through adapters exactly like text requests. The model sets
+`requires_raw_input_tokens = True` so vLLM passes the raw `input_ids` to the
+forward pass on the multimodal path; the switch then detects adapter control
+tokens as usual, and `embed_input_ids` applies the same token-exchange rewrite
+(control → substitute id) used for text — so an audio request that activates an
+adapter behaves identically to the text equivalent.
+
+## Tests
+
+- `tests/unit/test_asr.py` — CPU unit tests for the ASR backend (audio coercion,
+  resampling, transcription with a mocked pipeline). No GPU/vLLM required.
+- End-to-end (GPU): compose an `--enable-audio` checkpoint, then an audio request
+  through vLLM produces an answer and text-only requests are unaffected.

pyproject.toml

@@ -24,6 +24,10 @@ vllm = ["vllm>=0.19.1,<0.20.0"]
 vllm20 = ["vllm>=0.20.0,<0.21.0"]
 compose = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"]
 build = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"]  # Backward compatibility alias for compose
+# Audio (ASR) preprocessing — speech-to-text cascade in the vLLM backend.
+# transformers ships the ASR model itself (already a core dep); these add
+# audio decoding/resampling.
+audio = ["librosa", "soundfile"]
 tutorials = [
     "granite-switch[hf,vllm,compose]",
     "chromadb>=0.4.0",

src/granite_switch/composer/compose_granite_switch.py

@@ -60,7 +60,9 @@
     resolve_repo_path,
 )
 from granite_switch.composer.tokenizer_setup import (
+    add_audio_token,
     add_control_tokens,
+    configure_audio_chat_template,
     configure_chat_template,
     get_alora_first_invocation_token_id,
 )
@@ -570,6 +572,27 @@ def _compose_argparser():
         default=False,
         help="Include debug fields (original_path) in adapter_index.json",
     )
+    parser.add_argument(
+        "--enable-audio",
+        action="store_true",
+        default=False,
+        help="Enable the audio cascade: add the <|audio|> marker token and set "
+             "asr_enabled in the config so the vLLM backend transcribes audio.",
+    )
+    parser.add_argument(
+        "--asr-model",
+        type=str,
+        default=None,
+        help="HF id of the speech-to-text model the audio preprocessor loads. "
+             "Implies --enable-audio. Defaults to a small built-in model when unset.",
+    )
+    parser.add_argument(
+        "--asr-device",
+        type=str,
+        default="cpu",
+        help="Device the ASR model runs on (default: cpu). Use e.g. cuda:0 to "
+             "run transcription on GPU (watch vLLM's KV-cache memory budget).",
+    )
     return parser
 
 
@@ -773,12 +796,21 @@ def build():
 
     adapter_token_ids, special_tokens = add_control_tokens(tokenizer, all_discovered)
 
+    # Audio cascade: add the <|audio|> marker token before the embedding resize.
+    # Enabled by --enable-audio or by naming an --asr-model.
+    audio_enabled = args.enable_audio or args.asr_model is not None
+    audio_token_id = add_audio_token(tokenizer) if audio_enabled else None
+
     # Configure chat template with adapter mappings (Granite models only).
     # Non-Granite models preserve the upstream template verbatim because
     # the injection targets Granite-specific Jinja patterns.
     normalized_type = getattr(base_config, "model_type", "").replace("_switch", "")
     if normalized_type.startswith("granite"):
         configure_chat_template(tokenizer, all_discovered)
+        if audio_enabled:
+            # Make the chat template emit <|audio|> for audio content parts so
+            # the OpenAI server / chat() path works (the ASR processor replaces it).
+            configure_audio_chat_template(tokenizer)
     else:
         print("  Skipping chat template configuration (non-Granite model)")
 
@@ -834,6 +866,18 @@ def build():
         **optional_kwargs,
     )
 
+    # Record audio-cascade settings in the config so the checkpoint is
+    # self-describing and the vLLM backend gates audio on asr_enabled.
+    if audio_enabled:
+        model.config.asr_enabled = True
+        model.config.asr_model_id = args.asr_model
+        model.config.asr_device = args.asr_device
+        print(
+            f"  Audio cascade enabled "
+            f"(asr_model_id={args.asr_model or 'default'}, "
+            f"asr_device={args.asr_device}, audio_token_id={audio_token_id})"
+        )
+
     # Base model size (best effort)
     base_model_size_gb, _ = _get_directory_size(base_model_local_path)
     if base_model_size_gb is not None:

src/granite_switch/composer/tokenizer_setup.py

@@ -95,6 +95,62 @@ def add_control_tokens(
     return adapter_token_ids, special_tokens
 
 
+def add_audio_token(tokenizer, marker: str = "<|audio|>") -> int:
+    """Add the audio placeholder marker token to the tokenizer.
+
+    Used for the audio cascade: this single special token is placed in the
+    prompt and the vLLM ASR processor replaces it with the transcript tokens at
+    request time (see granite_switch.vllm.audio). Registering it as one special
+    token keeps the processor's prompt-replacement match clean.
+
+    Must be called before the model's embedding resize so the new row is sized
+    in. Returns the marker's token id.
+    """
+    print(f"\nAdding audio marker token: {marker}")
+    tokenizer.add_special_tokens({"additional_special_tokens": [marker]})
+    token_id = tokenizer.convert_tokens_to_ids(marker)
+    print(f"  {marker}: {token_id}")
+    return token_id
+
+
+def configure_audio_chat_template(tokenizer, marker: str = "<|audio|>") -> None:
+    """Make the chat template emit the audio marker for audio content parts.
+
+    The Granite content-part loop only handles ``entry.type == 'text'`` and
+    silently drops other parts. vLLM passes multimodal chat content to the
+    template as a *list of parts*, so without this the ``<|audio|>`` marker never
+    reaches the rendered prompt and the ASR processor's prompt replacement fails
+    (``Failed to apply prompt replacement for mm_items['audio'][0]``).
+
+    We inject an ``elif`` that appends the marker for any part whose ``type``
+    contains ``'audio'`` (covers ``audio`` / ``input_audio`` / ``audio_url``).
+    Call after :func:`configure_chat_template`, gated on audio being enabled.
+    """
+    template = tokenizer.chat_template
+    if template is None:
+        print("Warning: no chat template; skipping audio chat-template handling")
+        return
+
+    # The text-only branch of the Granite content-part loop:
+    old = (
+        "                    {%- set content.val = content.val + entry.text %}\n"
+        "                {%- endif %}"
+    )
+    if old not in template:
+        raise ValueError(
+            "Could not find the Granite content-part loop to inject audio "
+            "handling; the base chat template may have changed."
+        )
+    new = (
+        "                    {%- set content.val = content.val + entry.text %}\n"
+        "                {%- elif 'audio' in entry.type %}\n"
+        "                    {%- set content.val = content.val + '" + marker + "' %}\n"
+        "                {%- endif %}"
+    )
+    tokenizer.chat_template = template.replace(old, new, 1)
+    print(f"  Audio chat-template handling added (emits {marker} for audio parts)")
+
+
 def configure_chat_template(
     tokenizer,
     discovered_adapters: List[Tuple[Optional[str], str, str, Optional[str]]],

src/granite_switch/config.py

@@ -38,6 +38,20 @@ class GraniteSwitchConfig(GraniteMoeHybridConfig):
         lora_target_modules (List[str]): List of module GROUP names to apply LoRA to.
             Module groups: "qkv_proj", "o_proj", "shared_input_linear", "shared_output_linear".
             Default: all four groups
+
+        Audio (ASR) preprocessing parameters:
+            asr_enabled (bool): When True, the vLLM backend registers an audio
+                multimodal preprocessor that transcribes audio inputs to text and
+                splices the transcript tokens into the prompt before the decoder
+                runs. The decoder itself is unchanged and only ever sees text
+                tokens. Default: False.
+            asr_model_id (Optional[str]): HuggingFace id of the speech-to-text
+                model the preprocessor loads. When None and asr_enabled is True,
+                the backend falls back to a small built-in default. This makes the
+                checkpoint self-describing about which ASR front-end it expects.
+            asr_device (str): Device the ASR model runs on. Default "cpu" keeps
+                vLLM's GPU KV-cache budget clean; set to a CUDA device to trade GPU
+                memory for transcription latency.
         **kwargs: Additional arguments passed to GraniteConfig.
     """
 
@@ -56,6 +70,10 @@ def __init__(
         max_lora_rank: int = 8,
         adapter_ranks: List[int] = None,
         lora_target_modules: Optional[List[str]] = None,
+        # Audio (ASR) preprocessing parameters
+        asr_enabled: bool = False,
+        asr_model_id: Optional[str] = None,
+        asr_device: str = "cpu",
         # vLLM residual-norm convention (for bit-exact skinning equivalence)
         fused_add_norm: bool = False,
         # Parent class defaults (Granite 4 dense configuration)
@@ -136,6 +154,15 @@ def __init__(
         self.switch_head_dim = switch_head_dim
         self.fused_add_norm = fused_add_norm
 
+        # Audio (ASR) preprocessing parameters.
+        # The decoder is oblivious to audio: when enabled, the vLLM backend's
+        # multimodal preprocessor transcribes audio to text and injects the
+        # transcript tokens into the prompt before embedding. These fields make
+        # the checkpoint self-describing about its ASR front-end.
+        self.asr_enabled = asr_enabled
+        self.asr_model_id = asr_model_id
+        self.asr_device = asr_device
+
         # Adapter names
         self.adapter_names = adapter_names
 

src/granite_switch/vllm/audio/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Audio (ASR) preprocessing for the Granite Switch vLLM backend.
+
+This package implements the **alpha** audio pathway: a speech-to-text *cascade*.
+Audio is transcribed to text by a small ASR model and the transcript tokens are
+spliced into the prompt before the decoder runs. The Granite Switch decoder is
+unchanged and only ever sees text tokens — there is no in-model audio encoder yet
+(that is deferred future work: a trained projection from Granite Speech encoder
+embeddings into Granite token space).
+
+The ASR backend (:mod:`asr`) deliberately has no vLLM dependency so it can be
+unit-tested on CPU and reused by either the multimodal-processor integration or a
+fallback entrypoint wrapper.
+"""
+
+from .asr import ASRTranscriber, DEFAULT_ASR_MODEL_ID, transcribe
+
+__all__ = [
+    "ASRTranscriber",
+    "DEFAULT_ASR_MODEL_ID",
+    "transcribe",
+]