126 lines
4.0 KiB
Python
126 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from collections import namedtuple
|
|
from typing import List
|
|
import lark
|
|
|
|
schedule_parser = lark.Lark(r"""
|
|
!start: (prompt | /[][():]/+)*
|
|
prompt: (emphasized | scheduled | plain | WHITESPACE)*
|
|
!emphasized: "(" prompt ")"
|
|
| "(" prompt ":" prompt ")"
|
|
| "[" prompt "]"
|
|
| "{" prompt "}"
|
|
| "{" (prompt "|" prompt)* "}"
|
|
scheduled: "[" [prompt ":"] prompt ":" [WHITESPACE] NUMBER "]"
|
|
WHITESPACE: /\s+/
|
|
plain: /([^\\\[\]():|{}]|\\.)+/
|
|
%import common.SIGNED_NUMBER -> NUMBER
|
|
""")
|
|
|
|
def conditioning_set_timeset_range_prompt_mapper(clip):
|
|
previous_step = [0]
|
|
def wrapped(prompt):
|
|
start = previous_step[0]
|
|
end = prompt[0]
|
|
previous_step[0] = end
|
|
|
|
tokens = clip.tokenize(prompt[1])
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
|
conditioning = [[cond, {"pooled_output": pooled}]]
|
|
return set_range(conditioning, start, end)
|
|
|
|
def set_range(conditioning, start, end):
|
|
c = []
|
|
for t in conditioning:
|
|
d = t[1].copy()
|
|
d['start_percent'] = start
|
|
d['end_percent'] = end
|
|
n = [t[0], d]
|
|
c.append(n)
|
|
return c
|
|
return wrapped
|
|
|
|
def get_learned_conditioning_prompt_schedules(prompts, env='prod'):
|
|
"""
|
|
>>> g = lambda p: get_learned_conditioning_prompt_schedules([p], 10)[0]
|
|
>>> g("test")
|
|
[[10, 'test']]
|
|
>>> g("a [b:3]")
|
|
[[3, 'a '], [10, 'a b']]
|
|
>>> g("a [b: 3]")
|
|
[[3, 'a '], [10, 'a b']]
|
|
>>> g("a [[[b]]:2]")
|
|
[[2, 'a '], [10, 'a [[b]]']]
|
|
>>> g("[(a:2):3]")
|
|
[[3, ''], [10, '(a:2)']]
|
|
>>> g("a [b : c : 1] d")
|
|
[[1, 'a b d'], [10, 'a c d']]
|
|
>>> g("a[b:[c:d:2]:1]e")
|
|
[[1, 'abe'], [2, 'ace'], [10, 'ade']]
|
|
>>> g("a [unbalanced")
|
|
[[10, 'a [unbalanced']]
|
|
>>> g("a [b:.5] c")
|
|
[[5, 'a c'], [10, 'a b c']]
|
|
>>> g("a [{b|d{:.5] c") # not handling this right now
|
|
[[5, 'a c'], [10, 'a {b|d{ c']]
|
|
>>> g("((a][:b:c [d:3]")
|
|
[[3, '((a][:b:c '], [10, '((a][:b:c d']]
|
|
>>> g("[a|(b:1.1)]")
|
|
[[1, 'a'], [2, '(b:1.1)'], [3, 'a'], [4, '(b:1.1)'], [5, 'a'], [6, '(b:1.1)'], [7, 'a'], [8, '(b:1.1)'], [9, 'a'], [10, '(b:1.1)']]
|
|
"""
|
|
|
|
def collect_steps(tree):
|
|
res = [1]
|
|
|
|
class CollectSteps(lark.Visitor):
|
|
def scheduled(self, tree):
|
|
tree.children[-1] = float(tree.children[-1])
|
|
res.append(tree.children[-1])
|
|
|
|
CollectSteps().visit(tree)
|
|
return sorted(set(res))
|
|
|
|
def at_step(step, tree):
|
|
class AtStep(lark.Transformer):
|
|
def scheduled(self, args):
|
|
before, after, _, when = args
|
|
yield before or () if step <= when else after
|
|
def start(self, args):
|
|
def flatten(x):
|
|
if type(x) == str:
|
|
yield x
|
|
else:
|
|
for gen in x:
|
|
yield from flatten(gen)
|
|
return ''.join(flatten(args))
|
|
def plain(self, args):
|
|
yield args[0].value
|
|
def __default__(self, data, children, meta):
|
|
for child in children:
|
|
yield child
|
|
return AtStep().transform(tree)
|
|
|
|
def get_schedule(prompt):
|
|
try:
|
|
tree = schedule_parser.parse(prompt)
|
|
except lark.exceptions.LarkError as e:
|
|
if 0:
|
|
import traceback
|
|
traceback.print_exc()
|
|
if(env != 'prod'):
|
|
return {"error": "Illegal characters", "logs": str(e), "prompt": prompt}
|
|
else:
|
|
return [[1, prompt]]
|
|
return [[t, at_step(t, tree)] for t in collect_steps(tree)]
|
|
|
|
promptdict = {prompt: get_schedule(prompt) for prompt in set(prompts)}
|
|
return [promptdict[prompt] for prompt in prompts]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
|
|
else:
|
|
import torch # doctest faster |