Enforce the format of the output of a language model
__all__ = ['CharacterLevelParser', 
           'StringParser', 
           'RegexParser', 
           'JsonSchemaParser',
           'TokenEnforcer', 
           'generate_enforced',
           'build_transformers_prefix_allowed_tokens_fn']

from .characterlevelparser import CharacterLevelParser, StringParser
from .regexparser import RegexParser
from .jsonschemaparser import JsonSchemaParser
from .tokenenforcer import TokenEnforcer

try:
    from .transformerenforcer import generate_enforced, build_transformers_prefix_allowed_tokens_fn
except Exception as e:
    import logging
    logging.warning(f"Could not import transformerenforcer. Transformers-based functionality will not be available. Details: {e}")

lmformatenforcer/__init__.py
noamgat/lm-format-enforcer
import abc
from typing import Optional


class CharacterLevelParser(abc.ABC):
    """CharacterLevelParser is an interface for classes that can parse strings one character at a time, and determine which characters are allowed at any specific time"""
    @abc.abstractmethod
    def add_character(self, new_character: str) -> 'CharacterLevelParser':
        """Add a character to the parser, and return a new parser that represents the state of the parser after the character has been added. This has to be
        an immutable operation - the original CharacterLevelParser (self) must not be modified."""
        raise NotImplementedError()

    @abc.abstractmethod
    def get_allowed_characters(self) ->str:
        """Return a string containing all characters that are allowed at the current point in the parsing process."""
        raise NotImplementedError()
    
    @abc.abstractmethod
    def can_end(self) -> bool:
        """Return True if the parser is in a state where it can end (potentially finished parsing the desired structure), and False otherwise."""
        raise NotImplementedError()
    
    def shortcut_key(self) -> Optional[str]:
        """Optional. Return a string that denotes that this state is a repeating state, full tree traversal should be avoided."""
        return None


class StringParser(CharacterLevelParser):
    """RegexParser is an example CharacterLevelParser that only allows an exact string. It is a debugging / learning tool
    to show how CharacterLevelParser works together with TokenizerPrefixTree to filter the allowed tokens (some of whom may contain multiple characters)"""
    def __init__(self, string: str):
        self.target_str = string

    def add_character(self, new_character: str) -> CharacterLevelParser:
        if self.target_str.startswith(new_character):
            return StringParser(self.target_str[len(new_character):])
        else:
            raise ValueError(f"Expected '{self.target_str[0]}' but got '{new_character}'")

    def get_allowed_characters(self) -> str:
        return self.target_str[0] if self.target_str else ""

    def can_end(self) -> bool:
        return not self.target_str
    

class ForceStopParser(CharacterLevelParser):
    """A simple parser that forbids any characters except the stop token. Used to force stop LM operation"""
    def add_character(self, new_character: str) -> CharacterLevelParser:
        return self
    def get_allowed_characters(self) -> str:
        return ""
    def can_end(self) -> bool:
        return True

lmformatenforcer/characterlevelparser.py
noamgat/lm-format-enforcer
COMPLETE_ALPHABET = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!@#$%^&*()_+-=[]{};:,./<>? `'\""
lmformatenforcer/consts.py
noamgat/lm-format-enforcer
# https://github.com/koxudaxi/datamodel-code-generator/blob/master/datamodel_code_generator/parser/jsonschema.py
# MIT License

# Copyright (c) 2019 Koudai Aono

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from __future__ import annotations

import enum as _enum
from functools import lru_cache
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Set,
    Union,
)
from warnings import warn


from pydantic import (
    Field,
)

from .jsonschemaobjectutil import (
    PYDANTIC_V2,
    BaseModel,
    cached_property,
    field_validator,
    model_validator,
    is_url,
    Types,
    UnionIntFloat
)

if PYDANTIC_V2:
    from pydantic import ConfigDict


def get_model_by_path(
    schema: Union[Dict[str, Any], List[Any]], keys: Union[List[str], List[int]]
) -> Dict[Any, Any]:
    model: Union[Dict[Any, Any], List[Any]]
    if not keys:
        model = schema
    elif len(keys) == 1:
        if isinstance(schema, dict):
            model = schema.get(keys[0], {})  # type: ignore
        else:  # pragma: no cover
            model = schema[int(keys[0])]
    elif isinstance(schema, dict):
        model = get_model_by_path(schema[keys[0]], keys[1:])  # type: ignore
    else:
        model = get_model_by_path(schema[int(keys[0])], keys[1:])
    if isinstance(model, dict):
        return model
    raise NotImplementedError(  # pragma: no cover
        f'Does not support json pointer to array. schema={schema}, key={keys}'
    )


json_schema_data_formats: Dict[str, Dict[str, Types]] = {
    'integer': {
        'int32': Types.int32,
        'int64': Types.int64,
        'default': Types.integer,
        'date-time': Types.date_time,
        'unix-time': Types.int64,
    },
    'number': {
        'float': Types.float,
        'double': Types.double,
        'decimal': Types.decimal,
        'date-time': Types.date_time,
        'time': Types.time,
        'default': Types.number,
    },
    'string': {
        'default': Types.string,
        'byte': Types.byte,  # base64 encoded string
        'binary': Types.binary,
        'date': Types.date,
        'date-time': Types.date_time,
        'time': Types.time,
        'password': Types.password,
        'email': Types.email,
        'idn-email': Types.email,
        'uuid': Types.uuid,
        'uuid1': Types.uuid1,
        'uuid2': Types.uuid2,
        'uuid3': Types.uuid3,
        'uuid4': Types.uuid4,
        'uuid5': Types.uuid5,
        'uri': Types.uri,
        'uri-reference': Types.string,
        'hostname': Types.hostname,
        'ipv4': Types.ipv4,
        'ipv4-network': Types.ipv4_network,
        'ipv6': Types.ipv6,
        'ipv6-network': Types.ipv6_network,
        'decimal': Types.decimal,
        'integer': Types.integer,
    },
    'boolean': {'default': Types.boolean},
    'object': {'default': Types.object},
    'null': {'default': Types.null},
    'array': {'default': Types.array},
}


class JSONReference(_enum.Enum):
    LOCAL = 'LOCAL'
    REMOTE = 'REMOTE'
    URL = 'URL'


class Discriminator(BaseModel):
    propertyName: str
    mapping: Optional[Dict[str, str]] = None


class JsonSchemaObject(BaseModel):
    if not TYPE_CHECKING:
        if PYDANTIC_V2:

            @classmethod
            def get_fields(cls) -> Dict[str, Any]:
                return cls.model_fields

        else:

            @classmethod
            def get_fields(cls) -> Dict[str, Any]:
                return cls.__fields__

            @classmethod
            def model_rebuild(cls) -> None:
                cls.update_forward_refs()

    __constraint_fields__: Set[str] = {
        'exclusiveMinimum',
        'minimum',
        'exclusiveMaximum',
        'maximum',
        'multipleOf',
        'minItems',
        'maxItems',
        'minLength',
        'maxLength',
        'pattern',
        'uniqueItems',
    }
    __extra_key__: str = 'extras'

    @model_validator(mode='before')
    def validate_exclusive_maximum_and_exclusive_minimum(
        cls, values: Dict[str, Any]
    ) -> Any:
        exclusive_maximum: Union[float, bool, None] = values.get('exclusiveMaximum')
        exclusive_minimum: Union[float, bool, None] = values.get('exclusiveMinimum')

        if exclusive_maximum is True:
            values['exclusiveMaximum'] = values['maximum']
            del values['maximum']
        elif exclusive_maximum is False:
            del values['exclusiveMaximum']
        if exclusive_minimum is True:
            values['exclusiveMinimum'] = values['minimum']
            del values['minimum']
        elif exclusive_minimum is False:
            del values['exclusiveMinimum']
        return values

    @field_validator('ref')
    def validate_ref(cls, value: Any) -> Any:
        if isinstance(value, str) and '#' in value:
            if value.endswith('#/'):
                return value[:-1]
            elif '#/' in value or value[0] == '#' or value[-1] == '#':
                return value
            return value.replace('#', '#/')
        return value

    items: Union[List[JsonSchemaObject], JsonSchemaObject, bool, None] = None
    uniqueItems: Optional[bool] = None
    type: Union[str, List[str], None] = None
    format: Optional[str] = None
    pattern: Optional[str] = None
    minLength: Optional[int] = None
    maxLength: Optional[int] = None
    minimum: Optional[UnionIntFloat] = None
    maximum: Optional[UnionIntFloat] = None
    minItems: Optional[int] = None
    maxItems: Optional[int] = None
    multipleOf: Optional[float] = None
    exclusiveMaximum: Union[float, bool, None] = None
    exclusiveMinimum: Union[float, bool, None] = None
    additionalProperties: Union[JsonSchemaObject, bool, None] = None
    patternProperties: Optional[Dict[str, JsonSchemaObject]] = None
    oneOf: List[JsonSchemaObject] = []
    anyOf: List[JsonSchemaObject] = []
    allOf: List[JsonSchemaObject] = []
    enum: List[Any] = []
    writeOnly: Optional[bool] = None
    properties: Optional[Dict[str, Union[JsonSchemaObject, bool]]] = None
    required: List[str] = []
    ref: Optional[str] = Field(default=None, alias='$ref')
    nullable: Optional[bool] = False
    x_enum_varnames: List[str] = Field(default=[], alias='x-enum-varnames')
    description: Optional[str] = None
    title: Optional[str] = None
    example: Any = None
    examples: Any = None
    default: Any = None
    id: Optional[str] = Field(default=None, alias='$id')
    custom_type_path: Optional[str] = Field(default=None, alias='customTypePath')
    custom_base_path: Optional[str] = Field(default=None, alias='customBasePath')
    extras: Dict[str, Any] = Field(alias=__extra_key__, default_factory=dict)
    discriminator: Union[Discriminator, str, None] = None
    if PYDANTIC_V2:
        model_config = ConfigDict(
            arbitrary_types_allowed=True,
            ignored_types=(cached_property,),
        )
    else:

        class Config:
            arbitrary_types_allowed = True
            keep_untouched = (cached_property,)
            smart_casts = True

    if not TYPE_CHECKING:

        def __init__(self, **data: Any) -> None:
            super().__init__(**data)
            self.extras = {k: v for k, v in data.items() if k not in EXCLUDE_FIELD_KEYS}

    @cached_property
    def is_object(self) -> bool:
        return (
            self.properties is not None
            or self.type == 'object'
            and not self.allOf
            and not self.oneOf
            and not self.anyOf
            and not self.ref
        )

    @cached_property
    def is_array(self) -> bool:
        return self.items is not None or self.type == 'array'

    @cached_property
    def ref_object_name(self) -> str:  # pragma: no cover
        return self.ref.rsplit('/', 1)[-1]  # type: ignore

    @field_validator('items', mode='before')
    def validate_items(cls, values: Any) -> Any:
        # this condition expects empty dict
        return values or None

    @cached_property
    def has_default(self) -> bool:
        return 'default' in self.__fields_set__ or 'default_factory' in self.extras

    @cached_property
    def has_constraint(self) -> bool:
        return bool(self.__constraint_fields__ & self.__fields_set__)

    @cached_property
    def ref_type(self) -> Optional[JSONReference]:
        if self.ref:
            return get_ref_type(self.ref)
        return None  # pragma: no cover

    @cached_property
    def type_has_null(self) -> bool:
        return isinstance(self.type, list) and 'null' in self.type


@lru_cache()
def get_ref_type(ref: str) -> JSONReference:
    if ref[0] == '#':
        return JSONReference.LOCAL
    elif is_url(ref):
        return JSONReference.URL
    return JSONReference.REMOTE


def _get_type(type_: str, format__: Optional[str] = None) -> Types:
    if type_ not in json_schema_data_formats:
        return Types.any
    data_formats: Optional[Types] = json_schema_data_formats[type_].get(
        'default' if format__ is None else format__
    )
    if data_formats is not None:
        return data_formats

    warn(
        'format of {!r} not understood for {!r} - using default'
        ''.format(format__, type_)
    )
    return json_schema_data_formats[type_]['default']


JsonSchemaObject.model_rebuild()

DEFAULT_FIELD_KEYS: Set[str] = {
    'example',
    'examples',
    'description',
    'discriminator',
    'title',
    'const',
    'default_factory',
}

EXCLUDE_FIELD_KEYS = (set(JsonSchemaObject.get_fields()) - DEFAULT_FIELD_KEYS) | {
    '$id',
    '$ref',
    JsonSchemaObject.__extra_key__,
}

lmformatenforcer/external/jsonschemaobject.py
noamgat/lm-format-enforcer
# https://github.com/koxudaxi/datamodel-code-generator/blob/master/datamodel_code_generator/util.py
# MIT License

# Copyright (c) 2019 Koudai Aono

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from __future__ import annotations

from typing import TYPE_CHECKING, Any, Callable, TypeVar
from enum import Enum, auto

import pydantic
from packaging import version
from pydantic import BaseModel as _BaseModel

from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Iterator,
    TypeVar,
    Union,
)

PYDANTIC_VERSION = version.parse(
    pydantic.VERSION if isinstance(pydantic.VERSION, str) else str(pydantic.VERSION)
)

PYDANTIC_V2: bool = PYDANTIC_VERSION >= version.parse('2.0b3')

if PYDANTIC_V2:
    from pydantic import GetCoreSchemaHandler
    from pydantic_core import core_schema

if TYPE_CHECKING:
    cached_property = property
    from yaml import SafeLoader

    Protocol = object
    runtime_checkable: Callable[..., Any]

    from typing_extensions import Literal
else:
    try:
        from typing import Protocol
    except ImportError:
        from typing_extensions import Protocol  # noqa
    try:
        from typing import runtime_checkable
    except ImportError:
        from typing_extensions import runtime_checkable  # noqa
    try:
        from yaml import CSafeLoader as SafeLoader
    except ImportError:  # pragma: no cover
        from yaml import SafeLoader

    try:
        from functools import cached_property
    except ImportError:
        _NOT_FOUND = object()

        class cached_property:
            def __init__(self, func: Callable) -> None:
                self.func: Callable = func
                self.__doc__: Any = func.__doc__

            def __get__(self, instance: Any, owner: Any = None) -> Any:
                value = instance.__dict__.get(self.func.__name__, _NOT_FOUND)
                if value is _NOT_FOUND:  # pragma: no cover
                    value = instance.__dict__[self.func.__name__] = self.func(instance)
                return value


SafeLoader.yaml_constructors[
    'tag:yaml.org,2002:timestamp'
] = SafeLoader.yaml_constructors['tag:yaml.org,2002:str']


Model = TypeVar('Model', bound=_BaseModel)


def model_validator(
    mode: Literal['before', 'after'] = 'after',
) -> Callable[[Callable[[Model, Any], Any]], Callable[[Model, Any], Any]]:
    def inner(method: Callable[[Model, Any], Any]) -> Callable[[Model, Any], Any]:
        if PYDANTIC_V2:
            from pydantic import model_validator as model_validator_v2

            return model_validator_v2(mode=mode)(method)  # type: ignore
        else:
            from pydantic import root_validator

            return root_validator(method, pre=mode == 'before')  # type: ignore

    return inner


def field_validator(
    field_name: str,
    *fields: str,
    mode: Literal['before', 'after'] = 'after',
) -> Callable[[Any], Callable[[Model, Any], Any]]:
    def inner(method: Callable[[Model, Any], Any]) -> Callable[[Model, Any], Any]:
        if PYDANTIC_V2:
            from pydantic import field_validator as field_validator_v2

            return field_validator_v2(field_name, *fields, mode=mode)(method)  # type: ignore
        else:
            from pydantic import validator

            return validator(field_name, *fields, pre=mode == 'before')(method)  # type: ignore

    return inner


if PYDANTIC_V2:
    from pydantic import ConfigDict as ConfigDict
else:
    ConfigDict = dict  # type: ignore


class BaseModel(_BaseModel):
    if PYDANTIC_V2:
        model_config = ConfigDict(strict=False)


def is_url(ref: str) -> bool:
    return ref.startswith(('https://', 'http://'))


class Types(Enum):
    integer = auto()
    int32 = auto()
    int64 = auto()
    number = auto()
    float = auto()
    double = auto()
    decimal = auto()
    time = auto()
    string = auto()
    byte = auto()
    binary = auto()
    date = auto()
    date_time = auto()
    password = auto()
    email = auto()
    uuid = auto()
    uuid1 = auto()
    uuid2 = auto()
    uuid3 = auto()
    uuid4 = auto()
    uuid5 = auto()
    uri = auto()
    hostname = auto()
    ipv4 = auto()
    ipv4_network = auto()
    ipv6 = auto()
    ipv6_network = auto()
    boolean = auto()
    object = auto()
    null = auto()
    array = auto()
    any = auto()

class UnionIntFloat:
    def __init__(self, value: Union[int, float]) -> None:
        self.value: Union[int, float] = value

    def __int__(self) -> int:
        return int(self.value)

    def __float__(self) -> float:
        return float(self.value)

    def __str__(self) -> str:
        return str(self.value)

    @classmethod
    def __get_validators__(cls) -> Iterator[Callable[[Any], Any]]:
        yield cls.validate

    @classmethod
    def __get_pydantic_core_schema__(
        cls, _source_type: Any, _handler: GetCoreSchemaHandler
    ) -> core_schema.CoreSchema:
        from_int_schema = core_schema.chain_schema(
            [
                core_schema.union_schema(
                    [core_schema.int_schema(), core_schema.float_schema()]
                ),
                core_schema.no_info_plain_validator_function(cls.validate),
            ]
        )

        return core_schema.json_or_python_schema(
            json_schema=core_schema.no_info_plain_validator_function(cls.validate),
            python_schema=core_schema.union_schema(
                [
                    # check if it's an instance first before doing any further work
                    core_schema.is_instance_schema(UnionIntFloat),
                    from_int_schema,
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                lambda instance: instance.value
            ),
        )

    @classmethod
    def validate(cls, v: Any) -> UnionIntFloat:
        if isinstance(v, UnionIntFloat):
            return v
        elif not isinstance(v, (int, float)):  # pragma: no cover
            raise TypeError(f'{v} is not int or float')
        return cls(v)

lmformatenforcer/external/jsonschemaobjectutil.py
noamgat/lm-format-enforcer
from copy import deepcopy
from typing import Any, List, Optional, Union

from lmformatenforcer.consts import COMPLETE_ALPHABET

from .external.jsonschemaobject import JsonSchemaObject

from .characterlevelparser import CharacterLevelParser
from .consts import COMPLETE_ALPHABET

class JsonSchemaParser(CharacterLevelParser):
    def __init__(self, json_schema: Union[dict, JsonSchemaObject], existing_stack: List[Any] = None):
        self.model_class = json_schema if isinstance(json_schema, JsonSchemaObject) else JsonSchemaObject(**json_schema)
        self.object_stack: 'List[BaseParsingState]' = existing_stack or [ObjectParsingState(self.model_class, self)]

    def __deepcopy__(self, memo):
        # Avoid cloning the model class, since it is immutable
        for parser in self.object_stack:
            parser.root = None
        clone = JsonSchemaParser(self.model_class, deepcopy(self.object_stack, memo))
        for parser in self.object_stack:
            parser.root = self
        for parser in clone.object_stack:
            parser.root = clone 
        return clone
    
    def add_character(self, new_character: str) -> CharacterLevelParser:
        # The add_character contract requires immutability, therefore we clone before modifying.
        clone = deepcopy(self)
        clone.object_stack[-1].add_character(new_character)
        return clone

    def get_allowed_characters(self) -> str:
        return self.object_stack[-1].get_allowed_characters() if self.object_stack else ""

    def can_end(self) -> bool:
        return not self.object_stack
    
    def finish_parser(self, character_to_pass_to_parent: str = ""):
        self.object_stack.pop()
        if self.object_stack and character_to_pass_to_parent:
            self.object_stack[-1].add_character(character_to_pass_to_parent)

    def shortcut_key(self) -> Optional[str]:
        if self.object_stack:
            current_parser = self.object_stack[-1]
            if isinstance(current_parser, StringParsingState):
                if not current_parser.allowed_strings and current_parser.seen_opening_quote and not current_parser.seen_closing_quote:
                    # Performance optimization: When we are parsing a string that is not from a list of allowed strings, most tokens
                    # are legal. The exploration can be more costly than the LM itself for large tokenizers (because this is pure python),
                    # so we signal that we are in a "freetext" mode, and reuse the allowed token list throughout the run.
                    return 'json_freetext'
        return None


class BaseParsingState:
    def __init__(self, root: JsonSchemaParser):
        self.root = root

    def add_character(self, new_character: str):
        raise NotImplementedError()

    def get_allowed_characters(self) ->str:
        raise NotImplementedError()
    
    def can_end(self) -> bool:
        raise NotImplementedError()


class ObjectParsingStage:
    START_OBJECT = "StartObject"
    PARSING_KEY_OR_END = "ParsingKey"
    PARSING_VALUE = "ParsingValue"
    PARSING_SEPARATOR_OR_END = "ParsingSeparatorOrEnd"
    END_OBJECT = "EndObject"


def get_parser(
    parsing_state: JsonSchemaParser,
    value_schema: JsonSchemaObject,
    ending_characters: str,
) -> BaseParsingState:
    # Sometimes the schema is a union of a type and null, so we need to get the first type
    if value_schema.anyOf and len(value_schema.anyOf) == 2 and value_schema.anyOf[1].type == 'null':
        value_schema = value_schema.anyOf[0]
    if value_schema.type == "string":
        return StringParsingState(
            parsing_state,
            ending_characters,
            value_schema.enum,
            require_opening_quote=True,
        )
    elif value_schema.type == "object":
        return ObjectParsingState(value_schema, parsing_state)
    elif value_schema.type == None and value_schema.ref:
        value_class_name = value_schema.ref.split('/')[-1]
        extras = parsing_state.model_class.extras
        # Pydantic V1 and V2 have different names for the definitions field
        if 'definitions' in extras:
            definitions = extras['definitions']
        elif '$defs' in extras:
            definitions = extras['$defs']
        else:
            raise ValueError("No definitions found in schema")
        class_dict = definitions[value_class_name]
        value_schema = JsonSchemaObject(**class_dict)
        return get_parser(parsing_state, value_schema, ending_characters)
    elif value_schema.enum:
        is_numeric = all(isinstance(i, (int, float)) for i in value_schema.enum)
        is_string = all(isinstance(i, (str)) for i in value_schema.enum)
        if is_string:
            return StringParsingState(
            parsing_state,
            ending_characters,
            value_schema.enum,
            require_opening_quote=True,
        )
        elif is_numeric:
            return StringParsingState(
                parsing_state,
                ending_characters,
                [str(i) for i in value_schema.enum],
                require_opening_quote=False,
                require_closing_quote=False,
            )
        else:
            raise Exception("Unsupported enum type " + str(value_schema.enum))
    elif value_schema.type == "integer":
        return NumberParsingState(parsing_state, ending_characters, False)
    elif value_schema.type == "boolean":
        return StringParsingState(
            parsing_state,
            ending_characters,
            ["true", "false"],
            require_opening_quote=False,
            require_closing_quote=False,
        )
    elif value_schema.type == "number":
        return NumberParsingState(parsing_state, ending_characters, True)
    elif value_schema.type == "array":
        return ListParsingState(parsing_state, ending_characters, value_schema.items)
    else:
        raise Exception("Unsupported type " + str(value_schema.type))



class ObjectParsingState(BaseParsingState):
    schema_object: JsonSchemaObject
    current_stage: ObjectParsingStage
    existing_keys: List[str]
    current_key: str
    current_key_parser: Any  #  type: StringParsingState
    is_dictionary: bool

    def __init__(self, schema_object: JsonSchemaObject, root: JsonSchemaParser):
        super().__init__(root)
        self.schema_object = schema_object
        self.current_stage = ObjectParsingStage.START_OBJECT
        self.root = root
        self.existing_keys = []
        self.current_key = None
        self.current_key_parser = None
        # Javascript objects represent both classes and dictionaries, so we need to know which one we are parsing
        self.is_dictionary = self.schema_object.properties is None

    def add_character(self, new_character: str):
        if new_character.strip() == "":
            # In object scope, whitespaces can be ignored
            return
        if (
            self.current_stage == ObjectParsingStage.START_OBJECT
            and new_character == "{"
        ):
            self.current_stage = ObjectParsingStage.PARSING_KEY_OR_END
        elif self.current_stage == ObjectParsingStage.PARSING_KEY_OR_END:
            if new_character == "}":
                self.current_stage = ObjectParsingStage.END_OBJECT
                self.root.finish_parser()
            if new_character == '"':
                self.current_stage = ObjectParsingStage.PARSING_KEY_OR_END
                possible_keys = None
                if not self.is_dictionary:
                    possible_keys = list(self.schema_object.properties.keys())
                    possible_keys = list(
                        set(possible_keys).difference(self.existing_keys)
                    )
                ending_characters = ': '
                # We send require_opening_quote=True and then add_character('"') instead of require_opening_quote=False
                # Because there is a difference between "don't need a quote" and "received it before creating the parser"
                key_parser = StringParsingState(
                    self.root, ending_characters, possible_keys, require_opening_quote=True, require_closing_quote=True
                )
                key_parser.add_character('"')
                self.root.object_stack.append(key_parser)
                self.current_key_parser = key_parser
            if new_character == ":":
                self.current_stage = ObjectParsingStage.PARSING_VALUE
                self.current_key = self.current_key_parser.parsed_string
                self.existing_keys.append(self.current_key)
                if self.is_dictionary:
                    value_schema = self.schema_object.additionalProperties
                    can_continue = True
                    can_end = True
                else:
                    possible_keys = list(self.schema_object.properties.keys())
                    possible_keys = list(
                        set(possible_keys).difference(self.existing_keys)
                    )
                    value_schema = self.schema_object.properties[self.current_key]
                    can_continue = len(possible_keys) > 0
                    required_keys = self.schema_object.required or []
                    can_end = set(self.existing_keys).issuperset(required_keys)
                ending_characters = " "
                if can_continue:
                    ending_characters += ","
                if can_end:
                    ending_characters += "}"
                self.current_key_parser = get_parser(
                    self.root, value_schema, ending_characters
                )
                self.root.object_stack.append(self.current_key_parser)
                self.current_key_parser = None
        elif self.current_stage == ObjectParsingStage.PARSING_VALUE:
            # If we recieve a character during parsing value, it means that its the finishing character
            # of the value parser
            if new_character == '"':
                self.current_stage = ObjectParsingStage.PARSING_SEPARATOR_OR_END
            elif new_character == ",":
                self.current_stage = ObjectParsingStage.PARSING_KEY_OR_END
            elif new_character == "}":
                self.current_stage = ObjectParsingStage.END_OBJECT
                self.root.finish_parser()
        elif self.current_stage == ObjectParsingStage.PARSING_SEPARATOR_OR_END:
            if new_character == ",":
                self.current_stage = ObjectParsingStage.PARSING_KEY_OR_END
            elif new_character == "}":
                self.current_stage = ObjectParsingStage.END_OBJECT
                self.root.finish_parser()

    def get_allowed_characters(self) -> str:
        possible_keys = (
            list(self.schema_object.properties.keys())
            if not self.is_dictionary
            else None
        )
        required_keys = self.schema_object.required or []
        can_end = set(self.existing_keys).issuperset(required_keys)
        can_parse_key = self.is_dictionary or set(possible_keys).issuperset(
            self.existing_keys
        )

        possible_characters = [' ']
        if self.current_stage == ObjectParsingStage.START_OBJECT:
            possible_characters.append('{')
        elif self.current_stage == ObjectParsingStage.PARSING_KEY_OR_END:
            if can_end:
                possible_characters.append('}')
            if can_parse_key:
                possible_characters.append('"')
        elif self.current_stage == ObjectParsingStage.PARSING_VALUE:
            # Sometimes the value parser considers finishing, so it needs to know which continuations are possible
            if can_end:
                possible_characters.append('}')
            if can_parse_key:
                possible_characters.append(',')
        elif self.current_stage == ObjectParsingStage.PARSING_SEPARATOR_OR_END:
            if can_end:
                possible_characters.append('}')
            if can_parse_key:
                possible_characters.append(',')
        return "".join(possible_characters)


class StringParsingStage:
    START_TOKEN = "StartToken"
    PARSING_STRING = "ParsingString"
    END_TOKEN = "EndToken"


class PrimitiveParsingState(BaseParsingState):
    def __init__(self, root: JsonSchemaParser, ending_characters: str):
        super().__init__(root)
        self.stage = StringParsingStage.START_TOKEN
        self.parsed_string = ""
        self.ending_characters = ending_characters

    def add_character(self, new_character: str):
        if self.can_end() and new_character in self.ending_characters:
            self.root.finish_parser(new_character)
        else:
            self.parsed_string += new_character

    def get_allowed_characters(self) -> str:
        allowed_characters = self._get_allowed_primitive_characters()
        if self.can_end():
            allowed_characters += self.ending_characters
        return allowed_characters

    def _get_allowed_primitive_characters(self) -> str:
        return ''

    def can_end(self) -> bool:
        return True


class NumberParsingState(PrimitiveParsingState):
    def __init__(
        self,
        root: JsonSchemaParser,
        ending_characters: str,
        allow_floating_point: bool,
    ):
        super().__init__(root, ending_characters)
        self.allow_floating_point = allow_floating_point
        self.seen_decimal_point = False

    def add_character(self, new_character: str):
        if new_character == " ":
            return
        super().add_character(new_character)
        if new_character == ".":
            self.seen_decimal_point = True

    def _get_allowed_primitive_characters(self) -> str:
        allowed_characters = "0123456789"
        if not self.parsed_string:
            allowed_characters += "- "
        if self.allow_floating_point and not self.seen_decimal_point:
            allowed_characters += "."
        return allowed_characters

    def can_end(self) -> bool:
        return bool(self.parsed_string) and self.parsed_string[-1] in "0123456789"


class StringParsingState(PrimitiveParsingState):
    allowed_strings: List[str]
    parsed_string: str
    seen_closing_quote: bool
    seen_opening_quote: bool

    def __init__(
        self,
        root: JsonSchemaParser,
        ending_characters: str,
        allowed_strings: List[str],
        require_opening_quote: bool,
        require_closing_quote: bool = True,
    ):
        super().__init__(root, ending_characters)
        self.allowed_strings = allowed_strings
        self.seen_closing_quote = False
        self.seen_opening_quote = not require_opening_quote
        self.require_closing_quote = require_closing_quote
        self.require_opening_quote = require_opening_quote

    def add_character(self, new_character: str):
        if not self.parsed_string and new_character == ' ':
            return
        super().add_character(new_character)
        if new_character == '"':
            if not self.seen_opening_quote:
                self.seen_opening_quote = True
                self.parsed_string = ""
            else:
                self.seen_closing_quote = True
                self.parsed_string = self.parsed_string[:-1]

    def _get_allowed_primitive_characters(self) -> str:
        if not self.seen_opening_quote:
            return '" '
        if self.seen_closing_quote:
            return ''
        if self.allowed_strings:
            allowed_continuations = [
                s[len(self.parsed_string) :]
                for s in self.allowed_strings
                if s.startswith(self.parsed_string)
            ]
            allowed_next_characters = [allowed_continuation[0] for allowed_continuation in allowed_continuations if len(allowed_continuation) > 0]
            allowed_next_characters = list(set(allowed_next_characters))
            if self.parsed_string in self.allowed_strings and self.require_closing_quote:
                allowed_next_characters.append('"')
            if (not self.parsed_string) and (not self.seen_opening_quote or not self.require_opening_quote):
                allowed_next_characters.append(' ')
            return "".join(allowed_next_characters)
        else:
            return COMPLETE_ALPHABET

    def can_end(self) -> bool:
        if self.require_closing_quote:
            return self.seen_closing_quote
        else:
            if self.allowed_strings:
                return self.parsed_string in self.allowed_strings
            else:
                return bool(self.parsed_string)


class ListParsingState(PrimitiveParsingState):
    list_member_type: JsonSchemaObject
    seen_list_opener: bool = False
    seen_list_closer: bool = False

    def __init__(
        self,
        root: JsonSchemaParser,
        ending_characters: str,
        list_member_type: JsonSchemaObject,
    ):
        super().__init__(root, ending_characters)
        self.list_member_type = list_member_type

    def add_character(self, new_character: str):
        if self.seen_list_closer:
            super().add_character(new_character)
        if new_character == "[":
            self.seen_list_opener = True
            self.root.object_stack.append(
                get_parser(
                    self.root,
                    self.list_member_type,
                    "],",
                )
            )
        elif new_character == "]":
            self.seen_list_closer = True
        elif new_character == ",":
            if not self.seen_list_closer:
                self.root.object_stack.append(
                    get_parser(
                        self.root,
                        self.list_member_type,
                        "],",
                    )
                )

    def _get_allowed_primitive_characters(self) -> str:
        if not self.seen_list_opener:
            return "[ "
        elif not self.seen_list_closer:
            return "], "
        else:
            # The parent function will take care of allowing the ending tokens.
            return ""

    def can_end(self) -> bool:
        return self.seen_list_closer

lmformatenforcer/jsonschemaparser.py
noamgat/lm-format-enforcer
from typing import Dict, Union
import interegular
from interegular.fsm import anything_else

from .characterlevelparser import CharacterLevelParser
from .consts import COMPLETE_ALPHABET

class RegexParser(CharacterLevelParser):
    """RegexParser is an example CharacterLevelParser that only allows strings that match a given regular expression."""

    UNINITIALIZED_STATE = -1
    INVALID_STATE = -2

    class _Context:
        pattern: interegular.FSM
        anything_else_characters: str
        state_character_cache: Dict[int, str]
    
    context: _Context
    current_state: int

    def __init__(self, pattern: Union[str, _Context], current_state: int = UNINITIALIZED_STATE):
        if isinstance(pattern, str):
            self.context = RegexParser._Context()
            self.context.pattern = interegular.parse_pattern(pattern).to_fsm()
            self.context.state_character_cache = {}
            not_anything_else_characters = set([c for c in self.context.pattern.alphabet.keys() if c != anything_else])
            self.context.anything_else_characters = "".join([c for c in COMPLETE_ALPHABET if c not in not_anything_else_characters])
        else:
            self.context = pattern
        self.current_state: int = self.context.pattern.initial if current_state == RegexParser.UNINITIALIZED_STATE else current_state

    def add_character(self, new_character: str) -> 'RegexParser':
        if self.current_state == RegexParser.INVALID_STATE:
            return self
        
        state = self.current_state
        fsm = self.context.pattern
        # Mostly taken from FSM.accept()
        symbol = new_character
        if anything_else in fsm.alphabet and not symbol in fsm.alphabet:
            symbol = anything_else
        transition = fsm.alphabet[symbol]

        # Missing transition = transition to dead state
        if not (state in fsm.map and transition in fsm.map[state]):
            return RegexParser(self.context, RegexParser.INVALID_STATE)

        state = fsm.map[state][transition]

        return RegexParser(self.context, state)
    
    def can_end(self) -> bool:
        return self.current_state in self.context.pattern.finals
    
    def get_allowed_characters(self) -> str:
        if self.current_state not in self.context.pattern.map:
            return ''
        if self.current_state not in self.context.state_character_cache:
            allowed_characters = []
            state_map = self.context.pattern.map[self.current_state]
            for symbol, symbol_idx in self.context.pattern.alphabet.items():
                if symbol_idx in state_map:
                    if symbol == anything_else:
                        allowed_characters.append(self.context.anything_else_characters)
                    else:
                        allowed_characters.append(symbol)
            self.context.state_character_cache[self.current_state] = "".join(allowed_characters)
        return self.context.state_character_cache[self.current_state]


lmformatenforcer/regexparser.py
noamgat/lm-format-enforcer
from dataclasses import dataclass, field
from typing import Callable, Dict, Hashable, List, Optional, Tuple
import logging
from .characterlevelparser import CharacterLevelParser, ForceStopParser
from .tokenizerprefixtree import TokenizerPrefixTree, TokenizerPrefixTreeNode


class TokenEnforcer:
    """TokenEnforcer provides a token filtering mechanism, given a CharacterLevelParser and some information about the tokenizer.
    It is the main entry point for extending lm-format-enforcer to new inference libraries. See __init__() and get_allowed_tokens()"""
    @dataclass
    class OutputTensorState:
        str_so_far: str
        parser: CharacterLevelParser
        allowed_tokens: List[int] = field(default_factory=list)

    def __init__(self, regular_tokens: List[Tuple[int, str]], 
                 parser: CharacterLevelParser,
                 decoder: Callable[[List[int]], str],
                 eos_token_id: int):
        """
        Create a new TokenEnforcer.
        :param regular_tokens: A list of tuples (token_id, token_string) for all the regular (not special) tokens in the tokenizer vocabulary.
        Note that token_string is expected to include leading / trailing whitespaces if relevant.
        :param parser: A CharacterLevelParser that defines the allowed strings.
        :param decoder: A function that decodes a list of token ids into a string.
        :param eos_token_id: The token id of the end-of-string token.
        """
        self.prefix_states: Dict[Hashable, TokenEnforcer.OutputTensorState] = {}
        self.root_parser = parser
        self.tokenizer_tree = TokenizerPrefixTree(regular_tokens)
        self.decoder = decoder
        self.eos_token_id = eos_token_id

    def get_allowed_tokens(self, token_sequence: List[int]) -> List[int]:
        """
        Get a list of allowed tokens, given a list of tokens that were already generated.
        :param token_sequence: The tokens that were already generated, and the next token will be generated for.
        :return: A list of token ids that are allowed to be selected next.
        """
        # In order to elegantly support beam search and batching, we don't store per-batch information.
        # Instead, we store a hash of all the states (unique token tensors) we encountered so far.
        # When we encounter a new unique token tensor, we find the token tensor that led to it, and continue from there.
        sent_tuple = tuple(token_sequence)
        prev_step_tuple = sent_tuple[:-1]

        if sent_tuple in self.prefix_states:
            # We already calculated for this node, return cached list
            return self.prefix_states[sent_tuple].allowed_tokens
        elif prev_step_tuple not in self.prefix_states:
            # We have not encountered the tensor up to the before-last entry. This means that this is the first call - the instruction / prompt tensor.
            # Initialize the root node
            state = TokenEnforcer.OutputTensorState(str_so_far=self.decoder(token_sequence),
                                                    parser=self.root_parser)
            self.prefix_states[sent_tuple] = state
            self._compute_allowed_tokens(state)
            return state.allowed_tokens
        else:
            # Find the state that led to this node. We explicitly don't use the concept of "timestep" because of beam search        
            prev_step_state = self.prefix_states[prev_step_tuple]
            new_state = self._apply_new_characters(prev_step_state, token_sequence)
            self.prefix_states[sent_tuple] = new_state
            self._compute_allowed_tokens(new_state)
            return new_state.allowed_tokens

    def _compute_allowed_tokens(self, state: 'TokenEnforcer.OutputTensorState'):
        allowed_tokens: List[int] = []
        shortcut_key = state.parser.shortcut_key()
        self._collect_allowed_tokens(state.parser, self.tokenizer_tree.root, allowed_tokens, shortcut_key)
        if state.parser.can_end():
            allowed_tokens.append(self.eos_token_id)
        if not allowed_tokens:
            raise ValueError(f"Parser reached state with no allowed tokens")
        # root_state = next(state for state in self.prefix_states.values() if state.parser == self.root_parser)
        # print(f"Allowing {len(allowed_tokens)} tokens after {state.str_so_far[len(root_state.str_so_far):]}")
        state.allowed_tokens = allowed_tokens

    def _collect_allowed_tokens(self, parser: CharacterLevelParser, tree_node: TokenizerPrefixTreeNode, allowed_tokens: List[int], shortcut_key: Optional[str]):
        allowed_tokens.extend(tree_node.tokens)
        allowed_characters = parser.get_allowed_characters()
        relevant_characters = tree_node.children.keys()
        # This next line is the heart of the traversal algorithm. We only explore paths that are shared by both the parser and the tokenizer.
        characters_to_explore = set(relevant_characters).intersection(allowed_characters)
        
        # Performance optimization: If we are in JSON freetext, all of the tokens that don't contain quote, or end with quote, are legal, so we take
        # their cached list. If the quote character is allowed, we only need to dynamically explore the cases where the string starts with a quote.
        # This breaks the elegance of the API, but otherwise it is a huge performance hit.
        if shortcut_key == 'json_freetext':
            allowed_tokens.extend(self.tokenizer_tree.json_freetext_tokens)
            characters_to_explore = characters_to_explore.intersection(['"'])

        for character in characters_to_explore:
            next_parser = parser.add_character(character )
            next_tree_node = tree_node.children[character]
            self._collect_allowed_tokens(next_parser, next_tree_node, allowed_tokens, None)
            
    def _apply_new_characters(self, state: 'TokenEnforcer.OutputTensorState', token_sequence: List[int]):
        characters = self.decoder(token_sequence)
        new_state = TokenEnforcer.OutputTensorState(str_so_far=characters, parser=state.parser)
        new_characters = characters[len(state.str_so_far):]
        for character in new_characters:
            if character in new_state.parser.get_allowed_characters():
                new_state.parser = new_state.parser.add_character(character)
            else:
                # This can happen in beam / batch scenarios, when some of the batches finished but others are continuing.
                logging.debug(f"Received an invalid character '{character}', switching to ForceStopParser")
                new_state.parser = ForceStopParser()
        return new_state
        

    
lmformatenforcer/tokenenforcer.py
noamgat/lm-format-enforcer
from typing import Dict, List, Tuple


class TokenizerPrefixTreeNode:
    def __init__(self):
        self.tokens: List[int] = []
        self.children: Dict[str, TokenizerPrefixTreeNode] = {}


class TokenizerPrefixTree:
    def __init__(self, regular_tokens: List[Tuple[int, str]]):
        self.root = TokenizerPrefixTreeNode()
        self.json_freetext_tokens: List[int] = []
        for token_idx, decoded in regular_tokens:
            self._add_token_to_tree(decoded, token_idx, self.root)
            # Performance optimization - cache the tokens of all the strings that don't contain a quote in the middle.
            # When we are in a JSON freetext string field, they will all be permitted and this will save a lot of tree iterations.
            if '"' not in decoded or decoded.index('"') == len(decoded) - 1:
                self.json_freetext_tokens.append(token_idx)

    def _add_token_to_tree(self, token_str: str, token_idx: int, node: TokenizerPrefixTreeNode):
        for character in token_str:
            if character not in node.children:
                node.children[character] = TokenizerPrefixTreeNode()
            node = node.children[character]
        node.tokens.append(token_idx)

lmformatenforcer/tokenizerprefixtree.py
noamgat/lm-format-enforcer
from typing import Callable, List, Tuple, Union
from transformers import AutoModelForCausalLM

from .characterlevelparser import CharacterLevelParser
from transformers.generation.logits_process import LogitsWarper, PrefixConstrainedLogitsProcessor
from transformers.tokenization_utils import PreTrainedTokenizerBase

import torch

from .tokenenforcer import TokenEnforcer

class LogitsSaverWarper(LogitsWarper):
    def __init__(self) -> None:
        self.scores = []
        
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        cpu_scores = scores.clone().detach()
        self.scores.append(cpu_scores)
        return scores
    
class LogitsSaverManager:
    warper: LogitsSaverWarper

    def __init__(self, model: AutoModelForCausalLM):
        self.model = model
        self.warper = None
        self.old_warper = None

    def replace_logits_warper(self, filter_func = None):
        self.old_warper = self.model._get_logits_warper

        def new_logits_warper(generation_config):
            warpers = self.old_warper(generation_config)
            self.warper = LogitsSaverWarper()
            warpers.insert(0, self.warper)
            if filter_func is not None:
                processor = PrefixConstrainedLogitsProcessor(filter_func, 1)
                warpers.insert(1, processor)
            return warpers
        self.model._get_logits_warper = new_logits_warper

    def unreplace_logits_warper(self):
        self.model._get_logits_warper = self.old_warper

    def get_generated_scores(self, token_sequence):
        relevant_tokens = token_sequence[-len(self.warper.scores):]
        scores_matrix = torch.concat(self.warper.scores)  # n_tokens * vocab_size
        probs = torch.softmax(scores_matrix, dim=1)  # n_tokens * vocab_size
        token_probs = probs[torch.arange(probs.size(0)), relevant_tokens]  # n_tokens
        return token_probs.to('cpu').tolist()
    
    def get_leading_scores(self) -> Tuple[List[int], List[float]]:
        scores_matrix = torch.concat(self.warper.scores)  # n_tokens * vocab_size
        probs = torch.softmax(scores_matrix, dim=1)  # n_tokens * vocab_size
        best_tokens = torch.argmax(scores_matrix, dim=1)
        token_probs = probs[torch.arange(probs.size(0)), best_tokens]  # n_tokens
        token_probs_list = token_probs.to('cpu').tolist()
        return best_tokens.tolist(), token_probs_list


def _build_regular_tokens_list(tokenizer: PreTrainedTokenizerBase) -> List[Tuple[int, str]]:
    token_0 = tokenizer.encode("0")[-1]
    regular_tokens = []
    for token_idx in range(tokenizer.vocab_size):
        if token_idx in tokenizer.all_special_ids:
            continue
        # We prepend token 0 and skip the first letter of the result to get a space if the token is a start word.
        decoded = tokenizer.decode([token_0, token_idx])[1:]
        regular_tokens.append((token_idx, decoded))
    return regular_tokens


def build_transformers_prefix_allowed_tokens_fn(tokenizer: PreTrainedTokenizerBase, 
                                                character_level_parser: CharacterLevelParser) -> Callable[[int, torch.Tensor], List[int]]:
    """Build the prefix allowed tokens function that transformers will use to filter the tokens generated by the model. The result
    can be passed to the prefix_allowed_tokens_fn parameter of the generate() method of transformers models or pipeline configurations."""
    regular_tokens = _build_regular_tokens_list(tokenizer)
    token_enforcer = TokenEnforcer(regular_tokens, character_level_parser, tokenizer.decode, tokenizer.eos_token_id)

    def transformers_filter_allowed_tokens(batch_id: int, sent: torch.Tensor) -> List[int]:
        token_sequence = sent.tolist()
        return token_enforcer.get_allowed_tokens(token_sequence)
    
    return transformers_filter_allowed_tokens


def generate_enforced(model: AutoModelForCausalLM, 
                      tokenizer: PreTrainedTokenizerBase, 
                      character_level_parser: CharacterLevelParser, 
                      **kwargs: dict) -> Union[str, dict]:
    """Generate text from a model while enforcing a given format, generating enforcing diagnostic information. 
    This can be used instead of calling model.generate().
    If return_dict_in_generate and output_scores parameters are True, diagnostic information will be returned in the result.
    If you don't need this, consider using prefix_allowed_tokens_fn + build_transformers_prefix_allowed_tokens_fn() instead"""
    
    transformers_filter_allowed_tokens = build_transformers_prefix_allowed_tokens_fn(tokenizer, character_level_parser)
    
    is_multi_inputs = kwargs['input_ids'].shape[0] > 1
    is_multi_beams = kwargs.get('num_beams', 1) > 1
    support_diagnostics = not (is_multi_inputs or is_multi_beams)  # TODO: Support diagnostics in these cases as well.
    return_dict_in_generate = kwargs.get('return_dict_in_generate', False)
    output_scores = kwargs.get('output_scores', None)

    # We do some internals hacking in order to extract the data needed for diagnostics. If we weren't asked for them,
    # we are better off simply using prefix_allowed_tokens_fn parameter.
    should_run_in_advanced_mode = return_dict_in_generate and output_scores and support_diagnostics

    if should_run_in_advanced_mode:
        logits_saver = LogitsSaverManager(model)
        logits_saver.replace_logits_warper(transformers_filter_allowed_tokens)
        generate_kwargs = kwargs
        
        try:
            output = model.generate(**generate_kwargs)
        finally:
            logits_saver.unreplace_logits_warper()

        sequence = output.sequences[0] if return_dict_in_generate else output[0]
        sequence = output.sequences[0]
        generated_scores = logits_saver.get_generated_scores(sequence)
        generated_tokens = sequence[-len(generated_scores):].to('cpu').tolist()
        single_token_strs = [tokenizer.convert_ids_to_tokens([token], skip_special_tokens=False)[0] for token in generated_tokens]

        leading_tokens, leading_scores = logits_saver.get_leading_scores()
        leading_token_strs = [tokenizer.convert_ids_to_tokens([token], skip_special_tokens=False)[0] for token in leading_tokens]
        df_dict = {}  # In order to minimize the package's dependencies, we don't create a dataframe, but create a dataframe-like dictionary instead.
        df_dict['generated_token'] = single_token_strs
        df_dict['generated_token_idx'] = generated_tokens
        df_dict['generated_score'] = generated_scores
        df_dict['leading_token'] = leading_token_strs
        df_dict['leading_token_idx'] = leading_tokens
        df_dict['leading_score'] = leading_scores
        output.enforced_scores = df_dict
    else:
        output = model.generate(**kwargs, prefix_allowed_tokens_fn=transformers_filter_allowed_tokens)
    
    return output
lmformatenforcer/transformerenforcer.py
noamgat/lm-format-enforcer
from lmformatenforcer import CharacterLevelParser


def assert_parser_with_string(string: str, parser: CharacterLevelParser, expect_success: bool):
    for idx, character in enumerate(string):
        try:
            if character in parser.get_allowed_characters():
                parser = parser.add_character(character)
            else:
                if expect_success:
                    raise ValueError(f"Parser does not allow '{character}' at index {idx}")
                else:
                    return  # Success
        except Exception as e:
            raise Exception(f"Error parsing '{character}' at index {idx}: {e}", e)
    if parser.can_end() and not expect_success:
        raise ValueError("Parser succeeded when it should have failed")
    if not parser.can_end() and expect_success:
        raise ValueError("Parser did not reach end state when it should have")
tests/common.py
noamgat/lm-format-enforcer
from typing import Dict, List, Optional
from pydantic import BaseModel
from lmformatenforcer import JsonSchemaParser
from enum import Enum
 
from .common import assert_parser_with_string


def _test_json_schema_parsing_with_string(string: str, schema_dict: dict, expect_success: bool):
    parser = JsonSchemaParser(schema_dict)
    assert_parser_with_string(string, parser, expect_success)

class InnerModel(BaseModel):
    list_of_ints: List[int]

class IntegerEnum(Enum):
    ONE = 1
    TWO = 2
    THREE = 3
    FOUR = 4

class StringEnum(Enum):
    ONE = "One"
    TWO = "Two"
    THREE = "Three"
    FOUR = "Four"

class SampleModel(BaseModel):
    num: int
    dec: Optional[float] = None
    message: Optional[str] = None
    list_of_strings: Optional[List[str]] = None
    inner_dict: Optional[Dict[str, InnerModel]] = None
    simple_dict: Optional[Dict[str, int]] = None
    list_of_models: Optional[List[InnerModel]] = None
    enum: Optional[IntegerEnum] = None
    enum_dict: Optional[Dict[str, StringEnum]] = None
    true_or_false: Optional[bool] = None


def test_parsing_test_model():
    test_string = '{"num":1,"dec":1.1,"message":"ok","list_of_strings":["a","b","c"],"inner_dict":{"a":{"list_of_ints":[1,2,3]}}}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_invalid_key_in_json_string():
    test_string = '{"numa":1,"dec":1.1,"message":"ok","list_of_strings":["a","b","c"],"inner_dict":{"a":{"list_of_ints":[1,2,3]}}}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), False)

def test_incomplete_json():
    # Intentionally missing closing }
    test_string = '{"num":1,"dec":1.1,"message":"ok","list_of_strings":["a","b","c"],"inner_dict":{"a":{"list_of_ints":[1,2,3]}}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), False)

def test_invalid_value_type_in_json_string():
    test_string = '{"num:"1","dec":1.1,"message":"ok","list_of_strings":["a","b","c"],"inner_dict":{"a":{"list_of_ints":[1,2,3]}}}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), False)

def test_list_of_objects():
    test_string = '{"list_of_models":[{"list_of_ints":[1,2,3]},{"list_of_ints":[4,5,6]}],"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)
    test_string = '{"list_of_models": [{"list_of_ints":[1, 2, 3]} , {"list_of_ints":[4,5,6]}],"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_simple_dict():
    test_string = '{"simple_dict":{"a":1,"b":2,"c":3},"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_int_enum():
    test_string = '{"enum":4,"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_invalid_int_enum_value():
    test_string = '{"enum":5,"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), False)

def test_str_enum():
    test_string = '{"enum_dict":{"a":"One","b":"Two","c":"Three","d":"Four"},"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_invalid_str_enum_value():
    test_string = '{"enum_dict":{"a":"Onee"},"num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), False)

def test_whitespaces():
    test_string = '{ "message": "","num":1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_whitespace_before_number():
    test_string = '{"num": 1, "dec": 1.1}'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)

def test_whitespace_before_close():
    test_string = '{"num":1 }'
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), True)    

def test_required_field():
    test_string = '{"dec": 1.1}'  # num is a required field, doesn't exist, should fail.
    _test_json_schema_parsing_with_string(test_string, SampleModel.schema(), False)

def test_boolean_field():
    _test_json_schema_parsing_with_string('{"num":1,"true_or_false":false}', SampleModel.schema(), True)
    _test_json_schema_parsing_with_string('{"num":1,"true_or_false":true}', SampleModel.schema(), True)
    _test_json_schema_parsing_with_string('{"num":1,"true_or_false": true}', SampleModel.schema(), True)
    _test_json_schema_parsing_with_string('{"num":1,"true_or_false":falsy}', SampleModel.schema(), False)

    

tests/test_jsonschemaparser.py
noamgat/lm-format-enforcer
from lmformatenforcer import RegexParser
from .common import assert_parser_with_string

def _test_regex_parsing_with_string(string: str, regex: str, expect_success: bool):
    parser = RegexParser(regex)
    assert_parser_with_string(string, parser, expect_success)


def test_parsing_exact_string():
    _test_regex_parsing_with_string(
        string='abc123', 
        regex='abc123', 
        expect_success=True)

def test_parsing_exact_string_failure():
    _test_regex_parsing_with_string(
        string='abc124', 
        regex='abc123', 
        expect_success=False)
    
def test_parsing_exact_string_not_reaching_end():
    _test_regex_parsing_with_string(
        string='abc123', 
        regex='abc1234', 
        expect_success=False)


def test_parsing_letter_options():
    for letter in 'cdefghif':
        expect_success = letter in 'cdef'
        _test_regex_parsing_with_string(
            string=f'ab{letter}123', 
            regex='ab(c|d|e|f)123', 
            expect_success=expect_success)
        

def test_parsing_digits():
    for letter in '0123abcd':
        expect_success = letter.isnumeric()
        _test_regex_parsing_with_string(
            string=f'ab{letter}123', 
            regex='ab\d123', 
            expect_success=expect_success)
        

def test_parsing_repeat():
    for num_repeats in range(20):
        expect_success = num_repeats > 0
        _test_regex_parsing_with_string(
            string=f'ab{"c" * num_repeats}123', 
            regex='abc+123', 
            expect_success=expect_success)
        

def test_any_character():
    for num_repeats in range(20):
        for character in '0123abcd':
            expect_success = num_repeats > 0
            _test_regex_parsing_with_string(
                string=f'ab{character * num_repeats}123', 
                regex='ab.+123', 
                expect_success=expect_success)
            
def test_dates():
    # https://stackoverflow.com/q/15491894 , removed the ^ and $ because interegular doesn't support them
    date_regex = r'(0?[1-9]|[12][0-9]|3[01])[\/\-](0?[1-9]|1[012])[\/\-]\d{4}'
    _test_regex_parsing_with_string('01/01/2020', date_regex, True)
    _test_regex_parsing_with_string('29/04/1986', date_regex, True)
    _test_regex_parsing_with_string('001/01/2020', date_regex, False)
tests/test_regexparser.py
noamgat/lm-format-enforcer