Antonio Cordero Balcazar
e8643680f9
* Option to cleanup around extra network tags. * Improved cleanup. * Some refactoring.
904 lines
43 KiB
Python
904 lines
43 KiB
Python
from collections import namedtuple
|
|
import re
|
|
import math
|
|
import lark
|
|
|
|
|
|
class PromptPostProcessor: # pylint: disable=too-few-public-methods,too-many-instance-attributes
|
|
"""
|
|
The PromptPostProcessor class is responsible for processing and manipulating prompt strings.
|
|
|
|
The format for the negative tags is:
|
|
<!content!>
|
|
|
|
<!!x!content!>
|
|
|
|
with x being:
|
|
s - content is added at the start of the negative prompt. This is the default if no parameter exists.
|
|
e - content is added at the end of the negative prompt.
|
|
pN - content is added where the insertion point N is in the negative prompt or at the start if it does not exist. N can be 0 to 9.
|
|
iN - tags the position of insertion point N. Used only in the negative prompt and does not accept content. N can be 0 to 9.
|
|
|
|
Attributes:
|
|
NAME (str): The name of the prompt post-processor.
|
|
VERSION (str): The version of the prompt post-processor.
|
|
DEFAULT_STN_SEPARATOR (str): The default separator used for content sent to the negative prompt.
|
|
|
|
Methods:
|
|
__init__(self, script, opts=None, is_i2i=False): Initializes the PromptPostProcessor instance.
|
|
formatOutput(self, text: str) -> str: Formats the output text by encoding and decoding it.
|
|
STNTree: A nested class for interpreting and processing a tree generated by the prompt parser.
|
|
__find_tags(self, prompt): Finds tags in the given prompt and returns the modified prompt and tag information.
|
|
__add_to_insertion_points(self, negative_prompt, add_at_insertion_point): Adds the negative prompt to the insertion points.
|
|
__add_to_start(self, negative_prompt, add_at_start): Adds the elements in `add_at_start` list to the start of the `negative_prompt` string.
|
|
__add_to_end(self, negative_prompt, add_at_end): Adds the elements in `add_at_end` list to the end of `negative_prompt` string.
|
|
CleanupTree: A nested class for cleaning up a prompt parsed into a tree.
|
|
__cleanup(self, prompt, negative_prompt): Cleans up the prompt and negative prompt by removing extra spaces, empty constructs, and extra separators.
|
|
cleanup_text(self, text): Cleans up the given text by removing extra separators, breaks, and spaces.
|
|
trim_text(self, text): Trims the given text based on the specified cleanup options.
|
|
process_prompt(self, original_prompt, original_negative_prompt): Process the prompt and negative prompt by moving content to the negative prompt, and cleaning up.
|
|
"""
|
|
|
|
NAME = "Prompt Post-Processor"
|
|
VERSION = "2.3.0"
|
|
|
|
DEFAULT_STN_SEPARATOR = ", "
|
|
IFWILDCARDS_CHOICES = {
|
|
"ignore": "Ignore",
|
|
"remove": "Remove",
|
|
"warn": "Add visible warning",
|
|
"stop": "Stop the generation",
|
|
}
|
|
WILDCARD_WARNING = '(WARNING TEXT "INVALID WILDCARD" IN BRIGHT RED:1.5) BREAK\n'
|
|
WILDCARD_STOP = "INVALID WILDCARD! BREAK\n"
|
|
|
|
def __init__(
|
|
self,
|
|
script,
|
|
opts=None,
|
|
is_i2i=False,
|
|
):
|
|
"""
|
|
Initializes the PPP object.
|
|
|
|
Args:
|
|
script: The script object.
|
|
opts: Optional. The options object for configuring PPP behavior.
|
|
"""
|
|
self.script = script
|
|
self.logger = script.ppp_logger
|
|
self.opts = opts
|
|
self.is_i2i = is_i2i
|
|
self.debug = getattr(self.opts, "ppp_gen_debug", False) if opts is not None else False
|
|
if opts is not None and getattr(opts, "prompt_attention", "") == "Compel parser":
|
|
self.logger.warning("Compel parser is not supported!")
|
|
self.ifwildcards = (
|
|
getattr(opts, "ppp_gen_ifwildcards", self.IFWILDCARDS_CHOICES["ignore"])
|
|
if opts is not None
|
|
else self.IFWILDCARDS_CHOICES["ignore"]
|
|
)
|
|
|
|
self.stn_doi2i = getattr(opts, "ppp_stn_doi2i", False) if opts is not None else False
|
|
self.stn_ignore_repeats = getattr(opts, "ppp_stn_ignorerepeats", True) if opts is not None else True
|
|
self.stn_join_attention = getattr(opts, "ppp_stn_joinattention", True) if opts is not None else True
|
|
self.stn_separator = (
|
|
getattr(opts, "ppp_stn_separator", self.DEFAULT_STN_SEPARATOR)
|
|
if opts is not None
|
|
else self.DEFAULT_STN_SEPARATOR
|
|
)
|
|
|
|
self.cup_doi2i = getattr(opts, "ppp_cup_doi2i", False) if opts is not None else False
|
|
self.cup_extraspaces = getattr(opts, "ppp_cup_extraspaces", True) if opts is not None else True
|
|
self.cup_emptyconstructs = getattr(opts, "ppp_cup_emptyconstructs", True) if opts is not None else True
|
|
self.cup_extraseparators = getattr(opts, "ppp_cup_extraseparators", True) if opts is not None else True
|
|
self.cup_breaks = getattr(opts, "ppp_cup_breaks", True) if opts is not None else True
|
|
self.cup_ands = getattr(opts, "ppp_cup_ands", True) if opts is not None else True
|
|
self.cup_extranetworktags = getattr(opts, "ppp_cup_extranetworktags", False) if opts is not None else False
|
|
|
|
self.__insertion_point_tags = [f"<!!i{x}!!>" for x in range(10)]
|
|
# Process with lark (debug with https://www.lark-parser.org/ide/)
|
|
self.__parser_complete = lark.Lark(
|
|
r"""
|
|
start: (promptcomp | specialchars)* // BUG: sometimes it chooses specialchars instead of promptcomp and fails!!!
|
|
// prompt composition with AND
|
|
promptcomp: promptcomppart ([":" numpar] (/\bAND\b/ promptcomppart [":" numpar])+)?
|
|
promptcomppart: prompt
|
|
// prompt scheduling and alternation
|
|
alternate: "[" alternateoption ("|" alternateoption)+ "]"
|
|
alternateoption: prompt
|
|
scheduled: "[" [prompt ":"] prompt ":" numpar "]"
|
|
// wildcard extension support
|
|
wildcard: "__" /(?:(?!__)\S)+/ "__"
|
|
choices: "{" choice ("|" choice)* "}"
|
|
// we ignore weight and any other parameters in each choice
|
|
choice: prompt
|
|
// simple prompts
|
|
prompt: (emphasized | deemphasized | scheduled | alternate | extranetworktag | negtag | wildcard | choices | plain)*
|
|
nonegprompt: (emphasized | deemphasized | scheduled | alternate | extranetworktag | wildcard | choices | plain)*
|
|
// attention modifiers
|
|
emphasized: "(" prompt [":" numpar] ")"
|
|
deemphasized: "[" prompt "]"
|
|
// extra network tags
|
|
extranetworktag: "<" /(?!!)[^>]+/ ">"
|
|
// negative tags
|
|
negtag: "<!" [negtagparameters] nonegprompt "!>"
|
|
negtagparameters: "!" /s|e|[ip]\d/ "!"
|
|
// plain text and weights
|
|
numpar: WHITESPACE* NUMBER WHITESPACE*
|
|
WHITESPACE: /\s+/
|
|
plain: /((?!__|\bAND\b)[^\\[\]():|<>!{}]|\\.)+/s
|
|
specialchars: /[\]():|<>!{}]|\bAND\b/+
|
|
%import common.SIGNED_NUMBER -> NUMBER
|
|
""",
|
|
propagate_positions=True,
|
|
)
|
|
|
|
def formatOutput(self, text: str):
|
|
"""
|
|
Formats the output text by encoding it using unicode_escape and decoding it using utf-8.
|
|
|
|
Args:
|
|
text (str): The input text to be formatted.
|
|
|
|
Returns:
|
|
str: The formatted output text.
|
|
"""
|
|
return text.encode("unicode_escape").decode("utf-8")
|
|
|
|
class STNTree(lark.visitors.Interpreter):
|
|
"""
|
|
A class for interpreting and processing a tree generated by the prompt parser.
|
|
|
|
Attributes:
|
|
__ppp (object): The instance of the parent class.
|
|
__prompt (str): The prompt string.
|
|
__shell (list): The list of accumulated shell elements.
|
|
__negtags (list): The list of negative tags.
|
|
__already_processed (list): The list of already processed content.
|
|
add_at (dict): The dictionary to store the content to be added at different positions.
|
|
remove (list): The list of content to be removed from the prompt.
|
|
|
|
Methods:
|
|
__get_numpar_value(numpar): Get the numerical value from a numpar object.
|
|
scheduled(tree): Process a scheduling construct in the tree and add it to the accumulated shell.
|
|
alternate(tree): Process an alternation construct in the tree and add it to the accumulated shell.
|
|
emphasized(tree): Process an attention change construct in the tree and add it to the accumulated shell.
|
|
deemphasized(tree): Process a decrease attention construct in the tree and add it to the accumulated shell.
|
|
negtag(tree): Process a negative tag in the tree and add it to the list of negative tags.
|
|
start(tree): Process the given tree and perform necessary operations on the found negative tags.
|
|
"""
|
|
|
|
def __init__(self, ppp, prompt, add_at):
|
|
super().__init__()
|
|
self.__ppp = ppp
|
|
self.__prompt = prompt
|
|
self.AccumulatedShell = namedtuple("AccumulatedShell", ["type", "data", "position"])
|
|
AccumulatedShell = self.AccumulatedShell
|
|
self.__shell: list[AccumulatedShell] = []
|
|
self.NegTag = namedtuple("NegTag", ["start", "end", "content", "parameters", "shell"])
|
|
NegTag = self.NegTag
|
|
self.__negtags: list[NegTag] = []
|
|
self.__already_processed = []
|
|
self.add_at = add_at
|
|
self.remove = []
|
|
|
|
def __get_numpar_value(self, numpar):
|
|
"""
|
|
Get the numerical value from a numpar object.
|
|
|
|
Args:
|
|
numpar (object): The numpar object to extract the value from.
|
|
|
|
Returns:
|
|
float: The numerical value extracted from the numpar object.
|
|
"""
|
|
return float(next(x for x in numpar.children if x.type == "NUMBER").value)
|
|
|
|
def scheduled(self, tree):
|
|
"""
|
|
Process a scheduling construct in the tree and add it to the accumulated shell.
|
|
|
|
Args:
|
|
tree (Node): The tree node representing the scheduling construct.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
treemetaposition = (
|
|
[tree.meta.start_pos, tree.meta.end_pos] if hasattr(tree, "meta") and not tree.meta.empty else None
|
|
)
|
|
if len(tree.children) > 2: # before & after
|
|
before = tree.children[0]
|
|
else:
|
|
before = None
|
|
after = tree.children[-2]
|
|
numpar = tree.children[-1]
|
|
pos = self.__get_numpar_value(numpar)
|
|
if pos >= 1:
|
|
pos = int(pos)
|
|
# self.__shell.append(self.AccumulatedShell("sc", pos, treemetaposition))
|
|
if before is not None and hasattr(before, "data"):
|
|
if self.__ppp.debug:
|
|
before_metaposition = (
|
|
[before.meta.start_pos, before.meta.end_pos]
|
|
if hasattr(before, "meta") and not before.meta.empty
|
|
else "?"
|
|
)
|
|
self.__ppp.logger.info(f"Shell scheduled before at {before_metaposition} with position {pos}")
|
|
self.__shell.append(self.AccumulatedShell("scb", pos, treemetaposition))
|
|
self.visit(before)
|
|
self.__shell.pop()
|
|
if hasattr(after, "data"):
|
|
if self.__ppp.debug:
|
|
after_metaposition = (
|
|
[after.meta.start_pos, after.meta.end_pos]
|
|
if hasattr(after, "meta") and not after.meta.empty
|
|
else "?"
|
|
)
|
|
self.__ppp.logger.info(f"Shell scheduled after at {after_metaposition} with position {pos}")
|
|
self.__shell.append(self.AccumulatedShell("sca", pos, treemetaposition))
|
|
self.visit(after)
|
|
self.__shell.pop()
|
|
# self.__shell.pop()
|
|
|
|
def alternate(self, tree):
|
|
"""
|
|
Process an alternation construct in the tree and add it to the accumulated shell.
|
|
|
|
Args:
|
|
tree (Node): The tree node representing the alternation construct.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
treemetaposition = (
|
|
[tree.meta.start_pos, tree.meta.end_pos] if hasattr(tree, "meta") and not tree.meta.empty else None
|
|
)
|
|
# self.__shell.append(self.AccumulatedShell("al", len(tree.children), treemetaposition))
|
|
for i, opt in enumerate(tree.children):
|
|
if self.__ppp.debug:
|
|
metaposition = (
|
|
[opt.meta.start_pos, opt.meta.end_pos] if hasattr(opt, "meta") and not opt.meta.empty else "?"
|
|
)
|
|
self.__ppp.logger.info(f"Shell alternate at {metaposition} option {i+1}")
|
|
if hasattr(opt, "data"):
|
|
self.__shell.append(
|
|
self.AccumulatedShell("alo", {"pos": i + 1, "len": len(tree.children)}, treemetaposition)
|
|
)
|
|
self.visit(opt)
|
|
self.__shell.pop()
|
|
# self.__shell.pop()
|
|
|
|
def emphasized(self, tree):
|
|
"""
|
|
Process a attention change construct in the tree and add it to the accumulated shell.
|
|
|
|
Args:
|
|
tree (Node): The tree node representing the attention construct.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
treemetaposition = (
|
|
[tree.meta.start_pos, tree.meta.end_pos] if hasattr(tree, "meta") and not tree.meta.empty else None
|
|
)
|
|
numpar = tree.children[-1]
|
|
weight = self.__get_numpar_value(numpar) if numpar is not None else 1.1
|
|
if self.__ppp.debug:
|
|
self.__ppp.logger.info(f"Shell attention at {treemetaposition or '?'} with weight {weight}")
|
|
self.__shell.append(self.AccumulatedShell("at", weight, treemetaposition))
|
|
self.visit_children(tree)
|
|
self.__shell.pop()
|
|
|
|
def deemphasized(self, tree):
|
|
"""
|
|
Process a decrease attention construct in the tree and add it to the accumulated shell.
|
|
|
|
Args:
|
|
tree (Node): The tree node representing the decreased attention construct.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
weight = 0.9
|
|
treemetaposition = (
|
|
[tree.meta.start_pos, tree.meta.end_pos] if hasattr(tree, "meta") and not tree.meta.empty else None
|
|
)
|
|
if self.__ppp.debug:
|
|
self.__ppp.logger.info(f"Shell attention at {treemetaposition or '?'} with weight {weight}")
|
|
self.__shell.append(self.AccumulatedShell("at", weight, treemetaposition))
|
|
self.visit_children(tree)
|
|
self.__shell.pop()
|
|
|
|
def negtag(self, tree):
|
|
"""
|
|
Process a negative tag in the tree and add it to the list of negative tags.
|
|
|
|
Args:
|
|
tree (Node): The tree node representing the negative tag.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
treemetaposition = (
|
|
[tree.meta.start_pos, tree.meta.end_pos] if hasattr(tree, "meta") and not tree.meta.empty else None
|
|
)
|
|
negtagparameters = tree.children[0]
|
|
parameters = negtagparameters.children[0].value if negtagparameters is not None else ""
|
|
rest = []
|
|
for x in tree.children[1::]:
|
|
rest.append(
|
|
self.__prompt[x.meta.start_pos : x.meta.end_pos]
|
|
if hasattr(x, "meta") and not x.meta.empty
|
|
else x.value
|
|
)
|
|
content = "".join(rest)
|
|
self.__negtags.append(
|
|
self.NegTag(tree.meta.start_pos, tree.meta.end_pos, content, parameters, self.__shell.copy())
|
|
)
|
|
if self.__ppp.debug:
|
|
self.__ppp.logger.info(
|
|
f"Negative tag at {treemetaposition or '?'}: {parameters or 'with no parameters :'} {self.__ppp.formatOutput(content)}"
|
|
)
|
|
|
|
def start(self, tree):
|
|
"""
|
|
Process the given tree and perform necessary operations on the found negative tags.
|
|
|
|
Args:
|
|
tree: The tree to be processed.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
self.visit_children(tree)
|
|
# process the found negative tags
|
|
for negtag in self.__negtags:
|
|
if self.__ppp.stn_join_attention:
|
|
# join consecutive attention elements
|
|
for i in range(len(negtag.shell) - 1, 0, -1):
|
|
if negtag.shell[i].type == "at" and negtag.shell[i - 1].type == "at":
|
|
negtag.shell[i - 1] = self.AccumulatedShell(
|
|
"at",
|
|
math.floor(100 * negtag.shell[i - 1].data * negtag.shell[i].data)
|
|
/ 100, # we limit the new weight to two decimals
|
|
negtag.shell[i - 1].position,
|
|
)
|
|
negtag.shell.pop(i)
|
|
start = ""
|
|
end = ""
|
|
for s in negtag.shell:
|
|
match s.type:
|
|
case "at":
|
|
if s.data == 0.9:
|
|
start += "["
|
|
end = "]" + end
|
|
elif s.data == 1.1:
|
|
start += "("
|
|
end = ")" + end
|
|
else:
|
|
start += "("
|
|
end = f":{s.data})" + end
|
|
# case "sc":
|
|
case "scb":
|
|
start += "["
|
|
end = f"::{s.data}]" + end
|
|
case "sca":
|
|
start += "["
|
|
end = f":{s.data}]" + end
|
|
# case "al":
|
|
case "alo":
|
|
start += "[" + ("|" * int(s.data["pos"] - 1))
|
|
end = ("|" * int(s.data["len"] - s.data["pos"])) + "]" + end
|
|
content = start + negtag.content + end
|
|
position = negtag.parameters or "s"
|
|
if len(content) > 0:
|
|
if content not in self.__already_processed:
|
|
if self.__ppp.stn_ignore_repeats:
|
|
self.__already_processed.append(content)
|
|
if self.__ppp.debug:
|
|
self.__ppp.logger.info(
|
|
f"Adding content at position {position}: {self.__ppp.formatOutput(content)}"
|
|
)
|
|
if position == "e":
|
|
self.add_at["end"].append(content)
|
|
elif position.startswith("p"):
|
|
n = int(position[1])
|
|
self.add_at["insertion_point"][n].append(content)
|
|
else: # position == "s" or invalid
|
|
self.add_at["start"].append(content)
|
|
else:
|
|
self.__ppp.logger.warning(f"Ignoring repeated content: {self.__ppp.formatOutput(content)}")
|
|
# remove from prompt
|
|
self.remove.append([negtag.start, negtag.end])
|
|
|
|
def __find_tags(self, prompt):
|
|
"""
|
|
Finds tags in the given prompt and returns the modified prompt and tag information.
|
|
|
|
Args:
|
|
prompt (str): The input prompt.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the modified prompt and tag information.
|
|
|
|
"""
|
|
add_at = {"start": [], "insertion_point": [[] for x in range(10)], "end": []}
|
|
tree = self.__parser_complete.parse(prompt)
|
|
|
|
readtree = self.STNTree(self, prompt, add_at)
|
|
readtree.visit(tree)
|
|
|
|
for r in readtree.remove[::-1]:
|
|
prompt = prompt[: r[0]] + prompt[r[1] :]
|
|
|
|
add_at = readtree.add_at
|
|
if self.debug:
|
|
self.logger.info(f"New negative additions: {add_at}")
|
|
return prompt, add_at
|
|
|
|
def __add_to_insertion_points(self, negative_prompt, add_at_insertion_point):
|
|
"""
|
|
Adds the negative prompt to the insertion points.
|
|
|
|
Args:
|
|
negative_prompt (str): The negative prompt to be added.
|
|
add_at_insertion_point (list): A list of insertion points.
|
|
|
|
Returns:
|
|
str: The modified negative prompt.
|
|
"""
|
|
for n in range(10):
|
|
ipp = negative_prompt.find(self.__insertion_point_tags[n])
|
|
if ipp >= 0:
|
|
ipl = len(self.__insertion_point_tags[n])
|
|
if negative_prompt[ipp - len(self.stn_separator) : ipp] == self.stn_separator:
|
|
ipp -= len(self.stn_separator) # adjust for existing start separator
|
|
ipl += len(self.stn_separator)
|
|
add_at_insertion_point[n].insert(0, negative_prompt[:ipp])
|
|
if negative_prompt[ipp + ipl : ipp + ipl + len(self.stn_separator)] == self.stn_separator:
|
|
ipl += len(self.stn_separator) # adjust for existing end separator
|
|
endPart = negative_prompt[ipp + ipl :]
|
|
if len(endPart) > 0:
|
|
add_at_insertion_point[n].append(endPart)
|
|
negative_prompt = self.stn_separator.join(add_at_insertion_point[n])
|
|
else:
|
|
ipp = 0
|
|
if negative_prompt.startswith(self.stn_separator):
|
|
ipp = len(self.stn_separator)
|
|
add_at_insertion_point[n].append(negative_prompt[ipp:])
|
|
negative_prompt = self.stn_separator.join(add_at_insertion_point[n])
|
|
return negative_prompt
|
|
|
|
def __add_to_start(self, negative_prompt, add_at_start):
|
|
"""
|
|
Adds the elements in `add_at_start` list to the start of the `negative_prompt` string.
|
|
|
|
Args:
|
|
negative_prompt (str): The original negative prompt string.
|
|
add_at_start (list): The list of elements to be added at the start of the negative prompt.
|
|
|
|
Returns:
|
|
str: The updated negative prompt string with the elements added at the start.
|
|
"""
|
|
if len(negative_prompt) > 0:
|
|
ipp = 0
|
|
if negative_prompt.startswith(self.stn_separator):
|
|
ipp = len(self.stn_separator) # adjust for existing end separator
|
|
add_at_start.append(negative_prompt[ipp:])
|
|
negative_prompt = self.stn_separator.join(add_at_start)
|
|
return negative_prompt
|
|
|
|
def __add_to_end(self, negative_prompt, add_at_end):
|
|
"""
|
|
Adds the elements in `add_at_end` list to the end of `negative_prompt` string.
|
|
|
|
Args:
|
|
negative_prompt (str): The original negative prompt string.
|
|
add_at_end (list): The list of elements to be added at the end of `negative_prompt`.
|
|
|
|
Returns:
|
|
str: The updated negative prompt string with elements added at the end.
|
|
"""
|
|
if len(negative_prompt) > 0:
|
|
ipl = len(negative_prompt)
|
|
if negative_prompt.endswith(self.stn_separator):
|
|
ipl -= len(self.stn_separator) # adjust for existing start separator
|
|
add_at_end.insert(0, negative_prompt[:ipl])
|
|
negative_prompt = self.stn_separator.join(add_at_end)
|
|
return negative_prompt
|
|
|
|
def __sendtonegative(self, prompt, negative_prompt):
|
|
"""
|
|
Modifies the prompt and negative_prompt by moving content from the prompt to the negative prompt.
|
|
|
|
Args:
|
|
prompt (str): The prompt.
|
|
negative_prompt (str): The negative prompt.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the modified prompt and negative_prompt.
|
|
"""
|
|
if self.debug:
|
|
self.logger.info("Doing send-to-negative")
|
|
prompt, add_at = self.__find_tags(prompt)
|
|
negative_prompt = self.__add_to_insertion_points(negative_prompt, add_at["insertion_point"])
|
|
if len(add_at["start"]) > 0:
|
|
negative_prompt = self.__add_to_start(negative_prompt, add_at["start"])
|
|
if len(add_at["end"]) > 0:
|
|
negative_prompt = self.__add_to_end(negative_prompt, add_at["end"])
|
|
if self.debug:
|
|
self.logger.info(f"prompt after send-to-negative: {self.formatOutput(prompt)}")
|
|
self.logger.info(f"negative_prompt after send-to-negative: {self.formatOutput(negative_prompt)}")
|
|
return prompt, negative_prompt
|
|
|
|
class TransformerTree(lark.visitors.Transformer_NonRecursive):
|
|
"""
|
|
Transformer class for detecting wildcards and/or cleaning up a prompt parsed into a tree.
|
|
|
|
This class provides methods for transforming different constructs in a parse tree
|
|
based on certain conditions. It is used for detecting wildcards and cleaning up invalid
|
|
or empty constructs in the parse tree.
|
|
|
|
Args:
|
|
ppp (object): An instance of the parent class `ppp`.
|
|
|
|
Attributes:
|
|
__ppp (object): An instance of the parent class `ppp`.
|
|
|
|
Methods:
|
|
promptcomp(tree): Replicates prompt composition constructs.
|
|
scheduled(tree): Replicates or removes scheduling constructs based on conditions.
|
|
alternate(tree): Replicates or removes alternation constructs based on conditions.
|
|
emphasized(tree): Replicates or removes attention constructs based on conditions.
|
|
deemphasized(tree): Replicates or removes attention constructs based on conditions.
|
|
extranetworktag(tree): Replicates extra network constructs.
|
|
numpar(tree): Cleans up number parameter.
|
|
negtag(tree): Replicates or removes negative tag constructs based on conditions.
|
|
wildcard(tree): Replicates or removes wildcard constructs based on conditions.
|
|
choices(tree): Replicates or removes choices constructs based on conditions.
|
|
choice(tree): Replicates choices.
|
|
plain(tree): Cleans up plain text based on conditions.
|
|
__default__(data, children, meta): Default method for joining children and cleaning up text based on conditions.
|
|
"""
|
|
|
|
def __init__(self, ppp, phase="cleanup"):
|
|
super().__init__(visit_tokens=True)
|
|
self.__ppp = ppp
|
|
self.__phase = phase
|
|
self.detectedWildcards = []
|
|
|
|
def promptcomp(self, tree):
|
|
r = tree[0]
|
|
if len(tree) > 1:
|
|
if tree[1] is not None:
|
|
r += f":{tree[1]}"
|
|
for i in range(2, len(tree), 3):
|
|
if self.__phase == "cleanup" and self.__ppp.cup_ands:
|
|
r = re.sub(r"[, ]+$", " ", r)
|
|
if r[-1:].isalnum(): # add space if needed
|
|
r += " "
|
|
r += "AND"
|
|
t = tree[i + 1]
|
|
if self.__phase == "cleanup" and self.__ppp.cup_ands:
|
|
t = re.sub(r"^[, ]+", " ", t)
|
|
if t[0:1].isalnum(): # add space if needed
|
|
r += " "
|
|
r += t
|
|
if tree[i + 2] is not None:
|
|
r += f":{tree[i+2]}"
|
|
return r
|
|
|
|
def scheduled(self, tree):
|
|
if len(tree) == 0 and self.__phase == "cleanup" and self.__ppp.cup_emptyconstructs:
|
|
return "" # remove invalid scheduling construct (probably this is not reachable)
|
|
# replicate scheduling construct
|
|
if len(tree) > 0 and tree[0] is None:
|
|
return f"[{':'.join(tree[1:])}]"
|
|
return f"[{':'.join(tree)}]"
|
|
|
|
def alternate(self, tree):
|
|
if len(tree) == 0 and self.__phase == "cleanup" and self.__ppp.cup_emptyconstructs:
|
|
return "" # remove invalid alternation construct (probably this is not reachable)
|
|
return f"[{'|'.join(tree)}]" # replicate alternation construct
|
|
|
|
def emphasized(self, tree):
|
|
if (len(tree) == 0 or tree[0] == "") and self.__phase == "cleanup" and self.__ppp.cup_emptyconstructs:
|
|
return "" # remove empty attention construct
|
|
if len(tree) > 1 and tree[1] is not None:
|
|
return f"({tree[0]}:{tree[1]})" # replicate attention construct with weight
|
|
return f"({tree[0]})" # replicate attention construct without weight
|
|
|
|
def deemphasized(self, tree):
|
|
if (len(tree) == 0 or tree[0] == "") and self.__phase == "cleanup" and self.__ppp.cup_emptyconstructs:
|
|
return "" # remove empty attention construct (invalid scheduling or alternation constructs end up here too?)
|
|
return f"[{tree[0]}]" # replicate attention construct
|
|
|
|
def extranetworktag(self, tree):
|
|
return f"<{tree[0]}>" # replicate extra network construct
|
|
|
|
def numpar(self, tree):
|
|
return next(x for x in tree if x.type == "NUMBER").value.strip() # clean up number parameter
|
|
|
|
def negtag(self, tree):
|
|
if self.__phase == "cleanup":
|
|
return "" # remove negative tag construct (there shouldn't be any at this point)
|
|
parameters = "!" + tree[0] + "!" if tree[0] is not None else ""
|
|
content = "".join(tree[1::])
|
|
return f"<!{parameters}{content}!>" # replicate negative tag construct
|
|
|
|
def wildcard(self, tree):
|
|
content = f"__{tree[0]}__" # replicate wildcard construct
|
|
self.detectedWildcards.append(content)
|
|
if self.__phase == "wildcards" and self.__ppp.ifwildcards == self.__ppp.IFWILDCARDS_CHOICES["remove"]:
|
|
return ""
|
|
return content
|
|
|
|
def choices(self, tree):
|
|
content = "{" + "|".join(tree) + "}" # replicate wildcard choices construct
|
|
self.detectedWildcards.append(content)
|
|
if self.__phase == "wildcards" and self.__ppp.ifwildcards == self.__ppp.IFWILDCARDS_CHOICES["remove"]:
|
|
return ""
|
|
return content
|
|
|
|
def choice(self, tree):
|
|
return f"{tree[0]}" # replicate choice
|
|
|
|
def plain(self, tree):
|
|
if self.__phase == "cleanup":
|
|
return self.__ppp.cleanup_text(tree[0].value) # clean up plain text
|
|
return tree[0].value
|
|
|
|
def __default__(self, data, children, meta):
|
|
joined = "".join(children) # join all children
|
|
if self.__phase == "cleanup":
|
|
# take care of cleaning the joints only if there are no constructs that can be affected
|
|
if not re.match(r"[([<{]", joined):
|
|
joined = self.__ppp.cleanup_text(joined)
|
|
return joined
|
|
|
|
def __cleanup(self, prompt, negative_prompt):
|
|
"""
|
|
Cleans up the prompt and negative prompt by removing extra spaces, empty constructs, and extra separators.
|
|
|
|
Args:
|
|
prompt (str): The original prompt.
|
|
negative_prompt (str): The negative prompt.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the cleaned up prompt and negative prompt.
|
|
"""
|
|
if self.debug:
|
|
self.logger.info("Doing cleanup")
|
|
|
|
transformtree = self.TransformerTree(self, phase="cleanup")
|
|
try:
|
|
prompt_tree = self.__parser_complete.parse(prompt)
|
|
prompt = self.trim_text(transformtree.transform(prompt_tree))
|
|
except Exception as e: # pylint: disable=broad-except
|
|
self.logger.warning("Cleanup parsing failed on prompt!: %s", e)
|
|
try:
|
|
negativeprompt_tree = self.__parser_complete.parse(negative_prompt)
|
|
negative_prompt = self.trim_text(transformtree.transform(negativeprompt_tree))
|
|
except Exception as e: # pylint: disable=broad-except
|
|
self.logger.warning("Cleanup parsing failed on negative prompt!: %s", e)
|
|
|
|
if self.debug:
|
|
self.logger.info(f"prompt after cleanup: {self.formatOutput(prompt)}")
|
|
self.logger.info(f"negative_prompt after cleanup: {self.formatOutput(negative_prompt)}")
|
|
return prompt, negative_prompt
|
|
|
|
def cleanup_text(self, text):
|
|
"""
|
|
Cleans up the given text by removing extra separators, breaks, and spaces. This is called for plain text only or when there are no constructs.
|
|
|
|
Args:
|
|
text (str): The text to be cleaned up.
|
|
|
|
Returns:
|
|
str: The cleaned up text.
|
|
"""
|
|
# NOTE: we can't use start/end of line regex since the text might only be a part of a larger line due to the parser
|
|
if self.cup_extraseparators:
|
|
#
|
|
# sendtonegative separator
|
|
#
|
|
escapedSeparator = re.escape(self.stn_separator)
|
|
# collapse separators
|
|
text = re.sub(r"(?:\s*" + escapedSeparator + r"\s*){2,}", self.stn_separator, text)
|
|
#
|
|
# regular comma separator
|
|
#
|
|
# collapse separators
|
|
text = re.sub(r"(?:\s*,\s*){2,}", ", ", text)
|
|
if self.cup_breaks:
|
|
# collapse separators and commas before BREAK
|
|
text = re.sub(r"[, ]+BREAK\b", " BREAK", text)
|
|
# collapse separators and commas after BREAK
|
|
text = re.sub(r"\bBREAK[, ]+", "BREAK ", text)
|
|
# collapse separators and commas around BREAK
|
|
text = re.sub(r"[, ]+BREAK[, ]+", " BREAK ", text)
|
|
# collapse BREAKs
|
|
text = re.sub(r"\bBREAK(?:\s+BREAK)+\b", " BREAK ", text)
|
|
if self.cup_extraspaces:
|
|
# remove spaces before comma
|
|
text = re.sub(r"[ ]+,", ",", text)
|
|
# collapse spaces
|
|
text = re.sub(r"[ ]{2,}", " ", text)
|
|
return text
|
|
|
|
def trim_text(self, text):
|
|
"""
|
|
Trims the given text based on the specified cleanup options. This is only called for the reconstructed prompt.
|
|
|
|
Args:
|
|
text (str): The text to be trimmed.
|
|
|
|
Returns:
|
|
str: The trimmed text.
|
|
"""
|
|
# NOTE: here we can only do cleanups that can be done on the whole text, including inside constructs and around them
|
|
if self.cup_extraseparators:
|
|
#
|
|
# sendtonegative separator
|
|
#
|
|
escapedSeparator = re.escape(self.stn_separator)
|
|
# remove duplicate separator after starting parenthesis or bracket
|
|
text = re.sub(r"(\s*" + escapedSeparator + r"\s*[([])\s*" + escapedSeparator + r"\s*", r"\1", text)
|
|
# remove before colon or ending parenthesis or bracket
|
|
text = re.sub(r"\s*" + escapedSeparator + r"\s*([:)\]]\s*" + escapedSeparator + r"\s*)", r"\1", text)
|
|
# remove at start of prompt or line
|
|
text = re.sub(r"^(?:\s*" + escapedSeparator + r"\s*)", "", text, flags=re.MULTILINE)
|
|
# remove at end of prompt or line
|
|
text = re.sub(r"(?:\s*" + escapedSeparator + r"\s*)$", "", text, flags=re.MULTILINE)
|
|
#
|
|
# regular comma separator
|
|
#
|
|
# remove duplicate separators after starting parenthesis or bracket
|
|
text = re.sub(r"(\s*,\s*[([])\s*,\s*", r"\1", text)
|
|
# remove duplicate separators before colon or ending parenthesis or bracket
|
|
text = re.sub(r"\s*,\s*([:)\]]\s*,\s*)", r"\1", text)
|
|
# remove at start of prompt or line
|
|
text = re.sub(r"^\s*,\s*", "", text, flags=re.MULTILINE)
|
|
# remove at end of prompt or line
|
|
text = re.sub(r"\s*,\s*$", "", text, flags=re.MULTILINE)
|
|
if self.cup_breaks:
|
|
# remove spaces between start of line and BREAK
|
|
text = re.sub(r"^[ ]+BREAK\b", "BREAK", text, flags=re.MULTILINE)
|
|
# remove spaces between BREAK and end of line
|
|
text = re.sub(r"\bBREAK[ ]+$", "BREAK", text, flags=re.MULTILINE)
|
|
# remove at start of prompt
|
|
text = re.sub(r"\ABREAK\b", "", text)
|
|
# remove at end of prompt
|
|
text = re.sub(r"\bBREAK\Z", "", text)
|
|
if self.cup_ands:
|
|
# collapse ANDs with space after
|
|
text = re.sub(r"\bAND(?:\s+AND)+\s+", "AND ", text)
|
|
# collapse ANDs without space after
|
|
text = re.sub(r"\bAND(?:\s+AND)+\b", "AND", text)
|
|
# collapse separators and spaces before ANDs
|
|
text = re.sub(r"[, ]+AND\b", " AND", text)
|
|
# collapse separators and spaces after ANDs
|
|
text = re.sub(r"\bAND[, ]+", "AND ", text)
|
|
# remove at start of prompt
|
|
text = re.sub(r"\AAND\b", "", text)
|
|
# remove at end of prompt
|
|
text = re.sub(r"\bAND\Z", "", text)
|
|
if self.cup_extranetworktags:
|
|
#
|
|
# all cases since we can't find them inside plain text
|
|
#
|
|
# remove spaces before <
|
|
text = re.sub(r"\B\s+<(?!!)", "<", text)
|
|
# remove spaces after >
|
|
text = re.sub(r"(?<!!)>\s+\B", ">", text)
|
|
if self.cup_extraspaces:
|
|
# remove extra spaces after starting parenthesis or bracket
|
|
text = re.sub(r"([,\.;\s]+[([])\s+", r"\1", text)
|
|
# remove extra spaces before ending parenthesis or bracket
|
|
text = re.sub(r"\s+([)\]][,\.;\s]+)", r"\1", text)
|
|
# collapse spaces
|
|
# text = re.sub(r"[ ]{2,}", " ", text)
|
|
# remove spaces at start and end
|
|
text = text.strip()
|
|
return text
|
|
|
|
def __findwildcards(self, prompt, negative_prompt):
|
|
"""
|
|
Find and process wildcards in the prompt and negative_prompt strings.
|
|
|
|
Args:
|
|
prompt (str): The prompt string.
|
|
negative_prompt (str): The negative prompt string.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the processed prompt and negative_prompt strings.
|
|
"""
|
|
|
|
if self.debug:
|
|
self.logger.info("Doing wildcard processing")
|
|
|
|
p_transformtree = self.TransformerTree(self, phase="wildcards")
|
|
try:
|
|
p_tree = self.__parser_complete.parse(prompt)
|
|
prompt = p_transformtree.transform(p_tree)
|
|
except Exception as e: # pylint: disable=broad-except
|
|
self.logger.warning("Wildcards parsing failed in prompt!: %s", e)
|
|
|
|
np_transformtree = self.TransformerTree(self, phase="wildcards")
|
|
try:
|
|
np_tree = self.__parser_complete.parse(negative_prompt)
|
|
negative_prompt = np_transformtree.transform(np_tree)
|
|
except Exception as e: # pylint: disable=broad-except
|
|
self.logger.warning("Wildcards parsing failed in negative prompt!: %s", e)
|
|
|
|
foundP = False
|
|
foundNP = False
|
|
if len(p_transformtree.detectedWildcards) > 0:
|
|
foundP = True
|
|
self.logger.info(f"Found wildcards in prompt: {p_transformtree.detectedWildcards}")
|
|
if len(np_transformtree.detectedWildcards) > 0:
|
|
foundNP = True
|
|
self.logger.info(f"Found wildcards in negative prompt: {np_transformtree.detectedWildcards}")
|
|
|
|
if foundP or foundNP:
|
|
if self.ifwildcards == self.IFWILDCARDS_CHOICES["warn"]:
|
|
prompt = self.WILDCARD_WARNING + prompt
|
|
elif self.ifwildcards == self.IFWILDCARDS_CHOICES["stop"]:
|
|
self.logger.error("Found unprocessed wildcards! stopping the generation.")
|
|
if foundP:
|
|
prompt = self.WILDCARD_STOP + prompt
|
|
if foundNP:
|
|
negative_prompt = self.WILDCARD_STOP + negative_prompt
|
|
if hasattr(self.script, "ppp_interrupt"):
|
|
self.script.ppp_interrupt()
|
|
if self.debug:
|
|
self.logger.info(f"prompt after wildcards: {self.formatOutput(prompt)}")
|
|
self.logger.info(f"negative_prompt after wildcards: {self.formatOutput(negative_prompt)}")
|
|
return prompt, negative_prompt
|
|
|
|
def process_prompt(self, original_prompt, original_negative_prompt):
|
|
"""
|
|
Process the prompt and negative prompt by moving content to the negative prompt, and cleaning up.
|
|
|
|
Args:
|
|
original_prompt (str): The original prompt.
|
|
original_negative_prompt (str): The original negative prompt.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the processed prompt and negative prompt.
|
|
"""
|
|
try:
|
|
prompt = original_prompt
|
|
negative_prompt = original_negative_prompt
|
|
self.debug = getattr(self.opts, "ppp_gen_debug", False)
|
|
if not self.is_i2i or self.stn_doi2i or self.cup_doi2i:
|
|
if self.debug:
|
|
self.logger.info(f"Input prompt: {self.formatOutput(prompt)}")
|
|
p_tree = self.__parser_complete.parse(prompt)
|
|
self.logger.info(f"Tree from prompt:\n{p_tree.pretty()}")
|
|
self.logger.info(f"Input negative_prompt: {self.formatOutput(negative_prompt)}")
|
|
np_tree = self.__parser_complete.parse(negative_prompt)
|
|
self.logger.info(f"Tree from negative prompt:\n{np_tree.pretty()}")
|
|
|
|
if self.ifwildcards != self.IFWILDCARDS_CHOICES["ignore"]:
|
|
prompt, negative_prompt = self.__findwildcards(prompt, negative_prompt)
|
|
|
|
if not self.is_i2i or self.stn_doi2i:
|
|
prompt, negative_prompt = self.__sendtonegative(prompt, negative_prompt)
|
|
|
|
# pylint: disable-next=too-many-boolean-expressions
|
|
if (not self.is_i2i or self.cup_doi2i) and (
|
|
self.cup_extraspaces
|
|
or self.cup_emptyconstructs
|
|
or self.cup_extraseparators
|
|
or self.cup_breaks
|
|
or self.cup_ands
|
|
or self.cup_extranetworktags
|
|
):
|
|
prompt, negative_prompt = self.__cleanup(prompt, negative_prompt)
|
|
return prompt, negative_prompt
|
|
except Exception as e: # pylint: disable=broad-exception-caught
|
|
self.logger.exception(e)
|
|
return original_prompt, original_negative_prompt
|