Rewrite Parser Code

[wbclark]

In #55 we discussed a need to re-write the parser code to be more maintainable. I recommend that we could do this with https://github.com/lark-parser/lark

To demo what this could roughly look like, I tried out the below approach:

import abc
import dataclasses
import re
from enum import Enum
from typing import List, Tuple, Any, Optional
from lark import Lark, Transformer


class PromptKeyword(Enum):
    AND = 'AND'
    AND_PERP = 'AND_PERP'
    AND_SALT = 'AND_SALT'
    AND_TOPK = 'AND_TOPK'


prompt_keywords = [e.value for e in PromptKeyword]


class ConciliationStrategy(Enum):
    PERPENDICULAR = PromptKeyword.AND_PERP.value
    SALIENCE_MASK = PromptKeyword.AND_SALT.value
    SEMANTIC_GUIDANCE = PromptKeyword.AND_TOPK.value


conciliation_strategies = [e.value for e in ConciliationStrategy]


@dataclasses.dataclass
class PromptExpr(abc.ABC):
    weight: float
    conciliation: Optional[ConciliationStrategy]

    @abc.abstractmethod
    def accept(self, visitor, *args, **kwargs) -> Any:
        pass


@dataclasses.dataclass
class LeafPrompt(PromptExpr):
    prompt: str

    def accept(self, visitor, *args, **kwargs):
        return visitor.visit_leaf_prompt(self, *args, **kwargs)


@dataclasses.dataclass
class CompositePrompt(PromptExpr):
    children: List[PromptExpr]

    def accept(self, visitor, *args, **kwargs):
        return visitor.visit_composite_prompt(self, *args, **kwargs)


class PromptTransformer(Transformer):
    def composite_prompt(self, items):
        children = []
        conciliation = None
        for item in items:
            if isinstance(item, tuple):
                keyword, expr = item
                conciliation = ConciliationStrategy(keyword) if keyword != 'AND' else None
                if isinstance(expr, list):
                    children.extend(expr)
                else:
                    children.append(expr)
            else:
                children.append(item)
        return CompositePrompt(weight=1.0, conciliation=conciliation, children=children)

    def leaf_prompt(self, items):
        text, weight = items
        return LeafPrompt(weight=float(weight), prompt=text.strip(), conciliation=None)

    def prompt_keyword(self, items):
        return items[0]

    def AND(self, _):
        return 'AND'

    def AND_PERP(self, _):
        return 'AND_PERP'

    def AND_SALT(self, _):
        return 'AND_SALT'

    def AND_TOPK(self, _):
        return 'AND_TOPK'

    def TEXT(self, s):
        return s[0]

    def WEIGHT(self, w):
        return w[0]


grammar = """
?start: composite_prompt

composite_prompt: "[" prompt_expr (prompt_keyword prompt_expr)+ "]"

prompt_expr: leaf_prompt | composite_prompt
leaf_prompt: TEXT ":" WEIGHT
prompt_keyword: AND | AND_PERP | AND_SALT | AND_TOPK

AND: "AND"
AND_PERP: "AND_PERP"
AND_SALT: "AND_SALT"
AND_TOPK: "AND_TOPK"

WEIGHT: /\\d+(\\.\\d+)?/
TEXT: /[^:\\[\\]]+/

%import common.WS
%ignore WS
"""

parser = Lark(grammar, parser='lalr', transformer=PromptTransformer())

def parse_prompt(input_str: str) -> CompositePrompt:
    return parser.parse(input_str)

# Example usage
input_str = "[ rolling green hills :2 AND_PERP a golden retriever playing fetch :1.5 ]"
parsed_prompt = parse_prompt(input_str)
print(parsed_prompt)

It's not exactly correct yet! Right now this parses the example string as

CompositePrompt(weight=1.0, conciliation=None, children=[Tree(Token('RULE', 'prompt_expr'), [LeafPrompt(weight=2.0, conciliation=None, prompt='r')]), 'AND_PERP', Tree(Token('RULE', 'prompt_expr'), [LeafPrompt(weight=1.0, conciliation=None, prompt='a')])])

As you can see, this is incorrect in about a dozen different ways right now. But I thought it best to get your input before I try to go any further with it, in case you prefer going a different direction.

[ljleb]

I didn't see this issue until about now. My mail was flooded by notifications from sd-webui-forge.

This looks really good but it will be hard to make it work with lark. I appreciate the attempt at this, and if you ever get it to work it would be a very good contribution!

[wbclark]

Cool, I'll carve out some time to see if I can make an actually working implementation.