170 lines
6.5 KiB
Python
170 lines
6.5 KiB
Python
from dataclasses import replace
|
|
|
|
from ppp import PromptPostProcessor # type: ignore
|
|
from .base_tests import OutputTuple, InputTuple, TestPromptPostProcessorBase
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit("This script must not be run directly")
|
|
|
|
|
|
class TestChoices(TestPromptPostProcessorBase):
|
|
|
|
def setUp(self): # pylint: disable=arguments-differ
|
|
super().setUp(enable_file_logging=False)
|
|
|
|
# Choices tests
|
|
|
|
def test_ch_choices(self): # simple choices with weights
|
|
self.process(
|
|
InputTuple("the choices are: {3::choice1|2::choice2|choice3}", ""),
|
|
OutputTuple("the choices are: choice2", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_cyclical(self): # cyclical sampler cycles through all choices
|
|
ppp_instance = self.init_ppp("nocup")
|
|
self.process(
|
|
InputTuple("the choices are: {@choice1|choice2|choice3}", ""),
|
|
[
|
|
OutputTuple("the choices are: choice1", ""),
|
|
OutputTuple("the choices are: choice2", ""),
|
|
OutputTuple("the choices are: choice3", ""),
|
|
OutputTuple("the choices are: choice1", ""), # cycles back
|
|
],
|
|
ppp=ppp_instance,
|
|
)
|
|
|
|
def test_ch_cyclical_multiple_constructs(self): # two independent @ constructs cycle together
|
|
ppp_instance = self.init_ppp("nocup")
|
|
self.process(
|
|
InputTuple("{@a|b} {@c|d}", ""),
|
|
[
|
|
OutputTuple("a c", ""),
|
|
OutputTuple("a d", ""),
|
|
OutputTuple("b c", ""),
|
|
OutputTuple("b d", ""),
|
|
OutputTuple("a c", ""), # cycles back
|
|
],
|
|
ppp=ppp_instance,
|
|
)
|
|
|
|
def test_ch_cyclical_resets_on_prompt_change(self): # state resets when the prompt pair changes
|
|
ppp_instance = self.init_ppp("nocup")
|
|
# Advance the cycle to position 1 (choice2).
|
|
self.process(
|
|
InputTuple("the choices are: {@choice1|choice2|choice3}", ""),
|
|
[
|
|
OutputTuple("the choices are: choice1", ""),
|
|
OutputTuple("the choices are: choice2", ""),
|
|
],
|
|
ppp=ppp_instance,
|
|
)
|
|
# A different prompt must restart from position 0 (choice1).
|
|
self.process(
|
|
InputTuple("the choices are: {@choice1|choice2|choice3} different", ""),
|
|
OutputTuple("the choices are: choice1 different", ""),
|
|
ppp=ppp_instance,
|
|
)
|
|
|
|
def test_ch_cyclical_mixed_samplers(self): # @ construct cycles while a ~ construct alongside is unaffected
|
|
ppp_instance = self.init_ppp("nocup")
|
|
self.process(
|
|
InputTuple("{@a|b|c} {x|y}", ""),
|
|
[
|
|
OutputTuple("a y", ""),
|
|
OutputTuple("b x", ""),
|
|
OutputTuple("c x", ""),
|
|
OutputTuple("a y", ""), # @ cycles back
|
|
],
|
|
ppp=ppp_instance,
|
|
)
|
|
|
|
def test_ch_choices_withcomments(self): # choices with comments and multiline
|
|
self.process(
|
|
InputTuple(
|
|
"the choices are: {\n3::choice1 # this is option 1\n|2::choice2\n# this was option 2\n|choice3 # this is option 3\n}",
|
|
"",
|
|
),
|
|
OutputTuple("the choices are: choice2", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_choices_multiple(self): # choices with multiple selection
|
|
self.process(
|
|
InputTuple("the choices are: {~2$$, $$3::choice1|2:: choice2 |choice3}", ""),
|
|
OutputTuple("the choices are: choice2 , choice3", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_choices_if_multiple(self): # choices with if and multiple selection
|
|
self.process(
|
|
InputTuple("the choices are: {2$$, $$3::choice1|2 if _is_sd1::choice2|choice3}", ""),
|
|
OutputTuple("the choices are: choice1, choice3", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_choices_set_if_multiple(self): # choices with if user variable and multiple selection
|
|
self.process(
|
|
InputTuple("${var=test}the choices are: {2$$, $$3::choice1|2 if not var eq 'test'::choice2|choice3}", ""),
|
|
OutputTuple("the choices are: choice1, choice3", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_choices_set_if_nested(self): # nested choices with if user variable and multiple selection
|
|
self.process(
|
|
InputTuple(
|
|
"${var=test}the choices are: {2$$, $$3::choice1${var2=test2} {if var2 eq 'test2'::choice11|choice12}|2 if not var eq 'test'::choice2|choice3}",
|
|
"",
|
|
),
|
|
OutputTuple("the choices are: choice1 choice11, choice3", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_choicesinsidelora(self): # simple choices inside a lora
|
|
self.process(
|
|
InputTuple("<lora:test1:1><lora:test__other__name:1><lora:test2:{0.2|0.5|0.7|1}>", ""),
|
|
OutputTuple("<lora:test1:1><lora:test__other__name:1><lora:test2:0.7>", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
def test_ch_removelorawithchoices(self):
|
|
self.process(
|
|
InputTuple("<lora:test1:1><lora:test2:{0.2|0.5|0.7|1}>", ""),
|
|
OutputTuple("", ""),
|
|
ppp=PromptPostProcessor(
|
|
self.ppp_logger,
|
|
self.def_env_info,
|
|
replace(
|
|
self.defopts,
|
|
cup_remove_extranetwork_tags=True,
|
|
),
|
|
self.grammar_content,
|
|
self.interrupt,
|
|
self.wildcards_obj,
|
|
self.extranetwork_maps_obj,
|
|
),
|
|
)
|
|
|
|
def test_ch_cmd_includewildcard(self):
|
|
self.process(
|
|
InputTuple("{ch_one|ch_two|%0.5::include yaml/wildcard1}", ""),
|
|
OutputTuple("ch_two", ""),
|
|
ppp="nocup",
|
|
)
|
|
|
|
# Combinatorial
|
|
|
|
def test_ch_combinatorial(self):
|
|
self.process(
|
|
InputTuple("{choice1|choice2|choice3}, ${v:{option1|option2}}", ""),
|
|
[
|
|
OutputTuple("choice1, option1", ""),
|
|
OutputTuple("choice1, option2", ""),
|
|
OutputTuple("choice2, option1", ""),
|
|
OutputTuple("choice2, option2", ""),
|
|
OutputTuple("choice3, option1", ""),
|
|
OutputTuple("choice3, option2", "", {"v": "option2"}),
|
|
],
|
|
combinatorial=True,
|
|
)
|