102 lines
2.6 KiB
Python
102 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.metrics.utils import (
|
|
count_audio_chunk_frames,
|
|
count_audio_frames,
|
|
count_image_pixels,
|
|
count_tokens_from_outputs,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("audio_chunk", "expected"),
|
|
[
|
|
(SimpleNamespace(shape=(2, 320)), 320),
|
|
([0.0, 0.0, 0.0], 3),
|
|
(1.0, 1),
|
|
],
|
|
)
|
|
def test_count_audio_chunk_frames(audio_chunk: object, expected: int) -> None:
|
|
assert count_audio_chunk_frames(audio_chunk) == expected
|
|
|
|
|
|
def test_count_audio_frames_sums_audio_chunks() -> None:
|
|
mm_out = {
|
|
"audio": [
|
|
SimpleNamespace(shape=(1, 120)),
|
|
SimpleNamespace(shape=(2, 80)),
|
|
]
|
|
}
|
|
|
|
assert count_audio_frames(mm_out) == 200
|
|
|
|
|
|
def test_count_audio_frames_supports_model_outputs() -> None:
|
|
mm_out = {
|
|
"model_outputs": [
|
|
SimpleNamespace(shape=(1, 120)),
|
|
SimpleNamespace(shape=(2, 80)),
|
|
]
|
|
}
|
|
|
|
assert count_audio_frames(mm_out) == 200
|
|
|
|
|
|
@pytest.mark.parametrize("mm_out", [{}, {"audio": None}, {"model_outputs": None}])
|
|
def test_count_audio_frames_returns_zero_without_audio(mm_out: dict[str, object]) -> None:
|
|
assert count_audio_frames(mm_out) == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[
|
|
(None, 0),
|
|
(object(), 0),
|
|
(SimpleNamespace(size=(10, 20)), 200),
|
|
(SimpleNamespace(shape=(2, 3, 4, 5)), 40),
|
|
(SimpleNamespace(shape=(3, 4, 5)), 20),
|
|
(SimpleNamespace(shape=(6, 7, 3)), 42),
|
|
(SimpleNamespace(shape=(6, 7)), 42),
|
|
],
|
|
)
|
|
def test_count_image_pixels(value: object, expected: int) -> None:
|
|
assert count_image_pixels(value) == expected
|
|
|
|
|
|
def test_count_image_pixels_sums_nested_values() -> None:
|
|
images = [
|
|
SimpleNamespace(size=(10, 20)),
|
|
(
|
|
SimpleNamespace(shape=(3, 4, 5)),
|
|
SimpleNamespace(shape=(6, 7, 3)),
|
|
),
|
|
]
|
|
|
|
assert count_image_pixels(images) == 262
|
|
|
|
|
|
def test_count_tokens_from_outputs() -> None:
|
|
engine_outputs = [
|
|
SimpleNamespace(
|
|
outputs=[
|
|
SimpleNamespace(cumulative_token_ids=[1, 2, 3], token_ids=[3]),
|
|
SimpleNamespace(cumulative_token_ids=None, token_ids=[4, 5]),
|
|
SimpleNamespace(cumulative_token_ids=None, token_ids=None),
|
|
]
|
|
),
|
|
SimpleNamespace(outputs=None),
|
|
SimpleNamespace(
|
|
outputs=[
|
|
SimpleNamespace(cumulative_token_ids=[6], token_ids=[6]),
|
|
]
|
|
),
|
|
]
|
|
|
|
assert count_tokens_from_outputs(engine_outputs) == 6
|