WildAi
dc29843483
Vendored tree mocked at collection, so the suite runs green without perth, diffusers, s3tokenizer or conformer. Frozen widget snapshot pins node ids and widget order.
380 lines
16 KiB
Python
380 lines
16 KiB
Python
"""S8.3 -- the tag engine: five features, in one order, that order load-bearing.
|
|
|
|
A tag engine that runs its passes in the wrong order does not crash. It quietly
|
|
produces different speech: ``[Alice] <giggle>`` loses its character if the
|
|
angle-bracket pass runs first, ``[de:Alice]`` becomes a voice called
|
|
``"de:Alice"`` if the character pass claims it, and ``[pause:1.5]`` says the
|
|
word "pause" out loud. None of that is catchable except by asserting the *order*,
|
|
which is why the ordering tests below are about which stage produced a value and
|
|
not only about what the final text says.
|
|
|
|
Every documented tag form is parsed here, the three interchangeable pause words
|
|
are parsed with each documented unit, unsupported per-segment parameters are
|
|
shown to be ignored rather than fatal, and ``auto``'s chunk-method selection is
|
|
pinned at every boundary the corpus names -- including 500, where the strict
|
|
comparison makes the answer surprising and therefore worth a test.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from ComfyUI_Chatterbox.modules.text_processing import (
|
|
AUTO_CHUNK_THRESHOLDS,
|
|
CHUNK_COMBINATION_METHODS,
|
|
CRASH_PROTECTION_TEMPLATE,
|
|
MIN_SAFE_SEGMENT_CHARS,
|
|
PARAM_ALIASES,
|
|
TURBO_TAGS,
|
|
V2_SPECIAL_TOKENS,
|
|
apply_crash_protection,
|
|
chunk_text,
|
|
parse_duration,
|
|
plan_text,
|
|
resolve_v2_tokens,
|
|
select_chunk_method,
|
|
split_segments,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Durations and pause tags (stage 3)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDurations:
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[
|
|
("1.5", 1.5),
|
|
("2s", 2.0),
|
|
("500ms", 0.5),
|
|
("800MS", 0.8),
|
|
(" 1200 ms ", 1.2),
|
|
("3", 3.0),
|
|
("0.25s", 0.25),
|
|
(".5", 0.5),
|
|
],
|
|
)
|
|
def test_documented_spellings_parse(self, value: str, expected: float) -> None:
|
|
assert parse_duration(value) == pytest.approx(expected)
|
|
|
|
@pytest.mark.parametrize("value", ["", "abc", "5 parsecs", "-1", "1.2.3"])
|
|
def test_nonsense_is_refused(self, value: str) -> None:
|
|
"""A refusal is caught by the caller, which leaves the tag as literal text."""
|
|
with pytest.raises(ValueError):
|
|
parse_duration(value)
|
|
|
|
|
|
class TestPauseTagsAreInterchangeable:
|
|
"""``pause`` / ``wait`` / ``stop`` are three spellings of one tag."""
|
|
|
|
@pytest.mark.parametrize("word", ["pause", "wait", "stop"])
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[("1.5", 1.5), ("2s", 2.0), ("500ms", 0.5)],
|
|
)
|
|
def test_all_three_words_read_all_three_units(
|
|
self, word: str, value: str, expected: float
|
|
) -> None:
|
|
segments = split_segments(f"Hello there. [{word}:{value}] And more words.")
|
|
assert len(segments) == 2, "a pause opens a new segment"
|
|
assert segments[0].pause_after == pytest.approx(expected)
|
|
assert segments[0].text == "Hello there."
|
|
assert segments[1].text == "And more words."
|
|
assert segments[1].pause_after == 0.0
|
|
|
|
def test_the_word_is_case_insensitive(self) -> None:
|
|
assert split_segments("Hi. [PAUSE:2s] There.")[0].pause_after == 2.0
|
|
|
|
def test_an_unreadable_duration_is_kept_as_text_not_fatal(self) -> None:
|
|
segments = split_segments("Hi. [pause:soon] There.")
|
|
assert len(segments) == 1
|
|
assert "[pause:soon]" in segments[0].text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Character switching (stage 1) and language tags (stage 2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCharacterTags:
|
|
def test_a_bare_name_opens_a_segment(self) -> None:
|
|
segments = split_segments("[Alice] Hi there, Bob.")
|
|
assert len(segments) == 1
|
|
assert segments[0].character == "Alice"
|
|
assert segments[0].text == "Hi there, Bob."
|
|
|
|
def test_an_untagged_run_stays_in_the_same_segment(self) -> None:
|
|
assert len(split_segments("Just talking, with no tags at all.")) == 1
|
|
|
|
def test_two_characters_are_two_segments(self) -> None:
|
|
segments = split_segments("[Alice] One. [Bob] Two.")
|
|
assert [s.character for s in segments] == ["Alice", "Bob"]
|
|
assert [s.text for s in segments] == ["One.", "Two."]
|
|
|
|
def test_a_language_tag_is_not_a_character_named_de_alice(self) -> None:
|
|
"""The regression the ordering exists to prevent."""
|
|
segments = split_segments("[de:Alice] Hallo.")
|
|
assert len(segments) == 1
|
|
assert segments[0].character == "Alice"
|
|
assert segments[0].language == "de"
|
|
assert segments[0].text == "Hallo."
|
|
|
|
def test_a_display_name_resolves_to_its_code(self) -> None:
|
|
segments = split_segments("[German:Alice] Hallo.")
|
|
assert segments[0].language == "de"
|
|
assert segments[0].character == "Alice"
|
|
|
|
def test_a_language_tag_may_carry_no_character(self) -> None:
|
|
segments = split_segments("[de:] Hallo.")
|
|
assert segments[0].language == "de"
|
|
assert segments[0].character is None
|
|
|
|
def test_a_pause_tag_is_not_a_character_named_pause(self) -> None:
|
|
"""It is a pause, and a leading one has nothing to be silent after."""
|
|
segments = split_segments("[pause:1.5] Hallo.")
|
|
assert len(segments) == 1
|
|
assert segments[0].character is None
|
|
assert segments[0].text == "Hallo."
|
|
assert segments[0].pause_after == 0.0, (
|
|
"a pause with no text before it has no segment to attach to"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-segment overrides (stage 4)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPerSegmentOverrides:
|
|
def test_order_does_not_matter(self) -> None:
|
|
first = split_segments("[Alice|seed:42|temp:0.5] Hi.")[0]
|
|
second = split_segments("[Alice|temp:0.5|seed:42] Hi.")[0]
|
|
assert first.params == second.params
|
|
assert first.params == {"seed": "42", "temperature": "0.5"}
|
|
|
|
def test_case_does_not_matter(self) -> None:
|
|
segment = split_segments("[Alice|SEED:42|Temp:0.5] Hi.")[0]
|
|
assert segment.params == {"seed": "42", "temperature": "0.5"}
|
|
|
|
@pytest.mark.parametrize(
|
|
("alias", "canonical"),
|
|
[("temp", "temperature"), ("cfg", "cfg_weight"),
|
|
("cfg_weight", "cfg_weight"), ("exag", "exaggeration")],
|
|
)
|
|
def test_documented_aliases_resolve(self, alias: str, canonical: str) -> None:
|
|
segment = split_segments(f"[Alice|{alias}:0.5] Hi.")[0]
|
|
assert segment.params == {canonical: "0.5"}
|
|
assert canonical in PARAM_ALIASES.values()
|
|
|
|
def test_an_unsupported_parameter_is_ignored_not_fatal(self) -> None:
|
|
"""A script carried over from another tool must still run."""
|
|
segment = split_segments("[Alice|batch:4|temp:0.5] Hi.")[0]
|
|
assert segment.params == {"temperature": "0.5"}
|
|
assert segment.character == "Alice"
|
|
assert segment.text == "Hi."
|
|
|
|
def test_a_parameter_with_no_socket_at_all_is_ignored(self) -> None:
|
|
segment = split_segments("[Alice|batch_size:8] Hi.")[0]
|
|
assert segment.params == {}
|
|
assert segment.text == "Hi."
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# v2 angle-bracket tokens (stage 5) -- and the order it runs in
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestV2Tokens:
|
|
def test_all_thirty_documented_names_are_known(self) -> None:
|
|
assert len(V2_SPECIAL_TOKENS) == 30
|
|
assert len(set(V2_SPECIAL_TOKENS)) == 30
|
|
|
|
@pytest.mark.parametrize("token", V2_SPECIAL_TOKENS)
|
|
def test_every_documented_token_is_recognised(self, token: str) -> None:
|
|
resolved, seen = resolve_v2_tokens(f"Hello <{token}> there.")
|
|
assert seen == (token,), f"<{token}> was not recognised"
|
|
assert resolved == f"Hello <{token}> there."
|
|
|
|
def test_casing_is_normalised_except_where_it_is_significant(self) -> None:
|
|
assert resolve_v2_tokens("<GIGGLE>")[0] == "<giggle>"
|
|
assert resolve_v2_tokens("<UH>")[0] == "<UH>"
|
|
assert resolve_v2_tokens("<um>")[0] == "<UM>"
|
|
|
|
def test_an_unknown_angle_token_is_left_alone(self) -> None:
|
|
"""The upstream tokenizer carries 118 special tokens; only 30 are named.
|
|
|
|
Deleting a word the model can still speak is worse than passing it on.
|
|
"""
|
|
resolved, seen = resolve_v2_tokens("Hello <not_a_real_token> there.")
|
|
assert seen == ()
|
|
assert resolved == "Hello <not_a_real_token> there."
|
|
|
|
def test_a_giggle_after_a_character_resolves_after_the_character(self) -> None:
|
|
"""The load-bearing ordering assertion.
|
|
|
|
Stage 5 must not have run yet when the character is resolved, or
|
|
``[Alice] <giggle>`` loses its voice.
|
|
"""
|
|
segments = split_segments("[Alice] <giggle> Hello there, nice to meet you.")
|
|
assert len(segments) == 1
|
|
assert segments[0].character == "Alice", (
|
|
"the character was not resolved; stage 5 ran before stage 1"
|
|
)
|
|
assert "<giggle>" in segments[0].text, (
|
|
"the v2 token was rewritten during the character stage; it must still "
|
|
"be untouched text when the character resolves"
|
|
)
|
|
|
|
plan = plan_text("[Alice] <giggle> Hello there, nice to meet you.")
|
|
assert plan.segments[0].character == "Alice"
|
|
assert "<giggle>" in plan.text
|
|
assert "[" not in plan.text
|
|
|
|
def test_the_same_example_composes_with_a_pause(self) -> None:
|
|
"""The corpus's own worked example, end to end."""
|
|
plan = plan_text(
|
|
"[Alice] Hello! <giggle> Nice to meet you. [pause:0.5] How are you? "
|
|
"<whisper> I have a secret."
|
|
)
|
|
assert [s.character for s in plan.segments] == ["Alice", None]
|
|
assert plan.segments[0].pause_after == pytest.approx(0.5)
|
|
assert plan.segments[1].pause_after == 0.0
|
|
assert "<giggle>" in plan.text
|
|
assert "<whisper>" in plan.text
|
|
|
|
|
|
class TestTurboTagsAreADifferentVocabulary:
|
|
"""``[laugh]`` is a character name here, and a Turbo tag there."""
|
|
|
|
def test_square_brackets_are_a_character_on_the_classic_path(self) -> None:
|
|
segments = split_segments("[laugh] Hello.")
|
|
assert segments[0].character == "laugh"
|
|
|
|
def test_the_same_text_is_a_tag_on_the_turbo_path(self) -> None:
|
|
segments = split_segments("[laugh] Hello.", turbo_tags=True)
|
|
assert segments[0].character is None
|
|
assert segments[0].turbo_tags == ("laugh",)
|
|
|
|
def test_turbo_tags_are_recognised_case_insensitively(self) -> None:
|
|
assert split_segments("[Clear Throat] Hi.", turbo_tags=True)[0].turbo_tags == (
|
|
"clear throat",
|
|
)
|
|
|
|
def test_the_turbo_set_is_not_the_v2_set(self) -> None:
|
|
"""Nine square-bracket tokens against thirty angle-bracket ones.
|
|
|
|
Five names are shared between the two vocabularies (``sigh``, ``gasp``,
|
|
``cough``, ``sniff``, ``groan``); what differs is the *syntax* each
|
|
vocabulary uses, which is the point of this module.
|
|
"""
|
|
assert len(TURBO_TAGS) == 9
|
|
assert len(V2_SPECIAL_TOKENS) == 30
|
|
shared = set(TURBO_TAGS) & {token.lower() for token in V2_SPECIAL_TOKENS}
|
|
assert shared == {"sigh", "gasp", "cough", "sniff", "groan"}
|
|
assert "clear throat" not in V2_SPECIAL_TOKENS
|
|
assert "clear_throat" not in TURBO_TAGS
|
|
|
|
def test_an_angle_token_is_still_normalised_on_the_turbo_path(self) -> None:
|
|
segments = split_segments("[laugh] <giggle> Hi.", turbo_tags=True)
|
|
assert segments[0].turbo_tags == ("laugh",)
|
|
assert "<giggle>" in segments[0].text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Chunking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAutoChunkMethodBoundaries:
|
|
"""``auto`` selects by length, and both comparisons are strict.
|
|
|
|
The corpus phrases the rule as "<500 -> concatenate" while the thresholds are
|
|
written as ">1000" and ">500". Those disagree at exactly 500, so the strict
|
|
reading is implemented and 500 is pinned here -- a rule with an unstated
|
|
boundary is a rule with a bug in it.
|
|
"""
|
|
|
|
@pytest.mark.parametrize(
|
|
("chars", "expected"),
|
|
[
|
|
(499, "concatenate"),
|
|
(500, "concatenate"),
|
|
(501, "crossfade"),
|
|
(1000, "crossfade"),
|
|
(1001, "silence_padding"),
|
|
(5000, "silence_padding"),
|
|
],
|
|
)
|
|
def test_the_documented_boundaries(self, chars: int, expected: str) -> None:
|
|
assert select_chunk_method("auto", chars) == expected
|
|
assert plan_text("x" * chars).chunk_method == expected
|
|
|
|
def test_the_thresholds_are_published_not_implicit(self) -> None:
|
|
assert AUTO_CHUNK_THRESHOLDS == ((1000, "silence_padding"), (500, "crossfade"))
|
|
|
|
@pytest.mark.parametrize(
|
|
"method", [name for name in CHUNK_COMBINATION_METHODS if name != "auto"]
|
|
)
|
|
def test_an_explicit_method_passes_through(self, method: str) -> None:
|
|
"""Whatever the user picked is what runs, at any length.
|
|
|
|
``auto`` is the only value that gets *replaced*; the other three are
|
|
names, not hints, and a short script must not quietly switch the user
|
|
off ``silence_padding`` because it was under a thousand characters.
|
|
"""
|
|
assert select_chunk_method(method, 10) == method
|
|
assert select_chunk_method(method, 5000) == method
|
|
|
|
def test_the_resolved_method_is_never_auto(self) -> None:
|
|
for chars in (0, 1, 500, 5000):
|
|
assert select_chunk_method("auto", chars) != "auto"
|
|
|
|
|
|
class TestChunking:
|
|
def test_a_short_script_is_one_chunk(self) -> None:
|
|
assert chunk_text("Short enough.", max_chars_per_chunk=400) == ["Short enough."]
|
|
|
|
def test_a_long_script_is_cut(self) -> None:
|
|
text = "A sentence here. " * 100
|
|
chunks = chunk_text(text, 200)
|
|
assert len(chunks) > 1
|
|
assert all(len(chunk) <= 200 for chunk in chunks)
|
|
assert "".join(chunks).replace(" ", "") == text.replace(" ", "")
|
|
|
|
def test_cutting_is_disabled_by_the_flag(self) -> None:
|
|
text = "A sentence here. " * 100
|
|
assert chunk_text(text, 200, enabled=False) == [text]
|
|
|
|
def test_a_crash_protected_short_script_is_still_short(self) -> None:
|
|
plan = plan_text("Hi", enable_chunking=True, max_chars_per_chunk=400)
|
|
assert len(plan.chunks) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The whole pipeline
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPlanText:
|
|
def test_a_plain_script_yields_one_segment_and_one_chunk(self) -> None:
|
|
plan = plan_text("Hello there, this is a plain script.")
|
|
assert len(plan.segments) == 1
|
|
assert plan.chunks == ("Hello there, this is a plain script.",)
|
|
assert plan.char_count == len(plan.text)
|
|
|
|
def test_blank_input_is_an_empty_plan_not_an_error(self) -> None:
|
|
plan = plan_text(" ")
|
|
assert plan.segments == ()
|
|
assert plan.chunks == ("",)
|
|
|
|
def test_characters_are_listed_in_first_appearance_order(self) -> None:
|
|
plan = plan_text("[Bob] One. [Alice] Two. [Bob] Three.")
|
|
assert plan.characters == ("Bob", "Alice")
|
|
|
|
def test_the_default_template_is_the_documented_one(self) -> None:
|
|
assert CRASH_PROTECTION_TEMPLATE == "hmm ,, {seg} hmm ,,"
|
|
assert MIN_SAFE_SEGMENT_CHARS == 21
|
|
assert apply_crash_protection("Short.").startswith("hmm ,,")
|