| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Adapted from Qwen3-ForcedAligner inference: |
| 4 | # https://github.com/QwenLM/Qwen3-ASR |
| 5 | |
| 6 | """Offline forced alignment example using Qwen3-ForcedAligner-0.6B. |
| 7 | |
| 8 | Forced alignment takes audio and reference text as input and produces |
| 9 | word-level timestamps. The model predicts a time bin at each <timestamp> |
| 10 | token position; multiplying by ``timestamp_segment_time`` gives milliseconds. |
| 11 | |
| 12 | Usage:: |
| 13 | |
| 14 | python forced_alignment_offline.py \ |
| 15 | --model Qwen/Qwen3-ForcedAligner-0.6B |
| 16 | """ |
| 17 | |
| 18 | from argparse import Namespace |
| 19 | |
| 20 | import numpy as np |
| 21 | |
| 22 | from vllm import LLM, EngineArgs |
| 23 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 24 | |
| 25 | |
| 26 | def parse_args(): |
| 27 | parser = FlexibleArgumentParser() |
| 28 | parser = EngineArgs.add_cli_args(parser) |
| 29 | parser.set_defaults( |
| 30 | model="Qwen/Qwen3-ForcedAligner-0.6B", |
| 31 | runner="pooling", |
| 32 | enforce_eager=True, |
| 33 | hf_overrides={"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}, |
| 34 | ) |
| 35 | return parser.parse_args() |
| 36 | |
| 37 | |
| 38 | def build_prompt(words: list[str]) -> str: |
| 39 | """Build the forced alignment prompt from a word list. |
| 40 | |
| 41 | Format: <|audio_start|><|audio_pad|><|audio_end|> |
| 42 | word1<timestamp><timestamp>word2<timestamp><timestamp>... |
| 43 | """ |
| 44 | body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>" |
| 45 | return f"<|audio_start|><|audio_pad|><|audio_end|>{body}" |
| 46 | |
| 47 | |
| 48 | def main(args: Namespace): |
| 49 | llm = LLM(**vars(args)) |
| 50 | |
| 51 | config = llm.llm_engine.vllm_config.model_config.hf_config |
| 52 | timestamp_token_id = config.timestamp_token_id |
| 53 | timestamp_segment_time = config.timestamp_segment_time |
| 54 | |
| 55 | # Example: align these words against a 5-second audio clip |
| 56 | words = ["Hello", "world"] |
| 57 | prompt = build_prompt(words) |
| 58 | |
| 59 | # Use a 5-second silent audio as placeholder (replace with real audio) |
| 60 | sample_rate = 16000 |
| 61 | audio = np.zeros(sample_rate * 5, dtype=np.float32) |
| 62 | |
| 63 | outputs = llm.encode( |
| 64 | [{"prompt": prompt, "multi_modal_data": {"audio": audio}}], |
| 65 | pooling_task="token_classify", |
| 66 | ) |
| 67 | |
| 68 | for output in outputs: |
| 69 | logits = output.outputs.data # [num_tokens, classify_num] |
| 70 | predictions = logits.argmax(dim=-1) |
| 71 | token_ids = output.prompt_token_ids |
| 72 | |
| 73 | # Extract timestamps at <timestamp> positions |
| 74 | ts_predictions = [ |
| 75 | pred.item() * timestamp_segment_time |
| 76 | for tid, pred in zip(token_ids, predictions) |
| 77 | if tid == timestamp_token_id |
| 78 | ] |
| 79 | |
| 80 | # Pair up start/end times per word |
| 81 | for i, word in enumerate(words): |
| 82 | start_ms = ts_predictions[i * 2] |
| 83 | end_ms = ts_predictions[i * 2 + 1] |
| 84 | print(f"{word:15s} {start_ms / 1000:.3f}s - {end_ms / 1000:.3f}s") |
| 85 | |
| 86 | |
| 87 | if __name__ == "__main__": |
| 88 | args = parse_args() |
| 89 | main(args) |