Sy03andYueqian Lin
3bc3f1a7d0
Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _Processor:
|
|
"""Segment API of the model's processor with a deterministic tokenizer."""
|
|
|
|
def __init__(self):
|
|
self.model_config = SimpleNamespace(
|
|
n_vq=3,
|
|
im_start_token_id=1,
|
|
im_end_token_id=2,
|
|
audio_start_token_id=3,
|
|
audio_pad_token_id=1024,
|
|
audio_assistant_slot_token_id=5,
|
|
)
|
|
self.encoded = []
|
|
|
|
def _encode_text(self, text):
|
|
self.encoded.append(text)
|
|
return [100 + ord(c) % 50 for c in text]
|
|
|
|
def _user_prompt_prefix_ids(self):
|
|
return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
|
|
|
|
def _user_prompt_after_reference_ids(self, language, fields):
|
|
return self._encode_text(f"lang={fields.get('language', language)}")
|
|
|
|
def _assistant_prompt_prefix_ids(self):
|
|
return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
|
|
|
|
def reference(self, text, language, codes):
|
|
# Mirrors the processor's _build_continuation_codes + _pad (one item).
|
|
fields = {} if language is None else {"language": language}
|
|
ids = (
|
|
self._user_prompt_prefix_ids()
|
|
+ self._encode_text("None")
|
|
+ self._user_prompt_after_reference_ids(language, fields)
|
|
+ self._encode_text(text)
|
|
+ self._assistant_prompt_prefix_ids()
|
|
+ [3]
|
|
)
|
|
text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
|
|
text_rows[:, 0] = torch.tensor(ids)
|
|
audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
|
|
audio_rows[:, 0] = 5
|
|
audio_rows[:, 1:] = codes
|
|
unified = torch.cat([text_rows, audio_rows])
|
|
return unified[:, 0].tolist(), unified[:, 1:].contiguous()
|
|
|
|
|
|
@pytest.mark.parametrize("language", [None, "English"])
|
|
def test_matches_processor_and_caches_constant_segments(language):
|
|
proc = _Processor()
|
|
codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
|
|
kwargs = {"text": "ref words target words"}
|
|
if language is not None:
|
|
kwargs["language"] = language
|
|
ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
|
|
expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
|
|
assert ids == expected_ids
|
|
assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
|
|
|
|
proc.encoded.clear()
|
|
_local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
|
|
assert proc.encoded == ["other"]
|
|
|
|
|
|
def test_unsupported_layouts_fall_back():
|
|
proc = _Processor()
|
|
codes = torch.zeros((2, 3), dtype=torch.long)
|
|
assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
|
|
assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
|
|
assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
|
|
assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None
|