60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
from vllm_omni.inputs.preprocess import OmniInputPreprocessor
|
|
|
|
|
|
def _make_preprocessor(monkeypatch):
|
|
preprocessor = object.__new__(OmniInputPreprocessor)
|
|
monkeypatch.setattr(preprocessor, "_truncate_inputs", lambda tokens, tokenization_kwargs=None: tokens)
|
|
monkeypatch.setattr(
|
|
preprocessor,
|
|
"_process_multimodal",
|
|
lambda *args, **kwargs: {"prompt_token_ids": [1, 2, 3]},
|
|
)
|
|
monkeypatch.setattr(preprocessor, "_tokenize_prompt", lambda prompt_text, tokenization_kwargs=None: [9, 8, 7])
|
|
return preprocessor
|
|
|
|
|
|
def test_process_tokens_keeps_additional_information(monkeypatch):
|
|
preprocessor = _make_preprocessor(monkeypatch)
|
|
parsed = {
|
|
"prompt_token_ids": [1, 2, 3],
|
|
"prompt_embeds": "embeds",
|
|
"additional_information": {"task": ["tts"], "lang": ["auto"]},
|
|
}
|
|
|
|
inputs = OmniInputPreprocessor._process_tokens(preprocessor, parsed)
|
|
|
|
assert inputs["prompt_token_ids"] == [1, 2, 3]
|
|
assert inputs["prompt_embeds"] == "embeds"
|
|
assert inputs["additional_information"] == {"task": ["tts"], "lang": ["auto"]}
|
|
|
|
|
|
def test_process_text_keeps_additional_information(monkeypatch):
|
|
preprocessor = _make_preprocessor(monkeypatch)
|
|
parsed = {
|
|
"prompt": "hello",
|
|
"prompt_embeds": "embeds",
|
|
"additional_information": {"speaker": ["alice"]},
|
|
}
|
|
|
|
inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
|
|
|
|
assert inputs["prompt_token_ids"] == [9, 8, 7]
|
|
assert inputs["prompt_embeds"] == "embeds"
|
|
assert inputs["additional_information"] == {"speaker": ["alice"]}
|
|
|
|
|
|
def test_process_text_multimodal_skips_empty_payloads(monkeypatch):
|
|
preprocessor = _make_preprocessor(monkeypatch)
|
|
parsed = {
|
|
"prompt": "hello",
|
|
"multi_modal_data": {"image": "fake"},
|
|
"prompt_embeds": None,
|
|
"additional_information": None,
|
|
}
|
|
|
|
inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
|
|
|
|
assert inputs["prompt_token_ids"] == [1, 2, 3]
|
|
assert "prompt_embeds" not in inputs
|
|
assert "additional_information" not in inputs
|