feat: add hotwords support for vLLM ASR
This commit is contained in:
YingboHAO
+107
-91
@@ -1,14 +1,23 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test VibeVoice vLLM API with Streaming (Real-time output).
|
||||
Test VibeVoice vLLM API with Streaming and Optional Hotwords Support.
|
||||
|
||||
This script tests ASR transcription via the vLLM OpenAI-compatible API.
|
||||
By default, it runs standard transcription without hotwords.
|
||||
|
||||
Optionally, you can provide hotwords (context_info) to improve recognition
|
||||
of domain-specific content like proper nouns, technical terms, and speaker names.
|
||||
Hotwords are embedded in the prompt as "with extra info: {hotwords}".
|
||||
|
||||
Usage:
|
||||
python test_api.py [audio_path] [--url URL]
|
||||
python test_api_with_hotwords.py [audio_path] [--url URL] [--hotwords "word1,word2"]
|
||||
|
||||
Examples:
|
||||
python test_api.py # Use default audio
|
||||
python test_api.py /path/to/audio.wav # Specify audio file
|
||||
python test_api.py /path/to/audio.mp3 --url http://localhost:8000 # Custom URL
|
||||
# Standard transcription (no hotwords)
|
||||
python3 test_api.py audio.wav
|
||||
|
||||
# With hotwords for better recognition of specific terms
|
||||
python3 test_api.py audio.wav --hotwords "Microsoft,Azure,VibeVoice"
|
||||
"""
|
||||
import requests
|
||||
import json
|
||||
@@ -21,38 +30,38 @@ import argparse
|
||||
|
||||
|
||||
def _guess_mime_type(path: str) -> str:
|
||||
"""Guess MIME type from file extension."""
|
||||
ext = os.path.splitext(path)[1].lower()
|
||||
if ext == ".wav":
|
||||
return "audio/wav"
|
||||
if ext in (".mp3",):
|
||||
return "audio/mpeg"
|
||||
if ext in (".m4a",):
|
||||
return "audio/mp4"
|
||||
if ext in (".mp4", ".m4v", ".mov", ".webm"):
|
||||
return "video/mp4"
|
||||
if ext in (".flac",):
|
||||
return "audio/flac"
|
||||
if ext in (".ogg", ".opus"):
|
||||
return "audio/ogg"
|
||||
return "application/octet-stream"
|
||||
mime_map = {
|
||||
".wav": "audio/wav",
|
||||
".mp3": "audio/mpeg",
|
||||
".m4a": "audio/mp4",
|
||||
".mp4": "video/mp4",
|
||||
".flac": "audio/flac",
|
||||
".ogg": "audio/ogg",
|
||||
".opus": "audio/ogg",
|
||||
}
|
||||
return mime_map.get(ext, "application/octet-stream")
|
||||
|
||||
|
||||
def _get_duration_seconds_ffprobe(path: str) -> float:
|
||||
"""Get audio duration using ffprobe."""
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-show_entries",
|
||||
"format=duration",
|
||||
"-of",
|
||||
"default=noprint_wrappers=1:nokey=1",
|
||||
"ffprobe", "-v", "error",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "default=noprint_wrappers=1:nokey=1",
|
||||
path,
|
||||
]
|
||||
out = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8").strip()
|
||||
return float(out)
|
||||
|
||||
|
||||
def _is_video_file(path: str) -> bool:
|
||||
"""Check if the file is a video file that needs audio extraction."""
|
||||
ext = os.path.splitext(path)[1].lower()
|
||||
return ext in (".mp4", ".m4v", ".mov", ".webm", ".avi", ".mkv")
|
||||
|
||||
|
||||
def _extract_audio_from_video(video_path: str) -> str:
|
||||
"""
|
||||
Extract audio from video file (mp4/mov/webm) to a temporary mp3 file.
|
||||
@@ -74,26 +83,40 @@ def _extract_audio_from_video(video_path: str) -> str:
|
||||
return audio_path
|
||||
|
||||
|
||||
def _is_video_file(path: str) -> bool:
|
||||
"""Check if the file is a video file that needs audio extraction."""
|
||||
ext = os.path.splitext(path)[1].lower()
|
||||
return ext in (".mp4", ".m4v", ".mov", ".webm", ".avi", ".mkv")
|
||||
|
||||
|
||||
def test_transcription(audio_path: str, base_url: str = "http://localhost:8000"):
|
||||
"""Test ASR transcription with streaming output."""
|
||||
def test_transcription_with_hotwords(
|
||||
audio_path: str,
|
||||
context_info: str = None,
|
||||
base_url: str = "http://localhost:8000",
|
||||
):
|
||||
"""
|
||||
Test ASR transcription with customized hotwords.
|
||||
|
||||
print(f"Loading audio from: {audio_path}")
|
||||
Hotwords are embedded in the prompt text as "with extra info: {hotwords}".
|
||||
This helps the model recognize domain-specific terms more accurately.
|
||||
|
||||
Args:
|
||||
audio_path: Path to the audio file
|
||||
context_info: Hotwords string (e.g., "Microsoft,Azure,VibeVoice")
|
||||
base_url: vLLM server URL
|
||||
"""
|
||||
|
||||
print(f"=" * 70)
|
||||
print(f"Testing Customized Hotwords Support")
|
||||
print(f"=" * 70)
|
||||
print(f"Input file: {audio_path}")
|
||||
print(f"Hotwords: {context_info or '(none)'}")
|
||||
print()
|
||||
|
||||
# Handle video files: extract audio first
|
||||
temp_audio_path = None
|
||||
actual_audio_path = audio_path
|
||||
if _is_video_file(audio_path):
|
||||
print(f"Detected video file, extracting audio...")
|
||||
print(f"🎬 Detected video file, extracting audio...")
|
||||
temp_audio_path = _extract_audio_from_video(audio_path)
|
||||
actual_audio_path = temp_audio_path
|
||||
print(f"Audio extracted to: {temp_audio_path}")
|
||||
print(f"✅ Audio extracted to: {temp_audio_path}")
|
||||
|
||||
# Load audio
|
||||
try:
|
||||
duration = _get_duration_seconds_ffprobe(actual_audio_path)
|
||||
print(f"Audio duration: {duration:.2f} seconds")
|
||||
@@ -106,16 +129,30 @@ def test_transcription(audio_path: str, base_url: str = "http://localhost:8000")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error preparing audio: {e}")
|
||||
# Cleanup temp file if created
|
||||
if temp_audio_path and os.path.exists(temp_audio_path):
|
||||
os.remove(temp_audio_path)
|
||||
return
|
||||
|
||||
# Build the request
|
||||
url = f"{base_url}/v1/chat/completions"
|
||||
|
||||
show_keys = ["Start time", "End time", "Speaker ID", "Content"]
|
||||
prompt_text = (
|
||||
f"This is a {duration:.2f} seconds audio, please transcribe it with these keys: "
|
||||
+ ", ".join(show_keys)
|
||||
)
|
||||
|
||||
# Build prompt with optional hotwords
|
||||
# Hotwords are embedded as "with extra info: {hotwords}" in the prompt
|
||||
if context_info and context_info.strip():
|
||||
prompt_text = (
|
||||
f"This is a {duration:.2f} seconds audio, with extra info: {context_info.strip()}\n\n"
|
||||
f"Please transcribe it with these keys: " + ", ".join(show_keys)
|
||||
)
|
||||
print(f"\n📝 Hotwords embedded in prompt: '{context_info}'")
|
||||
else:
|
||||
prompt_text = (
|
||||
f"This is a {duration:.2f} seconds audio, please transcribe it with these keys: "
|
||||
+ ", ".join(show_keys)
|
||||
)
|
||||
print(f"\n📝 No hotwords provided")
|
||||
|
||||
mime = _guess_mime_type(actual_audio_path)
|
||||
data_url = f"data:{mime};base64,{audio_b64}"
|
||||
@@ -139,20 +176,19 @@ def test_transcription(audio_path: str, base_url: str = "http://localhost:8000")
|
||||
"temperature": 0.0,
|
||||
"stream": True,
|
||||
"top_p": 1.0,
|
||||
"repetition_penalty": 1.0,
|
||||
}
|
||||
|
||||
print(f"\nSending request to {url} (Streaming Mode)...")
|
||||
print(f"Prompt: {prompt_text}")
|
||||
print("-" * 60)
|
||||
print(f"\n{'=' * 70}")
|
||||
print(f"Sending request to {url}")
|
||||
print(f"{'=' * 70}")
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
|
||||
response = requests.post(url, json=payload, stream=True, timeout=12000)
|
||||
|
||||
if response.status_code == 200:
|
||||
print("Response received. Streaming content:\n")
|
||||
print("\n✅ Response received. Streaming content:\n")
|
||||
print("-" * 50)
|
||||
|
||||
printed = ""
|
||||
for line in response.iter_lines():
|
||||
@@ -162,92 +198,72 @@ def test_transcription(audio_path: str, base_url: str = "http://localhost:8000")
|
||||
if decoded_line.startswith("data: "):
|
||||
json_str = decoded_line[6:]
|
||||
if json_str.strip() == "[DONE]":
|
||||
print("\n\n[Finished]")
|
||||
print("\n" + "-" * 50)
|
||||
print("✅ [Finished]")
|
||||
break
|
||||
try:
|
||||
data = json.loads(json_str)
|
||||
|
||||
delta = data['choices'][0]['delta']
|
||||
content = delta.get('content', '')
|
||||
if content:
|
||||
|
||||
# vLLM/OpenAI-compatible streams may emit either
|
||||
# incremental deltas OR the full accumulated text.
|
||||
# Only print the newly-added part to avoid repeats.
|
||||
if content.startswith(printed):
|
||||
to_print = content[len(printed):]
|
||||
else:
|
||||
to_print = content
|
||||
|
||||
if to_print:
|
||||
print(to_print, end='', flush=True)
|
||||
printed += to_print
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
else:
|
||||
print(f"Error: {response.status_code}")
|
||||
print(f"❌ Error: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
print("\nRequest timed out!")
|
||||
print("❌ Request timed out!")
|
||||
except Exception as e:
|
||||
print(f"\nError: {e}")
|
||||
print(f"❌ Error: {e}")
|
||||
|
||||
print(f"\n{'-'*60}")
|
||||
print(f"Total time elapsed: {time.time() - t0:.2f}s")
|
||||
elapsed = time.time() - t0
|
||||
print(f"\n{'=' * 70}")
|
||||
print(f"⏱️ Total time elapsed: {elapsed:.2f}s")
|
||||
print(f"📊 RTF (Real-Time Factor): {elapsed / duration:.2f}x")
|
||||
print(f"{'=' * 70}")
|
||||
|
||||
# Cleanup temp audio file if created
|
||||
if temp_audio_path and os.path.exists(temp_audio_path):
|
||||
os.remove(temp_audio_path)
|
||||
print(f"Cleaned up temp file: {temp_audio_path}")
|
||||
print(f"🗑️ Cleaned up temp file: {temp_audio_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Test VibeVoice vLLM API with streaming output"
|
||||
description="Test VibeVoice vLLM API with Customized Hotwords"
|
||||
)
|
||||
parser.add_argument(
|
||||
"audio_path",
|
||||
nargs="?",
|
||||
default=None,
|
||||
help="Path to audio file (wav, mp3, flac, etc.) or video file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--url",
|
||||
default="http://localhost:8000",
|
||||
help="vLLM server base URL (default: http://localhost:8000)"
|
||||
help="vLLM server URL (default: http://localhost:8000)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hotwords",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Hotwords to improve recognition (e.g., 'Microsoft,Azure,VibeVoice')"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Find default audio if not specified
|
||||
audio_path = args.audio_path
|
||||
if audio_path is None:
|
||||
# Try to find a sample audio in common locations
|
||||
possible_paths = [
|
||||
# In VibeVoice demo folder
|
||||
os.path.join(os.path.dirname(__file__), "..", "..", "demo", "voices", "en-Carter_man.wav"),
|
||||
os.path.join(os.path.dirname(__file__), "..", "..", "demo", "voices", "zh-Anchen_man_bgm.wav"),
|
||||
# Relative to current directory
|
||||
"demo/voices/en-Carter_man.wav",
|
||||
"demo/voices/zh-Anchen_man_bgm.wav",
|
||||
]
|
||||
|
||||
for path in possible_paths:
|
||||
if os.path.exists(path):
|
||||
audio_path = path
|
||||
break
|
||||
|
||||
if audio_path is None:
|
||||
print("Error: No audio file specified and no default audio found.")
|
||||
print("Usage: python test_api.py <audio_path>")
|
||||
sys.exit(1)
|
||||
|
||||
if not os.path.exists(audio_path):
|
||||
print(f"Error: Audio file not found: {audio_path}")
|
||||
sys.exit(1)
|
||||
|
||||
test_transcription(audio_path, args.url)
|
||||
# Run test
|
||||
test_transcription_with_hotwords(
|
||||
audio_path=args.audio_path,
|
||||
context_info=args.hotwords,
|
||||
base_url=args.url,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user