diff --git a/docs/AUDIO.md b/docs/AUDIO.md index de64363..c05166d 100644 --- a/docs/AUDIO.md +++ b/docs/AUDIO.md @@ -33,6 +33,13 @@ settings into `config.json` so the checkpoint is self-describing: `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). +- `asr_dtype` — precision the ASR weights load in. Unset (default) derives it + from the device: `float16` on CUDA, `float32` on CPU. Half precision halves + the ASR weight footprint and is what the Whisper-family defaults expect, but + it is not universally safe — an encoder with **BatchNorm** layers raises + `Expected weight to have type Float but got Half`, since BatchNorm will not + promote a float16 weight against float32 features. Such a checkpoint needs + `--asr-dtype float32`. Accepted: `auto`, `float16`, `bfloat16`, `float32`. 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. diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index 3ec4587..426f5bb 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -67,6 +67,7 @@ configure_chat_template, get_alora_first_invocation_token_id, ) +from granite_switch.config import ASR_DTYPES # --------------------------------------------------------------------------- # Utility helpers (kept local — not worth a separate module) @@ -589,6 +590,16 @@ def _compose_argparser(): 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).", ) + parser.add_argument( + "--asr-dtype", + type=str, + default=None, + choices=ASR_DTYPES, + help="Precision the ASR weights load in. Default derives it from " + "--asr-device (float16 on CUDA, float32 on CPU); set float32 for an " + "encoder that cannot run in half precision (e.g. one with BatchNorm). " + "Implies --enable-audio.", + ) parser.add_argument( "--asr-pipeline-kwargs", type=json.loads, @@ -852,6 +863,7 @@ def build(): audio_enabled = ( args.enable_audio or args.asr_model is not None + or args.asr_dtype is not None or args.asr_pipeline_kwargs is not None or args.asr_generate_kwargs is not None or args.asr_max_audio_clips is not None @@ -934,6 +946,7 @@ def build(): model.config.asr_enabled = True model.config.asr_model_id = args.asr_model model.config.asr_device = args.asr_device + model.config.asr_dtype = args.asr_dtype # Optional pipeline-construction extras and default decode kwargs. Only # set when provided so the config stays minimal for the common case. if args.asr_pipeline_kwargs is not None: @@ -953,7 +966,9 @@ def build(): 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}, " + f"asr_device={args.asr_device}, " + f"asr_dtype={args.asr_dtype or 'auto'}, " + f"audio_token_id={audio_token_id}, " f"pipeline_kwargs={args.asr_pipeline_kwargs or {}}, " f"generate_kwargs={args.asr_generate_kwargs or {}}, " f"max_audio_clips={args.asr_max_audio_clips or 'default'}, " diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index 38c20cd..cc788d8 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -3,6 +3,9 @@ from transformers import GraniteMoeHybridConfig +# Accepted asr_dtype values. Keep in sync with vllm.audio.asr._ASR_DTYPE_NAMES. +ASR_DTYPES = ("auto", "float16", "bfloat16", "float32") + class GraniteSwitchConfig(GraniteMoeHybridConfig): """Configuration class for GraniteSwitch model. @@ -37,56 +40,38 @@ class GraniteSwitchConfig(GraniteMoeHybridConfig): 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. + Audio (ASR) preprocessing parameters (see docs/AUDIO.md): + asr_enabled (bool): Register the audio preprocessor that transcribes + audio and splices the transcript into the prompt. Default: False. + asr_model_id (Optional[str]): HF id of the speech-to-text model. None + falls back to a small built-in default. 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. - asr_pipeline_kwargs (Optional[dict]): Extra keyword arguments merged - into the ``transformers.pipeline(...)`` construction call for the - ASR model (e.g. ``{"chunk_length_s": 15}``). These affect how the - pipeline is built, so they are baked into the transcriber cache - key. None means "no extras". Default: None. - asr_generate_kwargs (Optional[dict]): Default decode-time keyword - arguments passed to the ASR model on every transcription (e.g. - ``{"language": "de", "task": "transcribe"}`` for a multilingual - Whisper). Applied at call time, so a single loaded pipeline can be - reused; per-request values (via ``mm_processor_kwargs``) override - these. Ignored by models that do not generate (e.g. CTC). Default: - None. - asr_max_audio_clips (int): Maximum number of audio clips accepted in a - single request. Each clip's transcript is spliced at its own - ``<|audio|>`` marker. vLLM enforces this as the modality ceiling - (``--limit-mm-per-prompt`` may lower it, not raise it). For the - cascade the clips cost no extra KV (transcripts are ordinary text - tokens bounded by the context window); the ceiling mainly guards - against one request triggering an unbounded number of synchronous - ASR transcriptions, and bounds the startup profiling pass. Default: - 32. - asr_chunk_length_s (float): Window length (seconds) our own long-audio - chunker splits a clip into before transcribing each window. Only - used for backends that do not self-chunk (see asr_self_chunks); the - default Whisper backend chunks internally and ignores this. Default: - 30.0. - asr_chunk_overlap_s (float): Overlap (seconds) between consecutive - chunker windows, so words straddling a boundary are whole in at - least one window; the transcript merge de-duplicates the overlap. - Only used when asr_self_chunks is False. Default: 5.0. - asr_self_chunks (bool): Whether the ASR backend handles long audio - itself. True for the Whisper pipeline (its internal - ``chunk_length_s`` does timestamp-based stitching, higher quality - than our text-level merge), so our chunker is bypassed. Set False - for a backend with a fixed input window (e.g. a future speech - encoder) to route audio through our encoder-agnostic - split/transcribe/merge chunker. Default: True. + vLLM's GPU KV-cache budget clean. + asr_dtype (Optional[str]): Precision the ASR weights load in, one of + ASR_DTYPES. None/"auto" derives it from asr_device (float16 on + CUDA). An encoder with BatchNorm layers must set "float32". + Default: None. + asr_pipeline_kwargs (Optional[dict]): Extra kwargs merged into the + ``transformers.pipeline(...)`` construction, e.g. + ``{"chunk_length_s": 15}``. Baked into the transcriber cache key. + Default: None. + asr_generate_kwargs (Optional[dict]): Default decode-time kwargs, e.g. + ``{"language": "de"}``. Applied per call, so one pipeline is + reused; per-request ``mm_processor_kwargs`` override them. Ignored + by non-generative backends. Default: None. + asr_max_audio_clips (int): Max audio clips per request. Bounds the + synchronous transcriptions one request can trigger and the startup + profiling pass; ``--limit-mm-per-prompt`` may lower it, not raise + it. Default: 32. + asr_chunk_length_s (float): Chunker window length in seconds. Only + used when asr_self_chunks is False. Default: 30.0. + asr_chunk_overlap_s (float): Overlap in seconds between chunker + windows, de-duplicated by the transcript merge. Only used when + asr_self_chunks is False. Default: 5.0. + asr_self_chunks (bool): True when the backend chunks long audio + itself (Whisper's timestamp stitching beats our text-level merge), + bypassing our chunker. False routes audio through the + split/transcribe/merge chunker instead. Default: True. **kwargs: Additional arguments passed to GraniteConfig. """ @@ -109,6 +94,7 @@ def __init__( asr_enabled: bool = False, asr_model_id: str | None = None, asr_device: str = "cpu", + asr_dtype: str | None = None, asr_pipeline_kwargs: dict | None = None, asr_generate_kwargs: dict | None = None, asr_max_audio_clips: int = 32, @@ -195,22 +181,19 @@ 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. + # Audio (ASR) preprocessing. The decoder is oblivious to audio; 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 - # Pipeline-construction extras (affect the built pipeline → cache key) - # and default decode kwargs (applied per call, per-request overridable). + # Validated here so a typo fails at compose time, not in a vLLM worker. + if asr_dtype is not None and asr_dtype not in ASR_DTYPES: + raise ValueError( + f"asr_dtype must be one of {ASR_DTYPES} or None, got {asr_dtype!r}" + ) + self.asr_dtype = asr_dtype self.asr_pipeline_kwargs = asr_pipeline_kwargs self.asr_generate_kwargs = asr_generate_kwargs - # Long-audio / multi-clip preprocessing. The transcript is spliced into - # the prompt as ordinary text tokens; a clip that makes the prompt exceed - # the context is rejected by vLLM's standard length check (no truncation). - # The chunker settings only apply to backends that do not self-chunk. if asr_max_audio_clips < 1: raise ValueError( f"asr_max_audio_clips must be >= 1, got {asr_max_audio_clips}" diff --git a/src/granite_switch/vllm/audio/__init__.py b/src/granite_switch/vllm/audio/__init__.py index 486bc25..feab7ce 100644 --- a/src/granite_switch/vllm/audio/__init__.py +++ b/src/granite_switch/vllm/audio/__init__.py @@ -1,16 +1,8 @@ # 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. +Speech-to-text cascade: audio is transcribed and the transcript tokens are +spliced into the prompt, so the decoder only ever sees text. See docs/AUDIO.md. """ from .asr import DEFAULT_ASR_MODEL_ID, ASRTranscriber, transcribe diff --git a/src/granite_switch/vllm/audio/asr.py b/src/granite_switch/vllm/audio/asr.py index d58ea8d..b95085a 100644 --- a/src/granite_switch/vllm/audio/asr.py +++ b/src/granite_switch/vllm/audio/asr.py @@ -1,23 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""Speech-to-text backend for the alpha audio cascade. - -A thin wrapper around a HuggingFace ASR model that turns raw audio into a -transcript string. Kept free of any vLLM import so it can be unit-tested on CPU -and reused by both the multimodal-processor integration and a fallback wrapper. - -Design notes: - * The model is loaded **lazily** on first use and cached per (model_id, - device) so engine startup stays cheap and a process loads each ASR model - at most once. - * Default device is CPU — this keeps vLLM's GPU KV-cache budget clean. Set a - CUDA device to trade GPU memory for transcription latency. - * The default model is a small, text-emitting open-source ASR model. The - alpha intentionally uses a complete audio->text model (encoder+decoder); - swapping in the Granite Speech 4.1 encoder is future work and only touches - this module. - * ``transcribe`` accepts the audio item shapes vLLM hands to multimodal - processors: a bare array, an ``(array, sampling_rate)`` tuple, a Python - list of floats, or a torch tensor. +"""Speech-to-text backend for the audio cascade. + +Wraps a HuggingFace ASR pipeline. Free of any vLLM import so it unit-tests on +CPU. The model loads lazily and is cached per (model_id, device, dtype, +pipeline_kwargs), so a process loads each ASR model at most once. + +Device defaults to CPU to keep vLLM's GPU KV-cache budget clean; dtype follows +the device unless a checkpoint sets ``asr_dtype``. See docs/AUDIO.md. """ from __future__ import annotations @@ -32,6 +21,43 @@ # checkpoint does not name its own (config.asr_model_id is None). DEFAULT_ASR_MODEL_ID = "distil-whisper/distil-small.en" +ASR_DTYPE_AUTO = "auto" + +# Keep in sync with config.ASR_DTYPES. +_ASR_DTYPE_NAMES = frozenset({"float16", "bfloat16", "float32"}) +_ASR_DTYPE_ALIASES = { + "fp16": "float16", + "half": "float16", + "bf16": "bfloat16", + "fp32": "float32", + "float": "float32", +} + + +def _resolve_torch_dtype(dtype: str | None, device: str) -> Any: + """Resolve an ``asr_dtype`` name to a ``torch.dtype``. + + None/"auto" derives it from the device: float16 on CUDA, float32 elsewhere + (CPU float16 is slow and partly unimplemented). Name a dtype explicitly for + an encoder that cannot run in half precision — BatchNorm raises on a float16 + weight against float32 features rather than promoting. + """ + import torch + + name = str(dtype or ASR_DTYPE_AUTO).lower() + name = _ASR_DTYPE_ALIASES.get(name, name) + if name == ASR_DTYPE_AUTO: + on_cuda = isinstance(device, str) and device.startswith("cuda") + return torch.float16 if on_cuda else torch.float32 + if name not in _ASR_DTYPE_NAMES: + raise ValueError( + f"Unsupported asr_dtype {dtype!r}. Expected {ASR_DTYPE_AUTO!r} (or " + f"None) to derive it from the device, or one of: " + f"{', '.join(sorted(_ASR_DTYPE_NAMES))}." + ) + return getattr(torch, name) + + _CHUNKING = None @@ -73,23 +99,18 @@ def _load_chunking(): class ASRTranscriber: - """Lazily-loaded ASR model wrapper exposing :meth:`transcribe`. - - Instances are cheap to construct; the underlying model is only materialized - on the first :meth:`transcribe` call (or an explicit :meth:`load`). - """ + """Lazily-loaded ASR model wrapper exposing :meth:`transcribe`.""" def __init__( self, model_id: str = DEFAULT_ASR_MODEL_ID, device: str = "cpu", pipeline_kwargs: Mapping[str, Any] | None = None, + dtype: str | None = None, ) -> None: self.model_id = model_id self.device = device - # Extra kwargs merged into the pipeline() construction. Because they - # change the built pipeline, get_transcriber folds them into the cache - # key so distinct construction options get distinct cached instances. + self.dtype = dtype self.pipeline_kwargs: dict[str, Any] = dict(pipeline_kwargs or {}) self._pipeline = None self._load_lock = threading.Lock() @@ -101,28 +122,17 @@ def load(self) -> None: with self._load_lock: if self._pipeline is not None: return - # Imported lazily so this module stays importable without - # transformers' heavy audio stack until it is actually used. - import torch + # Lazy: keeps this module importable without transformers' audio stack. from transformers import pipeline - torch_dtype = ( - torch.float16 - if isinstance(self.device, str) and self.device.startswith("cuda") - else torch.float32 - ) - # Built-in defaults, then let checkpoint-supplied pipeline_kwargs - # override any of them (e.g. a different chunk_length_s, or model - # kwargs a non-Whisper backend needs). kwargs: dict[str, Any] = { "task": "automatic-speech-recognition", "model": self.model_id, "device": self.device, - "torch_dtype": torch_dtype, - # Enable internal 30s chunking so audio longer than the model's - # native window is handled without us re-implementing it. + "torch_dtype": _resolve_torch_dtype(self.dtype, self.device), "chunk_length_s": 30, } + # pipeline_kwargs last: a checkpoint may override any default above. kwargs.update(self.pipeline_kwargs) self._pipeline = pipeline(**kwargs) @@ -135,32 +145,12 @@ def transcribe( chunk_length_s: float = 30.0, chunk_overlap_s: float = 5.0, ) -> str: - """Transcribe one audio clip to a text string. - - Args: - audio: The audio samples. Accepts a numpy array, a list of floats, a - torch tensor, or an ``(array, sampling_rate)`` tuple (vLLM's - ``AudioItem`` shape). - sampling_rate: Sample rate of ``audio`` in Hz. Required unless - ``audio`` is an ``(array, sampling_rate)`` tuple. Audio is - resampled to 16 kHz when needed. - generate_kwargs: Decode-time kwargs forwarded to the ASR model for - this call (e.g. ``{"language": "fr"}``). Applied per call so the - same loaded pipeline serves many languages. Passed only when - non-empty, so CTC/non-generative backends are unaffected. - self_chunks: True when the backend handles long audio itself (the - Whisper pipeline, via its internal ``chunk_length_s``); the whole - clip is passed in one call. False routes long audio through our - encoder-agnostic chunker (:mod:`.chunking`): split into - overlapping windows, transcribe each, and merge with overlap - de-duplication. - chunk_length_s: Window length for our chunker (only used when - ``self_chunks`` is False). - chunk_overlap_s: Window overlap for our chunker (only used when - ``self_chunks`` is False). - - Returns: - The transcript with surrounding whitespace stripped. + """Transcribe one audio clip, stripped. Resampled to 16 kHz as needed. + + ``sampling_rate`` is required unless ``audio`` is an ``(array, rate)`` + tuple. ``generate_kwargs`` is passed only when non-empty, so CTC backends + are unaffected. ``self_chunks=False`` routes long audio through + :mod:`.chunking` using ``chunk_length_s``/``chunk_overlap_s``. """ samples, sr = _coerce_audio(audio, sampling_rate) samples = _to_mono_float32(samples) @@ -168,12 +158,9 @@ def transcribe( self.load() - # Backend chunks internally (Whisper): one call over the whole clip. if self_chunks: return self._run_pipeline(samples, generate_kwargs) - # Encoder-agnostic path: split into overlapping windows, transcribe each, - # then stitch the per-window transcripts (de-duplicating the overlap). chunking = _load_chunking() segments = chunking.split_waveform( samples, _TARGET_SAMPLE_RATE, chunk_length_s, chunk_overlap_s @@ -187,7 +174,6 @@ def _run_pipeline( generate_kwargs: Mapping[str, Any] | None = None, ) -> str: """Run the loaded pipeline over an already-resampled mono waveform.""" - # Tell the pipeline the rate so it does not attempt its own resampling. call_kwargs: dict[str, Any] = {} if generate_kwargs: call_kwargs["generate_kwargs"] = dict(generate_kwargs) @@ -201,17 +187,14 @@ def _run_pipeline( # ── Module-level cache + convenience function ──────────────────────────────── -# Keyed on (model_id, device, frozen pipeline_kwargs). pipeline_kwargs are in -# the key because they change the constructed pipeline; generate_kwargs are NOT -# — they are applied per transcribe() call, so one cached pipeline serves them -# all (that is what makes per-request language selection cheap). +# Keyed on what changes the constructed pipeline. generate_kwargs are excluded: +# they apply per transcribe() call, so one cached pipeline serves every language. _TRANSCRIBERS: dict[tuple, ASRTranscriber] = {} _CACHE_LOCK = threading.Lock() -# Decode kwargs a client may set per request. Kept to language/task so a client -# cannot inject arbitrary (potentially expensive or unsafe) generation options; -# everything else is fixed by the checkpoint author in config.asr_generate_kwargs. +# Allowlisted so a client cannot inject arbitrary generation options; everything +# else is fixed by the checkpoint in config.asr_generate_kwargs. DEFAULT_ALLOWED_REQUEST_GENERATE_KEYS = frozenset({"language", "task"}) @@ -222,12 +205,9 @@ def resolve_generate_kwargs( ) -> dict[str, Any]: """Merge config-default decode kwargs with allowlisted per-request overrides. - ``config_defaults`` come from ``config.asr_generate_kwargs``. ``request`` is - the per-request mapping (vLLM's ``mm_processor_kwargs``); a top-level - ``language`` or a nested ``asr_generate_kwargs`` object may override the - defaults, but only keys in ``allowed_keys`` are honored. Request values win - so one deployed model can serve many languages. Pure and vLLM-free so it can - be unit-tested on CPU. + ``request`` is vLLM's ``mm_processor_kwargs``: a top-level ``language`` or a + nested ``asr_generate_kwargs`` may override the defaults, but only for keys + in ``allowed_keys``. Request values win, so one model serves many languages. """ merged: dict[str, Any] = dict(config_defaults or {}) if isinstance(request, Mapping): @@ -253,14 +233,15 @@ def get_transcriber( model_id: str | None = None, device: str = "cpu", pipeline_kwargs: Mapping[str, Any] | None = None, + dtype: str | None = None, ) -> ASRTranscriber: """Return a process-wide cached :class:`ASRTranscriber`. - Cached per ``(model_id, device, pipeline_kwargs)``. ``model_id`` of None - resolves to :data:`DEFAULT_ASR_MODEL_ID`. + Cached per ``(model_id, device, dtype, pipeline_kwargs)``. ``model_id`` of + None resolves to :data:`DEFAULT_ASR_MODEL_ID`. """ resolved = model_id or DEFAULT_ASR_MODEL_ID - key = (resolved, device, _freeze(pipeline_kwargs or {})) + key = (resolved, device, dtype, _freeze(pipeline_kwargs or {})) transcriber = _TRANSCRIBERS.get(key) if transcriber is None: with _CACHE_LOCK: @@ -270,6 +251,7 @@ def get_transcriber( model_id=resolved, device=device, pipeline_kwargs=pipeline_kwargs, + dtype=dtype, ) _TRANSCRIBERS[key] = transcriber return transcriber @@ -282,6 +264,7 @@ def transcribe( model_id: str | None = None, device: str = "cpu", pipeline_kwargs: Mapping[str, Any] | None = None, + dtype: str | None = None, generate_kwargs: Mapping[str, Any] | None = None, self_chunks: bool = True, chunk_length_s: float = 30.0, @@ -289,7 +272,10 @@ def transcribe( ) -> str: """Convenience wrapper: transcribe with the cached transcriber for the args.""" return get_transcriber( - model_id=model_id, device=device, pipeline_kwargs=pipeline_kwargs + model_id=model_id, + device=device, + pipeline_kwargs=pipeline_kwargs, + dtype=dtype, ).transcribe( audio, sampling_rate, diff --git a/src/granite_switch/vllm/audio/chunking.py b/src/granite_switch/vllm/audio/chunking.py index bcc8cb5..2dfa7bf 100644 --- a/src/granite_switch/vllm/audio/chunking.py +++ b/src/granite_switch/vllm/audio/chunking.py @@ -1,21 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 """Encoder-agnostic long-audio chunking for the cascade. -A backend with a fixed input window (e.g. a speech encoder that only accepts a -few seconds at a time) cannot ingest a long clip in one shot. This module splits -a long waveform into overlapping fixed-length windows and stitches the resulting -per-window transcripts back into one string — the same job Whisper's internal -``chunk_length_s`` does, but lifted **above** the backend so *any* transcriber -inherits long-audio support. - -Kept pure (numpy + stdlib, no vLLM / torch / transformers import) so it unit- -tests on CPU. The transcriber decides whether to use it via its ``self_chunks`` -flag: backends that already chunk internally (Whisper) bypass this entirely. - -Why overlap + dedup: cutting on hard boundaries can split a word across two -windows, so consecutive windows overlap and each boundary word is whole in at -least one of them. The overlap region is therefore transcribed twice; the merge -finds the repeated span at each seam and keeps one copy. +Splits a waveform into overlapping windows and stitches the per-window +transcripts, so a backend with a fixed input window inherits long-audio support. +Windows overlap because a hard cut can split a word; the overlap is therefore +transcribed twice and the merge keeps one copy. + +Pure numpy + stdlib (no torch/transformers) so it unit-tests on CPU. """ from __future__ import annotations @@ -24,9 +15,7 @@ import numpy as np -# Cap on how many trailing/leading words we search for a seam overlap. Comfortably -# larger than any plausible word count inside a few seconds of overlap, but bounded -# so the merge stays linear in transcript length. +# Bounded seam search so the merge stays linear in transcript length. _MAX_SEAM_WORDS = 60 @@ -36,21 +25,10 @@ def split_waveform( window_s: float, overlap_s: float, ) -> list[np.ndarray]: - """Split ``samples`` into overlapping windows of ``window_s`` seconds. - - Consecutive windows advance by ``window_s - overlap_s`` seconds, so each pair - shares ``overlap_s`` of audio. A clip already shorter than one window is - returned as a single segment (no copying/splitting). - - Args: - samples: 1-D mono waveform. - sr: Sample rate of ``samples`` in Hz. - window_s: Window length in seconds (must be > 0). - overlap_s: Overlap between consecutive windows in seconds (0 <= overlap_s - < window_s). + """Split a 1-D mono waveform into overlapping windows, in order. - Returns: - A list of 1-D numpy views/arrays, in order. + Windows advance by ``window_s - overlap_s``. A clip shorter than one window + is returned as a single segment. """ if window_s <= 0: raise ValueError(f"window_s must be > 0, got {window_s}") @@ -87,9 +65,8 @@ def _norm_word(word: str) -> str: def _seam_overlap(prev: list[str], nxt: list[str]) -> int: """Longest k such that the last k words of ``prev`` match the first k of ``nxt``. - Comparison is punctuation/case-insensitive because ASR often renders the - overlap region slightly differently on each side of the seam. Returns 0 when - there is no matching overlap. + Punctuation/case-insensitive: ASR renders the overlap slightly differently on + each side of a seam. Returns 0 when nothing matches. """ max_k = min(len(prev), len(nxt), _MAX_SEAM_WORDS) for k in range(max_k, 0, -1): @@ -101,12 +78,7 @@ def _seam_overlap(prev: list[str], nxt: list[str]) -> int: def merge_transcripts(transcripts: list[str]) -> str: - """Concatenate per-window transcripts, de-duplicating the overlap at each seam. - - For each new window, find the longest word-level overlap between the tail of - the text so far and the head of the new window, and drop that duplicated span - from the new window before appending. Empty windows are skipped. - """ + """Concatenate per-window transcripts, de-duplicating the overlap at each seam.""" merged: list[str] = [] for text in transcripts: words = text.split() diff --git a/src/granite_switch/vllm/audio/processor.py b/src/granite_switch/vllm/audio/processor.py index 17e39e2..61700c8 100644 --- a/src/granite_switch/vllm/audio/processor.py +++ b/src/granite_switch/vllm/audio/processor.py @@ -1,26 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -"""vLLM multimodal processor for the alpha audio cascade (Path A). - -Strategy (see also the design notes in :mod:`granite_switch.vllm.audio`): - -* Audio is declared as a multimodal input, so the standard vLLM server accepts - it and developers keep loading a single model. -* In :meth:`_call_hf_processor` we run ASR and tokenize the transcript. -* :meth:`_get_prompt_updates` replaces the single ``<|audio|>`` marker with the - **actual transcript token ids** via ``PromptReplacement``. After this the - transcript tokens are ordinary tokens in ``prompt_token_ids``; the scheduler - sizes KV for the real, runtime-determined length (no fixed audio window). -* The model's ``embed_multimodal`` then returns ``embed_tokens(transcript_ids)`` - for those placeholder positions — identical to what they'd get as normal text. - That redundant-but-consistent step is the exact seam the future projection - model reuses (swap the embedding source for a trained audio encoder). - -Modeled on vLLM 0.19.1's ``ultravox.py`` (the reference audio model), confirmed -against that version's API by the scratch probes. - -ALPHA SCOPE: audio is answered by the base model. We do not place adapter -control tokens for audio requests, so no token-exchange interaction with the -switch is needed here. +"""vLLM multimodal processor for the audio cascade. + +``_call_hf_processor`` runs ASR and tokenizes the transcript; +``_get_prompt_updates`` then replaces the ``<|audio|>`` marker with the real +transcript token ids via ``PromptReplacement``, so the scheduler sizes KV for the +runtime-determined length rather than a fixed audio window. + +Modeled on vLLM 0.19.1's ``ultravox.py``. Audio is answered by the base model — +no adapter control tokens are placed, so the switch is not involved. """ from __future__ import annotations @@ -53,17 +40,10 @@ resolve_generate_kwargs, ) -# The chat-template marker that stands in for an audio clip before replacement. AUDIO_MARKER = "<|audio|>" - -# ASR feature-extractor sample rate the audio is resampled to. _TARGET_SR = 16_000 - -# Fallback context length if the served max_model_len cannot be read (should not -# happen in practice; keeps the transcript budget finite either way). +# Keeps the transcript budget finite if max_model_len cannot be read. _FALLBACK_CONTEXT_LEN = 8192 - -# Dummy clip length (seconds) used during vLLM's profiling run. _DUMMY_AUDIO_SECONDS = 5 @@ -74,14 +54,11 @@ def _asr_enabled(self) -> bool: return bool(getattr(self.get_hf_config(), "asr_enabled", False)) def get_supported_mm_limits(self) -> Mapping[str, int | None]: - # Audio capability is gated per-checkpoint: a non-audio GraniteSwitch - # reports no modalities, so vLLM never profiles audio or loads ASR. + # No modalities on a non-audio checkpoint, so vLLM never loads ASR. if not self._asr_enabled(): return {} - # Configurable ceiling on audio clips per request; each clip's transcript - # is spliced at its own <|audio|> marker. Finite so vLLM can size KV for - # the worst case (clips x per-clip budget). --limit-mm-per-prompt may - # lower this, not raise it above the declared ceiling. + # Finite so vLLM can size KV for the worst case. --limit-mm-per-prompt + # may lower this ceiling, not raise it. return {"audio": self._asr_max_audio_clips()} def get_mm_max_tokens_per_item( @@ -91,16 +68,13 @@ def get_mm_max_tokens_per_item( ) -> Mapping[str, int] | None: if not self._asr_enabled(): return {} - # Worst-case transcript positions one clip can occupy, used only to size - # the encoder cache and profiling pass — NOT to bound requests. A clip - # cannot transcribe to more than the whole context (a longer prompt is - # rejected by vLLM's standard length check), so the per-clip share of the - # context window is the honest upper bound. vLLM calls this with count=1. + # Sizes the encoder cache and profiling pass only — does NOT bound + # requests (vLLM's prompt-length check does). A clip cannot exceed the + # whole context, so the per-clip share of it is the honest upper bound. count = mm_counts.get("audio", 1) or 1 return {"audio": max(1, seq_len // count)} def get_data_parser(self) -> MultiModalDataParser: - # Resample incoming audio to the ASR sample rate. return MultiModalDataParser(target_sr=_TARGET_SR) # --- ASR config resolved from the model's GraniteSwitchConfig --- @@ -113,6 +87,10 @@ def _asr_device(self) -> str: cfg = self.get_hf_config() return getattr(cfg, "asr_device", "cpu") or "cpu" + def _asr_dtype(self) -> str | None: + cfg = self.get_hf_config() + return getattr(cfg, "asr_dtype", None) + def _asr_pipeline_kwargs(self) -> Mapping[str, object]: cfg = self.get_hf_config() return getattr(cfg, "asr_pipeline_kwargs", None) or {} @@ -140,11 +118,8 @@ def _asr_chunk_overlap_s(self) -> float: def _max_model_len(self) -> int: """The served context window, or a safe fallback. - vLLM passes ``seq_len`` into ``get_mm_max_tokens_per_item`` for profiling, - but ``_call_hf_processor`` needs the served context at request time to size - the transcript budget, so read it from the processing context's model - config (falling back to the checkpoint's max_position_embeddings, then a - constant). + ``_call_hf_processor`` needs this at request time to size the transcript + budget; vLLM's profiling ``seq_len`` is not available there. """ model_config = getattr(self.ctx, "model_config", None) max_len = getattr(model_config, "max_model_len", None) @@ -183,19 +158,13 @@ def _transcribe( audio, generate_kwargs: Mapping[str, object] | None = None, ) -> list[int]: - """Transcribe one audio item to a list of token ids (full transcript). - - The transcript is spliced into the prompt as ordinary text tokens; it is - never truncated here. A clip whose transcript makes the prompt exceed the - context is rejected by vLLM's standard prompt-length check. - - Long audio is handled by the backend itself (Whisper) or by our - encoder-agnostic chunker, per the checkpoint's ``asr_self_chunks`` flag. - """ + """Transcribe one audio item to token ids. Never truncated here — an + oversized prompt is rejected by vLLM's own length check.""" transcriber = get_transcriber( model_id=self.info._asr_model_id(), device=self.info._asr_device(), pipeline_kwargs=self.info._asr_pipeline_kwargs(), + dtype=self.info._asr_dtype(), ) # The data parser already resampled to _TARGET_SR. text = transcriber.transcribe( @@ -219,13 +188,11 @@ def _call_hf_processor( tokenizer = self.info.get_tokenizer() audios = mm_data.get("audios", []) or [] - # Text-only request: just tokenize (mirrors Ultravox's text-only branch). if not audios: input_ids = tokenizer.encode(prompt, add_special_tokens=False) return BatchFeature(dict(input_ids=[input_ids]), tensor_type="pt") - # Resolve decode kwargs once (config defaults + allowlisted per-request - # overrides) and apply them to every audio item in this request. + # Resolved once, then applied to every audio item in this request. generate_kwargs = resolve_generate_kwargs( self.info._asr_generate_kwargs(), mm_kwargs, @@ -234,9 +201,7 @@ def _call_hf_processor( input_ids = tokenizer.encode(prompt, add_special_tokens=False) - # Transcribe each audio to token ids; concatenate flat with per-item sizes. - # Transcripts are spliced in full — an oversized request is rejected by - # vLLM's prompt-length check, not silently truncated here. + # Concatenated flat, with per-item sizes to split them back. per_item_ids = [self._transcribe(a, generate_kwargs) for a in audios] sizes = [len(ids) for ids in per_item_ids] flat_ids = [tid for ids in per_item_ids for tid in ids] @@ -257,7 +222,6 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: num_tokens = hf_inputs.get("audio_num_tokens", torch.zeros(0)) return dict( - # Flat transcript ids, split back per audio by audio_num_tokens. audio_token_ids=MultiModalFieldConfig.flat_from_sizes("audio", num_tokens), audio_num_tokens=MultiModalFieldConfig.batched("audio"), ) @@ -277,7 +241,6 @@ def _get_prompt_updates( def replacement(item_idx: int): s = int(starts[item_idx]) e = int(starts[item_idx + 1]) - # Replace <|audio|> with the real transcript token ids. return [int(t) for t in all_ids[s:e]] return [ diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index 9467846..e24466a 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -299,9 +299,8 @@ def forward( adapter_token_ids=self.adapter_token_ids, ) else: - # Either no switch, or the multimodal path where vLLM passes - # pre-merged inputs_embeds and input_ids is None. ALPHA: audio - # requests run on the base model (adapter_id = 0). + # No switch, or the multimodal path (input_ids is None because + # vLLM pre-merged inputs_embeds): run on base, adapter_id 0. if input_ids is not None: num_tokens = input_ids.shape[0] device = input_ids.device @@ -429,10 +428,8 @@ class GraniteSwitchForCausalLM( supports_multimodal = True - # Force vLLM to pass the raw input_ids to forward() on the multimodal path - # (not just precomputed inputs_embeds). The switch needs them to detect - # adapter control tokens, so audio requests route through adapters exactly - # like text requests. + # Without this, the multimodal path passes only inputs_embeds and the switch + # cannot see control tokens — audio requests would bypass adapters. requires_raw_input_tokens = True @classmethod diff --git a/tests/integration/test_adapter_routing_audio_enabled.py b/tests/integration/test_adapter_routing_audio_enabled.py index 131a01f..7bbc63b 100644 --- a/tests/integration/test_adapter_routing_audio_enabled.py +++ b/tests/integration/test_adapter_routing_audio_enabled.py @@ -2,14 +2,10 @@ """Adapters + audio (issue #47): every default adapter routes to its own index on an audio-enabled checkpoint. -Composes the default adapter set (RAG + Core + Guardian) with --enable-audio and -sweeps every adapter control token that resolved, asserting the switch maps -``adapter_token_ids[i]`` to index ``i + 1`` and leaves pre-control positions on -the base (``0``) — i.e. enabling audio does not perturb the token->index map. -For granite-4.1-3b that is 12 adapters: context_relevance ships no 4.1-3b flavor, -so 12 of the 13 defined adapters resolve. Complements the single-adapter sweep in -test_switch_e2e_compose and the serve-time answerability check in -test_answerability_over_audio. +Asserts the switch maps ``adapter_token_ids[i]`` to index ``i + 1`` and leaves +pre-control positions on base (``0``) — enabling audio must not perturb the +token->index map. For granite-4.1-3b, 12 of the 13 defined adapters resolve +(context_relevance ships no 4.1-3b flavor). Markers: slow + requires_model + gpu (opt-in via -m). """ @@ -26,8 +22,6 @@ pytest.skip("requires the HF backend ([hf] extra)", allow_module_level=True) -# The default adapter set is the union of the three granitelib libraries (the -# "12 adapters"): RAG + Core + Guardian. _DEFAULT_ADAPTER_LIBRARIES = [ "ibm-granite/granitelib-rag-r1.0", "ibm-granite/granitelib-core-r1.0", diff --git a/tests/integration/test_answerability_over_audio.py b/tests/integration/test_answerability_over_audio.py index e59f11b..538457b 100644 --- a/tests/integration/test_answerability_over_audio.py +++ b/tests/integration/test_answerability_over_audio.py @@ -2,18 +2,14 @@ """Adapters + audio: the RAG answerability adapter judging a query against a document delivered as audio (issue #47). -Covers the "adapters + audio" box. Where test_audio_serving_smoke.py only proves -an adapter control token doesn't crash the audio path, this checks the switch -reaches the *correct* verdict when the RAG document arrives as speech. - -Document-as-audio: ``documents=[{"text": "<|audio|>"}]`` — the Granite template -renders each document with ``doc | tojson``, so the marker lands in the -```` block where the ASR processor splices the transcript. The -``<|answerability|>`` control token is inserted by the same template (aLoRA -fallback path). Ground truth is one committed speech clip (fixtures/, generated -with SpeechT5 — MIT) driving both classes; the questions key on content that -transcribes cleanly, so this tests the decision, not transcription accuracy -(WER). +Unlike test_audio_serving_smoke.py, which only proves an adapter control token +doesn't crash the audio path, this checks the switch reaches the *correct* +verdict when the document arrives as speech. + +``documents=[{"text": "<|audio|>"}]`` works because the Granite template renders +each document with ``doc | tojson``, so the marker lands in the ```` +block for the ASR processor to splice. Questions key on words that transcribe +cleanly, so this tests the decision, not WER. Markers: slow + requires_model + gpu (opt-in via -m). """ @@ -37,10 +33,7 @@ def _load_experimental_pairs(): - """Extra (base, adapter) pairs from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS. - - JSON array of {"base": str, "adapter": str}; mirrors the other E2E files. - """ + """Extra (base, adapter) pairs from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS.""" raw = os.environ.get("GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS", "") if not raw: return [] @@ -134,8 +127,7 @@ def served(audio_rag_checkpoint): def _answerability_adapter_name(config): - # Discovery derives the name from the library layout — match rather than - # hard-code, and skip loudly if this checkpoint has no answerability adapter. + # Discovery derives the name from the library layout, so match, don't hard-code. names = list(getattr(config, "adapter_names", None) or []) for name in names: if "answerab" in name.lower(): @@ -155,8 +147,6 @@ def speech_clip(): def _build_answerability_prompt(tokenizer, adapter_name, question): - # documents=[{"text": "<|audio|>"}] places one audio marker in the template's - # block; adapter_name arms the answerability control-token insert. return tokenizer.apply_chat_template( [{"role": "user", "content": question}], documents=[{"text": _AUDIO_MARKER}], @@ -167,8 +157,7 @@ def _build_answerability_prompt(tokenizer, adapter_name, question): def _parse_label(text): - # Adapter emits the bare enum ("answerable"/"unanswerable"), maybe quoted. - # Match the longer label first so "answerable" doesn't shadow "unanswerable". + # Longer label first, so "answerable" doesn't shadow "unanswerable". norm = text.strip().strip('"').strip().lower() if "unanswerable" in norm: return "unanswerable" diff --git a/tests/integration/test_audio_serving_smoke.py b/tests/integration/test_audio_serving_smoke.py index b1eb3a6..3c0e9c7 100644 --- a/tests/integration/test_audio_serving_smoke.py +++ b/tests/integration/test_audio_serving_smoke.py @@ -1,29 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 """End-to-end vLLM serving smoke for the audio cascade (issue #47). -Boots a real composed, audio-enabled GraniteSwitch checkpoint under vLLM and -drives the three request shapes the #47 checklist calls out as the serving -smoke — **text + adapter + audio x1/x2/x3** — through one live engine. This is -the path the low-level tests deliberately bypass (CUDA graphs, torch.compile, -the ASR processor running inside vLLM's EngineCore subprocess, the multi-clip -splice, and generation), so it is the only test that proves those pieces -compose at serve time. - -Scope is a *smoke*: each request must complete and return well-formed output, -and multi-clip requests must be accepted up to the checkpoint's declared clip -ceiling. It intentionally does NOT assert transcript content — transcription -quality (WER) and adapter-routing correctness are separate #47 boxes covered by -the eval harness and `test_switch_e2e_compose.py` respectively. Synthetic audio -keeps the test asset-free; a silent/tonal clip exercises the serving plumbing -without shipping a speech fixture. - -Model construction goes through the compose CLI (CLAUDE.md gotcha #5): no test -hand-assembles a config. Audio is enabled with `--enable-audio`, which defaults -the ASR front-end to the small built-in model (distil-whisper/distil-small.en) -to keep CI cost down. - -Markers: @pytest.mark.slow + @pytest.mark.requires_model + @pytest.mark.gpu. -CI must opt in explicitly: `pytest -m "slow and requires_model and gpu"`. +Drives text + adapter + audio x1/x2/x3 through one live engine — the serve-time +path the low-level tests bypass (CUDA graphs, the ASR processor inside vLLM's +EngineCore subprocess, the multi-clip splice, generation). + +Deliberately does NOT assert transcript content: WER and adapter-routing +correctness are separate boxes, covered by the eval harness and +test_switch_e2e_compose.py. Synthetic tones keep the test asset-free. + +Opt in explicitly: `pytest -m "slow and requires_model and gpu"`. """ import importlib.util @@ -38,13 +24,8 @@ pytest.skip("requires vLLM installed", allow_module_level=True) -# ---------------------------------------------------------------------------- -# Base-model / adapter-library pairs — kept in lockstep with -# tests/integration/test_switch_e2e_compose.py so the two E2E files exercise -# the same model matrix. A fast CI profile can pin one pair via -k or the -# experimental env var below. -# ---------------------------------------------------------------------------- - +# Kept in lockstep with test_switch_e2e_compose.py so both E2E files exercise the +# same model matrix. _DEFAULT_BASE_MODEL_PAIRS = [ ("ibm-granite/granite-4.0-micro", "ibm-granite/granitelib-core-r1.0"), ("ibm-granite/granite-4.1-3b", "ibm-granite/granitelib-core-r1.0"), @@ -52,11 +33,10 @@ def _load_experimental_pairs(): - """Local/experimental pairings from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS. + """Local pairings from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS. - JSON array of {"base": str, "adapter": str}; base/adapter may be HF ids or - local paths. The mechanism is committed; the values are not. Mirrors - test_switch_e2e_compose.py so both E2E files share one extension knob. + JSON array of {"base": str, "adapter": str}, HF ids or local paths. The + mechanism is committed; the values are not. """ raw = os.environ.get("GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS", "") if not raw: @@ -85,9 +65,8 @@ def _load_experimental_pairs(): def audio_checkpoint(request, tmp_path_factory): """Compose one audio-enabled checkpoint per (base, adapter) pair. - Delegates to the compose CLI (same as test_switch_e2e_compose.py) with - `--enable-audio`, then returns the save dir. Module scope amortizes the - download-dominated first run across every smoke case for the pair. + Goes through the compose CLI, never a hand-assembled config (CLAUDE.md + gotcha #5). Module scope amortizes the download across the pair's cases. """ import subprocess import sys @@ -123,8 +102,7 @@ def served(audio_checkpoint): """Boot vLLM once for the checkpoint and share it across smoke cases. Tokenizer init stays ON (unlike the argmax-equivalence test): the ASR - processor needs a tokenizer to encode both the prompt and the transcript, - and the adapter case needs it to tokenize text around the control token. + processor needs it to encode the prompt and the transcript. """ import gc @@ -177,9 +155,8 @@ def test_text_only_serving(served): def test_adapter_control_token_serving(served): """An adapter control token routes through the switch under serving. - Correctness of the routing is `test_switch_e2e_compose.py`'s job; here we - only prove the switch path runs end-to-end in the live engine without - crashing and still generates. + Routing correctness is test_switch_e2e_compose.py's job; the bar here is + that the switch path runs in the live engine and still generates. """ from vllm import SamplingParams from vllm.inputs import TokensPrompt @@ -203,9 +180,7 @@ def test_adapter_control_token_serving(served): def test_audio_clip_serving(served, num_clips): """Audio x1/x2/x3: N markers + N clips transcribe, splice, and generate. - Exercises the multi-clip long-audio serving path through the ASR processor - running inside vLLM's engine subprocess. Content is not asserted (WER is a - separate #47 box); the bar is a completed request with well-formed output. + Content is not asserted; the bar is a completed, well-formed request. """ from vllm import SamplingParams diff --git a/tests/unit/test_asr.py b/tests/unit/test_asr.py index d5e00e0..69187ce 100644 --- a/tests/unit/test_asr.py +++ b/tests/unit/test_asr.py @@ -7,6 +7,7 @@ leaf module directly by file path rather than through the package. """ +import contextlib import importlib.util import pathlib from unittest import mock @@ -282,15 +283,92 @@ def test_config_not_mutated(self): assert cfg == {"language": "de"} +@contextlib.contextmanager +def _patched_pipeline(factory): + """Patch both lookup paths: `from transformers import pipeline` re-resolves to + transformers.pipelines.pipeline, but transformers caches it on the top-level + module after the first access, so a later test would get the real one.""" + with ( + mock.patch("transformers.pipelines.pipeline", factory), + mock.patch("transformers.pipeline", factory), + ): + yield + + +class TestResolveTorchDtype: + """asr_dtype resolution. float16-on-CUDA is the default, but overridable.""" + + def test_auto_on_cuda_is_float16(self): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype(None, "cuda:0") is torch.float16 + assert asr._resolve_torch_dtype("auto", "cuda") is torch.float16 + + def test_auto_on_cpu_is_float32(self): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype(None, "cpu") is torch.float32 + + def test_explicit_float32_overrides_cuda_default(self): + # The BatchNorm-encoder case: CUDA must not force half precision. + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype("float32", "cuda:0") is torch.float32 + + def test_explicit_bfloat16(self): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype("bfloat16", "cpu") is torch.bfloat16 + + @pytest.mark.parametrize( + "name,attr", + [ + ("fp16", "float16"), + ("half", "float16"), + ("bf16", "bfloat16"), + ("fp32", "float32"), + ("FLOAT32", "float32"), + ], + ) + def test_aliases_and_case(self, name, attr): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype(name, "cpu") is getattr(torch, attr) + + def test_unknown_name_raises(self): + pytest.importorskip("torch") + with pytest.raises(ValueError, match="Unsupported asr_dtype"): + asr._resolve_torch_dtype("int8", "cpu") + + def test_dtype_is_part_of_cache_key(self): + a = asr.get_transcriber("dt", "cuda:0", dtype="float32") + b = asr.get_transcriber("dt", "cuda:0", dtype="float16") + c = asr.get_transcriber("dt", "cuda:0", dtype="float32") + assert a is not b + assert a is c + + def test_load_passes_resolved_dtype(self): + torch = pytest.importorskip("torch") + factory = mock.Mock(return_value=mock.Mock()) + with _patched_pipeline(factory): + asr.ASRTranscriber(model_id="m", device="cuda:0", dtype="float32").load() + assert factory.call_args.kwargs["torch_dtype"] is torch.float32 + + def test_pipeline_kwargs_torch_dtype_still_wins(self): + torch = pytest.importorskip("torch") + factory = mock.Mock(return_value=mock.Mock()) + with _patched_pipeline(factory): + asr.ASRTranscriber( + model_id="m", + device="cpu", + dtype="float32", + pipeline_kwargs={"torch_dtype": torch.bfloat16}, + ).load() + assert factory.call_args.kwargs["torch_dtype"] is torch.bfloat16 + + class TestLoadMergesPipelineKwargs: def test_pipeline_kwargs_override_defaults(self): # load() must merge config-supplied pipeline_kwargs over the built-in # defaults (e.g. override chunk_length_s, add extra kwargs). pytest.importorskip("torch") fake_pipe_factory = mock.Mock(return_value=mock.Mock()) - # transformers is a lazy module: `from transformers import pipeline` - # re-resolves to transformers.pipelines.pipeline, so patch there. - with mock.patch("transformers.pipelines.pipeline", fake_pipe_factory): + with _patched_pipeline(fake_pipe_factory): t = asr.ASRTranscriber( model_id="m", device="cpu", diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index f25e2a7..6252544 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -106,9 +106,21 @@ def test_asr_defaults_off(self): assert cfg.asr_enabled is False assert cfg.asr_model_id is None assert cfg.asr_device == "cpu" + assert cfg.asr_dtype is None assert cfg.asr_pipeline_kwargs is None assert cfg.asr_generate_kwargs is None + def test_asr_dtype_round_trip(self, tmp_path): + cfg = GraniteSwitchConfig( + num_adapters=0, asr_enabled=True, asr_device="cuda:0", asr_dtype="float32" + ) + cfg.save_pretrained(tmp_path) + assert GraniteSwitchConfig.from_pretrained(tmp_path).asr_dtype == "float32" + + def test_invalid_asr_dtype_raises(self): + with pytest.raises(ValueError, match="asr_dtype"): + GraniteSwitchConfig(num_adapters=0, asr_dtype="fp8") + def test_longaudio_defaults(self): cfg = GraniteSwitchConfig(num_adapters=0) assert cfg.asr_max_audio_clips == 32 diff --git a/tests/vllm/test_audio_processor.py b/tests/vllm/test_audio_processor.py index 2d81f3d..b158b27 100644 --- a/tests/vllm/test_audio_processor.py +++ b/tests/vllm/test_audio_processor.py @@ -1,21 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 """vLLM-tier tests for the audio ASR multimodal processor. -These exercise the code that lives in ``granite_switch.vllm.audio.processor`` -and therefore needs vLLM importable (the base classes come from vLLM). They do -NOT need a GPU or a real ASR model: the transcriber is faked, so what is under -test is our plumbing — - -- per-checkpoint gating of the audio modality on ``asr_enabled``, -- the config accessors for ``asr_model_id`` / ``asr_device`` / - ``asr_pipeline_kwargs`` / ``asr_generate_kwargs``, -- and the Level-2 seam: config-default decode kwargs merged with an allowlisted - per-request override (vLLM's ``mm_processor_kwargs``) reaching the transcriber. - -The pure merge/kwargs logic itself is unit-tested on CPU in -``tests/unit/test_asr.py``; this file checks it is wired through the processor. - -Requires vLLM installed. Skipped otherwise (e.g. CPU-only dev machines). +Needs vLLM importable (the base classes come from it) but no GPU and no real ASR +model — the transcriber is faked, so what is under test is the plumbing: modality +gating on ``asr_enabled``, the config accessors, and decode kwargs reaching the +transcriber. The merge logic itself is unit-tested in tests/unit/test_asr.py. """ import importlib.util @@ -43,9 +32,7 @@ def _make_info(*, max_model_len=131072, **cfg_attrs): """A ProcessingInfo whose get_hf_config() returns a stub config. - Bypasses __init__ (which needs a full vLLM InputProcessingContext) — we only - exercise methods that read the config. A stub ``ctx`` supplies the served - ``max_model_len`` used to size the transcript budget. + Bypasses __init__, which needs a full vLLM InputProcessingContext. """ info = object.__new__(GraniteSwitchASRProcessingInfo) cfg = SimpleNamespace(**cfg_attrs) @@ -167,10 +154,13 @@ def transcribe( capture["chunk_overlap_s"] = chunk_overlap_s return "hello world" - def fake_get_transcriber(model_id=None, device="cpu", pipeline_kwargs=None): + def fake_get_transcriber( + model_id=None, device="cpu", pipeline_kwargs=None, dtype=None + ): capture["model_id"] = model_id capture["device"] = device capture["pipeline_kwargs"] = pipeline_kwargs + capture["dtype"] = dtype return FakeTranscriber() monkeypatch.setattr(proc_mod, "get_transcriber", fake_get_transcriber) @@ -196,6 +186,19 @@ def test_transcribe_forwards_pipeline_and_generate_kwargs(self, monkeypatch): assert capture["generate_kwargs"] == {"language": "fr"} assert capture["sampling_rate"] == proc_mod._TARGET_SR + def test_dtype_forwarded_from_config(self, monkeypatch): + capture = {} + info = _make_info(asr_enabled=True, asr_device="cuda:0", asr_dtype="float32") + proc = _make_processor(info, monkeypatch, capture) + proc._transcribe(np.zeros(1600, dtype=np.float32), {}) + assert capture["dtype"] == "float32" + + def test_dtype_defaults_to_none(self, monkeypatch): + capture = {} + proc = _make_processor(_make_info(asr_enabled=True), monkeypatch, capture) + proc._transcribe(np.zeros(1600, dtype=np.float32), {}) + assert capture["dtype"] is None + def test_empty_generate_kwargs_becomes_none(self, monkeypatch): capture = {} info = _make_info(asr_enabled=True, asr_model_id="w", asr_device="cpu") @@ -340,7 +343,9 @@ def transcribe( monkeypatch.setattr( proc_mod, "get_transcriber", - lambda model_id=None, device="cpu", pipeline_kwargs=None: FakeTranscriber(), + lambda model_id=None, device="cpu", pipeline_kwargs=None, dtype=None: ( + FakeTranscriber() + ), ) return proc @@ -410,7 +415,9 @@ def transcribe(self, audio, **kw): monkeypatch.setattr( proc_mod, "get_transcriber", - lambda model_id=None, device="cpu", pipeline_kwargs=None: FakeTranscriber(), + lambda model_id=None, device="cpu", pipeline_kwargs=None, dtype=None: ( + FakeTranscriber() + ), ) bf = proc._call_hf_processor(