from typing import Tuple, List, Dict, Optional, Set
from extensions.output_template.symbols import Symbol, Terminal, NonTerminal, Sequence, Alternative, Repeat, RegExp, \
    AnyToken
from extensions.output_template.state_machine import Advance, Matcher
from extensions.output_template.utils import shared, AllowedTokens
import torch, re

RE_RULE = re.compile(r'\s*([-a-z]+)\s*::=\s*(.*)', re.MULTILINE | re.DOTALL)
RE_NEWLINE = re.compile(r'[ \t]*\n[ \t\n]*(.*)', re.MULTILINE | re.DOTALL)
RE_TERMINAL = re.compile(r'[ \t]*([-a-z]+)[ \t]*(.*)', re.DOTALL)
RE_ANYTOKEN = re.compile(r'[ \t]*\.[ \t]*\*[ \t]*(.*)', re.DOTALL)
RE_OR = re.compile(r'[ \t\n]*\|[ \t]*(.*)', re.MULTILINE | re.DOTALL)
RE_COMMENT = re.compile(r'([^#]*)#[^\n]*(.*)', re.MULTILINE | re.DOTALL)


class Grammar:
    """ Grammar and also state machine used to match LLM output """

    def __init__(self, definition: str):
        self.rules: Dict[str, Symbol] = {}
        self.active_matcher: Optional[Matcher] = None
        self.reset(definition)

    def stop(self):
        self.active_matcher = None

    def reset(self, definition: str = None):
        self.stop()
        if definition:
            text = definition
            self.rules = {}

            # Strip comments
            m = RE_COMMENT.match(text)
            while m:
                text = m.group(1) + m.group(2)
                m = RE_COMMENT.match(text)

            while text:
                m = RE_RULE.match(text)
                if not m:
                    raise GrammarError("expected rule")
                rule_name, text = m.groups()
                if rule_name in self.rules:
                    raise ValidationError(f"duplicate rule '{rule_name}'")
                self.rules[rule_name], text = parse_rule(text)

            if "root" not in self.rules:
                raise ValidationError("missing 'root' rule")
            for rule in self.rules.values():
                rule.validate(self)

        self.enter_rule("root")

    def resolve(self, symbol: "Symbol") -> "Symbol":
        # Resolves NonTerminal into rule and returns Symbol it represents
        dont_loop: Set[Terminal] = set([symbol])
        while isinstance(symbol, NonTerminal):
            dont_loop.add(symbol)
            if symbol.name not in self.rules:
                raise ValidationError(f"invalid rule name: '{symbol.name}'")
            if self.rules[symbol.name] in dont_loop:
                raise ValidationError(f"infinite loop detected at symbol '{symbol.name}'")
            symbol = self.rules[symbol.name]
        return symbol

    def enter_rule(self, name: str):
        """
        Sets active symbol to specific rule.
        Used for testing.
        """
        if name not in self.rules:
            raise ValueError(f"invalid rule name: '{name}'")
        self.active_matcher = self.resolve(self.rules[name]).enter(self)

    def get_rule_name(self, symbol: "Symbol"):
        for (name, s) in self.rules.items():
            if symbol is s:
                return name
        return None

    def get_effective_matcher(self) -> Optional["Matcher"]:
        """
        Recursively descends rule hierarchy and returns symbol
        that will effectively decide on next token
        """
        return self.active_matcher.get_effective_matcher() if self.active_matcher else None

    def update_scores(self, scores: torch.FloatTensor) -> torch.FloatTensor:
        """
        Calculates probability scores of next token according to current state.
        May update and return same object as one that was passed as argument.
        """
        # TODO: how to cache or optimize this?
        if self.active_matcher:
            allowed = self.active_matcher.get_allowed_tokens(self)
            if allowed.look_ahead:
                allowed.allow_eos = True

            allowed.apply(scores)
        else:
            # Grammar reached terminal token. Force EOS
            AllowedTokens(allowed={int(shared.tokenizer.eos_token_id)}, allow_eos=True).apply(scores)
        return scores

    def advance(self, token_id: int):
        try:
            if self.active_matcher:
                # from extensions.output_template.script import logger
                # logger.warning(f"Feeding {token_id} into {self.active_matcher}.")
                a = self.active_matcher.advance(self, token_id)
                if a == Advance.Reject:
                    if token_id == shared.tokenizer.eos_token_id:
                        self.active_matcher = None
                    else:
                        raise GenerationError
                elif a == Advance.Done:
                    self.active_matcher = None
        except GenerationError as e:
            from extensions.output_template.script import logger
            logger.warning("LLM failed to generate token conforming to grammar")
            self.active_matcher = None


class GrammarError(ValueError):
    pass


class ValidationError(GrammarError):
    pass


class GenerationError(Exception):
    pass


def find_unescaped_index(haystack: str, needle: str, start=0) -> int:
    index = start
    while True:
        index, index2 = haystack.find(needle, index), haystack.find("\\", index)
        if index < 0:
            return len(haystack)
        if index2 >= 0 and index2 < index:
            index = index2 + 2
        else:
            return index


def parse_sequence(text: str, parentheses=False) -> Tuple[Sequence, str]:
    seq = []
    while text:
        if text[0] in '"\'':
            # Terminal rule
            try:
                end_index = find_unescaped_index(text, text[0], 1)
                t = text[1:end_index].encode("utf-8").decode("unicode_escape")
                seq.append(Terminal(t))
                text = text[end_index+1:]
            except ValueError:
                raise GrammarError(f"unmatched {text[0]}")
        elif RE_TERMINAL.match(text):
            # Non-terminal rule
            t, text = RE_TERMINAL.match(text).groups()
            seq.append(NonTerminal(t))
        elif text[0] in " \t":
            # Whitespace
            text = text[1:]
        elif text[0] == "[":
            # Regexp rule
            try:
                end_index = find_unescaped_index(text, "]", 1)
            except ValueError:
                raise GrammarError(f"unmatched {text[0]}")
            try:
                seq.append(RegExp(text[0:end_index+1]))
                text = text[end_index+1:]
            except ValueError:
                raise GrammarError(f"invalid pattern {text[0:end_index]}")
        elif text[0] == "(":
            # Parenthesized rule
            text = text[1:]
            t, text = parse_rule(text, parentheses=True)
            seq.append(t)
            pass
        elif parentheses and text[0] == ")":
            text = text[1:]
            break
        elif RE_ANYTOKEN.match(text):
            # '.*' rule
            text, = RE_ANYTOKEN.match(text).groups()
            seq.append(AnyToken())
        elif text[0] in "*?+":
            # Repeat rule
            if not seq:
                raise GrammarError(f"unexpected '{text[0]}'")
            left = seq.pop()
            if text[0] == "+":
                # A+ is converted into sequence (A A*)
                if text[0] in "+" and isinstance(left, RegExp):
                    # If child is regexp, extend its rule so multi-character tokens are matched
                    left = RegExp(left.value + text[0])
                seq.append(Sequence([left, Repeat("*", left)]))
            else:
                seq.append(Repeat(text[0], left))
            text = text[1:]
        elif RE_OR.match(text):
            text, = RE_OR.match(text).groups()
            if not seq:
                raise GrammarError(f"unexpected '|'")
            left = seq.pop()
            right, text = parse_rule(text, parentheses=parentheses)
            seq.append(Alternative([left, right]))
            break
        elif RE_NEWLINE.match(text):
            # Newline
            text, = RE_NEWLINE.match(text).groups()
            if not parentheses:
                break
        else:
            raise GrammarError(f"unexpected '{text[0:5]}'...")

    return Sequence(seq), text


def parse_rule(text: str, parentheses=False) -> Tuple[Symbol, str]:
    rv, text = parse_sequence(text, parentheses=parentheses)
    if len(rv.items) == 1:
        rv = rv.items[0]
    return rv, text

grammar.py
im-not-tom/text-generation-webui-output-template-extension
from typing import Optional
from extensions.output_template.grammar import Grammar
from extensions.output_template.utils import shared
from functools import partial
import torch, transformers
try:
    from modules.logging_colors import logger
except ModuleNotFoundError:
    # Allows testing by running script outside text-generation-ui
    from logging import Logger
    logger = Logger(__name__)
    transformers = None
    class LogitsProcessor:
        pass


EMPTY_GRAMMAR = "root ::= .*"

params = {
    "grammar": Grammar(EMPTY_GRAMMAR),
    "enabled": False,
    "template": "",
    "token_dictionary": None,
    "used_tokenizer": None,
    "scores_size": 0,
}


class TemplatingLogitsProcessor(transformers.LogitsProcessor):

    def __init__(self):
        super().__init__()
        self.last_input_size = 0

    def __call__(self, input_ids: Optional[torch.LongTensor], scores: torch.FloatTensor):
        if params["enabled"]:
            params["scores_size"] = len(scores[0])
            grammar: Grammar = params["grammar"]

            if input_ids is not None:
                # input_ids are None when running from tests.
                input_size = len(input_ids[0])
                if input_size <= self.last_input_size:
                    logger.warning("output_template: input size unexpectedly decreased. Restarting grammar (except wrong output)")
                    grammar.reset()
                elif self.last_input_size != 0:
                    for token_id in input_ids[0][self.last_input_size:]:
                        grammar.advance(int(token_id))
                self.last_input_size = input_size

            return grammar.update_scores(scores)
        return scores


def logits_processor_modifier(processor_list, input_ids):
    """
    Adds logits processors to the list, allowing you to access and modify
    the next token probabilities.
    Only used by loaders that use the transformers library for sampling.
    """
    processor_list.append(TemplatingLogitsProcessor())
    return processor_list


def token_generated_callback(input_ids, scores):
    if params["enabled"]:
        grammar: Grammar = params["grammar"]
        new_token = input_ids[0][-1]
        grammar.advance(int(new_token))


def input_modifier(string, state, is_chat=False):
    """
    Initializes template and appends initial simple text to input.
    Note: In chat_mode, this extension does nothing.
    """
    if not is_chat:
        if "grammar" in state or params["template"]:
            grammar: Grammar = params["grammar"]
            if "grammar" in state:
                grammar.reset(state["grammar"] or EMPTY_GRAMMAR)
            else:
                grammar.reset(params["template"])
            params["enabled"] = True
        else:
            params["enabled"] = False
    return string


def ui():
    import gradio as gr
    output_template = gr.Textbox(value="", placeholder="Enter output template", label="Output Template",
                               info='Output Template to use', lines=5)

    def update_output_template(x):
        logger.info("output_template updated")
        params.update({"template": x})

    output_template.change(update_output_template, output_template, None)

script.py
im-not-tom/text-generation-webui-output-template-extension
from typing import List, Optional, Set
from extensions.output_template.utils import get_token_dictionary, AllowedTokens
from enum import IntEnum


class Advance(IntEnum):
    Again = 0
    Done = 1
    Reject = 2
    TryNext = 3


class Matcher:
    ''' Piece of state machine used to match the Symbol '''

    def __init__(self, symbol: "Symbol"):
        self.symbol = symbol

    def __repr__(self):
        return f"<{self.__class__.__name__} {self.debug()}>"

    def debug(self) -> str:
        """
        Returns string describing current state.
        There's no logic here, just what seem'd most readable at the time
        """
        return repr(self.symbol)

    def get_effective_matcher(self) -> Optional["Matcher"]:
        raise NotImplementedError(f"get_effective_matcher on {self.__class__.__name__}") 

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        """
        Returns Advance.Again if given token matches symbol only partially and more tokens are expected.
        Returns Advance.Done if given token (at current state) matches rest of this symbol.
        Method should not be called again after returning Advance.Done.
        Returns Advance.Reject if token doesn't match
        Returns Advance.TryNext if not matching token is expected and next symbol should be tried
        """
        raise NotImplementedError(f"advance on {self.__class__.__name__}")

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        raise NotImplementedError(f"get_allowed_tokens on {self.__class__.__name__}")


class TerminalMatcher(Matcher):
    symbol: "Terminal"

    def __init__(self, t: "Terminal"):
        super().__init__(t)
        self.index = 0

    def debug(self) -> str:
        if self.index <= 0 or self.index >= len(self.symbol.value):
            return repr(self.symbol)
        return f"""t'{repr(self.symbol.value[:self.index]).strip("'")}[{repr(self.symbol.value[self.index:]).strip("'")}]'"""

    def get_effective_matcher(self) -> "Matcher":
        return self

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        if self.index not in self.symbol.allowed_cache:
            d = get_token_dictionary()
            allowed = set()
            for token_id in d:
                if d[token_id] and d[token_id] == self.symbol.value[self.index:self.index+len(d[token_id])]:
                    allowed.add(token_id)
            self.symbol.allowed_cache[self.index] = allowed
        return AllowedTokens(allowed=self.symbol.allowed_cache[self.index])

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        d = get_token_dictionary()
        if d[token_id] != self.symbol.value[self.index:self.index+len(d[token_id])]:
            return Advance.Reject
        self.index += len(get_token_dictionary()[token_id])
        if self.index >= len(self.symbol.value):
            return Advance.Done
        return Advance.Again


class RegExpMatcher(Matcher):
    symbol: "RegExp"

    def __init__(self, t: "RegExp"):
        super().__init__(t)

    def debug(self) -> str:
        return f"""r{self.symbol.value}"""

    def get_effective_matcher(self) -> "Matcher":
        return self

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        if self.symbol.negative:
            if self.symbol.value not in self.symbol.banned_cache:
                d = get_token_dictionary()
                banned = set()
                for token_id in d:
                    if self.symbol.re.search(d[token_id]):
                        banned.add(token_id)
                self.symbol.banned_cache[self.symbol.value] = banned
            return AllowedTokens(banned=self.symbol.banned_cache[self.symbol.value])
        else:
            if self.symbol.value not in self.symbol.allowed_cache:
                d = get_token_dictionary()
                allowed = set()
                for token_id in d:
                    if self.symbol.re.match(d[token_id]):
                        allowed.add(token_id)
                self.symbol.allowed_cache[self.symbol.value] = allowed
            return AllowedTokens(allowed=self.symbol.allowed_cache[self.symbol.value])

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        d = get_token_dictionary()
        if self.symbol.negative:
            if self.symbol.re.search(d[token_id]):
                return Advance.Reject
        elif not self.symbol.re.match(d[token_id]):
            return Advance.Reject
        # TODO: use index? How to deal with tokens that match partially?
        return Advance.Done


class AnyTokenMatcher(Matcher):
    symbol: "AnyToken"

    def debug(self) -> str:
        return ".*"

    def get_effective_matcher(self) -> "Matcher":
        return self

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        return AllowedTokens(allow_eos=True)

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        return Advance.Again


class SequenceMatcher(Matcher):
    symbol: "Sequence"

    def __init__(self, symbol: "Sequence", items: List[Optional[Matcher]]):
        super().__init__(symbol)
        self.items = items
        self.index = 0

    def debug(self) -> str:
        return f'''({" ".join([
            f"[{self.items[i].debug()}]" if i == self.index and self.items[i]
            else repr(self.symbol.items[i])
            for i in range(len(self.symbol.items))
        ])})'''

    def get_effective_matcher(self) -> Optional[Matcher]:
        assert self.index < len(self.items)
        if self.items[self.index]:
            return self.items[self.index].get_effective_matcher()
        return None

    def ensure_matcher(self, g: "Grammar", i=0) -> "SequenceMatcher":
        if not self.items[i]:
            self.items[i] = g.resolve(self.symbol.items[i]).enter(g)
        return self

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        assert self.index < len(self.items)
        rv = self.items[self.index].get_allowed_tokens(g)
        if rv.look_ahead:
            i = self.index
            ahead = rv
            while i < len(self.symbol.items) - 1:
                i += 1
                self.ensure_matcher(g, i)
                ahead = self.items[i].get_allowed_tokens(g)
                rv = rv.combine(ahead)
                if not ahead.look_ahead:
                    break
            if not ahead.look_ahead:
                rv.look_ahead = False
        return rv

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        a = self.items[self.index].advance(g, token_id)
        if a in (Advance.Done, Advance.TryNext):
            if self.index < len(self.symbol.items) - 1:
                self.index += 1
                self.ensure_matcher(g, self.index)
                if a == Advance.TryNext:
                    return self.advance(g, token_id)
                a = Advance.Again
        return a


class AlternativeMatcher(Matcher):
    symbol: "Alternative"

    def __init__(self, symbol: "Alternative", items: Set[Matcher]):
        super().__init__(symbol)
        self.items = items

    def debug(self):
        return f'({" | ".join([x.debug() for x in self.items])})'

    def get_effective_matcher(self) -> "Matcher":
        if len(self.items) == 1:
            return list(self.items)[0].get_effective_matcher()
        return self

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        rv = None
        for i in self.items:
            a = i.get_allowed_tokens(g)
            rv = rv.combine(a) if rv else a
        # TODO: should this return 'ban everything' if no alternative is left?
        # TODO: should such state be even possible?
        return rv or AllowedTokens()

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        best_a = Advance.Reject
        for i in list(self.items):
            a = i.advance(g, token_id)
            if a in (Advance.Reject, Advance.TryNext):
                if a == Advance.TryNext and best_a == Advance.Reject:
                    best_a = Advance.TryNext
                self.items.remove(i)
            else:
                best_a = Advance.Done
                if a == Advance.Done:
                    self.items.remove(i)
        if len(self.items) == 0:
            return best_a
        return Advance.Again


class RepeatMatcher(Matcher):
    symbol: "Repeat"

    def __init__(self, symbol: "Symbol", effective_item: Matcher):
        super().__init__(symbol)
        self.effective_item = effective_item
        self.inside = False

    def get_effective_matcher(self) -> "Matcher":
        return self.effective_item.get_effective_matcher() if self.inside else self

    def get_allowed_tokens(self, g: "Grammar") -> AllowedTokens:
        rv = self.effective_item.get_allowed_tokens(g)
        if not self.inside:
            return rv.set_ahead()
        return rv

    def advance(self, g: "Grammar", token_id: int) -> Advance:
        a = self.effective_item.advance(g, token_id)
        if a == Advance.Reject:
            if self.inside:
                return a
            return Advance.TryNext
        elif a == Advance.Done:
            if self.symbol.mode == "*":
                self.effective_item = g.resolve(self.symbol.item).enter(g)
                self.inside = False
                return Advance.Again
            else:   # mode == "?"
                return Advance.Done
        elif a == Advance.Again:
            self.inside = True
        return a

state_machine.py
im-not-tom/text-generation-webui-output-template-extension
from extensions.output_template.state_machine import *
import re


class Symbol:
    def validate(self, g: "Grammar"):
        """
        Validates against grammar.
        Currently just checks that for each Terminal there is rule defined.
        """
        pass

    def enter(self, g: "Grammar") -> Matcher:
        raise NotImplementedError(f"enter on {self.__class__.__name__}")


class NonTerminal(Symbol):
    def __init__(self, name: str):
        self.name = name

    def __repr__(self):
        return self.name

    def validate(self, g: "Grammar"):
        g.resolve(self).validate(g)

    def enter(self, g: "Grammar") -> Matcher:
        return g.resolve(self).enter(g)


class Terminal(Symbol):
    def __init__(self, value: str):
        self.value = value
        self.allowed_cache = {}

    def __repr__(self):
        return f't{repr(self.value)}'

    def validate(self, g: "Grammar"):
        if not self.value:
            raise ValueError("empty terminal")

    def enter(self, g: "Grammar") -> Matcher:
        return TerminalMatcher(self)


class RegExp(Symbol):
    allowed_cache = {}
    banned_cache = {}

    def __init__(self, value: str):
        self.value = value
        if self.value.startswith("[^"):
            # Special handling to optimize.
            # Instead of allowing all tokens but these, rest is banned
            self.negative = True
            self.re = re.compile("[" + value[2:], re.MULTILINE | re.DOTALL)
        else:
            self.negative = False
            self.re = re.compile("^" + value + "$", re.MULTILINE | re.DOTALL)

    def __repr__(self):
        return f'r{self.value}'

    def enter(self, g: "Grammar") -> Matcher:
        return RegExpMatcher(self)


class AnyToken(Symbol):
    """
    Special symbol to which '.*' is translated.
    Just matches anything, basically turning grammar off.
    """

    def __repr__(self):
        return f'.*'

    def enter(self, g: "Grammar") -> Matcher:
        return AnyTokenMatcher(self)


class Collection(Symbol):
    def __init__(self, items: List[Symbol]):
        self.items = items


class Sequence(Collection):
    def __init__(self, items: List[Symbol]):
        super().__init__(items)
        self.effective = []
        self.index = 0

    def __repr__(self):
        return f'({" ".join([repr(x) for x in self.items])})'

    def enter(self, g: "Grammar") -> Matcher:
        return SequenceMatcher(self, [
            None
            for m in self.items
        ]).ensure_matcher(g)


class Alternative(Collection):
    def __init__(self, items: List[Matcher]):
        super().__init__([])
        for i in items:
            if isinstance(i, Alternative):
                self.items += i.items
            else:
                self.items.append(i)
        self.possible = set()

    def __repr__(self):
        return f'({" | ".join([repr(x) for x in self.items])})'

    def validate(self, g: "Grammar"):
        for item in self.items:
            g.resolve(item).validate(g)

    def enter(self, g: "Grammar") -> "Matcher":
        return AlternativeMatcher(self, {
            g.resolve(m).enter(g)
            for m in self.items
        })


class Repeat(Symbol):
    def __init__(self, mode: str, item: Symbol):
        assert mode in "*?"
        self.item = item
        self.mode = mode

    def __repr__(self):
        if isinstance(self.item, (Terminal, NonTerminal)):
            return f'({repr(self.item)}){self.mode}'
        else:
            return f'{repr(self.item)}{self.mode}'

    def validate(self, g: "Grammar"):
        self.item.validate(g)

    def enter(self, g: "Grammar") -> Matcher:
        return RepeatMatcher(self, g.resolve(self.item).enter(g))

symbols.py
im-not-tom/text-generation-webui-output-template-extension
"""
Fake tokenizer implementation used for testing.
Uses only 127 tokens with few IDs reserved for multi-character tokens
to simulate possible issues.
"""
from typing import List

tokens = {
    0:  "\x00",  # EOS
    10: "\n",    # newline
    # Randomly chosen strings
    # Ordered from longest to shortest and order is important.
    1:  "Universe",
    2:  "token",
    3:  "world",
    4:  "stock",
    5:  '..."',
    6:  "hall",
    7:  "the",
    8:  "...",
    9:  "and",
    11: "com",
    12: "Neg",
    13: "   ",
    14: "end",
    15: "six",
    16: "tab",
    17: "- [",
    18: "gg",
    19: "He",
    20: "- ",
    21: "ni",
    22: "oo",
    23: "[]",
    24: "or",
    25: "ro",
    26: "),",
    27: "of",
    28: "to",
    29: "by",
    30: "++",
    31: "],",
    # From space above, all printable characters
    **{i: chr(i) for i in range(32, 127)},
}


def encode(text) -> List[int]:
    rv = []
    while text:
        for i in tokens:
            if text.startswith(tokens[i]):
                rv.append(i)
                text = text[len(tokens[i]):]
                break
        else:
            text = text[1:]
    return rv


def decode(token_ids: List[int]) -> str:
    return "".join([tokens[i] for i in token_ids])

test_tokenizer.py
im-not-tom/text-generation-webui-output-template-extension
import os
from typing import Union

os.environ["OT_TESTING"] = "1"
from extensions.output_template.script import TemplatingLogitsProcessor, params
from extensions.output_template.utils import encode, decode, shared, MINUS_INF
from extensions.output_template.grammar import Grammar, Repeat, RegExp
from extensions.output_template.state_machine import AnyTokenMatcher
from torch import Tensor
import math, random, json

EOS = shared.tokenizer.eos_token_id
TEMPLATE = """
root ::= "Alice:" space action
greeting ::= "hello \\"world\\" \\U00003041"
action ::= speech | bullet | command
space ::= " "
bullet ::= ("- " "I'll go" space "to" space location space "and do") [^\n]+
command ::= "/go " location
location ::= ("hall" | "kitchen")
speech ::= "\\"" [^"\n]+ "\\""
"""


def random_scores():
    return Tensor([[ 0.0001 + (math.floor(random.random() * 100) / 100.0) for _ in range(127) ]])


def set_score(token_id: Union[str, int, list], scores, value=1000.0):
    if type(token_id) is list:
        for i in token_id:
            scores = set_score(i, scores, value)
        return scores
    if type(token_id) is str:
        token_id = encode(token_id)[0]
    scores[..., token_id] = value
    return scores


def scores_to_text(scores):
    scores = scores[0]
    return " ".join([
        f"{repr(chr(i))}:{scores[i]}"
        for i in range(len(scores))
        if scores[i] > 0
    ])


def sample_test(scores) -> int:
    # Returns single generated token
    TemplatingLogitsProcessor()(None, scores)
    best = int(scores.argmax())
    grammar: Grammar = params["grammar"]
    grammar.advance(best)
    return best


def get_text(until=EOS, score_fn=random_scores) -> str:
    tokens = []
    while True:
        t = sample_test(score_fn())
        if t != EOS:
            tokens.append(t)
        if type(until) is str and until in decode([t]):
            break
        elif t == until:
            break

    return decode(tokens)


def test_grammar_parser():
    g = Grammar(TEMPLATE)
    assert len(g.rules) == 8

    g.reset("""
        root ::= "hi"
        # Testing case when grammar ends with non-terminated line with comment""")


def test_terminal():
    grammar: Grammar = params["grammar"]
    grammar.reset("""root ::= 'Hello world' [\n]+""")
    EOL = encode("\n")[0]
    t = get_text(EOL)
    assert t == "Hello world\n"
    assert grammar.active_matcher

    grammar.reset()
    scores = grammar.update_scores(random_scores())
    assert len(encode("He")) == 1
    assert scores[..., encode("He")] > MINUS_INF
    assert scores[..., encode("H")] > MINUS_INF
    grammar.advance(encode("H")[0])
    assert ord('e') == sample_test(random_scores())
    matcher = grammar.get_effective_matcher()
    while grammar.get_effective_matcher() is matcher:
        sample_test(random_scores())
    assert isinstance(grammar.get_effective_matcher().symbol, (RegExp, Repeat))


def test_alternate():
    grammar: Grammar = params["grammar"]
    grammar.reset(TEMPLATE)
    grammar.enter_rule("action")
    sample_test(set_score("/", random_scores()))
    text = get_text()
    assert text in ("go hall", "go kitchen")


def test_sequence():
    grammar: Grammar = params["grammar"]
    grammar.reset(TEMPLATE)
    get_text(encode(" ")[0])
    sample_test(set_score("-", random_scores()))
    assert 32 == sample_test(random_scores())
    tokens = []
    while not isinstance(grammar.get_effective_matcher().symbol, (Repeat, RegExp)):
        t = sample_test(random_scores())
        if t == EOS:
            break
        tokens.append(t)

    assert decode(tokens) in (
        "I'll go to hall and do",
        "I'll go to kitchen and do",
    )


def test_regexp():
    grammar: Grammar = params["grammar"]
    grammar.reset("root ::= [a-z]")
    scores = TemplatingLogitsProcessor()(None, random_scores())
    assert len([x for x in scores[0] if x > MINUS_INF]) == 26
    assert scores[..., EOS] == MINUS_INF        # Also make sure that EOS is banned in any case

    # Tests that combination of repetition and regexp allows multi-char tokens
    grammar.reset("root ::= [a-z]+")
    scores = TemplatingLogitsProcessor()(None, random_scores())
    assert len([x for x in scores[0] if x > MINUS_INF]) > 26
    assert scores[..., encode("world")] > MINUS_INF
    assert scores[..., EOS] == MINUS_INF

    # Tests banning specific character
    grammar.reset("root ::= [^\n]+")
    scores = TemplatingLogitsProcessor()(None, random_scores())
    assert scores[..., 10] < 0
    assert scores[..., EOS] == MINUS_INF


def test_repeat():
    # Tests "speech" rule in TEMPLATE grammar
    grammar: Grammar = params["grammar"]
    grammar.reset(TEMPLATE)
    get_text(encode(" ")[0])
    sample_test(set_score('"', random_scores()))
    q = encode('"')[0]
    for _ in range(10):
        # Just generate few random chars
        assert sample_test(set_score('"', random_scores(), MINUS_INF)) not in (EOS, q)
    # Test that " is still allowed to be generated
    scores = TemplatingLogitsProcessor()(None, set_score('"', random_scores()))
    assert scores[..., q] > MINUS_INF
    assert q == sample_test(scores)
    assert EOS == sample_test(random_scores())

    # Tests grammar that generates repeated combination of foobars
    grammar: Grammar = params["grammar"]
    grammar.reset("""
        root ::= many
        many ::= one one one+
        one ::= foo | bar
        foo ::= "foo"
        bar ::= "b" "a"+ "r"
    """)
    text = get_text()
    assert "foo" in text or "ba" in text

    # Actual case that was broken originally
    grammar: Grammar = params["grammar"]
    grammar.reset("""
        root ::= line line line
        line ::= "Alice:" space action newline
        action ::= speech | bullet | command
        space ::= " "
        newline ::= "\n"
        bullet ::= ("- " "I'll go" space "to" space location space "and do") gibberish
        command ::= "/go" space location
        location ::= ("hall" | "out" | "room")
        speech ::= "\\"" ">" gibberish "\\""
        gibberish ::= [^-/">\n]+
    """)
    text = get_text()
    commands = len(text.split("/go")) - 1
    bullets = len(text.split("-")) - 1
    speech = len(text.split(">")) - 1
    assert 3 == commands + bullets + speech


def test_json():
    grammar: Grammar = params["grammar"]
    grammar.reset("""
        root   ::= object
        value  ::= object | array | string | number | ("true" | "false" | "null") ws

        object ::=
          "{" ws (
                    string ":" ws value
            ("," ws string ":" ws value)*
          )? "}" ws

        array  ::=
          "[" ws (
                    value
            ("," ws value)*
          )? "]" ws

        string ::=
          "\\"" (
            [^"\\\\\n] |
            "\\\\" (["\\\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])        # escapes
          )* "\\"" ws

        number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

        # Optional space: by convention, applied in this grammar after literal chars when allowed
        ws ::= ([ \\t\\n] ws)?
    """)

    random.seed(2342343231)
    # 1st token has to be {
    assert ord("{") == sample_test(random_scores())
    # Any number of whitespace has to be allowed
    assert ord(" ") == sample_test(set_score(32, random_scores(), 1000.0))
    assert ord(" ") == sample_test(set_score(32, random_scores(), 1000.0))
    assert ord(" ") == sample_test(set_score(32, random_scores(), 1000.0))
    assert ord("\n") == sample_test(set_score(10, random_scores(), 1000.0))
    # Ban }. If not whitespace, only '"' may follow
    assert ord('"') == sample_test(set_score(list("} \n"), random_scores(), MINUS_INF))
    # Grab some tokens and force string to finish
    for _ in range(5):
        sample_test(set_score('"', random_scores(), MINUS_INF))
    assert ord('"') == sample_test(set_score('"', random_scores(), 1000))
    # Now only whitespace and ':' should be allowed
    scores = random_scores()
    TemplatingLogitsProcessor()(None, scores)
    assert scores[..., ord(':')] > 0
    assert len([x for x in scores[0] if x > 0]) == 3    # ':', space and newline
    # Go over ':' and ban whitespace. Now start-of-value tokens should be allowed
    sample_test(set_score(':', random_scores()))
    scores = set_score(list(' \n'), random_scores(), MINUS_INF)
    TemplatingLogitsProcessor()(None, scores)
    # 14 characters =  0-9, minus sign, {, [ and '"'.
    # Additionally, 3 characters for n, t and f of null, true and false  
    assert len([ x for x in scores[0] if x > 0 ]) == 17
    # Force 'true' and verify repetition
    sample_test(set_score('t', random_scores()))
    assert "true" == 't' + get_text('e')
    assert ord(',') == sample_test(set_score(list('} \n'), random_scores(), MINUS_INF))
    # Force another key and make sure it generates quotes properly (this was failing before)
    sample_test(set_score('"', random_scores()))

    a = get_text('"')
    assert '"' not in a[:-1]
    scores = random_scores()
    TemplatingLogitsProcessor()(None, scores)
    assert len([x for x in scores[0] if x > 0]) == 3    # ':', space and newline
    assert ord(':') == sample_test(set_score(list(' \n'), random_scores(), MINUS_INF))

    # Now just restart grammar and test it generates some proper json
    for _ in range(100):
        # (and do it few times, cos it tends to generate just {})
        grammar.reset()
        a = get_text()
        json.loads(a)


def test_any_token():
    """ Test nonstandard rule to disable grammar """
    grammar: Grammar = params["grammar"]
    grammar.reset("""
        root ::= (donotend)
        donotend ::= (.*)
    """)
    for i in range(255):
        z = sample_test(random_scores())
    assert isinstance(grammar.get_effective_matcher(), AnyTokenMatcher)


if __name__ == "__main__":
    params["scores_size"] = 127
    params["enabled"] = True
    test_grammar_parser()
    test_terminal()
    test_alternate()
    test_sequence()
    test_regexp()
    test_repeat()
    test_json()
    test_any_token()

tests.py
im-not-tom/text-generation-webui-output-template-extension
from typing import List, Set, Dict
import torch, os
MINUS_INF = -float("inf")


if "OT_TESTING" in os.environ:
    from collections import namedtuple
    shared = namedtuple("Shared", ['tokenizer'])(namedtuple("Tokenizer", ['eos_token_id'])(0))
    from extensions.output_template.test_tokenizer import encode, decode 
else:
    from modules import shared

    def encode(text) -> List[int]:
        return shared.tokenizer.encode(str(text), add_special_tokens=False)

    def decode(token_ids: List[int]) -> str:
        return shared.tokenizer.decode(token_ids)


class AllowedTokens:
    """
    Utility class used to combine and eventually allow (or ban) generation of those tokens that may match
    (or for sure don't match) template.

    When applied on scores, following order is done:
      1. If 'banned' is not empty, all but banned tokens are allowed.
         If token is both allowed and banned, it's allowed. 
      2. if 'allowed' is not empty (and banned is), all but allowed tokens are banned.
      3. if 'allow_eos' is False, end-of-string token is banned in any case.

    'look_ahead' is used by Repeat symbol to signal that next symbol should also be considered.
    """
    def __init__(self, *, allowed=None, banned=None, look_ahead=False, allow_eos=False):
        self.allowed: Set[int] = allowed or set()
        self.banned: Set[int] = banned or set()
        assert (self.allowed and not self.banned) or (self.banned and not self.allowed) or not (self.allowed and self.banned)
        self.look_ahead = look_ahead
        self.allow_eos = allow_eos

    def combine(self, other: "AllowedTokens") -> "AllowedTokens":
        """ Returns new instance which is combination of self and other """
        allowed = set()
        banned = set()
        if (not self.allowed and not self.banned) or (not other.allowed and not other.banned):
            # One of self/other is 'allow all'
            pass
        elif self.allowed and other.allowed:
            # Both are 'allow only these'
            assert not self.banned and not other.banned
            allowed = self.allowed.union(other.allowed)
        elif self.banned and other.banned:
            # Both are 'ban only these'
            assert not self.allowed and not other.allowed
            banned = self.banned.intersection(other.banned)
        elif self.allowed and other.banned:
            # I have allowed tokens, other has banned tokens.
            # Allow everything but those we both banned
            assert not self.banned and not other.allowed
            banned = other.banned - self.allowed
        elif other.allowed and self.banned:
            # As above but reversed
            return other.combine(self)
        else:
            assert False, "impossible combination"

        return AllowedTokens(
            allow_eos=self.allow_eos or other.allow_eos,
            look_ahead=self.look_ahead or other.look_ahead,
            allowed=allowed,
            banned=banned,
        )

    def set_ahead(self):
        """ Returns copy of self with 'look_ahead' set to True """
        return AllowedTokens(
            allow_eos=self.allow_eos,
            allowed=self.allowed,
            banned=self.banned,
            look_ahead=True,
        )

    def __repr__(self):
        data = []
        if self.look_ahead or self.allow_eos:
            data.append(",".join([
                "ahead" if self.look_ahead else "",
                "eos" if self.allow_eos else ""
            ]).strip(","))
        data.append(f"allowed={self.allowed}")
        data.append(f"banned={self.banned}")
        return f"<AllowedTokens {' '.join(data)}>"

    def apply(self, scores: torch.FloatTensor):
        if self.allowed and not self.banned:
            s = scores.new_full(scores.shape, False, dtype=torch.bool)
            for a in self.allowed:
                s[..., a] = True
            if self.allow_eos:
                s[..., shared.tokenizer.eos_token_id] = True
            scores[~s] = MINUS_INF
        if self.banned or not self.allow_eos:
            s = scores.new_full(scores.shape, True, dtype=torch.bool)
            for a in self.banned:
                if a not in self.allowed:
                    s[..., a] = False
            if not self.allow_eos:
                s[..., shared.tokenizer.eos_token_id] = False
            scores[~s] = MINUS_INF


def get_token_dictionary() -> Dict[int, str]:
    from extensions.output_template.script import params, logger
    if not params["token_dictionary"] or params["used_tokenizer"] is not shared.tokenizer:
        assert params["scores_size"]
        logger.info("output_template: Creating token dictionary. This takes few seconds, but is done only once.")
        if "OT_TESTING" in os.environ:
            params["token_dictionary"] = {i: decode([i]) for i in range(params["scores_size"])}
        else:
            params["token_dictionary"] = {
                token_id: (
                    shared.tokenizer.decode([token_id])
                    if "▁" not in tmp
                    else tmp.replace("▁", " ")
                )
                for (token_id, tmp) in (
                    (i, shared.tokenizer.convert_ids_to_tokens(i))
                    for i in range(params["scores_size"])
                )
            }
        params["used_tokenizer"] = shared.tokenizer
        logger.info("output_template: Done creating token dictionary.")
    return params["token_dictionary"]

utils.py
im-not-tom/text-generation-webui-output-template-extension