wehub-resource-sync
221 lines
7.5 KiB
Python
221 lines
7.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Unit tests for speech token-usage accounting (issue #4646).
|
|
|
|
These exercise the pure usage logic without a model/server: input-token
|
|
breakdown (text + reference-audio gating) and output-token accumulation.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.entrypoints.openai.speech_usage import (
|
|
SpeechOutputTokenCounter,
|
|
build_speech_usage,
|
|
gate_audio_tokens,
|
|
qwen3_tts_input_token_details,
|
|
)
|
|
|
|
|
|
# A deterministic stand-in for a real tokenizer: 1 token per whitespace word.
|
|
def word_count(text: str) -> int:
|
|
return len(text.split())
|
|
|
|
|
|
# --- input breakdown: gating ------------------------------------------------
|
|
|
|
|
|
def test_gate_audio_tokens_base_in_context_counts_ref_frames():
|
|
assert (
|
|
gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120)
|
|
== 120
|
|
)
|
|
|
|
|
|
def test_gate_audio_tokens_customvoice_is_zero():
|
|
assert (
|
|
gate_audio_tokens(
|
|
task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120
|
|
)
|
|
== 0
|
|
)
|
|
|
|
|
|
def test_gate_audio_tokens_x_vector_only_is_zero_even_with_ref():
|
|
# x-vector cloning inserts NO codec frames, so audio_tokens must be 0
|
|
# even though ref_code_length is populated.
|
|
assert (
|
|
gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=None, ref_code_length=120) == 0
|
|
)
|
|
|
|
|
|
def test_gate_audio_tokens_icl_override_wins():
|
|
assert (
|
|
gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=False, ref_code_length=120) == 0
|
|
)
|
|
assert (
|
|
gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=True, ref_code_length=80) == 80
|
|
)
|
|
|
|
|
|
def test_gate_audio_tokens_icl_override_ignored_for_non_base():
|
|
# icl_mode_override is honored only for the Base task; CustomVoice stays 0
|
|
# even if such a flag leaks into its params (mirrors the estimator).
|
|
assert (
|
|
gate_audio_tokens(
|
|
task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=True, ref_code_length=120
|
|
)
|
|
== 0
|
|
)
|
|
|
|
|
|
def test_gate_audio_tokens_bad_ref_length_is_zero():
|
|
assert (
|
|
gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length="oops")
|
|
== 0
|
|
)
|
|
|
|
|
|
# --- input breakdown: Qwen3-TTS reproduces issue #4646 observation ----------
|
|
|
|
|
|
def test_customvoice_text_scales_with_input_no_audio():
|
|
short = qwen3_tts_input_token_details(
|
|
input_text="hi there",
|
|
instructions=None,
|
|
tts_params={"task_type": ["CustomVoice"]},
|
|
count_text_tokens=word_count,
|
|
)
|
|
long = qwen3_tts_input_token_details(
|
|
input_text="hi there this is a much longer sentence",
|
|
instructions=None,
|
|
tts_params={"task_type": ["CustomVoice"]},
|
|
count_text_tokens=word_count,
|
|
)
|
|
assert short.audio_tokens == 0 and long.audio_tokens == 0
|
|
assert long.text_tokens > short.text_tokens # scales with input
|
|
|
|
|
|
def test_base_icl_text_tracks_input_audio_tracks_ref():
|
|
# The core #4646 fix: text_tokens reflect the input (not dropped), and
|
|
# audio_tokens track the reference audio independently.
|
|
params = {"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]}
|
|
short = qwen3_tts_input_token_details(
|
|
input_text="hi", instructions=None, tts_params=params, count_text_tokens=word_count
|
|
)
|
|
long = qwen3_tts_input_token_details(
|
|
input_text="hi there friend", instructions=None, tts_params=params, count_text_tokens=word_count
|
|
)
|
|
# text now varies with input (the bug was that it was dropped for Base)
|
|
assert long.text_tokens > short.text_tokens
|
|
# audio fixed by the reference, independent of input text
|
|
assert short.audio_tokens == long.audio_tokens == 100
|
|
|
|
|
|
def test_instructions_count_toward_text():
|
|
without = qwen3_tts_input_token_details(
|
|
input_text="hello world",
|
|
instructions=None,
|
|
tts_params={"task_type": ["CustomVoice"]},
|
|
count_text_tokens=word_count,
|
|
)
|
|
with_instr = qwen3_tts_input_token_details(
|
|
input_text="hello world",
|
|
instructions="speak slowly and warmly",
|
|
tts_params={"task_type": ["CustomVoice"]},
|
|
count_text_tokens=word_count,
|
|
)
|
|
assert with_instr.text_tokens == without.text_tokens + word_count("speak slowly and warmly")
|
|
|
|
|
|
# --- build usage ------------------------------------------------------------
|
|
|
|
|
|
def test_build_speech_usage_aggregates_and_totals():
|
|
details = qwen3_tts_input_token_details(
|
|
input_text="hi there",
|
|
instructions=None,
|
|
tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]},
|
|
count_text_tokens=word_count,
|
|
)
|
|
usage = build_speech_usage(details, output_tokens=250)
|
|
assert usage.input_tokens == details.text_tokens + 100
|
|
assert usage.output_tokens == 250
|
|
assert usage.total_tokens == usage.input_tokens + 250
|
|
assert usage.input_token_details.text_tokens == 2
|
|
assert usage.input_token_details.audio_tokens == 100
|
|
|
|
|
|
def test_batch_item_result_serializes_usage():
|
|
# Guards the batch per-item usage carrier (and that it is optional/None on error).
|
|
from vllm_omni.entrypoints.openai.protocol.audio import SpeechBatchItemResult
|
|
|
|
details = qwen3_tts_input_token_details(
|
|
input_text="hello world",
|
|
instructions=None,
|
|
tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [50]},
|
|
count_text_tokens=word_count,
|
|
)
|
|
item = SpeechBatchItemResult(
|
|
index=0,
|
|
status="success",
|
|
audio_data="x",
|
|
media_type="audio/wav",
|
|
usage=build_speech_usage(details, output_tokens=30),
|
|
)
|
|
dumped = item.model_dump()
|
|
assert dumped["usage"]["input_tokens"] == 2 + 50
|
|
assert dumped["usage"]["output_tokens"] == 30
|
|
assert dumped["usage"]["input_token_details"] == {"text_tokens": 2, "audio_tokens": 50}
|
|
# Errored items carry no usage.
|
|
assert SpeechBatchItemResult(index=1, status="error", error="boom").model_dump()["usage"] is None
|
|
|
|
|
|
# --- output token accumulation ----------------------------------------------
|
|
# The generated codec-token count is read from per-stage pipeline metrics on the
|
|
# final output: res.metrics["stage_metrics"][<stage>]["num_tokens_out"]. Only the
|
|
# AR stage (0) is non-zero; code2wav (stage 1) reports 0.
|
|
|
|
|
|
@dataclass
|
|
class _FakeRes:
|
|
metrics: dict
|
|
|
|
|
|
def _final_res(stage0_out, stage1_out=0):
|
|
return _FakeRes(
|
|
metrics={
|
|
"stage_metrics": {
|
|
"0": {"num_tokens_in": 110, "num_tokens_out": stage0_out},
|
|
"1": {"num_tokens_in": 0, "num_tokens_out": stage1_out},
|
|
}
|
|
}
|
|
)
|
|
|
|
|
|
def test_output_tokens_from_stage_metrics():
|
|
acc = SpeechOutputTokenCounter()
|
|
# Intermediate outputs carry empty metrics -> contribute nothing.
|
|
acc.observe(_FakeRes(metrics={}))
|
|
acc.observe(_FakeRes(metrics={}))
|
|
# Final output carries stage metrics; stage-0 num_tokens_out is the count.
|
|
acc.observe(_final_res(stage0_out=77))
|
|
assert acc.total() == 77
|
|
|
|
|
|
def test_output_tokens_takes_max_across_stages():
|
|
acc = SpeechOutputTokenCounter()
|
|
acc.observe(_final_res(stage0_out=42, stage1_out=0))
|
|
assert acc.total() == 42
|
|
|
|
|
|
def test_output_tokens_zero_without_metrics():
|
|
acc = SpeechOutputTokenCounter()
|
|
acc.observe(_FakeRes(metrics={}))
|
|
acc.observe(_FakeRes(metrics={"stage_metrics": {}}))
|
|
assert acc.total() == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(pytest.main([__file__, "-v"]))
|