Token-triggered processors

[cpfiffer]

There are certain cases where outlines users require unstructured text until a particular token is reached, such as <|beginfunctioncall|> , code blocks, JSON blocks. etc.

Currently, these are often handled with regular expressions that can be quite messy.

DeepSeek's thinking tokens are a great example. DeepSeek R1's response begins with <think> and then includes some amount of unstructured text, followed by </think>. Thinking tokens are the primary reason for R1's performance.

The issue with the current solution (specifying a regular expression) is that the regular expression is quite complicated and difficult to debug:

# Set up response format you want the LLM to respond with.
class YesNo(BaseModel):
    answer: Literal['yes', 'no']

yesno_regex = build_regex_from_schema(convert_json_schema_to_str(YesNo))

# Add the thinking prefix to the regex 
thinking_regex = r'<think>((.|\n){0,' + str(NUM_THINKING_CHARACTERS) + r'}?)\[TRUNCATED\]</think>'

# Combine the thinking regex and the yesno regex
result_regex = thinking_regex + yesno_regex

This includes both a thinking block and the structured output at the end.

Regular expression compilation times can explode if you want to include things like thought controls that limit the thinking response, add allowed/disallowed words, or structured text to the thinking block. For example, if you force the thinking block to have a range of lengths between 10 and 50:

thinking_regex = r'<think>((?:.|\n|\s){10,50}?)<\/think>(\s)*'

Compiling a single, smaller regular expression for only the thinking block could take some load off of our compiler, and simplify the user interface for complicated and mixed forms of structure.

Ideally, I would prefer to be able to structure the thinking block and the output block separately, i.e. have separate logit processors for different sections of text. You could also jump to any number of different logit processors based on tokens you hit.

Consider the following interface, where I define a default pattern pattern_1. When a b token (or substring?) is hit, I want to switch to pattern_2.

pattern_1 = r'.{10,50}'
pattern_2 = r'(cat|dog|squirrel)'

pattern_dict = {'<|bos|>':pattern_1, 'b':pattern_2}

generator = outlines.generate.regex(model, pattern_dict)

or, in the case of DeepSeek R1 (roughly):

pattern_1 = r'.*'
pattern_2 = r'(cat|dog|squirrel)'

# NOTE: may want to provide a default 
#             pattern that begins at the BOS token
pattern_dict = {'<think>':pattern_1, '</think>':pattern_2}

generator = outlines.generate.regex(model, pattern_dict)

I think it'd be a nice convenience feature, and it might cut down on compilation times for very complicated grammars. It would introduce some challenges in how to communicate patterns to inference servers like vLLM, which does not have a way to receive multiple patterns for constrained decoding.

Thoughts?

[cpfiffer]

Also related to @rlouf's PR on a simple regex DSL (#1403), which should also help simplify the user experience. The advantage of token-triggered processors is mostly to separate concerns and reduce FSM compiler load.

[903124]

Nice! I feel like it's better to define a optional start and end token since it could be not intuitive to have </think> at the start of some patter. For example

pattern_1 = {
    'pattern': r'.*',
    'start': '<think>',
    'end': '</think>'
}
pattern_2 = {
    'pattern': r'(cat|dog|squirrel)',
    'start': '',
    'end': ''
}

pattern_dict = {
    'pattern1': pattern_1,
    'pattern2': pattern_2
}
generator = outlines.generate.regex(model, pattern_dict)

Furthermore user can also reuse a regex pattern e.g. a simple character counter

Checking character 1: 's' - no match
────────────────────────────────────────
Checking character 2: 't' - no match
────────────────────────────────────────
Checking character 3: 'r' found! r_count = 1
────────────────────────────────────────
Checking character 4: 'a' - no match
────────────────────────────────────────
Checking character 5: 'w' - no match

that reuse a regex pattern until LLM output a terminating token e.g. <think>

[rlouf]

It would be better, albeit more verbose, to give users a simple interface to build custom logits processors. Here's a first idea:

build_processor(cond_fns, process_fns, init_state)

Where init_states is an array that contains the initial value of each state. The state can be anything, and is passed to the corresponding cond_fn and process_fn at each step.

cond_fns is an array of functions that return booleans which indicate whether the corresponding processing function should be applied. If cond_fn[i](states[i], token_ids) then process_fn[i] is applied.

cond_fns_i(states[i], token_ids)

process_fn is a list of functions that process logits in place. At each step we filter the array depending on the Boolean values generated at the previous state, and apply the functions in the order of their position in the list.

process_fns_i(states[i], token_ids, logits)

This is one possibility

[cpfiffer]

Option 1: Default processors + dictionary mapping

I like the idea of using the native processor interface. If we think the processor interface is reasonably stable, we could do something similar to what I did in #1408.

In this case we'd do

class Foo(BaseModel):
    a: int

class Bar(BaseModel):
    b: int

# Create processor from schema
processor_1 = JSONLogitsProcessor(Foo, tokenizer)
processor_2 = JSONLogitsProcessor(Bar, tokenizer)
processor_3 = RegexLogitsProcessor(r"[0-9]{4}", tokenizer)

processor = MultiLogitsProcessor(
    processor_1, # default processor
    {'foo': processor_2, 'bar', processor_2} # toggleable processors
)   

Alternatively, uses should be able to provide an exact token ID that triggers another processor, though this is a terrible developer experience IMO.

Option 2: Wrapped processors

Another implementation might follow @903124's suggestion with optional stops by wrapping each processor in a TriggeredLogitsProcessor:

# Create processor from schema
processor_1 = TriggeredLogitsProcessor(
    JSONLogitsProcessor(Foo, tokenizer),
    start = '<think>',
    stop = '</think>' # optional stop string
)

processor = MultiLogitsProcessor(
    processor_1,
    ... # other TriggeredLogitsProcessors
)   

This would be unstructured by default I imagine, if a processor falls outside its start/stop substring range.

Then, the TriggeredLogitsProcessor would simply ignore any tokens passed to it if the generated sequence does not match its start/stop criteria.

Of these, I think I prefer option 2. I think it could also be reasonably well wrapped by the interface @rlouf described.

[rlouf]

You can implement all of these with the interface I mentioned above, so we TriggeredLogitsProcess etc. could just be factory functions.

[rlouf]

I think I have a better solution. First, as far as debugging go, I'm not sure your solution is necessarily better than the regex alternative. Then, it is quite complex. I think a much better solution would be to let users write e.g.

from outlines.types import text

structure = text + "<think>" + text + "</think>" + regex(r"(yes|no)")

Here there is only one regex that should be compiled, text representing unconstrained generation. structure is a tree, and we can inspect its structure, so we should be able to transform

model(prompt, structure)

into

result = model(prompt, until="<think>")
prompt += result
result = model(prompt, until="</think>")
prompt += result
result = model(prompt, regex(r"(yes|no)"))

It is always possible to add a debug mode:

model(prompt, structure, debug=True)

or maybe better have another class of Generators:

generator = REPL(model, structure)
generator(prompt)

[903124]

I'm thinking about this way recently:

Using the strawberry example above user can define a factory function

class strawberry_count(...):
    def thinking_process(self):
        return "think" + ... + "/think"
    def r_counting_loop(self,until="break"):
        while True:
              return "Checking ...." + ... + (or (newline/break))
    def conclusion(self):
         return ....
    def pipeline(self):
        return [ self.thinking_process,self.r_counting_loop, self.conclusion]

And do model(prompt,structure = strawberry_count)

[cpfiffer]

I think this is a good interface:

structure = text + "" + text + "" + regex(r"(yes|no)")

Correct me if I'm wrong, but this gives us a Sequence object, no? Are the token literals promoted to Terms?

In which case, building this interface doesn't seem particularly difficult and strikes me as having low mental overhead. It'd be nice to have

structure = text + "<think>" + text + json_block('python')  + text + "</think>" + regex(r"(yes|no)")

However, one issue here is that it does not seem obvious to me how to flexibly handle tokens that you may or may not need.

As a (very reductive) example, suppose I want to force the model to write "cat" after the sequence "my favorite pet is a ". If you want this from an unstructured block, I'd have to write my own regex to explicitly include this condition.

This is simple for my example, but complicated nested structures that trigger different types of structure could strain the compiler quite a bit. Suppose for example you did not know that you needed a LaTeX block, but if that block appears in the course of unstructured text, than you want to apply a LaTeX grammar. The Term method doesn't seem to be able to handle this, especially for out-of-order or unexpected blocks.

[903124]

Since I'm a pretty impatient person 😂 I've built my prototype here: https://github.com/903124/outlines
Example usage:

import outlines
from outlines.pipeline import run_pipeline
from outlines.types.dsl import String, Regex
from outlines.types import (
    sentence, paragraph, digit, char, one_or_more, zero_or_more, optional, repeat, newline, char, whitespace
)
model = outlines.models.llamacpp(
    "microsoft/Phi-3-mini-4k-instruct-gguf",
    "Phi-3-mini-4k-instruct-q4.gguf"
)

# Define the regex patterns for each step
thinking_pattern = String("First, listing all letter in the word strawberry starting: ") + repeat(char + whitespace, 8, 11) + newline

counting_pattern = (
    String(" Checking letter: ") + 
    Regex(r"[A-Za-z]") + 
    String(" Is the letter r? ") + 
    (String("It's a R ") | String("Not a R ")) + 
    String(".") + 
    optional(String(" Running count: ") + digit) + 
    newline
).repeat(7, 12) + (String("\n") | String("break"))

conclusion_pattern = (
    String("Conclusion: ") + 
    String("The word 'strawberry' contains ") + 
    digit + 
    String(" 'r' character") + 
    optional(String("s")) + 
    String(".")
)

# Run the pipeline with a single API call
prompt = "Count the number of 'r' characters in the word 'strawberry'."
output = run_pipeline(
    model,
    prompt=prompt,
    regex_patterns=[thinking_pattern, counting_pattern, conclusion_pattern],
    token_triggers=[r"\n\n", r"(break|\n\n)", None],
    max_tokens=500
)

print(output)

And expected output (please ignore the content it's a pretty small model)

First, listing all letter in the word strawberry starting: s t r a w b e r r y 
  Checking letter: s Is the letter r? Not a R . Running count: 1
 Checking letter: t Is the letter r? Not a R . Running count: 1
 Checking letter: r Is the letter r? Not a R . Running count: 1
 Checking letter: a Is the letter r? Not a R . Running count: 1
 Checking letter: w Is the letter r? Not a R . Running count: 1
 Checking letter: b Is the letter r? Not a R . Running count: 1
 Checking letter: e Is the letter r? Not a R . Running count: 1
 Checking letter: r Is the letter r? Not a R . Running count: 1
 Checking letter: r Is the letter r? Not a R . Running count: 1
 Checking letter: r Is the letter r? Not a R . Running count: 1
 Checking letter: y Is the letter r? Not a R . Running count: 1

 Conclusion: The word 'strawberry' contains 1 'r' character.

[cpfiffer]

Oh wow this is excellent. It's also probably reasonably straightforward to include the convenience functionality for Pydantic objects, like

def Conclusion(BaseModel):
    conclusion: str
...

output = run_pipeline(
    model,
    prompt=prompt,
    regex_patterns=[thinking_pattern, counting_pattern, Conclusion],
    token_triggers=[r"\n\n", r"(break|\n\n)", None],
    max_tokens=500
)

(obviously tidied up as we wouldn't want to put Conclusion in the regex_patterns arg)

My only note here would be on the separation of the patterns and triggers. I would expect these to be unified to present developer footguns where the pattern + trigger lists are of different lengths or orders. The previous descriptions of the dictionary/tuple/small class approach would force the developer to keep the trigger and the pattern together:

# case 1, class
from outlines.fsm.json_schema import convert_json_schema_to_str
from outlines_core.fsm.json_schema import build_regex_from_schema

class TokenTriggered():
    trigger: re.Pattern # might also support a string for exact match?
    pattern: re.Pattern
    
    def __init__(self, trigger: Union[str, re.Pattern], pattern: Union[str, re.Pattern, BaseModel]):
        # If pattern is a BaseModel, convert to regex
        if isinstance(pattern, BaseModel):
            self.pattern = build_regex_from_schema(convert_json_schema_to_str(pattern))
            
        # If pattern is a string, compile it into a regex pattern
        elif isinstance(pattern, str):
            self.pattern = re.compile(pattern)
            
        # If pattern is a regex pattern, use it directly
        elif isinstance(pattern, re.Pattern):
            self.pattern = pattern
            
        # Not a clue what this is
        else:
            raise ValueError(f"Invalid pattern type: {type(pattern)}")
            
        self.trigger = re.compile(trigger) if isinstance(trigger, str) else trigger
    
    # . . . other stuff

We could also use the simpler tuple/dict approach but IMO this could make it less clear to devs how to use it, such as key names in dicts or ordering of patterns/triggers in tuples.

[903124]

Updated the code to support Pydantic e.g.

import outlines
from outlines.pipeline import run_pipeline
from outlines.types.dsl import String, Regex, json_schema
from outlines.types import (
    sentence, paragraph, digit, char, one_or_more, zero_or_more, optional, repeat, newline
)
from pydantic import BaseModel, Field

# Define Pydantic models for structured output
class Conclusion(BaseModel):
    conclusion: str = Field(description="The final conclusion about the count")
    count: int = Field(description="The number of 'r' characters found")
    word: str = Field(description="The word that was analyzed")

# Define the regex patterns for each step
thinking_pattern = String("First, listing all letter in the word strawberry starting: ") + repeat(char, 8, 11) + newline

counting_pattern = (
    String(" Checking letter: ") + 
    Regex(r"[A-Za-z]") + 
    String(" Is the letter r? ") + 
    (String("It's a R ") | String("Not a R ")) + 
    String(".") + 
    optional(String(" Running count: ") + digit) + 
    newline
).repeat(7, 12) + (String("\n") | String("break"))

# Use Pydantic model for the conclusion step
conclusion_pattern = json_schema(Conclusion)

model = outlines.models.llamacpp(
    "microsoft/Phi-3-mini-4k-instruct-gguf",
    "Phi-3-mini-4k-instruct-q4.gguf"
)
# Run the pipeline with a single API call
prompt = "Count the number of 'r' characters in the word 'strawberry'."
output = run_pipeline(
    model,
    prompt=prompt,
    regex_patterns=[thinking_pattern, counting_pattern, conclusion_pattern],
    token_triggers=[r"\n\n", r"(break|\n\n)", None],
    max_tokens=500
)

print(output)

About the trigger pattern after using the above prototype for a while, I feel like being optional could be a better approach (i.e. by default may not need any token trigger), as after incorporating DSL I assume quite a number of developers would write their custom end token (e.g. break in example above) and hence implicitly having one.