@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 191 lines 5.82 KB
Raw
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 """Online forced alignment example using Qwen3-ForcedAligner-0.6B-hf.
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 This example uses STEP pooling to return logits only for <timestamp> tokens.
13 The ``step_tag_id`` value 151705 is this checkpoint's ``timestamp_token_id``
14 from its config.json.
15
16 Start the server with:
17
18 vllm serve Qwen/Qwen3-ForcedAligner-0.6B-hf \
19 --runner pooling \
20 --convert classify \
21 --pooler-config \
22 '{"tok_pooling_type":"STEP","step_tag_id":151705,"use_activation":false}' \
23 --enforce-eager \
24 --trust-request-chat-template
25
26 Then run:
27
28 python forced_alignment_online.py
29 """
30
31 import argparse
32 import json
33 import mimetypes
34 import wave
35 from io import BytesIO
36 from pathlib import Path
37 from typing import Any
38
39 import numpy as np
40 import pybase64 as base64
41 import requests
42 import torch
43 from huggingface_hub import hf_hub_download
44
45 RAW_CONTENT_CHAT_TEMPLATE = "{{ messages[0]['content'] }}"
46
47
48 def build_prompt(words: list[str]) -> str:
49 """Build the forced alignment prompt from a word list.
50
51 Format: <|audio_start|><|audio_pad|><|audio_end|>
52 word1<timestamp><timestamp>word2<timestamp><timestamp>...
53 """
54 body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>"
55 return f"<|audio_start|><|audio_pad|><|audio_end|>{body}"
56
57
58 def encode_audio_data_uri(audio_path: Path) -> str:
59 mime_type = mimetypes.guess_type(audio_path)[0] or "audio/wav"
60 audio_base64 = base64.b64encode(audio_path.read_bytes()).decode("utf-8")
61 return f"data:{mime_type};base64,{audio_base64}"
62
63
64 def encode_silent_wav_data_uri(sample_rate: int = 16000, duration_s: int = 5) -> str:
65 audio = np.zeros(sample_rate * duration_s, dtype=np.int16)
66
67 with BytesIO() as audio_buffer:
68 with wave.open(audio_buffer, "wb") as wav_file:
69 wav_file.setnchannels(1)
70 wav_file.setsampwidth(np.dtype(np.int16).itemsize)
71 wav_file.setframerate(sample_rate)
72 wav_file.writeframes(audio.tobytes())
73
74 audio_base64 = base64.b64encode(audio_buffer.getvalue()).decode("utf-8")
75
76 return f"data:audio/wav;base64,{audio_base64}"
77
78
79 def build_payload(model: str, prompt: str, audio_uri: str) -> dict[str, Any]:
80 return {
81 "model": model,
82 "messages": [
83 {
84 "role": "user",
85 "content": [
86 {"type": "text", "text": prompt},
87 {"type": "audio_url", "audio_url": {"url": audio_uri}},
88 ],
89 }
90 ],
91 "task": "token_classify",
92 "chat_template": RAW_CONTENT_CHAT_TEMPLATE,
93 }
94
95
96 def post_http_request(payload: dict[str, Any], api_url: str) -> requests.Response:
97 headers = {"User-Agent": "Test Client"}
98 return requests.post(api_url, headers=headers, json=payload)
99
100
101 def parse_response(response: requests.Response) -> dict[str, Any]:
102 try:
103 result = response.json()
104 except ValueError as exc:
105 raise RuntimeError(
106 f"Server returned non-JSON response: {response.text}"
107 ) from exc
108
109 if response.status_code != 200 or "data" not in result:
110 raise RuntimeError(f"Server error ({response.status_code}): {result}")
111
112 return result
113
114
115 def load_timestamp_segment_time(model: str) -> float:
116 model_path = Path(model)
117 config_path = (
118 model_path / "config.json"
119 if model_path.exists()
120 else Path(hf_hub_download(repo_id=model, filename="config.json"))
121 )
122
123 with config_path.open() as f:
124 config = json.load(f)
125
126 return config["timestamp_segment_time"]
127
128
129 def parse_args():
130 parser = argparse.ArgumentParser()
131 parser.add_argument("--host", type=str, default="localhost")
132 parser.add_argument("--port", type=int, default=8000)
133 parser.add_argument(
134 "--model",
135 type=str,
136 default="Qwen/Qwen3-ForcedAligner-0.6B-hf",
137 )
138 parser.add_argument(
139 "--audio-path",
140 type=Path,
141 default=None,
142 help="Optional audio file. Defaults to a 5-second silent WAV.",
143 )
144 parser.add_argument(
145 "--words",
146 nargs="+",
147 default=["Hello", "world"],
148 help="Reference words to align against the audio.",
149 )
150 return parser.parse_args()
151
152
153 def main(args):
154 api_url = f"http://{args.host}:{args.port}/pooling"
155 prompt = build_prompt(args.words)
156 audio_uri = (
157 encode_audio_data_uri(args.audio_path)
158 if args.audio_path
159 else encode_silent_wav_data_uri()
160 )
161 payload = build_payload(args.model, prompt, audio_uri)
162
163 pooling_response = post_http_request(payload=payload, api_url=api_url)
164 result = parse_response(pooling_response)
165
166 timestamp_segment_time = load_timestamp_segment_time(args.model)
167 output = result["data"][0]
168 logits = torch.tensor(output["data"])
169 predictions = logits.argmax(dim=-1)
170
171 expected_timestamps = len(args.words) * 2
172 if len(predictions) != expected_timestamps:
173 raise RuntimeError(
174 f"Expected {expected_timestamps} timestamp predictions, "
175 f"but received {len(predictions)}. Check that the server was started "
176 "with STEP pooling."
177 )
178
179 ts_predictions = [
180 prediction.item() * timestamp_segment_time for prediction in predictions
181 ]
182
183 for i, word in enumerate(args.words):
184 start_ms = ts_predictions[i * 2]
185 end_ms = ts_predictions[i * 2 + 1]
186 print(f"{word:15s} {start_ms / 1000:.3f}s - {end_ms / 1000:.3f}s")
187
188
189 if __name__ == "__main__":
190 args = parse_args()
191 main(args)