GenLM Control

A TokenSampler is an object responsible for proposing new tokens during the generation process. Given a context sequence, a sampler generates an individual token $x$ and attaches a log importance weight $w$ to it.

Note: Samplers also return a log-probability (often denoted as logp) which corresponds to the log-probability of all random choices made by the sampler. This is intended for testing purposes and is not used during the actual generation process.

You can combine the constraints of multiple language models by creating a Product potential. Use the .spawn() method on a PromptedLLM to create a shallow copy that shares the same underlying model but allows for a different prompt. The * operator between two language models (or potentials) defines a Product potential, which represents the element-wise product of their distributions.

# Spawn a new LLM with a different prompt
bos_llm = mtl_llm.spawn()
bos_llm.set_prompt_from_str("Boston is")

# Create a Product potential
product = mtl_llm * bos_llm

# Sample from the product
token_sampler = direct_token_sampler(product)
sequences = await token_sampler.smc(n_particles=5, max_tokens=25, ess_threshold=0.5)

The Product class allows you to combine two potentials using the multiplication operator (*). A Product is itself a potential and can be chained.

Behavior:

• Vocabulary: The resulting vocabulary is the intersection of the two input vocabularies ($\mathcal{A} = \mathcal{A}_1 \cap \mathcal{A}_2$).
• Weights: Weights for prefix, complete, and next-token are the sum of the individual weights in log space.

Warning: If the vocabulary overlap is less than 10% of either potential's vocabulary, a warning will be raised. Be cautious when combining potentials with minimal overlap.

# Example: Prompt intersection
mtl_llm = PromptedLLM.from_name("gpt2")
mtl_llm.set_prompt_from_str("Montreal is")

bos_llm = mtl_llm.spawn()
bos_llm.set_prompt_from_str("Boston is")

# Create product using multiplication operator
product = mtl_llm * bos_llm

Deciding which optimization to use depends on the nature of your Potential's workload:

• Auto-batching: Use when the potential's batch operations (e.g., batch_complete) are more efficient than sequential individual operations. Note that many built-in classes like PromptedLLM have auto-batching enabled by default.
• Multiprocessing: Use when the potential's operations are compute-intensive and can benefit from parallel processing.

Warning: Multiprocessing incurs significant overhead. Do not use it if the potential's operations are not compute-intensive.

A Potential is the core object in genlm-control. It encodes constraints or preferences by assigning non-negative weights to sequences of tokens. Potentials serve two primary roles:

1. Samplers: They act as components of samplers to propose new tokens during the generation process.
2. Critics: They act as critics to reweight sequences based on whether they satisfy specific constraints at each step of generation.

Potentials assign non-negative real weights to token sequences (computed in log space for stability). A potential is defined by two core weighting functions:

• complete: Assigns weights to sequences considered "finished" or "complete" (e.g., a grammatically valid sentence).
• prefix: Assigns weights to partial sequences that could potentially be extended into valid complete sequences.

Crucial Requirement: If a prefix has zero weight, all of its possible extensions and completions must also have zero weight. A prefix cannot rule out a sequence that could later become valid.

To set up a development environment for genlm-control, ensure you have Python >= 3.11 installed. It is recommended to use a dedicated virtual environment. You can install the package in editable mode along with testing and documentation dependencies using pip or uv.

# Using pip
pip install -e ".[test,docs]"

# Using uv (recommended for speed)
uv pip install -e ".[test,docs]"

To constrain generation with regular expressions, use BoolFSA.from_regex. Since BoolFSA typically operates on individual bytes, you must use the .coerce() method to map the FSA to the language model's vocabulary. Once coerced, use the AWRS token sampler, which efficiently uses the FSA to restrict the sampling space.

from genlm.control import BoolFSA, AWRS

# Create regex constraint
best_fsa = BoolFSA.from_regex(r"\sthe\s(best|worst).*")

# Coerce FSA to the product vocabulary
coerced_fsa = best_fsa.coerce(product, f=b"".join)

# Use AWRS sampler
token_sampler = AWRS(product, coerced_fsa)

sequences = await token_sampler.smc(n_particles=5, max_tokens=25, ess_threshold=0.5)

The token_sampler.smc(...) method is asynchronous and returns a coroutine. How you call it depends on your execution environment:

1. In a Jupyter Notebook or an async def function: Use await directly.
2. In a standard Python script (.py file) at the top level: Use asyncio.run() to execute the coroutine within an event loop.

Warning: Calling asyncio.run() from within an already-running event loop will raise a RuntimeError.

# In Jupyter or async function:
sequences = await token_sampler.smc(...)

# In a regular Python script:
import asyncio
sequences = asyncio.run(token_sampler.smc(...))

Follow these steps to create a clean environment for development:

Using Conda:

conda create -n genlm python=3.11
conda activate genlm
uv pip install -e ".[test,docs]"

Using uv:

uv venv .venv --python 3.11
source .venv/bin/activate
uv pip install -e ".[test,docs]"

# Conda workflow
conda create -n genlm python=3.11
conda activate genlm
uv pip install -e ".[test,docs]"

# uv workflow
uv venv .venv --python 3.11
source .venv/bin/activate
uv pip install -e ".[test,docs]"