Decorators ​

Variable decorators offer a simple mechanism to execute custom variable value transformations during the decoding process. They enable custom pre-processing, post-processing or streaming of model output.

Post-Processing Decorators ​

The standard use of decorators is to post-process a generated variable value:

def screaming(value):
    """Decorator to convert a string to uppercase"""
    return value.upper()

"Say 'this is a test':[@screaming TEST]"

Say 'this is a test': TESTTHIS IS A TEST

The example above shows a simple decorator that converts the generated variable value to uppercase. By default, decorators are applied to the variable value after the variable has finished decoding.

Value and Prompt Representation Next to simple use as shown above, the function signature of a decorator can also be extended to obtain more information about the variable context:

from lmql.runtime.program_state import ProgramState
from typing import Any

def aslist(value: Any, prompt_value: str, context: ProgramState):
    """Decorator to convert a comma-separated string into a List[str]"""
    return value.split(", "), prompt_value

"A (comma-separated) list of pancake ingredients: [@aslist ANSWER]"

A (comma-separated) list of pancake ingredients: ANSWER Flour, Eggs, Milk, Baking Powder, Salt, Butter

Here, the decorator aslist() has three arguments:

• value: any: the post-processed value of ANSWER after it has been fully generated
• prompt_value: str, the text representation of ANSWER as it was generated by the model
• context: ProgramState, the current program state.

The decorator returns a tuple of (value, prompt_value), where value is the updated program value for ANSWER and prompt_value is the updated text representation.

Differentiating between value and prompt_value allows decorators to modify the program value of ANSWER as well as its text representation as used in the prompt. For instance, in the snippet above, the remaining program can operate on ANSWER as the Python List[str] of ingredients, while the model still sees the comma-separated string representation of the list.

If a decorator does not provide value and prompt_value as a tuple, the prompt value is assumed to be str(value), i.e. the string representation of the converted value.

Streaming Decorators ​

Another form of decorators are streaming decorators. Streaming decorators are applied to the variable value during the decoding process, i.e. the decorator is called for intermediate value of a variable as it is being decoded. This allows for variable-specific streaming of the output:

from lmql.runtime.program_state import ProgramState

@lmql.decorators.streaming
def stream(value: str, context: ProgramState):
    """Decorator to stream the variable value"""
    print("VALUE", [value])

"Enumerate the alphabet without spaces:[@stream TEST]"

During execution of this query, stream() is called for each intermediate value of TEST, leading to the output shown above. This allows you to stream output as it is being generated, e.g. to show partial responses in a chat application or on the command line.

VALUE ['']
VALUE ['']
VALUE ['\n']
VALUE ['\n\n']
VALUE ['\n\nABC']
VALUE ['\n\nABCDEF']
...

Pre-Processing Decorators ​

Lastly, a decorator can also hook into query execution right before the generation of a variable begins. For instance, consider the following caching decorator function:

from lmql.runtime.program_state import ProgramState
from lmql.language.qstrings import TemplateVariable

cached_values = {
    "ITEM0": "A shopping cart"
}

@lmql.decorators.pre
def cache(variable: TemplateVariable, context: ProgramState):
    """Decorator to cache variable values by name"""
    return cached_values.get(variable.name, variable)

"""A list of things not to forget when going to the supermarket:
-[@cache ITEM0]
-[@cache ITEM1]
-[@cache ITEM0]
""" where STOPS_BEFORE(ITEM0, "\n") and STOPS_BEFORE(ITEM1, "\n")

Given a pre-determined list of fixed values per variable name like ITEM0, query execution only actually invokes the model for ITEM1. For all occurences of ITEM0, the cached value is used instead.

For this, a pre-processing decorator either returns the provided variable: TemplateVariable object to indicate that the variable should be generated as usual, or it returns a string value to indicate that the variable should be replaced by a fixed value instead.

Advanced Decorator Behavior ​

Class-Based Decorators ​

More advanced decorator behavior may require hooking into the query execution process at multiple stages. For this, the lmql.runtime.decorators.LMQLDecorator class can be implemented, an instance of which can then act as a decorator function that is invoked at multiple stages of the query execution process, e.g. you can annotate the variable as "Hello [@DecoratorCls(<args>) WORLD]".

Multiple Decorators ​

Multiple decorators can also be chained on a single variable, e.g.:

"Say 'this is a test':[@a @b @c TEST]"

During query execution, the pre-stage (if defined), invokes a(), b() and c() in order, passing the result of each decorator to the next one. The post-stage (if defined) then invokes c(), b() and a() in reverse order, passing the result of each decorator to the next one. Streaming-stage decorators do not have a return value, meaning that the order of execution does not have a particular effect.

Decorator Arguments ​

Decorators can also be provided with arguments. The implementation of such operators is achieved similarly to the implementation of decorators in standard Python via local variable capture:

def prefix(prefix: str):
    # actual decorator function
    def decorator(value: str):
        return prefix + value.strip()
    
    return decorator

"Say 'this is a test':[@prefix('PREFIX: ') TEST]"

model-output::

Say 'this is a test': [TEST PREFIX: This is a test]