| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import numpy as np |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | MODEL = "Qwen/Qwen3-ForcedAligner-0.6B" |
| 9 | CLASSIFY_NUM = 5000 |
| 10 | TIMESTAMP_TOKEN_ID = 151705 |
| 11 | |
| 12 | |
| 13 | def build_prompt(words: list[str]) -> str: |
| 14 | body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>" |
| 15 | return f"<|audio_start|><|audio_pad|><|audio_end|>{body}" |
| 16 | |
| 17 | |
| 18 | @pytest.mark.parametrize("model", [MODEL]) |
| 19 | @pytest.mark.parametrize("dtype", ["bfloat16"]) |
| 20 | @torch.inference_mode() |
| 21 | def test_qwen3_forced_aligner( |
| 22 | vllm_runner, |
| 23 | model: str, |
| 24 | dtype: str, |
| 25 | ) -> None: |
| 26 | words = ["Hello", "world"] |
| 27 | prompt = build_prompt(words) |
| 28 | |
| 29 | # 5-second silent audio at 16kHz |
| 30 | audio = np.zeros(16000 * 5, dtype=np.float32) |
| 31 | |
| 32 | with vllm_runner( |
| 33 | model, |
| 34 | runner="pooling", |
| 35 | dtype=dtype, |
| 36 | enforce_eager=True, |
| 37 | max_model_len=512, |
| 38 | hf_overrides={ |
| 39 | "architectures": [ |
| 40 | "Qwen3ASRForcedAlignerForTokenClassification", |
| 41 | ], |
| 42 | }, |
| 43 | ) as vllm_model: |
| 44 | outputs = vllm_model.llm.encode( |
| 45 | [{"prompt": prompt, "multi_modal_data": {"audio": audio}}], |
| 46 | pooling_task="token_classify", |
| 47 | ) |
| 48 | |
| 49 | # Validate output structure |
| 50 | assert len(outputs) == 1 |
| 51 | logits = outputs[0].outputs.data |
| 52 | assert logits.dim() == 2 |
| 53 | assert logits.shape[1] == CLASSIFY_NUM |
| 54 | |
| 55 | # Validate timestamp extraction |
| 56 | token_ids = outputs[0].prompt_token_ids |
| 57 | predictions = logits.argmax(dim=-1) |
| 58 | ts_indices = [i for i, t in enumerate(token_ids) if t == TIMESTAMP_TOKEN_ID] |
| 59 | |
| 60 | # 2 words x 2 timestamps each (start + end) = 4 |
| 61 | assert len(ts_indices) == 4 |
| 62 | |
| 63 | ts_preds = [predictions[i].item() for i in ts_indices] |
| 64 | assert all(p >= 0 for p in ts_preds) |
| 65 | # end >= start for each word |
| 66 | assert ts_preds[1] >= ts_preds[0] # Hello |
| 67 | assert ts_preds[3] >= ts_preds[2] # world |