diff --git a/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py b/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py new file mode 100644 index 000000000..fbf5ac0b0 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import pytest +from huggingface_hub import hf_hub_download +from tokenizers import Tokenizer as HFTokenizer + +from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import ( + _tokenize_with_nonverbal_tags, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _ids(tok: HFTokenizer, s: str) -> list[int]: + return tok.encode(s).ids + + +class TestNonVerbalTags: + @classmethod + def setup_class(cls): + tokenizer_path = hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json") + cls.tokenizer = HFTokenizer.from_file(tokenizer_path) + + def test_plain_text_fallback(self): + tok = self.tokenizer + text = "hello world" + assert _tokenize_with_nonverbal_tags(text, tok) == _ids(tok, text) + + def test_single_nonverbal_tag_at_beginning(self): + tok = self.tokenizer + text = "[laughter]HelloWorld" + expected = _ids(tok, "[laughter]") + _ids(tok, "Hello") + _ids(tok, "World") + assert _tokenize_with_nonverbal_tags(text, tok) == expected + + def test_single_nonverbal_tag_in_middle(self): + tok = self.tokenizer + text = "Hello[laughter]world" + expected = _ids(tok, "Hello") + _ids(tok, "[laughter]") + _ids(tok, "world") + assert _tokenize_with_nonverbal_tags(text, tok) == expected + + def test_multiple_nonverbal_tags(self): + tok = self.tokenizer + text = "[sigh]a[question-en]b[surprise-oh]" + expected = ( + _ids(tok, "[sigh]") + + _ids(tok, "a") + + _ids(tok, "[question-en]") + + _ids(tok, "b") + + _ids(tok, "[surprise-oh]") + ) + assert _tokenize_with_nonverbal_tags(text, tok) == expected + + def test_tag_only(self): + tok = self.tokenizer + text = "[question-en]" + assert _tokenize_with_nonverbal_tags(text, tok) == _ids(tok, text) + + def test_unknown_tag_not_special_cased(self): + tok = self.tokenizer + text = "x[unknown-tag]y" + assert _tokenize_with_nonverbal_tags(text, tok) == _ids(tok, text) diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index 745baae13..9671e2758 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import os +import re from collections.abc import Iterable from typing import ClassVar @@ -53,6 +54,76 @@ def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig): return post_process_func +def _combine_text(text, ref_text: str | None = None) -> str: + # combine with reference text if not None + if ref_text: + full_text = ref_text.strip() + " " + text.strip() + else: + full_text = text.strip() + + # filter out newline / carriage-return characters + full_text = re.sub(r"[\r\n]+", "", full_text) + + # replace Chinese parentheses with English ones + full_text = full_text.replace("\uff08", "(").replace("\uff09", ")") + + # collapse consecutive spaces / tabs into a single space + full_text = re.sub(r"[ \t]+", " ", full_text) + + # remove spaces around chinese characters + chinese_range = r"[\u4e00-\u9fff]" + pattern = rf"(?<={chinese_range})\s+|\s+(?={chinese_range})" + full_text = re.sub(pattern, "", full_text) + + return full_text + + +_NONVERBAL_PATTERN = re.compile( + r"\[(laughter|sigh|confirmation-en|question-en|question-ah|question-oh|" + r"question-ei|question-yi|surprise-ah|surprise-oh|surprise-wa|" + r"surprise-yo|dissatisfaction-hnn)\]" +) + + +def _tokenize_with_nonverbal_tags(text: str, tokenizer) -> list[int]: + """Tokenize text containing non-verbal tags, handling each tag independently. + + Non-verbal tags are tokenized standalone to guarantee consistent token + IDs regardless of surrounding language context (Chinese, English, etc.). + + Args: + text: Full text string potentially containing non-verbal tags. + tokenizer: HuggingFace text tokenizer instance. + Returns: + Token IDs list of length seq_len. + """ + parts = [] + last_end = 0 + for m in _NONVERBAL_PATTERN.finditer(text): + if m.start() > last_end: + segment = text[last_end : m.start()] + ids = tokenizer.encode(segment) + if ids: + parts.append(ids) + tag_ids = tokenizer.encode(m.group()) + if tag_ids: + parts.append(tag_ids) + last_end = m.end() + if last_end < len(text): + segment = text[last_end:] + ids = tokenizer.encode(segment) + if ids: + parts.append(ids) + + if not parts: + return tokenizer.encode(text).ids + else: + combined = [] + for p in parts: + combined.extend(p.ids) + return combined + + class OmniVoicePipeline(nn.Module, SupportAudioOutput): """OmniVoice text-to-speech pipeline for the diffusion engine. @@ -212,14 +283,13 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): target_len = max(1, int(target_len)) # Build text prompt with control tokens - style = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" - if ref_text: - full_text = f"{ref_text} {text}" - else: - full_text = text - full_prompt = f"{style}<|text_start|>{full_text}<|text_end|>" - encoding = self.tokenizer.encode(full_prompt) - text_tokens = torch.tensor(encoding.ids, dtype=torch.long, device=device) + style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" + full_text = _combine_text(ref_text=ref_text, text=text) + wrapped_text = f"<|text_start|>{full_text}<|text_end|>" + style_tokens = self.tokenizer.encode(style_text).ids + text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) + encoding_ids = style_tokens + text_tokens + text_tokens = torch.tensor(encoding_ids, dtype=torch.long, device=device) text_len = text_tokens.shape[0] # Encode reference audio tokens if provided (with voice caching)