Creating a composite FSM class - best practice?

[pj-ml]

Hi everyone! I have a math/statistics background, and I have just started to learn more about best practice software design patterns and habits when writing code in Python. I need a class of FSM (for a NER use case) that I can easily use to run multiple FSMs as a single FSM. Here is some code I wrote, and I was hoping to get some feedback on how to improve it further (or whether I should go an entirely different route):

from outlines.fsm.fsm import FSM, FSMState


class FSMRecursive(FSM):
    """
    Implements a finite state machine (FSM) with support for hierarchical states through recursion. This class allows
    representing complex state relationships by nesting multiple `FSMRecursive` instances, enabling simultaneous 
    representation of multiple states as a single composite state. Each instance can have its own states, transitions, 
    and children FSMs, allowing for the construction of intricate state machines with layered state logic.

    Attributes
    ----------
        name (str): Identifier for the FSM instance.
        num_states (int | None): The number of states the FSM can have. If None, the FSM is considered to have a dynamic number of states.
        disabled_state (FSMState | None): A specific state marked as disabled. If the FSM is in this state, it will return False instead of a list of allowed token IDs when calling the `allowed_token_ids` method.
        children (List["FSMRecursive"]): Nested FSMRecursive instances representing sub-states of this FSM.

    The class handles state transitions, state validation, and token acceptance, with mechanisms to integrate the state
    and behavior of child FSMs into the parent's operations, facilitating the modular construction of state machines.    """

    def __init__(
        self,
        name: str,
        num_states: int | None = None,
        disabled_state: FSMState | None = None,
        children: List["FSMRecursive"] | None = None,
        returned_token_ids_if_child_fsms_disabled = List[int] | None = None,
    ):
        if children is not None & returned_token_ids_if_child_fsms_disabled is None:
            raise ValueError("`returned_token_ids_if_child_fsms_disabled` cannot be None if argument `children` is specified. When all child fsms are in a disabled state the method `allowed_token_ids_if_has_child_fsms` must have a default set of tokens to return")
        
        self.name = name
        self.num_states = num_states
        self.disabled_state = disabled_state
        self.start_bit = None
        self.end_bit = None
        self.children = [] if children is None else children
        self.returned_token_ids_if_child_fsms_disabled = returned_token_ids_if_child_fsms_disabled
        self.set_children_bit_ranges()
        if self.end_bit > 62:
            raise ValueError(
                f"Error in `FSMRecursive` class or subclass with name {self.name}: End bit cannot be greater than 62 (63 bits available for positive ints and 1 bit for sign). This is due to too many states being recursively nested requiring more than 64 bits to represent."
            )

    def set_children_bit_ranges(self, start_bit: int = 0) -> None:
        """Sets the start and end bits for the sub-state. This is needed when extracting the sub-state from the FSMState.

        Parameters
        ----------
            start_bit
                The start bit for the sub-state. This is the bit where the sub-state starts in the full state.
        """
        next_start_bit = start_bit
        for child in self.children:
            child.set_children_bit_ranges(start_bit=next_start_bit)
            next_start_bit = child.end_bit + 1

        # next_start_bit - 1 is the end bit of the last child
        self.set_bit_range(start_bit, next_start_bit - 1)

    def set_bit_range(self, start_bit: int, end_bit: int | None = None) -> None:
        """Sets the start and end bits for the sub-state. This is needed when extracting the sub-state from the FSMState.
        """
        if start_bit < 0:
            raise ValueError(
                f"Error in `FSMRecursive` class or subclass with name {self.name}: Start bit cannot be negative."
            )

        if self.num_states is None & end_bit is None:
            raise ValueError(
                f"Error in `FSMRecursive` class or subclass with name {self.name}: Number of states must be set during class initialization using the `num_states` parameter to utilize set_bit_range() method with `end_bit` parameter set to None."
            )

        self.start_bit = start_bit
        self.end_bit = end_bit if end_bit is not None else start_bit + (self.num_states - 1).bit_length() - 1

    def is_disabled_state(self, state: FSMState) -> bool:
        """
        Returns
        ----------
        True if the FSM is disabled
        """
        return state == self.disabled_state

    def get_state(self, state: FSMState) -> int | None:
        """Extracts the FSMSubState state from the FSMState using the FSMSubState's start and end bits.

        Parameters
        ----------
            state: The integer representing the full state.

        Returns
        The integer value of the sub-state extracted from the specified bit range.
        """
        if self.start_bit is None or self.end_bit is None:
            raise ValueError(
                "Start and end bits must be set using the set_bit_range() method before extracting the state."
            )
        # Calculate the mask by shifting 1's left to cover the range of bits, then subtracting 1 to turn them into 1's.
        # Example: For a range 4-7, (1 << (7-4+1)) = 10000 in binary, subtracting 1 gives 1111.
        mask = (1 << (self.end_bit - self.start_bit + 1)) - 1
        # Shift the state right by the start bit, then apply the mask to isolate the bits of interest.
        sub_state = (state >> self.start_bit) & mask
        return FSMState(sub_state)

    def allowed_token_ids_if_has_child_fsms(self, state: FSMState) -> List[int] | bool:
        """
        Returns
        ----------
        False if this FSM is in a disabled state, otherwise a list of allowed token IDs
        """
        result_ids = set()
        all_disabled = True
        for child in self.children:
            child_state = child.get_state(state)

            if child.is_disabled_state(child_state):
                continue

            # if at least one child makes it past the `is_disabled_state` check set `all_disabeld` to False
            all_disabled = False

            if isinstance(child_state, List):
                child_ids = child.allowed_token_ids(child_state)
                result_ids.update(child_ids)
            else:
                raise ValueError(f"Sub-state for FSM '{self.name}' is not a list of states.")

        if all_disabled:
            return self.returned_token_ids_if_child_fsms_disabled

        if len(result_ids) == 0:
            raise ValueError(f"No allowed token IDs found for FSM '{self.name}'.")

        return list(result_ids)

    def allowed_token_ids_if_has_no_child_fsms(self, state: FSMState) -> List[int] | bool:
        """
        Returns
        ----------
            False if this FSM is not enabled.
            List of allowed token IDs if a subset of tokens are allowed.
        """
        ...

    def allowed_token_ids(self, state: FSMState) -> List[int] | bool:
        """
        Returns
        ----------
            False if this FSM is not enabled.
            List of allowed token IDs if a subset of tokens are allowed.
        """
        return (
            self.allowed_token_ids_if_has_child_fsms(state)
            if len(self.child_fsms) > 0
            else self.allowed_token_ids_if_has_no_child_fsms(state)
        )

    def is_final_state(self, state: FSMState) -> bool:
        # if any of the child FSMs are in a final state, then this FSM is in a final state
        for child in self.children:
            child_state = child.get_state(state)
            if child.is_final_state(child_state):
                return True

        # after checking all child FSMs, if none are in a final state, then check if this FSM is in a final state
        return self.get_state(state) == self.final_state

    def next_state_if_has_child_fsms(self, state: FSMState, token_id: int) -> FSMState:
        """
        Returns
        ----------
        Next FSM state
        """
        next_state = 0
        for child in self.children:
            child_state = child.get_state(state)

            if child.is_final_state(child_state):
                return self.final_state

            child_next_state = child.next_state(child_state, token_id)
            next_state |= child_next_state << child.start_bit

        return FSMState(next_state)

    def next_state_if_has_no_child_fsms(self, state: FSMState, token_id: int) -> FSMState:
        """
        Returns
        ----------
        Next FSM state
        """
        ...

    def next_state(self, state: FSMState, token_id: int) -> FSMState:
        """
        Returns
        ----------
        Next FSM state
        """
        return (
            self.next_state_if_has_child_fsms(state, token_id)
            if len(self.child_fsms) > 0
            else self.next_state_if_has_no_child_fsms(state, token_id)
        )

    def copy(self) -> "FSMRecursive":
        """Creates a deep copy of this FSMRecursive instance, including all attributes and nested FSMRecursive instances.

        Returns
        -------
        FSMRecursive
            A new FSMRecursive instance that is a deep copy of this instance.
        """
        # Create a copy of the current instance without the children
        copied_instance = FSMRecursive(
            name=self.name,
            num_states=self.num_states,
            disabled_state=self.disabled_state,
            children=[],  # Start with an empty list for children
        )

        # Deep copy the children separately to avoid shared references
        for child in self.children:
            copied_child = child.copy()  # Recursively call copy on each child
            copied_instance.children.append(copied_child)

        # Copy other attributes if necessary (e.g., start_bit, end_bit)
        copied_instance.set_bit_range(self.start_bit, self.end_bit)

        # Return the deep copied instance
        return copied_instance

[rlouf]

Hi! As someone who also has a stats background I'm very intrigued; why would you need several FSMs? How is it supposed to work?

[pj-ml]

Thanks for showing interest!

I hoped it could help with enforcing the named entity recognition (NER) methodology I previously mentioned here explosion/spacy-llm#346. I am not very experienced with regex, but I tried my hand at solving the problem with regex anyway. Between myself, google, and ChatGPT, we couldn't figure out a good solution.

In the proposed NER methodology, I need to force the LLM to generate an output sequence containing all of the tokens from the input sequence (excluding the instructions) with some special labels (e.g. "<|entity_end|>") in between to indicate and label the entities.

The start of an entity could be indicated by a string like "<|entity_start|>". Since my task is a zero-shot NER problem, "<|entity_start|>" is not represented by a single token but by multiple tokens. Suppose I have an input sequence with the string "<"; I need to allow both my input_sequence_enforcer FSM and my entity_sequence_enforcer FSM to suggest the next allowed tokens because it isn't clear whether the LLM has started labeling an entity or is continuing the input sequence.

I find it easier to split this NER task's constrained decoding as multiple FSMs that I later combine, e.g., input_sequence_enforcer , entity_sequence_enforcer and entity_label_enforcer.

This way, I create

1. one FSM that keeps track of where we are in the input sequence
2. one FSM enforcing the formatting when capturing a set of tokens belonging to an entity
3. one FSM enforcing the formatting and appropriate selection of one of a predefined set of entity labels

[lapp0]

It sounds like you're describing a union of FSM. This is currently done in outlines by just merging the patterns, e.g.

FSM(pattern='aabb') & FSM('pattern=cccc') == FSM(pattern='aabb|cccc')

Compiling the combined pattern to FSM can be extremely expensive (and in some cases impossible for interregular) so I think your approach of evaluating FSMs in parallel rather than compiling to a single FSM via a union operation makes sense.

[pj-ml]

Oh, that's interesting; I didn't know that one could merge patterns like that.

Exactly, yes. The union of the token IDs produced by the allowed_token_ids method for each of the (enabled) child FSMs is taken.

Side note:
The disabled state was added so that I would have an easy way of indicating that a particular FSM was not going to impose any constraints on the tokens produced (i.e., the full set of token ids would be allowed).

On second thought, the cost of returning all token ids is probably not that high, even if you have to perform set operations... Which means allowing the allowed_token_ids method to return False probably adds unnecessary complexity...

Curious to hear your thoughts on the side note.

[lapp0]

Oh, that's interesting; I didn't know that one could merge patterns like that.

It's pseudocode. interegular allows FSM.union(fsm0, fsm1) but outlines doesn't allow it, you would need to union directly within the pattern. However in my profiling results there are a few problems with this

• 1. It has superlinear time complexity, making unioned patterns extremely expensive in some cases.
• 2. It is illegal to union some lookaheads with interegular

For this reason, I think we should introduce something along the lines of your class.

Curious to hear your thoughts on the side note.

It's slightly expensive, for a single union it allows for 160 operations per second for a vocabulary of 30,000

>>> import time
>>> set1 = {f"string_{i}" for i in range(0, 20000)} 
>>> set2 = {f"string_{i}" for i in range(10000, 30000)}
>>> 
>>> start_time = time.time(); union_set = set1.union(set2); end_time = time.time()
>>> 
>>> end_time - start_time
0.006146430969238281
>>> 1/(end_time - start_time)
162.69604344453063

To optimize, we might represent tokens in a numpy array:

>>> import numpy as np
>>>
>>> array1 = np.zeros(30000, dtype=bool)
>>> array2 = np.zeros(30000, dtype=bool)
>>>
>>> array1[0:20000] = True
>>> array2[10000:30000] = True
>>> 
>>> start_time = time.time(); result_array = np.logical_or(array1, array2); end_time = time.time()
>>>
>>> end_time - start_time
2.6464462280273438e-05
>>> 1/(end_time - start_time)
37786.52252252252

[pj-ml]

@lapp0, this sounded good. Did you guys find a more elegant implementation than what I originally suggested?

[rlouf]

I think that should work:

https://outlines-dev.github.io/outlines/reference/custom_fsm_ops/