wangyu
e338f0ece3
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] <your.email@example.com>
130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""
|
|
E2E expansion tests for MiMo-Audio online serving (nightly Omni · Function Test with H100 · Single-GPU).
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.media import generate_synthetic_audio
|
|
from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
|
|
from tests.helpers.stage_config import get_deploy_config_path
|
|
from vllm_omni.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf_specific,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.slow, pytest.mark.omni]
|
|
|
|
MIMO_AUDIO_TOKENIZER_REPO = "XiaomiMiMo/MiMo-Audio-Tokenizer"
|
|
CHAT_TEMPLATE_PATH = str(
|
|
Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "mimo_audio" / "chat_template.jinja"
|
|
)
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
|
|
models = ["XiaomiMiMo/MiMo-Audio-7B-Instruct"]
|
|
|
|
|
|
def download_tokenizer():
|
|
tokenizer_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", MIMO_AUDIO_TOKENIZER_REPO)
|
|
if os.path.exists(tokenizer_path):
|
|
return tokenizer_path
|
|
local_path = download_weights_from_hf_specific(
|
|
model_name_or_path=MIMO_AUDIO_TOKENIZER_REPO,
|
|
cache_dir=None,
|
|
allow_patterns=["*"],
|
|
require_all=True,
|
|
)
|
|
return local_path
|
|
|
|
|
|
try:
|
|
stage_configs = [get_deploy_config_path("mimo_audio.yaml")]
|
|
tokenizer_path = download_tokenizer()
|
|
os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
|
|
|
|
test_params = [
|
|
OmniServerParams(
|
|
model=model,
|
|
stage_config_path=stage_config,
|
|
server_args=["--chat-template", CHAT_TEMPLATE_PATH],
|
|
)
|
|
for model in models
|
|
for stage_config in stage_configs
|
|
]
|
|
except Exception as exc:
|
|
pytest.skip(
|
|
f"MiMo-Audio expansion tests skipped: module setup failed ({type(exc).__name__}: {exc})",
|
|
allow_module_level=True,
|
|
)
|
|
|
|
|
|
def get_prompt(prompt_type="text_only"):
|
|
prompts = {
|
|
"text_only": "What is the capital of China? Answer in 20 words.",
|
|
"audio": "What one English word is repeated in the audio?",
|
|
}
|
|
return prompts.get(prompt_type, prompts["text_only"])
|
|
|
|
|
|
def get_max_batch_size(size_type="few"):
|
|
batch_sizes = {"few": 5, "medium": 100, "large": 256}
|
|
return batch_sizes.get(size_type, 5)
|
|
|
|
|
|
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
|
def test_text_to_text_001(omni_server, online_client) -> None:
|
|
"""
|
|
Test text input processing and text-only output generation via OpenAI API.
|
|
Deploy Setting: default yaml
|
|
Input Modal: text
|
|
Output Modal: text
|
|
Datasets: few requests
|
|
"""
|
|
messages = dummy_messages_from_mix_data(content_text=get_prompt())
|
|
|
|
request_config = {
|
|
"model": omni_server.model,
|
|
"messages": messages,
|
|
"stream": False,
|
|
"modalities": ["text"],
|
|
"key_words": {"text": ["beijing"]},
|
|
}
|
|
|
|
online_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
|
|
|
|
|
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
|
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
|
@pytest.mark.skip(reason="CI failed 8571")
|
|
def test_audio_to_text_audio_001(omni_server, online_client) -> None:
|
|
"""
|
|
Test audio and text input processing and text/audio output generation via OpenAI API.
|
|
Deploy Setting: default yaml
|
|
Input Modal: text + audio
|
|
Output Modal: text + audio
|
|
Input Setting: stream=True
|
|
Datasets: single request
|
|
"""
|
|
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1, sample_rate=24000)['base64']}"
|
|
messages = dummy_messages_from_mix_data(
|
|
audio_data_url=audio_data_url,
|
|
content_text=get_prompt("audio"),
|
|
)
|
|
|
|
request_config = {
|
|
"model": omni_server.model,
|
|
"messages": messages,
|
|
"stream": True,
|
|
"sampling_params_list": [{"max_tokens": 64}, {"max_tokens": 64}],
|
|
"key_words": {
|
|
"audio": ["test"],
|
|
},
|
|
}
|
|
|
|
online_client.send_omni_request(request_config)
|