Skip to content

Common library

steerability.algorithms.output_control.common

Output control component library.

Factors the output category into reusable components: candidate policies, per-candidate value functions, full-vocabulary logit sources, sequence scorers, a segment-search driver, a phased driver, composable stopping criteria, KV-cache utilities, and the PrefixKeyedProcessor base for stateful logits processors.

candidate_forward

Same-model forward of candidate continuations with prefix KV-cache reuse.

Some candidate values require forwarding the pipeline's own model mid-step to read candidate hidden states; CandidateForward performs those forwards. The states it reports lie at the raw output boundary of the final decoder layer (location="layer_output" in the capture utilities), before the model's final norm. These passes are marked as auxiliary via auxiliary_pass(aligned=True), so state-control accounting keeps them out of condition scoring and gate updates while transforms still apply at the candidates' true positions (prefix and candidate positions lie on the generation's own coordinate axis). At hook points where the state runtime cannot read pass positions, it skips transforming these passes and warns once. Values scored by an auxiliary model are unaffected.

CandidateForward

Batched forward of (prefix + candidate_k) for all K candidates on a given model.

Maintains an internal prefix KV cache keyed on the prefix ids. When a call's prefix extends the cached one, only the delta tokens are forwarded to extend the cache (one 1-token forward per decode step in the common case); any non-extension (rewind, restart, new generation, scoring replay from an unrelated prefix) rebuilds from scratch. The candidate evaluation then repeats a fresh copy of the cache across the K candidates, so the incremental cache survives the call. Explicit cache_position is passed on every forward, including the prefix rebuild (this is what generate does internally), and the model derives token positions from it, so the prefix must be unpadded: an attention mask containing zeros is rejected. Supports batch size 1 only.

Parameters:

Name Type Description Default
model PreTrainedModel

The model to forward through (the pipeline's own model for same-model values).

 required
Source code in steerability/algorithms/output_control/common/candidate_forward.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
class CandidateForward:
    """Batched forward of `(prefix + candidate_k)` for all K candidates on a given model.

    Maintains an internal prefix KV cache keyed on the prefix ids. When a call's prefix extends the
    cached one, only the delta tokens are forwarded to extend the cache (one 1-token forward per
    decode step in the common case); any non-extension (rewind, restart, new generation, scoring
    replay from an unrelated prefix) rebuilds from scratch. The candidate evaluation then repeats a
    fresh copy of the cache across the K candidates, so the incremental cache survives the call.
    Explicit `cache_position` is passed on every forward, including the prefix rebuild (this is what
    `generate` does internally), and the model derives token positions from it, so the prefix must
    be unpadded: an attention mask containing zeros is rejected. Supports batch size 1 only.

    Args:
        model: The model to forward through (the pipeline's own model for same-model values).
    """

    def __init__(self, model: PreTrainedModel):
        self.model = model
        self._final_layer = model.get_submodule(resolve_model_layout(model).layer_names[-1])
        self._cached_ids: torch.Tensor | None = None  # [1, T_c]
        self._cached_mask: torch.Tensor | None = None  # [1, T_c]
        self._cache = None  # past_key_values covering _cached_ids

    def _sync_cache(self, prefix_ids: torch.Tensor, full_mask: torch.Tensor) -> None:
        """Bring the internal cache up to `prefix_ids` (extend by the delta, or rebuild)."""
        if not extends_prefix(self._cached_ids, prefix_ids):
            out = self.model(
                input_ids=prefix_ids, attention_mask=full_mask, use_cache=True, return_dict=True,
                cache_position=torch.arange(prefix_ids.size(1), device=prefix_ids.device),
            )
            self._cache = out.past_key_values
        else:
            cached_len = self._cached_ids.size(1)
            if prefix_ids.size(1) > cached_len:
                delta = prefix_ids[:, cached_len:]
                positions = torch.arange(cached_len, prefix_ids.size(1), device=prefix_ids.device)
                out = self.model(
                    input_ids=delta, attention_mask=full_mask, past_key_values=self._cache,
                    use_cache=True, cache_position=positions, return_dict=True,
                )
                self._cache = out.past_key_values
            # equal length + extends -> cache already current; nothing to do
        self._cached_ids = prefix_ids.detach()
        self._cached_mask = full_mask

    @torch.no_grad()
    def last_hidden_states(
        self,
        prefix_ids: torch.Tensor,
        candidate_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Return final-layer hidden states at the candidate position for each candidate.

        The returned states lie at the raw output boundary of the final decoder layer
        (`location="layer_output"` in the capture utilities), before the model's final norm,
        recovered with a forward hook on that layer. The hook is registered per call, around the
        candidate forward only, so it observes the output after any session-registered
        state-control hooks on the same module.

        Args:
            prefix_ids: `[1, T]` prefix (batch size 1).
            candidate_ids: `[1, K]` candidate next tokens.
            attention_mask: Optional prefix mask; right-extended with ones to the prefix length.
                Must contain no zeros.

        Returns:
            A tensor `[K, H]` of candidate-position hidden states, one per candidate.

        Raises:
            ValueError: If the prefix batch size is not 1, or the attention mask contains zeros.
            RuntimeError: If the candidate forward does not pass through the final decoder layer
                exactly once.
        """
        if prefix_ids.size(0) != 1:
            raise ValueError("CandidateForward supports batch size 1 only.")
        if attention_mask is not None and not bool(attention_mask.all()):
            raise ValueError(
                "CandidateForward positions tokens by sequence index and requires an unpadded prefix; "
                "the attention mask contains zeros."
            )

        num = candidate_ids.size(1)
        device = prefix_ids.device
        full_mask = full_prefix_mask(prefix_ids, attention_mask)
        with auxiliary_pass(aligned=True):
            self._sync_cache(prefix_ids, full_mask)

            prefix_len = prefix_ids.size(1)
            repeated = repeat_cache(self._cache, num, preserve_input=True)
            cand_tokens = candidate_ids.reshape(num, 1)
            cand_mask = torch.cat(
                [full_mask.repeat(num, 1), torch.ones(num, 1, device=device, dtype=full_mask.dtype)],
                dim=1,
            )
            positions = torch.arange(prefix_len, prefix_len + 1, device=device)

            final_boundary: list[torch.Tensor] = []

            def _grab_final(module, args, output):
                final_boundary.append(output[0] if isinstance(output, tuple) else output)

            handle = self._final_layer.register_forward_hook(_grab_final)
            try:
                self.model(
                    input_ids=cand_tokens,
                    attention_mask=cand_mask,
                    past_key_values=repeated,
                    use_cache=True,
                    cache_position=positions,
                    return_dict=True,
                )
            finally:
                handle.remove()
        if len(final_boundary) != 1:
            raise RuntimeError(
                f"Expected exactly one final-layer forward for the candidate batch, "
                f"observed {len(final_boundary)}."
            )
        return final_boundary[0][:, -1, :]  # [K, H]
_cache = None instance-attribute
_cached_ids = None instance-attribute
_cached_mask = None instance-attribute
_final_layer = model.get_submodule(resolve_model_layout(model).layer_names[-1]) instance-attribute
model = model instance-attribute
__init__(model)
Source code in steerability/algorithms/output_control/common/candidate_forward.py 
39
40
41
42
43
44
def __init__(self, model: PreTrainedModel):
    self.model = model
    self._final_layer = model.get_submodule(resolve_model_layout(model).layer_names[-1])
    self._cached_ids: torch.Tensor | None = None  # [1, T_c]
    self._cached_mask: torch.Tensor | None = None  # [1, T_c]
    self._cache = None  # past_key_values covering _cached_ids
_sync_cache(prefix_ids, full_mask)

Bring the internal cache up to prefix_ids (extend by the delta, or rebuild).

Source code in steerability/algorithms/output_control/common/candidate_forward.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def _sync_cache(self, prefix_ids: torch.Tensor, full_mask: torch.Tensor) -> None:
    """Bring the internal cache up to `prefix_ids` (extend by the delta, or rebuild)."""
    if not extends_prefix(self._cached_ids, prefix_ids):
        out = self.model(
            input_ids=prefix_ids, attention_mask=full_mask, use_cache=True, return_dict=True,
            cache_position=torch.arange(prefix_ids.size(1), device=prefix_ids.device),
        )
        self._cache = out.past_key_values
    else:
        cached_len = self._cached_ids.size(1)
        if prefix_ids.size(1) > cached_len:
            delta = prefix_ids[:, cached_len:]
            positions = torch.arange(cached_len, prefix_ids.size(1), device=prefix_ids.device)
            out = self.model(
                input_ids=delta, attention_mask=full_mask, past_key_values=self._cache,
                use_cache=True, cache_position=positions, return_dict=True,
            )
            self._cache = out.past_key_values
        # equal length + extends -> cache already current; nothing to do
    self._cached_ids = prefix_ids.detach()
    self._cached_mask = full_mask
last_hidden_states(prefix_ids, candidate_ids, attention_mask=None)

Return final-layer hidden states at the candidate position for each candidate.

The returned states lie at the raw output boundary of the final decoder layer (location="layer_output" in the capture utilities), before the model's final norm, recovered with a forward hook on that layer. The hook is registered per call, around the candidate forward only, so it observes the output after any session-registered state-control hooks on the same module.

Parameters:

Name Type Description Default
prefix_ids Tensor

[1, T] prefix (batch size 1).

 required
candidate_ids Tensor

[1, K] candidate next tokens.

 required
attention_mask Tensor | None

Optional prefix mask; right-extended with ones to the prefix length. Must contain no zeros.

 None

Returns:

Type Description
Tensor

A tensor [K, H] of candidate-position hidden states, one per candidate.

Raises:

Type Description
ValueError

If the prefix batch size is not 1, or the attention mask contains zeros.
RuntimeError

If the candidate forward does not pass through the final decoder layer exactly once.
Source code in steerability/algorithms/output_control/common/candidate_forward.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
@torch.no_grad()
def last_hidden_states(
    self,
    prefix_ids: torch.Tensor,
    candidate_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Return final-layer hidden states at the candidate position for each candidate.

    The returned states lie at the raw output boundary of the final decoder layer
    (`location="layer_output"` in the capture utilities), before the model's final norm,
    recovered with a forward hook on that layer. The hook is registered per call, around the
    candidate forward only, so it observes the output after any session-registered
    state-control hooks on the same module.

    Args:
        prefix_ids: `[1, T]` prefix (batch size 1).
        candidate_ids: `[1, K]` candidate next tokens.
        attention_mask: Optional prefix mask; right-extended with ones to the prefix length.
            Must contain no zeros.

    Returns:
        A tensor `[K, H]` of candidate-position hidden states, one per candidate.

    Raises:
        ValueError: If the prefix batch size is not 1, or the attention mask contains zeros.
        RuntimeError: If the candidate forward does not pass through the final decoder layer
            exactly once.
    """
    if prefix_ids.size(0) != 1:
        raise ValueError("CandidateForward supports batch size 1 only.")
    if attention_mask is not None and not bool(attention_mask.all()):
        raise ValueError(
            "CandidateForward positions tokens by sequence index and requires an unpadded prefix; "
            "the attention mask contains zeros."
        )

    num = candidate_ids.size(1)
    device = prefix_ids.device
    full_mask = full_prefix_mask(prefix_ids, attention_mask)
    with auxiliary_pass(aligned=True):
        self._sync_cache(prefix_ids, full_mask)

        prefix_len = prefix_ids.size(1)
        repeated = repeat_cache(self._cache, num, preserve_input=True)
        cand_tokens = candidate_ids.reshape(num, 1)
        cand_mask = torch.cat(
            [full_mask.repeat(num, 1), torch.ones(num, 1, device=device, dtype=full_mask.dtype)],
            dim=1,
        )
        positions = torch.arange(prefix_len, prefix_len + 1, device=device)

        final_boundary: list[torch.Tensor] = []

        def _grab_final(module, args, output):
            final_boundary.append(output[0] if isinstance(output, tuple) else output)

        handle = self._final_layer.register_forward_hook(_grab_final)
        try:
            self.model(
                input_ids=cand_tokens,
                attention_mask=cand_mask,
                past_key_values=repeated,
                use_cache=True,
                cache_position=positions,
                return_dict=True,
            )
        finally:
            handle.remove()
    if len(final_boundary) != 1:
        raise RuntimeError(
            f"Expected exactly one final-layer forward for the candidate batch, "
            f"observed {len(final_boundary)}."
        )
    return final_boundary[0][:, -1, :]  # [K, H]

candidates

Candidate token policies for step-shaped output controls.

A candidate policy turns a full-vocabulary score row into the small set of tokens a per-candidate value function will evaluate. Composition position matters and is a feature: a value-guided processor with policy="surviving" placed after other processors steers exactly the distribution they left behind (SASA's behavior with a composed processor stack).

CandidatePolicy = Literal['top_k', 'top_p', 'surviving'] module-attribute

select_candidates(scores, policy, *, k=None, p=None)

Return (candidate_ids [B, K], candidate_scores [B, K]) under policy.

Policies:

- `top_k`: the `k` highest-scoring tokens per row (batch-safe).
- `top_p`: the smallest nucleus with cumulative probability >= `p`. Batch size 1 only
  (K varies per row); raises on `B > 1`.
- `surviving`: every token with a finite score, i.e. whatever earlier processors in the
  composed stack left finite. Batch size 1 only, for the same reason.

Parameters:

Name Type Description Default
scores Tensor

Next-token logits [B, V].

 required
policy CandidatePolicy

One of top_k, top_p, surviving.

 required
k int | None

Number of candidates for top_k.

 None
p float | None

Nucleus threshold for top_p.

 None

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple (candidate_ids, candidate_scores), each shaped [B, K].

Raises:

Type Description
ValueError

On an unknown policy, a missing k/p, or B > 1 for top_p / surviving.
Source code in steerability/algorithms/output_control/common/candidates.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def select_candidates(
    scores: torch.Tensor,
    policy: CandidatePolicy,
    *,
    k: int | None = None,
    p: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return `(candidate_ids [B, K], candidate_scores [B, K])` under `policy`.

    Policies:

        - `top_k`: the `k` highest-scoring tokens per row (batch-safe).
        - `top_p`: the smallest nucleus with cumulative probability >= `p`. Batch size 1 only
          (K varies per row); raises on `B > 1`.
        - `surviving`: every token with a finite score, i.e. whatever earlier processors in the
          composed stack left finite. Batch size 1 only, for the same reason.

    Args:
        scores: Next-token logits `[B, V]`.
        policy: One of `top_k`, `top_p`, `surviving`.
        k: Number of candidates for `top_k`.
        p: Nucleus threshold for `top_p`.

    Returns:
        A tuple `(candidate_ids, candidate_scores)`, each shaped `[B, K]`.

    Raises:
        ValueError: On an unknown policy, a missing `k`/`p`, or `B > 1` for `top_p` / `surviving`.
    """
    batch_size = scores.size(0)

    if policy == "top_k":
        if k is None or k <= 0:
            raise ValueError("top_k policy requires a positive `k`.")
        k = min(k, scores.size(-1))
        candidate_scores, candidate_ids = torch.topk(scores, k, dim=-1)
        return candidate_ids, candidate_scores

    if policy == "top_p":
        if p is None:
            raise ValueError("top_p policy requires `p`.")
        if batch_size != 1:
            raise ValueError("top_p candidate policy supports batch size 1 only (K varies per row).")
        sorted_logits, sorted_indices = torch.sort(scores, descending=False)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
        keep = cumulative_probs > (1 - p)
        keep = keep.scatter(1, sorted_indices, keep)
        cand_ids = torch.nonzero(keep)[:, 1].unsqueeze(0)
        return cand_ids, scores.gather(1, cand_ids)

    if policy == "surviving":
        if batch_size != 1:
            raise ValueError("surviving candidate policy supports batch size 1 only (K varies per row).")
        cand_ids = torch.nonzero(scores > -torch.inf)[:, 1].unsqueeze(0)
        return cand_ids, scores.gather(1, cand_ids)

    raise ValueError(f"Unknown candidate policy: {policy!r}.")

criteria

Composable stopping criteria.

Plain StoppingCriteria usable both as pipeline-composed criteria (via OutputControl.get_stopping_criteria) and inside phased-driver phases. Each returns a per-row boolean tensor: True where that row should stop.

BudgetTokens

Bases: StoppingCriteria

Stop a row once it has generated n tokens past prompt_len.

Parameters:

Name Type Description Default
n int

Token budget (new tokens) before stopping.

 required
prompt_len int

Number of prompt tokens (the budget counts tokens beyond this length).

 required
Source code in steerability/algorithms/output_control/common/criteria.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class BudgetTokens(StoppingCriteria):
    """Stop a row once it has generated `n` tokens past `prompt_len`.

    Args:
        n: Token budget (new tokens) before stopping.
        prompt_len: Number of prompt tokens (the budget counts tokens beyond this length).
    """

    def __init__(self, n: int, prompt_len: int):
        self._n = n
        self._prompt_len = prompt_len

    def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> torch.Tensor:
        fired = input_ids.size(1) - self._prompt_len >= self._n
        return torch.full((input_ids.size(0),), fired, dtype=torch.bool, device=input_ids.device)
_n = n instance-attribute
_prompt_len = prompt_len instance-attribute
__call__(input_ids, scores, **kwargs)
Source code in steerability/algorithms/output_control/common/criteria.py 
75
76
77
def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> torch.Tensor:
    fired = input_ids.size(1) - self._prompt_len >= self._n
    return torch.full((input_ids.size(0),), fired, dtype=torch.bool, device=input_ids.device)
__init__(n, prompt_len)
Source code in steerability/algorithms/output_control/common/criteria.py 
71
72
73
def __init__(self, n: int, prompt_len: int):
    self._n = n
    self._prompt_len = prompt_len

StopOnSubstring

Bases: StoppingCriteria

Stop a row once the newly generated text contains text.

Decodes only the tokens generated after prompt_len per check (so the substring test runs on the continuation, not the prompt). Decoding per step is the cost of a text-level stop condition; for token-level stops prefer StopOnTokens.

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

Tokenizer used to decode the continuation.

 required
text str

The substring that triggers stopping.

 required
prompt_len int

Number of prompt tokens to skip when decoding the continuation.

 required
Source code in steerability/algorithms/output_control/common/criteria.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class StopOnSubstring(StoppingCriteria):
    """Stop a row once the newly generated text contains `text`.

    Decodes only the tokens generated after `prompt_len` per check (so the substring test runs on
    the continuation, not the prompt). Decoding per step is the cost of a text-level stop condition;
    for token-level stops prefer `StopOnTokens`.

    Args:
        tokenizer: Tokenizer used to decode the continuation.
        text: The substring that triggers stopping.
        prompt_len: Number of prompt tokens to skip when decoding the continuation.
    """

    def __init__(self, tokenizer: PreTrainedTokenizerBase, text: str, prompt_len: int):
        self._tokenizer = tokenizer
        self._text = text
        self._prompt_len = prompt_len

    def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> torch.Tensor:
        fired = []
        for row in input_ids:
            continuation = self._tokenizer.decode(
                row[self._prompt_len:], skip_special_tokens=False
            )
            fired.append(self._text in continuation)
        return torch.tensor(fired, dtype=torch.bool, device=input_ids.device)
_prompt_len = prompt_len instance-attribute
_text = text instance-attribute
_tokenizer = tokenizer instance-attribute
__call__(input_ids, scores, **kwargs)
Source code in steerability/algorithms/output_control/common/criteria.py 
33
34
35
36
37
38
39
40
def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> torch.Tensor:
    fired = []
    for row in input_ids:
        continuation = self._tokenizer.decode(
            row[self._prompt_len:], skip_special_tokens=False
        )
        fired.append(self._text in continuation)
    return torch.tensor(fired, dtype=torch.bool, device=input_ids.device)
__init__(tokenizer, text, prompt_len)
Source code in steerability/algorithms/output_control/common/criteria.py 
28
29
30
31
def __init__(self, tokenizer: PreTrainedTokenizerBase, text: str, prompt_len: int):
    self._tokenizer = tokenizer
    self._text = text
    self._prompt_len = prompt_len

StopOnTokens

Bases: StoppingCriteria

Stop a row once its last token is one of ids.

Parameters:

Name Type Description Default
ids Iterable[int]

Token ids that trigger stopping.

 required
Source code in steerability/algorithms/output_control/common/criteria.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class StopOnTokens(StoppingCriteria):
    """Stop a row once its last token is one of `ids`.

    Args:
        ids: Token ids that trigger stopping.
    """

    def __init__(self, ids: Iterable[int]):
        self._ids = set(int(i) for i in ids)

    def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> torch.Tensor:
        last = input_ids[:, -1]
        fired = torch.tensor(
            [int(t.item()) in self._ids for t in last],
            dtype=torch.bool,
            device=input_ids.device,
        )
        return fired
_ids = set(int(i) for i in ids) instance-attribute
__call__(input_ids, scores, **kwargs)
Source code in steerability/algorithms/output_control/common/criteria.py 
53
54
55
56
57
58
59
60
def __call__(self, input_ids: torch.Tensor, scores, **kwargs) -> torch.Tensor:
    last = input_ids[:, -1]
    fired = torch.tensor(
        [int(t.item()) in self._ids for t in last],
        dtype=torch.bool,
        device=input_ids.device,
    )
    return fired
__init__(ids)
Source code in steerability/algorithms/output_control/common/criteria.py 
50
51
def __init__(self, ids: Iterable[int]):
    self._ids = set(int(i) for i in ids)

drivers

Decoding drivers (segment search and phased splicing) and their support components.

frontier

Top-k, finished, and best-so-far bookkeeping for the segment search in SearchDriver.

Frontier

Keep top-k beams by score, track per-beam finished flags and best-so-far.

Because beams arrive right-padded to a common length, a kept beam's continuation is first stripped of trailing pad_token_id positions to recover its true length. The beam is finished when the stripped continuation ends in a token of the eos set, or when its stripped length reaches max_new_tokens. A beam cut by a caller-supplied stopping criterion classifies as unfinished.

Parameters:

Name Type Description Default
keep_k int

Number of beams to retain each iteration.

 required
eos_token_id int | list[int] | None

Finished-beam token id(s), as an int, a list of ints, or None (disables the EOS check).

 required
input_length int

Prompt length, used to slice each beam's continuation.

 required
max_new_tokens int | None

Global token budget (None disables the budget check).

 required
pad_token_id int | None

Padding token id stripped from continuation tails, or None (no stripping).

 None
Source code in steerability/algorithms/output_control/common/drivers/frontier.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class Frontier:
    """Keep top-k beams by score, track per-beam finished flags and best-so-far.

    Because beams arrive right-padded to a common length, a kept beam's continuation is first
    stripped of trailing `pad_token_id` positions to recover its true length. The beam is
    finished when the stripped continuation ends in a token of the eos set, or when its
    stripped length reaches `max_new_tokens`. A beam cut by a caller-supplied stopping
    criterion classifies as unfinished.

    Args:
        keep_k: Number of beams to retain each iteration.
        eos_token_id: Finished-beam token id(s), as an int, a list of ints, or None (disables
            the EOS check).
        input_length: Prompt length, used to slice each beam's continuation.
        max_new_tokens: Global token budget (None disables the budget check).
        pad_token_id: Padding token id stripped from continuation tails, or None (no stripping).
    """

    def __init__(self, keep_k: int, eos_token_id: int | list[int] | None, input_length: int,
                 max_new_tokens: int | None, pad_token_id: int | None = None):
        self.keep_k = keep_k
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.input_length = input_length
        self.max_new_tokens = max_new_tokens
        self.best_ids: torch.Tensor | None = None
        self.best_score = float("-inf")

    def keep(self, beams: torch.Tensor, scores: list[float]) -> FrontierStep:
        """Select the top-k beams, update best-so-far, and compute finished flags.

        Args:
            beams: Candidate beam sequences `[N, T]`, right-padded to a common length.
            scores: One score per beam.

        Returns:
            A `FrontierStep` with the kept beams, their scores, and their finished flags.
        """
        score_tensor = torch.tensor(scores, device=beams.device)
        k = min(self.keep_k, score_tensor.numel())
        top_idx = torch.topk(score_tensor, k).indices
        kept = beams[top_idx]
        kept_scores = score_tensor[top_idx].tolist()

        reasons = infer_finish_reasons(
            kept[:, self.input_length:],
            {"max_new_tokens": self.max_new_tokens},
            eos_token_id=self.eos_token_id,
            pad_token_id=self.pad_token_id,
        )
        finished_flags = [reason is not None for reason in reasons]

        best_local = int(torch.argmax(torch.tensor(kept_scores)))
        if kept_scores[best_local] > self.best_score:
            self.best_score = kept_scores[best_local]
            self.best_ids = kept[best_local]

        return FrontierStep(kept_ids=kept, kept_scores=kept_scores, finished_flags=finished_flags)
best_ids = None instance-attribute
best_score = float('-inf') instance-attribute
eos_token_id = eos_token_id instance-attribute
input_length = input_length instance-attribute
keep_k = keep_k instance-attribute
max_new_tokens = max_new_tokens instance-attribute
pad_token_id = pad_token_id instance-attribute
__init__(keep_k, eos_token_id, input_length, max_new_tokens, pad_token_id=None)
Source code in steerability/algorithms/output_control/common/drivers/frontier.py 
44
45
46
47
48
49
50
51
52
def __init__(self, keep_k: int, eos_token_id: int | list[int] | None, input_length: int,
             max_new_tokens: int | None, pad_token_id: int | None = None):
    self.keep_k = keep_k
    self.eos_token_id = eos_token_id
    self.pad_token_id = pad_token_id
    self.input_length = input_length
    self.max_new_tokens = max_new_tokens
    self.best_ids: torch.Tensor | None = None
    self.best_score = float("-inf")
keep(beams, scores)

Select the top-k beams, update best-so-far, and compute finished flags.

Parameters:

Name Type Description Default
beams Tensor

Candidate beam sequences [N, T], right-padded to a common length.

 required
scores list[float]

One score per beam.

 required

Returns:

Type Description
FrontierStep

A FrontierStep with the kept beams, their scores, and their finished flags.
Source code in steerability/algorithms/output_control/common/drivers/frontier.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def keep(self, beams: torch.Tensor, scores: list[float]) -> FrontierStep:
    """Select the top-k beams, update best-so-far, and compute finished flags.

    Args:
        beams: Candidate beam sequences `[N, T]`, right-padded to a common length.
        scores: One score per beam.

    Returns:
        A `FrontierStep` with the kept beams, their scores, and their finished flags.
    """
    score_tensor = torch.tensor(scores, device=beams.device)
    k = min(self.keep_k, score_tensor.numel())
    top_idx = torch.topk(score_tensor, k).indices
    kept = beams[top_idx]
    kept_scores = score_tensor[top_idx].tolist()

    reasons = infer_finish_reasons(
        kept[:, self.input_length:],
        {"max_new_tokens": self.max_new_tokens},
        eos_token_id=self.eos_token_id,
        pad_token_id=self.pad_token_id,
    )
    finished_flags = [reason is not None for reason in reasons]

    best_local = int(torch.argmax(torch.tensor(kept_scores)))
    if kept_scores[best_local] > self.best_score:
        self.best_score = kept_scores[best_local]
        self.best_ids = kept[best_local]

    return FrontierStep(kept_ids=kept, kept_scores=kept_scores, finished_flags=finished_flags)
FrontierStep dataclass

Result of keeping the top-k beams for one search iteration.

Attributes:

Name Type Description
kept_ids Tensor

The top-k beam sequences [k, T].
kept_scores list[float]

Their scores.
finished_flags list[bool]

Per-kept-beam finished flag (EOS hit or budget reached).
Source code in steerability/algorithms/output_control/common/drivers/frontier.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
@dataclass
class FrontierStep:
    """Result of keeping the top-k beams for one search iteration.

    Attributes:
        kept_ids: The top-k beam sequences `[k, T]`.
        kept_scores: Their scores.
        finished_flags: Per-kept-beam finished flag (EOS hit or budget reached).
    """

    kept_ids: torch.Tensor
    kept_scores: list[float]
    finished_flags: list[bool]
finished_flags instance-attribute
kept_ids instance-attribute
kept_scores instance-attribute
__init__(kept_ids, kept_scores, finished_flags)

phased

Phase-plan types (Fixed, Generated) and PhasedDriver, which executes a plan by splicing token streams.

A phase plan is a list of Fixed (append text without generating) and Generated (generate until a boundary) phases. Every Generated phase delegates to model.generate with the received stacks (the driver contract satisfied per phase). An extract_after output rule reproduces ThinkingIntervention's tail extraction.

Fixed dataclass

Splice this text's tokens into the stream without generating.

Attributes:

Name Type Description
text str | Callable[[str, dict], str]

A literal string, or a callable (prompt_text, params) -> str.
replace bool

When True, this phase's tokens replace the current stream rather than being appended. Used to model a prompt-rewriting transform (ThinkingIntervention) whose text already contains the prompt.
add_special_tokens bool

Whether to add special tokens when tokenizing (True for a replacing phase that becomes the new prompt; False for an appended snippet).
Source code in steerability/algorithms/output_control/common/drivers/phased.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
@dataclass(frozen=True)
class Fixed:
    """Splice this text's tokens into the stream without generating.

    Attributes:
        text: A literal string, or a callable `(prompt_text, params) -> str`.
        replace: When True, this phase's tokens replace the current stream rather than being
            appended. Used to model a prompt-rewriting transform (ThinkingIntervention) whose text
            already contains the prompt.
        add_special_tokens: Whether to add special tokens when tokenizing (True for a replacing
            phase that becomes the new prompt; False for an appended snippet).
    """

    text: str | Callable[[str, dict], str]
    replace: bool = False
    add_special_tokens: bool = False
add_special_tokens = False class-attribute instance-attribute
replace = False class-attribute instance-attribute
text instance-attribute
__init__(text, replace=False, add_special_tokens=False)
Generated dataclass

Generate until a boundary: the until substring, any token in until_token_ids, or the token budget, whichever occurs first.

All three compose with the pipeline's criteria for the phase. until_token_ids is the backend-portable form of a delimiter that tokenizes to a special token, which a stop string cannot express (skip_special_tokens=True strips it before a vLLM stop-string match).

Attributes:

Name Type Description
until str | None

Substring that ends the phase (via StopOnSubstring). None disables.
until_token_ids tuple[int, ...]

Token ids any one of which ends the phase once it is the last generated token (via StopOnTokens). On the session path the ids extend stop_token_ids the way until extends stop_strings. Empty disables.
budget int | None

Max new tokens for the phase (via BudgetTokens). None disables.
Source code in steerability/algorithms/output_control/common/drivers/phased.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
@dataclass(frozen=True)
class Generated:
    """Generate until a boundary: the `until` substring, any token in `until_token_ids`, or the
    token `budget`, whichever occurs first.

    All three compose with the pipeline's criteria for the phase. `until_token_ids` is the
    backend-portable form of a delimiter that tokenizes to a special token, which a stop string
    cannot express (`skip_special_tokens=True` strips it before a vLLM stop-string match).

    Attributes:
        until: Substring that ends the phase (via `StopOnSubstring`). None disables.
        until_token_ids: Token ids any one of which ends the phase once it is the last generated
            token (via `StopOnTokens`). On the session path the ids extend `stop_token_ids` the
            way `until` extends `stop_strings`. Empty disables.
        budget: Max new tokens for the phase (via `BudgetTokens`). None disables.
    """

    until: str | None = None
    until_token_ids: tuple[int, ...] = ()
    budget: int | None = None

    def __post_init__(self) -> None:
        object.__setattr__(self, "until_token_ids", tuple(int(i) for i in self.until_token_ids))
budget = None class-attribute instance-attribute
until = None class-attribute instance-attribute
until_token_ids = () class-attribute instance-attribute
__init__(until=None, until_token_ids=(), budget=None)
__post_init__()
Source code in steerability/algorithms/output_control/common/drivers/phased.py 
61
62
def __post_init__(self) -> None:
    object.__setattr__(self, "until_token_ids", tuple(int(i) for i in self.until_token_ids))
PhasedDriver

Bases: DecodingDriver

Execute a per-example phase plan by splicing token streams.

Subclasses / callers provide plan(prompt_text, params) -> list[Fixed | Generated]. Every Generated phase delegates to model.generate with the received stacks. An extract_after output rule reproduces ThinkingIntervention's tail extraction: decode, rsplit on the marker, keep the original prompt's token prefix + the re-tokenized remainder.

Plans are per example; batched inputs are handled by looping over rows.

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
class PhasedDriver(DecodingDriver):
    """Execute a per-example phase plan by splicing token streams.

    Subclasses / callers provide `plan(prompt_text, params) -> list[Fixed | Generated]`. Every
    `Generated` phase delegates to `model.generate` with the received stacks. An `extract_after`
    output rule reproduces ThinkingIntervention's tail extraction: decode, rsplit on the marker, keep
    the original prompt's token prefix + the re-tokenized remainder.

    Plans are per example; batched inputs are handled by looping over rows.
    """

    RUNTIME_KWARGS_SCHEMA = [
        {
            "name": "params",
            "type": "dict",
            "scope": "call",
            "help": (
                "Phase-plan parameters for this call: a mapping whose scalar values apply to every prompt row and "
                "whose list values carry one entry per row, of batch length."
            ),
        },
    ]

    def __init__(self, extract_after: str | None = None):
        self.extract_after = extract_after
        self.tokenizer = None  # injected by the pipeline

    def max_rollouts_per_query(self) -> int | None:
        """None: the phase plan is per example, so the base class declares no static bound.
        Subclasses with a fixed plan (e.g. `PhasedDecoding`, `BudgetForcing`) override it."""
        return None

    def requirements(self) -> Requirements:
        """Phase splicing is client-side and generated phases run through the session, so no
        phase requires anything beyond the session contract."""
        return Requirements()

    def plan(self, prompt_text: str, params: dict) -> list:
        """Return the phase plan for one example. Subclasses override."""
        raise NotImplementedError

    def _params_per_example(self, runtime_kwargs: dict, batch_size: int) -> list[dict]:
        """Slice a dict-of-lists `params` per example (ThinkingIntervention's convention)."""
        params_agg = runtime_kwargs.get("params", None)
        if params_agg is None:
            return [{} for _ in range(batch_size)]
        if isinstance(params_agg, dict) and any(isinstance(v, (list, tuple)) for v in params_agg.values()):
            out = []
            for i in range(batch_size):
                p_i = {}
                for k, v in params_agg.items():
                    if isinstance(v, (list, tuple)):
                        if len(v) != batch_size:
                            raise ValueError(
                                f"params['{k}'] has length {len(v)}, but batch size is {batch_size}."
                            )
                        p_i[k] = v[i]
                    else:
                        p_i[k] = v
                out.append(p_i)
            return out
        return [params_agg] * batch_size

    def decode(self, input_ids, attention_mask, model: PreTrainedModel | None, logits_processors,
               stopping_criteria, runtime_kwargs, session=None, **gen_kwargs) -> torch.Tensor:
        if self.tokenizer is None:
            raise RuntimeError("PhasedDriver requires a tokenizer; steer() must run first.")

        runtime_kwargs = runtime_kwargs or {}
        via_session = session is not None
        base_generate = resolve_generate_callable(model, runtime_kwargs, session=session)

        if input_ids.dim() == 1:
            input_ids = input_ids.unsqueeze(0)
        batch_size = input_ids.size(0)
        params_per_example = self._params_per_example(runtime_kwargs, batch_size)
        original_prompts = self.tokenizer.decode(input_ids, skip_special_tokens=True)
        original_lengths = [row.size(0) for row in input_ids]

        final_sequences: list[torch.Tensor] = []
        for i in range(batch_size):
            row_ids = input_ids[i:i + 1]
            prompt_text = original_prompts[i]
            params = params_per_example[i]
            plan = self.plan(prompt_text, params)
            full = self._run_plan(
                plan, row_ids, prompt_text, params, base_generate, via_session,
                logits_processors, stopping_criteria, gen_kwargs,
            )
            final_sequences.append(self._finalize(full[0], original_lengths[i]))

        padded = self.tokenizer.pad(
            {"input_ids": [seq.tolist() for seq in final_sequences]},
            padding=True,
            return_tensors="pt",
        ).to(input_ids.device)
        return padded["input_ids"]

    def _run_plan(self, plan, row_ids, prompt_text, params, base_generate, via_session,
                  logits_processors, stopping_criteria, gen_kwargs) -> torch.Tensor:
        """Execute one example's plan; return the full spliced sequence `[1, L]`."""
        current = row_ids
        for phase in plan:
            if isinstance(phase, Fixed):
                text = phase.text(prompt_text, params) if callable(phase.text) else phase.text
                fixed_ids = self.tokenizer(
                    text, add_special_tokens=phase.add_special_tokens, return_tensors="pt"
                )["input_ids"].to(current.device)
                current = fixed_ids if phase.replace else torch.cat([current, fixed_ids], dim=1)
            elif isinstance(phase, Generated):
                current = self._generate_phase(
                    phase, current, base_generate, via_session,
                    logits_processors, stopping_criteria, gen_kwargs,
                )
            else:
                raise TypeError(f"Unknown phase type: {type(phase).__name__}")
        return current

    def _generate_phase(self, phase: Generated, current, base_generate, via_session,
                        logits_processors, stopping_criteria, gen_kwargs) -> torch.Tensor:
        """Run one Generated phase, composing its boundary with the pipeline's stop rules.

        On the session path the boundary lowers to normalized parameters (`until` as a stop
        string, `until_token_ids` as extra stop token ids, `budget` as a tightened
        `max_new_tokens`), so the phase runs on any backend; a raw generate callable receives the
        boundary as prompt-anchored criteria instead.
        """
        criteria = list(stopping_criteria) if stopping_criteria is not None else []
        kwargs = dict(gen_kwargs)

        if via_session:
            if phase.until is not None:
                existing = kwargs.get("stop_strings") or ()
                if isinstance(existing, str):
                    existing = (existing,)
                kwargs["stop_strings"] = (*existing, phase.until)
            if phase.until_token_ids:
                existing_ids = tuple(kwargs.get("stop_token_ids") or ())
                kwargs["stop_token_ids"] = (*existing_ids, *phase.until_token_ids)
            if phase.budget is not None:
                cap = kwargs.get("max_new_tokens")
                kwargs["max_new_tokens"] = phase.budget if cap is None else min(cap, phase.budget)
        else:
            current_len = current.size(1)
            if phase.until is not None:
                criteria.append(StopOnSubstring(self.tokenizer, phase.until, current_len))
            if phase.until_token_ids:
                criteria.append(StopOnTokens(phase.until_token_ids))
            if phase.budget is not None:
                criteria.append(BudgetTokens(phase.budget, current_len))
                if "max_new_tokens" not in kwargs:
                    kwargs["max_new_tokens"] = phase.budget

        extra = stack_generate_kwargs(logits_processors, StoppingCriteriaList(criteria))

        attention_mask = torch.ones_like(current)
        outputs = base_generate(input_ids=current, attention_mask=attention_mask, **extra, **kwargs)
        if not isinstance(outputs, torch.Tensor):
            outputs = outputs[0] if isinstance(outputs, (list, tuple)) else outputs
        return outputs

    def _finalize(self, out_ids: torch.Tensor, original_length: int) -> torch.Tensor:
        """Apply the extract_after output rule (or return the full spliced sequence)."""
        if self.extract_after is None:
            return out_ids
        keep_prefix = out_ids[:original_length]
        decoded = self.tokenizer.decode(out_ids, skip_special_tokens=False)
        remainder_txt = decoded.rsplit(self.extract_after, 1)[-1].lstrip()
        remainder_ids = (
            self.tokenizer(remainder_txt, add_special_tokens=False, return_tensors="pt")["input_ids"]
            .to(out_ids.device)
            .squeeze(0)
        )
        return torch.cat([keep_prefix, remainder_ids], dim=0)
RUNTIME_KWARGS_SCHEMA = [{'name': 'params', 'type': 'dict', 'scope': 'call', 'help': 'Phase-plan parameters for this call: a mapping whose scalar values apply to every prompt row and whose list values carry one entry per row, of batch length.'}] class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
enabled = True class-attribute instance-attribute
extract_after = extract_after instance-attribute
include_in_scoring = True class-attribute instance-attribute
same_model_forwards = False class-attribute instance-attribute
supports_batching = False class-attribute instance-attribute
tokenizer = None instance-attribute
__init__(extract_after=None)
Source code in steerability/algorithms/output_control/common/drivers/phased.py 
88
89
90
def __init__(self, extract_after: str | None = None):
    self.extract_after = extract_after
    self.tokenizer = None  # injected by the pipeline
_configure()

Post-construction hook, called after Args fields are mirrored onto the instance.

Driver presets override this to map their mirrored args onto the fields their generic base reads, so subclasses never bypass a parent __init__. Default no-op.

Source code in steerability/algorithms/core/base_control.py 
67
68
69
70
71
72
73
def _configure(self) -> None:
    """Post-construction hook, called after `Args` fields are mirrored onto the instance.

    Driver presets override this to map their mirrored args onto the fields their generic base
    reads, so subclasses never bypass a parent `__init__`. Default no-op.
    """
    pass
_finalize(out_ids, original_length)

Apply the extract_after output rule (or return the full spliced sequence).

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
def _finalize(self, out_ids: torch.Tensor, original_length: int) -> torch.Tensor:
    """Apply the extract_after output rule (or return the full spliced sequence)."""
    if self.extract_after is None:
        return out_ids
    keep_prefix = out_ids[:original_length]
    decoded = self.tokenizer.decode(out_ids, skip_special_tokens=False)
    remainder_txt = decoded.rsplit(self.extract_after, 1)[-1].lstrip()
    remainder_ids = (
        self.tokenizer(remainder_txt, add_special_tokens=False, return_tensors="pt")["input_ids"]
        .to(out_ids.device)
        .squeeze(0)
    )
    return torch.cat([keep_prefix, remainder_ids], dim=0)
_generate_phase(phase, current, base_generate, via_session, logits_processors, stopping_criteria, gen_kwargs)

Run one Generated phase, composing its boundary with the pipeline's stop rules.

On the session path the boundary lowers to normalized parameters (until as a stop string, until_token_ids as extra stop token ids, budget as a tightened max_new_tokens), so the phase runs on any backend; a raw generate callable receives the boundary as prompt-anchored criteria instead.

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def _generate_phase(self, phase: Generated, current, base_generate, via_session,
                    logits_processors, stopping_criteria, gen_kwargs) -> torch.Tensor:
    """Run one Generated phase, composing its boundary with the pipeline's stop rules.

    On the session path the boundary lowers to normalized parameters (`until` as a stop
    string, `until_token_ids` as extra stop token ids, `budget` as a tightened
    `max_new_tokens`), so the phase runs on any backend; a raw generate callable receives the
    boundary as prompt-anchored criteria instead.
    """
    criteria = list(stopping_criteria) if stopping_criteria is not None else []
    kwargs = dict(gen_kwargs)

    if via_session:
        if phase.until is not None:
            existing = kwargs.get("stop_strings") or ()
            if isinstance(existing, str):
                existing = (existing,)
            kwargs["stop_strings"] = (*existing, phase.until)
        if phase.until_token_ids:
            existing_ids = tuple(kwargs.get("stop_token_ids") or ())
            kwargs["stop_token_ids"] = (*existing_ids, *phase.until_token_ids)
        if phase.budget is not None:
            cap = kwargs.get("max_new_tokens")
            kwargs["max_new_tokens"] = phase.budget if cap is None else min(cap, phase.budget)
    else:
        current_len = current.size(1)
        if phase.until is not None:
            criteria.append(StopOnSubstring(self.tokenizer, phase.until, current_len))
        if phase.until_token_ids:
            criteria.append(StopOnTokens(phase.until_token_ids))
        if phase.budget is not None:
            criteria.append(BudgetTokens(phase.budget, current_len))
            if "max_new_tokens" not in kwargs:
                kwargs["max_new_tokens"] = phase.budget

    extra = stack_generate_kwargs(logits_processors, StoppingCriteriaList(criteria))

    attention_mask = torch.ones_like(current)
    outputs = base_generate(input_ids=current, attention_mask=attention_mask, **extra, **kwargs)
    if not isinstance(outputs, torch.Tensor):
        outputs = outputs[0] if isinstance(outputs, (list, tuple)) else outputs
    return outputs
_params_per_example(runtime_kwargs, batch_size)

Slice a dict-of-lists params per example (ThinkingIntervention's convention).

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def _params_per_example(self, runtime_kwargs: dict, batch_size: int) -> list[dict]:
    """Slice a dict-of-lists `params` per example (ThinkingIntervention's convention)."""
    params_agg = runtime_kwargs.get("params", None)
    if params_agg is None:
        return [{} for _ in range(batch_size)]
    if isinstance(params_agg, dict) and any(isinstance(v, (list, tuple)) for v in params_agg.values()):
        out = []
        for i in range(batch_size):
            p_i = {}
            for k, v in params_agg.items():
                if isinstance(v, (list, tuple)):
                    if len(v) != batch_size:
                        raise ValueError(
                            f"params['{k}'] has length {len(v)}, but batch size is {batch_size}."
                        )
                    p_i[k] = v[i]
                else:
                    p_i[k] = v
            out.append(p_i)
        return out
    return [params_agg] * batch_size
_run_plan(plan, row_ids, prompt_text, params, base_generate, via_session, logits_processors, stopping_criteria, gen_kwargs)

Execute one example's plan; return the full spliced sequence [1, L].

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def _run_plan(self, plan, row_ids, prompt_text, params, base_generate, via_session,
              logits_processors, stopping_criteria, gen_kwargs) -> torch.Tensor:
    """Execute one example's plan; return the full spliced sequence `[1, L]`."""
    current = row_ids
    for phase in plan:
        if isinstance(phase, Fixed):
            text = phase.text(prompt_text, params) if callable(phase.text) else phase.text
            fixed_ids = self.tokenizer(
                text, add_special_tokens=phase.add_special_tokens, return_tensors="pt"
            )["input_ids"].to(current.device)
            current = fixed_ids if phase.replace else torch.cat([current, fixed_ids], dim=1)
        elif isinstance(phase, Generated):
            current = self._generate_phase(
                phase, current, base_generate, via_session,
                logits_processors, stopping_criteria, gen_kwargs,
            )
        else:
            raise TypeError(f"Unknown phase type: {type(phase).__name__}")
    return current
cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py 
195
196
197
198
199
200
201
def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass
clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name Type Description Default
seed int | None

Optional seed forwarded to the clone's reseed().

 None

Returns:

Type Description
BaseControl

The clone.
Source code in steerability/algorithms/core/base_control.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone
decode(input_ids, attention_mask, model, logits_processors, stopping_criteria, runtime_kwargs, session=None, **gen_kwargs)

Run the decoding procedure; return full sequence ids (prompt + continuation).

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def decode(self, input_ids, attention_mask, model: PreTrainedModel | None, logits_processors,
           stopping_criteria, runtime_kwargs, session=None, **gen_kwargs) -> torch.Tensor:
    if self.tokenizer is None:
        raise RuntimeError("PhasedDriver requires a tokenizer; steer() must run first.")

    runtime_kwargs = runtime_kwargs or {}
    via_session = session is not None
    base_generate = resolve_generate_callable(model, runtime_kwargs, session=session)

    if input_ids.dim() == 1:
        input_ids = input_ids.unsqueeze(0)
    batch_size = input_ids.size(0)
    params_per_example = self._params_per_example(runtime_kwargs, batch_size)
    original_prompts = self.tokenizer.decode(input_ids, skip_special_tokens=True)
    original_lengths = [row.size(0) for row in input_ids]

    final_sequences: list[torch.Tensor] = []
    for i in range(batch_size):
        row_ids = input_ids[i:i + 1]
        prompt_text = original_prompts[i]
        params = params_per_example[i]
        plan = self.plan(prompt_text, params)
        full = self._run_plan(
            plan, row_ids, prompt_text, params, base_generate, via_session,
            logits_processors, stopping_criteria, gen_kwargs,
        )
        final_sequences.append(self._finalize(full[0], original_lengths[i]))

    padded = self.tokenizer.pad(
        {"input_ids": [seq.tolist() for seq in final_sequences]},
        padding=True,
        return_tensors="pt",
    ).to(input_ids.device)
    return padded["input_ids"]
export_constraint(runtime_kwargs=None)

The control's declarative constrained-decoding source, or None.

A control whose per-step masking compiles from a declarative source returns a ConstraintSource; on a backend advertising Capability.GUIDED_DECODING the pipeline renders it onto the engine's native structured-output parameters in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name Type Description Default
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 None

Returns:

Type Description
ConstraintSource | None

The constraint source, or None.
Source code in steerability/algorithms/output_control/base.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def export_constraint(self, runtime_kwargs: dict | None = None) -> ConstraintSource | None:
    """The control's declarative constrained-decoding source, or None.

    A control whose per-step masking compiles from a declarative source returns a
    `ConstraintSource`; on a backend advertising `Capability.GUIDED_DECODING` the pipeline
    renders it onto the engine's native structured-output parameters in place of the
    control's live processor. The default returns None, which keeps the control on the live
    processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The constraint source, or None.
    """
    return None
export_generation_params(runtime_kwargs=None)

The control's sampling-expressible contribution, or None.

A control whose behavior is expressible as normalized generation parameters returns a mapping over a subset of stop_strings, stop_token_ids, max_new_tokens, and min_new_tokens; the pipeline merges it into the call's GenerationParams (stop rules union with the caller's; token bounds only tighten) and does not additionally collect the control's live processors and criteria for that call, so the control executes on every backend through the session's composed stop rules. The default returns None, which keeps the control on the live processor/criteria mechanism.

Parameters:

Name Type Description Default
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 None

Returns:

Type Description
Mapping[str, Any] | None

The parameter contribution, or None.
Source code in steerability/algorithms/output_control/base.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def export_generation_params(self, runtime_kwargs: dict | None = None) -> Mapping[str, Any] | None:
    """The control's sampling-expressible contribution, or None.

    A control whose behavior is expressible as normalized generation parameters returns a
    mapping over a subset of `stop_strings`, `stop_token_ids`, `max_new_tokens`, and
    `min_new_tokens`; the pipeline merges it into the call's `GenerationParams` (stop rules
    union with the caller's; token bounds only tighten) and does not additionally collect
    the control's live processors and criteria for that call, so the control executes on
    every backend through the session's composed stop rules. The default returns None, which
    keeps the control on the live processor/criteria mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The parameter contribution, or None.
    """
    return None
export_processor_spec(runtime_kwargs=None)

The control's engine-hosted processor form, or None.

A control whose per-step logit math is expressible in an engine's served processor vocabulary returns a ProcessorSpec; on a backend advertising Capability.PER_STEP_LOGIT_SPECS with the spec's kind, the pipeline submits it as a ProcessorSpecEntry in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name Type Description Default
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 None

Returns:

Type Description
ProcessorSpec | None

The processor spec, or None.
Source code in steerability/algorithms/output_control/base.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
def export_processor_spec(self, runtime_kwargs: dict | None = None) -> ProcessorSpec | None:
    """The control's engine-hosted processor form, or None.

    A control whose per-step logit math is expressible in an engine's served processor
    vocabulary returns a `ProcessorSpec`; on a backend advertising
    `Capability.PER_STEP_LOGIT_SPECS` with the spec's kind, the pipeline submits it as a
    `ProcessorSpecEntry` in place of the control's live processor. The default returns
    None, which keeps the control on the live processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The processor spec, or None.
    """
    return None
export_state()

Steer-time state to persist when freezing, keyed by logical name.

Values are typed artifacts (SteeringVector, Probe, ProbeSet, a Memory implementation, CheckpointArtifact, LoRAArtifact, torch.Tensor, or a Path to an on-disk product). Must be called after steer(). The default returns an empty mapping (nothing to persist).

Returns:

Type Description
dict[str, Any]

Mapping from logical name to artifact value.
Source code in steerability/algorithms/core/base_control.py 
117
118
119
120
121
122
123
124
125
126
127
128
def export_state(self) -> dict[str, Any]:
    """Steer-time state to persist when freezing, keyed by logical name.

    Values are typed artifacts (`SteeringVector`, `Probe`, `ProbeSet`, a `Memory`
    implementation, `CheckpointArtifact`, `LoRAArtifact`, `torch.Tensor`, or a `Path` to an
    on-disk product). Must be called after `steer()`. The default returns an empty mapping
    (nothing to persist).

    Returns:
        Mapping from logical name to artifact value.
    """
    return {}
fit_identity()

The object whose canonical form defines this control's fit digest, or None.

The digest detects recipe edits that invalidate frozen artifacts, so the returned object should cover exactly the inputs a re-fit would consume (training data, fit specs) and exclude inert application parameters (strengths, multipliers). The default returns None (no fits).

Returns:

Type Description
Any | None

The fit-identity object, or None.
Source code in steerability/algorithms/core/base_control.py 
161
162
163
164
165
166
167
168
169
170
171
172
def fit_identity(self) -> Any | None:
    """The object whose canonical form defines this control's fit digest, or None.

    The digest detects recipe edits that invalidate frozen artifacts, so the returned
    object should cover exactly the inputs a re-fit would consume (training data, fit
    specs) and exclude inert application parameters (strengths, multipliers). The default
    returns None (no fits).

    Returns:
        The fit-identity object, or None.
    """
    return None
frozen_form(state)

The (registry method key, constructor kwargs) of this control's frozen form.

state is the output of export_state(); artifact values inside the returned kwargs are replaced by store references at encode time. When steer_fits() is empty and state is empty, the recipe is the frozen form and the control's own key and recipe args are returned unchanged. A control whose freeze expands into several entries may return a list of (method key, kwargs) pairs instead.

Parameters:

Name Type Description Default
state dict[str, Any]

The output of export_state().

 required

Returns:

Type Description
tuple[str, dict[str, Any]]

The frozen method key and its constructor kwargs.

Raises:

Type Description
NotFreezableError

If the control produces fits or state but declares no frozen form.
Source code in steerability/algorithms/core/base_control.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The `(registry method key, constructor kwargs)` of this control's frozen form.

    `state` is the output of `export_state()`; artifact values inside the returned kwargs
    are replaced by store references at encode time. When `steer_fits()` is empty and
    `state` is empty, the recipe is the frozen form and the control's own key and recipe
    args are returned unchanged. A control whose freeze expands into several entries may
    return a list of `(method key, kwargs)` pairs instead.

    Args:
        state: The output of `export_state()`.

    Returns:
        The frozen method key and its constructor kwargs.

    Raises:
        NotFreezableError: If the control produces fits or state but declares no frozen
            form.
    """
    if not state and not self.steer_fits():
        from steerability.algorithms.core.registry import method_key_for

        args = getattr(self, "args", None)
        kwargs = {field.name: getattr(args, field.name) for field in fields(args) if field.init} \
            if args is not None else {}
        return method_key_for(type(self)), kwargs
    raise NotFreezableError(
        f"{type(self).__name__} produces steer-time state but declares no frozen form; "
        "override frozen_form() (and export_state()) to support freezing."
    )
get_logits_processors(input_ids, runtime_kwargs, **kwargs)

The control's logits processors for the current generation.

Called once per generate() / compute_logprobs() call, after input and state controls have prepared the prompt (mirrors StateControl.get_hooks). **kwargs carries attention_mask and the caller's generation kwargs. Returned objects follow the HF LogitsProcessor convention; in-list order is preserved by the composition.

A processor must behave as a function of (prefix_ids, scores). Internal state is permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch, since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes teacher-forced (subclass common.processors.base.PrefixKeyedProcessor to satisfy this mechanically). Return fresh processor instances from this hook; it is invoked once per call precisely so that per-generation state is isolated.

Parameters:

Name Type Description Default
input_ids Tensor

The steered prompt token ids [batch, seq_len].

 required
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 required

Returns:

Type Description
list

A list of HF LogitsProcessor-style objects.
Source code in steerability/algorithms/output_control/base.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def get_logits_processors(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's logits processors for the current generation.

    Called once per `generate()` / `compute_logprobs()` call, after input and state
    controls have prepared the prompt (mirrors `StateControl.get_hooks`). `**kwargs`
    carries `attention_mask` and the caller's generation kwargs. Returned objects
    follow the HF `LogitsProcessor` convention; in-list order is preserved by the composition.

    A processor must behave as a function of `(prefix_ids, scores)`. Internal state is
    permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch,
    since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes
    teacher-forced (subclass `common.processors.base.PrefixKeyedProcessor` to satisfy this
    mechanically). Return fresh processor instances from this hook; it is invoked once per call
    precisely so that per-generation state is isolated.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `LogitsProcessor`-style objects.
    """
    return []
get_stopping_criteria(input_ids, runtime_kwargs, **kwargs)

The control's stopping criteria.

Not applied during scoring (there is no loop to stop). Same call convention as get_logits_processors.

Parameters:

Name Type Description Default
input_ids Tensor

The steered prompt token ids [batch, seq_len].

 required
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 required

Returns:

Type Description
list

A list of HF StoppingCriteria-style objects.
Source code in steerability/algorithms/output_control/base.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def get_stopping_criteria(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's stopping criteria.

    Not applied during scoring (there is no loop to stop). Same call convention as
    `get_logits_processors`.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `StoppingCriteria`-style objects.
    """
    return []
max_rollouts_per_query()

None: the phase plan is per example, so the base class declares no static bound. Subclasses with a fixed plan (e.g. PhasedDecoding, BudgetForcing) override it.

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
92
93
94
95
def max_rollouts_per_query(self) -> int | None:
    """None: the phase plan is per example, so the base class declares no static bound.
    Subclasses with a fixed plan (e.g. `PhasedDecoding`, `BudgetForcing`) override it."""
    return None
plan(prompt_text, params)

Return the phase plan for one example. Subclasses override.

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
102
103
104
def plan(self, prompt_text: str, params: dict) -> list:
    """Return the phase plan for one example. Subclasses override."""
    raise NotImplementedError
requirements()

Phase splicing is client-side and generated phases run through the session, so no phase requires anything beyond the session contract.

Source code in steerability/algorithms/output_control/common/drivers/phased.py 
 97
 98
 99
100
def requirements(self) -> Requirements:
    """Phase splicing is client-side and generated phases run through the session, so no
    phase requires anything beyond the session contract."""
    return Requirements()
steer(model, tokenizer=None, session=None, **kwargs)

Optional one-time preparation (e.g., load a reward model, fit a probe).

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/output_control/base.py 
207
208
209
210
211
212
def steer(self, model: PreTrainedModel, tokenizer=None, session=None, **kwargs) -> None:
    """Optional one-time preparation (e.g., load a reward model, fit a probe).

    `session` is a `SteeringSession` on the steering backend, provided by the pipeline.
    """
    pass
steer_access()

The model access this instance's steer step requires, on the ModelAccess ladder.

The default is ModelAccess.FACTS (layout and tokenizer only). A control whose steer step generates or scores through the session declares ROLLOUTS, one that captures hidden states declares CAPTURE, and one that touches the model as a live torch.nn.Module declares MODULE. Configuration determines the result. The pipeline hands steer() a session scoped to the declared rung, and the live model only at MODULE. A control may retain the pipeline model beyond steer() only if its generate phase requires Capability.IN_PROCESS_TORCH.

Returns:

Type Description
ModelAccess

The declared access rung.
Source code in steerability/algorithms/core/base_control.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def steer_access(self) -> ModelAccess:
    """The model access this instance's steer step requires, on the `ModelAccess` ladder.

    The default is `ModelAccess.FACTS` (layout and tokenizer only). A control whose steer
    step generates or scores through the session declares `ROLLOUTS`, one that captures
    hidden states declares `CAPTURE`, and one that touches the model as a live
    `torch.nn.Module` declares `MODULE`. Configuration determines the result. The pipeline
    hands `steer()` a session scoped to the declared rung, and the live model only at
    `MODULE`. A control may retain the pipeline model beyond `steer()` only if its
    generate phase requires `Capability.IN_PROCESS_TORCH`.

    Returns:
        The declared access rung.
    """
    return ModelAccess.FACTS
steer_fits()

The fit artifacts this instance's steer step will produce, for the steer plan.

Each entry is (artifact, artifact_class), where artifact is the fit source or recipe class name and artifact_class is "direction" or "calibrated". The default is an empty tuple (no fits).

Returns:

Type Description
tuple[tuple[str, str], ...]

The declared fit artifacts, in declaration order.
Source code in steerability/algorithms/core/base_control.py 
105
106
107
108
109
110
111
112
113
114
115
def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

    Each entry is `(artifact, artifact_class)`, where `artifact` is the fit source or
    recipe class name and `artifact_class` is `"direction"` or `"calibrated"`. The default
    is an empty tuple (no fits).

    Returns:
        The declared fit artifacts, in declaration order.
    """
    return ()

proposer

Rollout proposal for SearchDriver.

The composed logits and stopping stacks are forwarded into every rollout (each only when non-empty, following HFGenerateDriver's pattern), so a step-level control such as RAD steers every proposed continuation.

SegmentProposer

Propose n continuations of up to segment_len tokens from a frontier of sequences.

Parameters:

Name Type Description Default
mode str

"beam" (num_beams=n, num_return_sequences=n, DeAL's rollout) or "sample" (do_sample=True, num_return_sequences=n).

 'beam'
Source code in steerability/algorithms/output_control/common/drivers/proposer.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class SegmentProposer:
    """Propose `n` continuations of up to `segment_len` tokens from a frontier of sequences.

    Args:
        mode: `"beam"` (`num_beams=n, num_return_sequences=n`, DeAL's rollout) or `"sample"`
            (`do_sample=True, num_return_sequences=n`).
    """

    def __init__(self, mode: str = "beam"):
        if mode not in ("beam", "sample"):
            raise ValueError(f"mode must be 'beam' or 'sample', got {mode!r}.")
        self.mode = mode

    def propose(
        self,
        frontier_ids: torch.Tensor,
        *,
        n: int,
        segment_len: int,
        processors: LogitsProcessorList,
        criteria: StoppingCriteriaList,
        model: PreTrainedModel,
        base_generate: Callable[..., torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,
        **gen_kwargs,
    ) -> torch.Tensor:
        """Roll out `n` continuations of `segment_len` tokens from `frontier_ids`.

        Args:
            frontier_ids: Current context `[F, T]` to continue from.
            n: Number of continuations to produce.
            segment_len: Max new tokens per rollout.
            processors: The composed `LogitsProcessorList` (applied to every rollout when non-empty).
            criteria: The composed `StoppingCriteriaList` (applied to every rollout when non-empty).
            model: The language model.
            base_generate: Optional override for the generate callable (else `model.generate`).
            attention_mask: Optional attention mask matching `frontier_ids`.

        Returns:
            Full sequences `[F * n, T + segment_len]` (or shorter if a rollout stopped early).
        """
        generate = base_generate if base_generate is not None else model.generate

        extra = stack_generate_kwargs(processors, criteria)
        if attention_mask is not None:
            extra["attention_mask"] = attention_mask

        kwargs = dict(gen_kwargs)
        kwargs.update(
            {
                "max_new_tokens": segment_len,
                "num_return_sequences": n,
            }
        )
        if self.mode == "beam":
            kwargs["num_beams"] = n
        else:
            kwargs["do_sample"] = True

        return generate(input_ids=frontier_ids, **extra, **kwargs)
mode = mode instance-attribute
__init__(mode='beam')
Source code in steerability/algorithms/output_control/common/drivers/proposer.py 
25
26
27
28
def __init__(self, mode: str = "beam"):
    if mode not in ("beam", "sample"):
        raise ValueError(f"mode must be 'beam' or 'sample', got {mode!r}.")
    self.mode = mode
propose(frontier_ids, *, n, segment_len, processors, criteria, model, base_generate=None, attention_mask=None, **gen_kwargs)

Roll out n continuations of segment_len tokens from frontier_ids.

Parameters:

Name Type Description Default
frontier_ids Tensor

Current context [F, T] to continue from.

 required
n int

Number of continuations to produce.

 required
segment_len int

Max new tokens per rollout.

 required
processors LogitsProcessorList

The composed LogitsProcessorList (applied to every rollout when non-empty).

 required
criteria StoppingCriteriaList

The composed StoppingCriteriaList (applied to every rollout when non-empty).

 required
model PreTrainedModel

The language model.

 required
base_generate Callable[..., Tensor] | None

Optional override for the generate callable (else model.generate).

 None
attention_mask Tensor | None

Optional attention mask matching frontier_ids.

 None

Returns:

Type Description
Tensor

Full sequences [F * n, T + segment_len] (or shorter if a rollout stopped early).
Source code in steerability/algorithms/output_control/common/drivers/proposer.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def propose(
    self,
    frontier_ids: torch.Tensor,
    *,
    n: int,
    segment_len: int,
    processors: LogitsProcessorList,
    criteria: StoppingCriteriaList,
    model: PreTrainedModel,
    base_generate: Callable[..., torch.Tensor] | None = None,
    attention_mask: torch.Tensor | None = None,
    **gen_kwargs,
) -> torch.Tensor:
    """Roll out `n` continuations of `segment_len` tokens from `frontier_ids`.

    Args:
        frontier_ids: Current context `[F, T]` to continue from.
        n: Number of continuations to produce.
        segment_len: Max new tokens per rollout.
        processors: The composed `LogitsProcessorList` (applied to every rollout when non-empty).
        criteria: The composed `StoppingCriteriaList` (applied to every rollout when non-empty).
        model: The language model.
        base_generate: Optional override for the generate callable (else `model.generate`).
        attention_mask: Optional attention mask matching `frontier_ids`.

    Returns:
        Full sequences `[F * n, T + segment_len]` (or shorter if a rollout stopped early).
    """
    generate = base_generate if base_generate is not None else model.generate

    extra = stack_generate_kwargs(processors, criteria)
    if attention_mask is not None:
        extra["attention_mask"] = attention_mask

    kwargs = dict(gen_kwargs)
    kwargs.update(
        {
            "max_new_tokens": segment_len,
            "num_return_sequences": n,
        }
    )
    if self.mode == "beam":
        kwargs["num_beams"] = n
    else:
        kwargs["do_sample"] = True

    return generate(input_ids=frontier_ids, **extra, **kwargs)

search

Segment-search decoding driver: propose continuations, score them, keep the top k, iterate.

Every rollout receives the composed stacks, so step-level controls steer every lookahead of every segment-search driver.

SearchDriver

Bases: DecodingDriver

Segment-search decoding driver: propose, sequence-score, keep top-k, iterate.

Supports batch size 1 only (raises otherwise). decode() pops max_new_tokens as the global budget, builds a SegmentProposer with the received stacks, and runs the loop. A kept beam is finished when its continuation, with trailing pad tokens stripped, ends in a token of the eos set (the tokenizer's eos plus any ids on the model's generation config) or when its stripped length reaches the global budget. Finished beams leave the frontier, and a beam cut by a caller-supplied stopping criterion classifies as unfinished. The reward_params runtime kwarg is row-scoped, so its per-row form is one mapping, merged into the scorer's params on every scoring call, and a batched delivery is a one-element sequence holding that mapping (the driver handles one prompt per call).

Can be constructed directly (its positional constructor below) or as a preset: a subclass with an Args dataclass maps its mirrored args onto these fields in _configure() (see DeAL), so it never bypasses this __init__.

Parameters:

Name Type Description Default
scorer SequenceScorer

A SequenceScorer (prompt, continuations, params) -> list[float].

 required
segment_len int

Max new tokens per rollout.

 required
num_candidates int

Number of continuations proposed per iteration.

 required
keep_k int

Number of beams retained each iteration.

 required
max_iterations int

Maximum search iterations.

 required
propose_mode str

"beam" or "sample".

 'beam'
Source code in steerability/algorithms/output_control/common/drivers/search.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
class SearchDriver(DecodingDriver):
    """Segment-search decoding driver: propose, sequence-score, keep top-k, iterate.

    Supports batch size 1 only (raises otherwise). `decode()` pops `max_new_tokens` as the global
    budget, builds a `SegmentProposer` with the received stacks, and runs the loop. A kept beam is
    finished when its continuation, with trailing pad tokens stripped, ends in a token of the eos
    set (the tokenizer's eos plus any ids on the model's generation config) or when its stripped
    length reaches the global budget. Finished beams leave the frontier, and a beam cut by a
    caller-supplied stopping criterion classifies as unfinished. The `reward_params` runtime kwarg
    is row-scoped, so its per-row form is one mapping, merged into the scorer's params on every
    scoring call, and a batched delivery is a one-element sequence holding that mapping (the driver
    handles one prompt per call).

    Can be constructed directly (its positional constructor below) or as a preset: a subclass with an
    `Args` dataclass maps its mirrored args onto these fields in `_configure()` (see DeAL), so it never
    bypasses this `__init__`.

    Args:
        scorer: A `SequenceScorer` `(prompt, continuations, params) -> list[float]`.
        segment_len: Max new tokens per rollout.
        num_candidates: Number of continuations proposed per iteration.
        keep_k: Number of beams retained each iteration.
        max_iterations: Maximum search iterations.
        propose_mode: `"beam"` or `"sample"`.
    """

    RUNTIME_KWARGS_SCHEMA = [
        {
            "name": "reward_params",
            "type": "dict",
            "scope": "row",
            "help": (
                "Entries merged into the scorer's params mapping on every scoring call of this "
                "generation. The per-row form is one mapping; the driver handles one prompt per "
                "call, so a batched delivery is a one-element sequence holding that mapping. A "
                "per-sample mapping (for example a reference answer under 'reference') reaches "
                "the scorer's row through SampleSequenceScorer."
            ),
        },
    ]

    def __init__(
        self,
        scorer: SequenceScorer,
        segment_len: int,
        num_candidates: int,
        keep_k: int,
        max_iterations: int,
        propose_mode: str = "beam",
    ):
        self.scorer = scorer
        self.segment_len = segment_len
        self.num_candidates = num_candidates
        self.keep_k = keep_k
        self.max_iterations = max_iterations
        self.propose_mode = propose_mode
        self.tokenizer = None  # injected by the pipeline

    def max_rollouts_per_query(self) -> int:
        """`num_candidates * (1 + (max_iterations - 1) * keep_k)`.

        The first iteration proposes `num_candidates` continuations from the single-row prompt;
        each later iteration proposes `num_candidates` from each of up to `keep_k` retained
        beams.
        """
        return self.num_candidates * (1 + (self.max_iterations - 1) * self.keep_k)

    def requirements(self) -> Requirements:
        """Rollouts run through the session, so sampled proposals require nothing beyond the
        session contract; beam proposals require `Capability.BEAM_PROPOSALS`."""
        if getattr(self, "propose_mode", "sample") == "beam":
            return Requirements(generate=needs(
                Capability.BEAM_PROPOSALS,
                hint="use propose_mode='sample' or run this pipeline on the huggingface backend",
            ))
        return Requirements()

    def decode(self, input_ids, attention_mask, model: PreTrainedModel | None, logits_processors,
               stopping_criteria, runtime_kwargs, session=None, **gen_kwargs) -> torch.Tensor:
        if input_ids.dim() != 2 or input_ids.size(0) != 1:
            raise NotImplementedError("SearchDriver handles one prompt at a time (batch size 1).")
        if self.tokenizer is None:
            raise RuntimeError("SearchDriver requires a tokenizer; steer() must run first.")

        runtime_kwargs = runtime_kwargs or {}
        base_generate = resolve_generate_callable(model, runtime_kwargs, session=session)

        prompt_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
        input_length = input_ids.size(1)
        global_budget = gen_kwargs.pop("max_new_tokens", None)

        # segment_len=None means "one full-budget segment" (best-of-N as a config); each rollout then
        # spans the whole generation budget in a single iteration
        segment_len = self.segment_len if self.segment_len is not None else global_budget
        if segment_len is None:
            raise ValueError(
                "SearchDriver requires a segment length: set `segment_len` or pass `max_new_tokens` "
                "as the generation budget (segment_len=None uses the budget as the segment length)."
            )

        reward_params = {
            **_resolve_reward_params(runtime_kwargs),
            "segment_len": segment_len,
            "num_candidates": self.num_candidates,
            "keep_k": self.keep_k,
            "max_iterations": self.max_iterations,
        }
        tokenizer_eos = getattr(self.tokenizer, "eos_token_id", None)
        eos_ids = {tokenizer_eos} if tokenizer_eos is not None else set()
        if model is not None:
            configured = getattr(model.generation_config, "eos_token_id", None)
            eos_ids.update([configured] if isinstance(configured, int) else (configured or []))

        proposer = SegmentProposer(mode=self.propose_mode)
        frontier = Frontier(
            keep_k=self.keep_k,
            eos_token_id=sorted(eos_ids) or None,
            input_length=input_length,
            max_new_tokens=global_budget,
            pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
        )

        current_ids = input_ids
        kept = None
        for _ in range(self.max_iterations):
            # safe to deepcopy: the composed stacks travel as explicit decode() parameters, never inside gen_kwargs
            rollout_kwargs = copy.deepcopy(gen_kwargs)
            frontier_mask = infer_attention_mask_from_ids(current_ids, self.tokenizer.pad_token_id)
            beams = proposer.propose(
                current_ids,
                n=self.num_candidates,
                segment_len=segment_len,
                processors=logits_processors,
                criteria=stopping_criteria,
                model=model,
                base_generate=base_generate,
                attention_mask=frontier_mask,
                **rollout_kwargs,
            )
            continuations = self.tokenizer.decode(
                beams[:, input_length:], skip_special_tokens=True
            )
            scores = self.scorer(prompt_text, continuations, reward_params)
            if len(scores) != beams.size(0):
                raise RuntimeError(f"Scorer returned {len(scores)} scores for {beams.size(0)} beams.")

            step = frontier.keep(beams, scores)
            kept = step
            if all(step.finished_flags):
                break

            unfinished = [i for i, f in enumerate(step.finished_flags) if not f]
            if not unfinished:
                break
            current_ids = step.kept_ids[unfinished]

        final = frontier.best_ids if frontier.best_ids is not None else kept.kept_ids[0]
        return final.unsqueeze(0) if final.dim() == 1 else final
RUNTIME_KWARGS_SCHEMA = [{'name': 'reward_params', 'type': 'dict', 'scope': 'row', 'help': "Entries merged into the scorer's params mapping on every scoring call of this generation. The per-row form is one mapping; the driver handles one prompt per call, so a batched delivery is a one-element sequence holding that mapping. A per-sample mapping (for example a reference answer under 'reference') reaches the scorer's row through SampleSequenceScorer."}] class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
enabled = True class-attribute instance-attribute
include_in_scoring = True class-attribute instance-attribute
keep_k = keep_k instance-attribute
max_iterations = max_iterations instance-attribute
num_candidates = num_candidates instance-attribute
propose_mode = propose_mode instance-attribute
same_model_forwards = False class-attribute instance-attribute
scorer = scorer instance-attribute
segment_len = segment_len instance-attribute
supports_batching = False class-attribute instance-attribute
tokenizer = None instance-attribute
__init__(scorer, segment_len, num_candidates, keep_k, max_iterations, propose_mode='beam')
Source code in steerability/algorithms/output_control/common/drivers/search.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def __init__(
    self,
    scorer: SequenceScorer,
    segment_len: int,
    num_candidates: int,
    keep_k: int,
    max_iterations: int,
    propose_mode: str = "beam",
):
    self.scorer = scorer
    self.segment_len = segment_len
    self.num_candidates = num_candidates
    self.keep_k = keep_k
    self.max_iterations = max_iterations
    self.propose_mode = propose_mode
    self.tokenizer = None  # injected by the pipeline
_configure()

Post-construction hook, called after Args fields are mirrored onto the instance.

Driver presets override this to map their mirrored args onto the fields their generic base reads, so subclasses never bypass a parent __init__. Default no-op.

Source code in steerability/algorithms/core/base_control.py 
67
68
69
70
71
72
73
def _configure(self) -> None:
    """Post-construction hook, called after `Args` fields are mirrored onto the instance.

    Driver presets override this to map their mirrored args onto the fields their generic base
    reads, so subclasses never bypass a parent `__init__`. Default no-op.
    """
    pass
cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py 
195
196
197
198
199
200
201
def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass
clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name Type Description Default
seed int | None

Optional seed forwarded to the clone's reseed().

 None

Returns:

Type Description
BaseControl

The clone.
Source code in steerability/algorithms/core/base_control.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone
decode(input_ids, attention_mask, model, logits_processors, stopping_criteria, runtime_kwargs, session=None, **gen_kwargs)

Run the decoding procedure; return full sequence ids (prompt + continuation).

Source code in steerability/algorithms/output_control/common/drivers/search.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def decode(self, input_ids, attention_mask, model: PreTrainedModel | None, logits_processors,
           stopping_criteria, runtime_kwargs, session=None, **gen_kwargs) -> torch.Tensor:
    if input_ids.dim() != 2 or input_ids.size(0) != 1:
        raise NotImplementedError("SearchDriver handles one prompt at a time (batch size 1).")
    if self.tokenizer is None:
        raise RuntimeError("SearchDriver requires a tokenizer; steer() must run first.")

    runtime_kwargs = runtime_kwargs or {}
    base_generate = resolve_generate_callable(model, runtime_kwargs, session=session)

    prompt_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
    input_length = input_ids.size(1)
    global_budget = gen_kwargs.pop("max_new_tokens", None)

    # segment_len=None means "one full-budget segment" (best-of-N as a config); each rollout then
    # spans the whole generation budget in a single iteration
    segment_len = self.segment_len if self.segment_len is not None else global_budget
    if segment_len is None:
        raise ValueError(
            "SearchDriver requires a segment length: set `segment_len` or pass `max_new_tokens` "
            "as the generation budget (segment_len=None uses the budget as the segment length)."
        )

    reward_params = {
        **_resolve_reward_params(runtime_kwargs),
        "segment_len": segment_len,
        "num_candidates": self.num_candidates,
        "keep_k": self.keep_k,
        "max_iterations": self.max_iterations,
    }
    tokenizer_eos = getattr(self.tokenizer, "eos_token_id", None)
    eos_ids = {tokenizer_eos} if tokenizer_eos is not None else set()
    if model is not None:
        configured = getattr(model.generation_config, "eos_token_id", None)
        eos_ids.update([configured] if isinstance(configured, int) else (configured or []))

    proposer = SegmentProposer(mode=self.propose_mode)
    frontier = Frontier(
        keep_k=self.keep_k,
        eos_token_id=sorted(eos_ids) or None,
        input_length=input_length,
        max_new_tokens=global_budget,
        pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
    )

    current_ids = input_ids
    kept = None
    for _ in range(self.max_iterations):
        # safe to deepcopy: the composed stacks travel as explicit decode() parameters, never inside gen_kwargs
        rollout_kwargs = copy.deepcopy(gen_kwargs)
        frontier_mask = infer_attention_mask_from_ids(current_ids, self.tokenizer.pad_token_id)
        beams = proposer.propose(
            current_ids,
            n=self.num_candidates,
            segment_len=segment_len,
            processors=logits_processors,
            criteria=stopping_criteria,
            model=model,
            base_generate=base_generate,
            attention_mask=frontier_mask,
            **rollout_kwargs,
        )
        continuations = self.tokenizer.decode(
            beams[:, input_length:], skip_special_tokens=True
        )
        scores = self.scorer(prompt_text, continuations, reward_params)
        if len(scores) != beams.size(0):
            raise RuntimeError(f"Scorer returned {len(scores)} scores for {beams.size(0)} beams.")

        step = frontier.keep(beams, scores)
        kept = step
        if all(step.finished_flags):
            break

        unfinished = [i for i, f in enumerate(step.finished_flags) if not f]
        if not unfinished:
            break
        current_ids = step.kept_ids[unfinished]

    final = frontier.best_ids if frontier.best_ids is not None else kept.kept_ids[0]
    return final.unsqueeze(0) if final.dim() == 1 else final
export_constraint(runtime_kwargs=None)

The control's declarative constrained-decoding source, or None.

A control whose per-step masking compiles from a declarative source returns a ConstraintSource; on a backend advertising Capability.GUIDED_DECODING the pipeline renders it onto the engine's native structured-output parameters in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name Type Description Default
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 None

Returns:

Type Description
ConstraintSource | None

The constraint source, or None.
Source code in steerability/algorithms/output_control/base.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def export_constraint(self, runtime_kwargs: dict | None = None) -> ConstraintSource | None:
    """The control's declarative constrained-decoding source, or None.

    A control whose per-step masking compiles from a declarative source returns a
    `ConstraintSource`; on a backend advertising `Capability.GUIDED_DECODING` the pipeline
    renders it onto the engine's native structured-output parameters in place of the
    control's live processor. The default returns None, which keeps the control on the live
    processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The constraint source, or None.
    """
    return None
export_generation_params(runtime_kwargs=None)

The control's sampling-expressible contribution, or None.

A control whose behavior is expressible as normalized generation parameters returns a mapping over a subset of stop_strings, stop_token_ids, max_new_tokens, and min_new_tokens; the pipeline merges it into the call's GenerationParams (stop rules union with the caller's; token bounds only tighten) and does not additionally collect the control's live processors and criteria for that call, so the control executes on every backend through the session's composed stop rules. The default returns None, which keeps the control on the live processor/criteria mechanism.

Parameters:

Name Type Description Default
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 None

Returns:

Type Description
Mapping[str, Any] | None

The parameter contribution, or None.
Source code in steerability/algorithms/output_control/base.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def export_generation_params(self, runtime_kwargs: dict | None = None) -> Mapping[str, Any] | None:
    """The control's sampling-expressible contribution, or None.

    A control whose behavior is expressible as normalized generation parameters returns a
    mapping over a subset of `stop_strings`, `stop_token_ids`, `max_new_tokens`, and
    `min_new_tokens`; the pipeline merges it into the call's `GenerationParams` (stop rules
    union with the caller's; token bounds only tighten) and does not additionally collect
    the control's live processors and criteria for that call, so the control executes on
    every backend through the session's composed stop rules. The default returns None, which
    keeps the control on the live processor/criteria mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The parameter contribution, or None.
    """
    return None
export_processor_spec(runtime_kwargs=None)

The control's engine-hosted processor form, or None.

A control whose per-step logit math is expressible in an engine's served processor vocabulary returns a ProcessorSpec; on a backend advertising Capability.PER_STEP_LOGIT_SPECS with the spec's kind, the pipeline submits it as a ProcessorSpecEntry in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name Type Description Default
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 None

Returns:

Type Description
ProcessorSpec | None

The processor spec, or None.
Source code in steerability/algorithms/output_control/base.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
def export_processor_spec(self, runtime_kwargs: dict | None = None) -> ProcessorSpec | None:
    """The control's engine-hosted processor form, or None.

    A control whose per-step logit math is expressible in an engine's served processor
    vocabulary returns a `ProcessorSpec`; on a backend advertising
    `Capability.PER_STEP_LOGIT_SPECS` with the spec's kind, the pipeline submits it as a
    `ProcessorSpecEntry` in place of the control's live processor. The default returns
    None, which keeps the control on the live processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The processor spec, or None.
    """
    return None
export_state()

Steer-time state to persist when freezing, keyed by logical name.

Values are typed artifacts (SteeringVector, Probe, ProbeSet, a Memory implementation, CheckpointArtifact, LoRAArtifact, torch.Tensor, or a Path to an on-disk product). Must be called after steer(). The default returns an empty mapping (nothing to persist).

Returns:

Type Description
dict[str, Any]

Mapping from logical name to artifact value.
Source code in steerability/algorithms/core/base_control.py 
117
118
119
120
121
122
123
124
125
126
127
128
def export_state(self) -> dict[str, Any]:
    """Steer-time state to persist when freezing, keyed by logical name.

    Values are typed artifacts (`SteeringVector`, `Probe`, `ProbeSet`, a `Memory`
    implementation, `CheckpointArtifact`, `LoRAArtifact`, `torch.Tensor`, or a `Path` to an
    on-disk product). Must be called after `steer()`. The default returns an empty mapping
    (nothing to persist).

    Returns:
        Mapping from logical name to artifact value.
    """
    return {}
fit_identity()

The object whose canonical form defines this control's fit digest, or None.

The digest detects recipe edits that invalidate frozen artifacts, so the returned object should cover exactly the inputs a re-fit would consume (training data, fit specs) and exclude inert application parameters (strengths, multipliers). The default returns None (no fits).

Returns:

Type Description
Any | None

The fit-identity object, or None.
Source code in steerability/algorithms/core/base_control.py 
161
162
163
164
165
166
167
168
169
170
171
172
def fit_identity(self) -> Any | None:
    """The object whose canonical form defines this control's fit digest, or None.

    The digest detects recipe edits that invalidate frozen artifacts, so the returned
    object should cover exactly the inputs a re-fit would consume (training data, fit
    specs) and exclude inert application parameters (strengths, multipliers). The default
    returns None (no fits).

    Returns:
        The fit-identity object, or None.
    """
    return None
frozen_form(state)

The (registry method key, constructor kwargs) of this control's frozen form.

state is the output of export_state(); artifact values inside the returned kwargs are replaced by store references at encode time. When steer_fits() is empty and state is empty, the recipe is the frozen form and the control's own key and recipe args are returned unchanged. A control whose freeze expands into several entries may return a list of (method key, kwargs) pairs instead.

Parameters:

Name Type Description Default
state dict[str, Any]

The output of export_state().

 required

Returns:

Type Description
tuple[str, dict[str, Any]]

The frozen method key and its constructor kwargs.

Raises:

Type Description
NotFreezableError

If the control produces fits or state but declares no frozen form.
Source code in steerability/algorithms/core/base_control.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The `(registry method key, constructor kwargs)` of this control's frozen form.

    `state` is the output of `export_state()`; artifact values inside the returned kwargs
    are replaced by store references at encode time. When `steer_fits()` is empty and
    `state` is empty, the recipe is the frozen form and the control's own key and recipe
    args are returned unchanged. A control whose freeze expands into several entries may
    return a list of `(method key, kwargs)` pairs instead.

    Args:
        state: The output of `export_state()`.

    Returns:
        The frozen method key and its constructor kwargs.

    Raises:
        NotFreezableError: If the control produces fits or state but declares no frozen
            form.
    """
    if not state and not self.steer_fits():
        from steerability.algorithms.core.registry import method_key_for

        args = getattr(self, "args", None)
        kwargs = {field.name: getattr(args, field.name) for field in fields(args) if field.init} \
            if args is not None else {}
        return method_key_for(type(self)), kwargs
    raise NotFreezableError(
        f"{type(self).__name__} produces steer-time state but declares no frozen form; "
        "override frozen_form() (and export_state()) to support freezing."
    )
get_logits_processors(input_ids, runtime_kwargs, **kwargs)

The control's logits processors for the current generation.

Called once per generate() / compute_logprobs() call, after input and state controls have prepared the prompt (mirrors StateControl.get_hooks). **kwargs carries attention_mask and the caller's generation kwargs. Returned objects follow the HF LogitsProcessor convention; in-list order is preserved by the composition.

A processor must behave as a function of (prefix_ids, scores). Internal state is permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch, since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes teacher-forced (subclass common.processors.base.PrefixKeyedProcessor to satisfy this mechanically). Return fresh processor instances from this hook; it is invoked once per call precisely so that per-generation state is isolated.

Parameters:

Name Type Description Default
input_ids Tensor

The steered prompt token ids [batch, seq_len].

 required
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 required

Returns:

Type Description
list

A list of HF LogitsProcessor-style objects.
Source code in steerability/algorithms/output_control/base.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def get_logits_processors(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's logits processors for the current generation.

    Called once per `generate()` / `compute_logprobs()` call, after input and state
    controls have prepared the prompt (mirrors `StateControl.get_hooks`). `**kwargs`
    carries `attention_mask` and the caller's generation kwargs. Returned objects
    follow the HF `LogitsProcessor` convention; in-list order is preserved by the composition.

    A processor must behave as a function of `(prefix_ids, scores)`. Internal state is
    permitted only as memoization keyed on the prefix and must re-derive on a prefix mismatch,
    since drivers may restart, rewind, or reorder sequences, and scoring replays prefixes
    teacher-forced (subclass `common.processors.base.PrefixKeyedProcessor` to satisfy this
    mechanically). Return fresh processor instances from this hook; it is invoked once per call
    precisely so that per-generation state is isolated.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `LogitsProcessor`-style objects.
    """
    return []
get_stopping_criteria(input_ids, runtime_kwargs, **kwargs)

The control's stopping criteria.

Not applied during scoring (there is no loop to stop). Same call convention as get_logits_processors.

Parameters:

Name Type Description Default
input_ids Tensor

The steered prompt token ids [batch, seq_len].

 required
runtime_kwargs dict | None

Per-call parameters supplied to generate().

 required

Returns:

Type Description
list

A list of HF StoppingCriteria-style objects.
Source code in steerability/algorithms/output_control/base.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def get_stopping_criteria(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's stopping criteria.

    Not applied during scoring (there is no loop to stop). Same call convention as
    `get_logits_processors`.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `StoppingCriteria`-style objects.
    """
    return []
max_rollouts_per_query()

num_candidates * (1 + (max_iterations - 1) * keep_k).

The first iteration proposes num_candidates continuations from the single-row prompt; each later iteration proposes num_candidates from each of up to keep_k retained beams.

Source code in steerability/algorithms/output_control/common/drivers/search.py 
119
120
121
122
123
124
125
126
def max_rollouts_per_query(self) -> int:
    """`num_candidates * (1 + (max_iterations - 1) * keep_k)`.

    The first iteration proposes `num_candidates` continuations from the single-row prompt;
    each later iteration proposes `num_candidates` from each of up to `keep_k` retained
    beams.
    """
    return self.num_candidates * (1 + (self.max_iterations - 1) * self.keep_k)
requirements()

Rollouts run through the session, so sampled proposals require nothing beyond the session contract; beam proposals require Capability.BEAM_PROPOSALS.

Source code in steerability/algorithms/output_control/common/drivers/search.py 
128
129
130
131
132
133
134
135
136
def requirements(self) -> Requirements:
    """Rollouts run through the session, so sampled proposals require nothing beyond the
    session contract; beam proposals require `Capability.BEAM_PROPOSALS`."""
    if getattr(self, "propose_mode", "sample") == "beam":
        return Requirements(generate=needs(
            Capability.BEAM_PROPOSALS,
            hint="use propose_mode='sample' or run this pipeline on the huggingface backend",
        ))
    return Requirements()
steer(model, tokenizer=None, session=None, **kwargs)

Optional one-time preparation (e.g., load a reward model, fit a probe).

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/output_control/base.py 
207
208
209
210
211
212
def steer(self, model: PreTrainedModel, tokenizer=None, session=None, **kwargs) -> None:
    """Optional one-time preparation (e.g., load a reward model, fit a probe).

    `session` is a `SteeringSession` on the steering backend, provided by the pipeline.
    """
    pass
steer_access()

The model access this instance's steer step requires, on the ModelAccess ladder.

The default is ModelAccess.FACTS (layout and tokenizer only). A control whose steer step generates or scores through the session declares ROLLOUTS, one that captures hidden states declares CAPTURE, and one that touches the model as a live torch.nn.Module declares MODULE. Configuration determines the result. The pipeline hands steer() a session scoped to the declared rung, and the live model only at MODULE. A control may retain the pipeline model beyond steer() only if its generate phase requires Capability.IN_PROCESS_TORCH.

Returns:

Type Description
ModelAccess

The declared access rung.
Source code in steerability/algorithms/core/base_control.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def steer_access(self) -> ModelAccess:
    """The model access this instance's steer step requires, on the `ModelAccess` ladder.

    The default is `ModelAccess.FACTS` (layout and tokenizer only). A control whose steer
    step generates or scores through the session declares `ROLLOUTS`, one that captures
    hidden states declares `CAPTURE`, and one that touches the model as a live
    `torch.nn.Module` declares `MODULE`. Configuration determines the result. The pipeline
    hands `steer()` a session scoped to the declared rung, and the live model only at
    `MODULE`. A control may retain the pipeline model beyond `steer()` only if its
    generate phase requires `Capability.IN_PROCESS_TORCH`.

    Returns:
        The declared access rung.
    """
    return ModelAccess.FACTS
steer_fits()

The fit artifacts this instance's steer step will produce, for the steer plan.

Each entry is (artifact, artifact_class), where artifact is the fit source or recipe class name and artifact_class is "direction" or "calibrated". The default is an empty tuple (no fits).

Returns:

Type Description
tuple[tuple[str, str], ...]

The declared fit artifacts, in declaration order.
Source code in steerability/algorithms/core/base_control.py 
105
106
107
108
109
110
111
112
113
114
115
def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

    Each entry is `(artifact, artifact_class)`, where `artifact` is the fit source or
    recipe class name and `artifact_class` is `"direction"` or `"calibrated"`. The default
    is an empty tuple (no fits).

    Returns:
        The declared fit artifacts, in declaration order.
    """
    return ()
_resolve_reward_params(runtime_kwargs)

The reward_params mapping for this call, from either delivery form.

A mapping is one row's value (a direct call); a sequence is the row-aligned form (a batched caller or the evaluation collator) and must hold exactly one element, since the driver handles one prompt per call. A missing key or a None value gives an empty mapping.

Parameters:

Name Type Description Default
runtime_kwargs Mapping[str, Any]

The call's runtime kwargs.

 required

Returns:

Type Description
dict[str, Any]

A new mapping of the row's reward params, empty when the key is absent or None.

Raises:

Type Description
ValueError

If a sequence does not hold exactly one element.
TypeError

If the value, or the sequence's element, is not a mapping.
Source code in steerability/algorithms/output_control/common/drivers/search.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def _resolve_reward_params(runtime_kwargs: Mapping[str, Any]) -> dict[str, Any]:
    """The `reward_params` mapping for this call, from either delivery form.

    A mapping is one row's value (a direct call); a sequence is the row-aligned form (a batched
    caller or the evaluation collator) and must hold exactly one element, since the driver handles
    one prompt per call. A missing key or a None value gives an empty mapping.

    Args:
        runtime_kwargs: The call's runtime kwargs.

    Returns:
        A new mapping of the row's reward params, empty when the key is absent or None.

    Raises:
        ValueError: If a sequence does not hold exactly one element.
        TypeError: If the value, or the sequence's element, is not a mapping.
    """
    value = runtime_kwargs.get("reward_params")
    if value is None:
        return {}
    if isinstance(value, Mapping):
        return dict(value)
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        if len(value) != 1:
            raise ValueError(
                f"reward_params is row-scoped and the driver handles one prompt per call; "
                f"got a sequence of length {len(value)}."
            )
        row = value[0]
        if row is None:
            return {}
        if not isinstance(row, Mapping):
            raise TypeError(f"reward_params rows must be mappings; got {type(row).__name__}.")
        return dict(row)
    raise TypeError(
        f"reward_params must be a mapping or a one-element sequence of mappings; got {type(value).__name__}."
    )

granite_heads

Sequence-classification heads for the Granite and GraniteMoeHybrid architectures.

transformers ships no AutoModelForSequenceClassification head for the Granite families, so a Granite causal-LM checkpoint cannot be loaded as a scalar-head reward model out of the box. These two classes supply the head by mixing GenericForSequenceClassification (the pooled-last-token classifier the Llama and Qwen heads use) with each family's PreTrainedModel. Callers select the loading class with sequence_classifier_class, which reads the checkpoint config and returns the toolkit head for a Granite family and AutoModelForSequenceClassification otherwise, so a head shipped by a future transformers version wins.

GRANITE_SEQUENCE_CLASSIFIERS = {GraniteConfig: GraniteForSequenceClassification, GraniteMoeHybridConfig: GraniteMoeHybridForSequenceClassification} module-attribute

_HUB_KWARGS = ('cache_dir', 'force_download', 'local_files_only', 'proxies', 'revision', 'subfolder', 'token', 'trust_remote_code') module-attribute

GraniteForSequenceClassification

Bases: GenericForSequenceClassification, GranitePreTrainedModel

Sequence-classification head for the granite architecture.

Source code in steerability/algorithms/output_control/common/granite_heads.py 
26
27
class GraniteForSequenceClassification(GenericForSequenceClassification, GranitePreTrainedModel):
    """Sequence-classification head for the `granite` architecture."""

GraniteMoeHybridForSequenceClassification

Bases: GenericForSequenceClassification, GraniteMoeHybridPreTrainedModel

Sequence-classification head for the granitemoehybrid architecture.

Source code in steerability/algorithms/output_control/common/granite_heads.py 
30
31
class GraniteMoeHybridForSequenceClassification(GenericForSequenceClassification, GraniteMoeHybridPreTrainedModel):
    """Sequence-classification head for the `granitemoehybrid` architecture."""

sequence_classifier_class(model_id, **from_pretrained_kwargs)

The class that loads model_id as a sequence classifier.

Reads the checkpoint config (the hub-related entries of from_pretrained_kwargs are forwarded) and returns AutoModelForSequenceClassification when its mapping covers the config class, the toolkit head from GRANITE_SEQUENCE_CLASSIFIERS when the config is a Granite family the mapping does not cover, and AutoModelForSequenceClassification otherwise (its from_pretrained then raises the standard unrecognized-config error). The resolution is explicit rather than through AutoModelForSequenceClassification.register, which transformers 5.13 and later ignore for native config classes.

Parameters:

Name Type Description Default
model_id str

HF hub id or local path.

 required
**from_pretrained_kwargs

The kwargs the caller will pass to from_pretrained.

 {}

Returns:

Type Description
type

The class to call from_pretrained on.
Source code in steerability/algorithms/output_control/common/granite_heads.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def sequence_classifier_class(model_id: str, **from_pretrained_kwargs) -> type:
    """The class that loads `model_id` as a sequence classifier.

    Reads the checkpoint config (the hub-related entries of `from_pretrained_kwargs` are forwarded) and
    returns `AutoModelForSequenceClassification` when its mapping covers the config class, the toolkit
    head from `GRANITE_SEQUENCE_CLASSIFIERS` when the config is a Granite family the mapping does not
    cover, and `AutoModelForSequenceClassification` otherwise (its `from_pretrained` then raises the
    standard unrecognized-config error). The resolution is explicit rather than through
    `AutoModelForSequenceClassification.register`, which transformers 5.13 and later ignore for native
    config classes.

    Args:
        model_id: HF hub id or local path.
        **from_pretrained_kwargs: The kwargs the caller will pass to `from_pretrained`.

    Returns:
        The class to call `from_pretrained` on.
    """
    hub_kwargs = {name: from_pretrained_kwargs[name] for name in _HUB_KWARGS if name in from_pretrained_kwargs}
    config = AutoConfig.from_pretrained(model_id, **hub_kwargs)
    if type(config) in MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING:
        return AutoModelForSequenceClassification
    return GRANITE_SEQUENCE_CLASSIFIERS.get(type(config), AutoModelForSequenceClassification)

kv_cache

KV-cache utilities for same-model candidate evaluation.

repeat_cache / select_cache handle the cache-format compatibility that same-model value functions need: repeat a prefix cache across K candidates, then select the chosen candidate's slice. extends_prefix / full_prefix_mask are the pure tensor helpers an incremental prefix cache needs: whether a new prefix extends the cached one, and the full attention mask spanning a prefix.

Mutation contract: both functions may mutate the input cache in-place on some paths (batch_repeat_interleave, in-place per-layer tensor reassignment) and return a fresh cache on others (the raw-tuple path). Treat the input cache as consumed after the call and use only the returned handle. The exception is repeat_cache(..., preserve_input=True), which never mutates the input (it builds a fresh DynamicCache from the input's per-layer tensors); use it when the input cache must survive the call (e.g. an incremental prefix cache repeated across candidates).

Cache layouts handled, in dispatch order: Cache objects exposing per-layer tensors through cache.layers[i].keys / cache.layers[i].values (the transformers v5 layout), objects exposing key_cache / value_cache tensor lists, and raw per-layer (key, value) tuples.

_layer_tensors(cache)

The per-layer (keys, values) tensors of a layer-based Cache, or None.

Reads the transformers v5 layout (cache.layers[i].keys / cache.layers[i].values). Returns None when cache exposes no layers sequence of that shape.

Source code in steerability/algorithms/output_control/common/kv_cache.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
def _layer_tensors(cache: Cache | tuple) -> list[tuple[torch.Tensor | None, torch.Tensor | None]] | None:
    """The per-layer `(keys, values)` tensors of a layer-based `Cache`, or None.

    Reads the transformers v5 layout (`cache.layers[i].keys` / `cache.layers[i].values`). Returns
    None when `cache` exposes no `layers` sequence of that shape.
    """
    layers = getattr(cache, "layers", None)
    if layers is None:
        return None
    try:
        return [(layer.keys, layer.values) for layer in layers]
    except AttributeError:
        return None

extends_prefix(cached_ids, ids)

Whether ids extends cached_ids row-for-row (a shared, no-shorter prefix).

Parameters:

Name Type Description Default
cached_ids Tensor | None

The previously cached prefix ids [B, T_c], or None when nothing is cached.

 required
ids Tensor

The candidate new prefix ids [B, T].

 required

Returns:

Type Description
bool

True when cached_ids is not None and ids is at least as long and matches cached_ids
bool

over its first T_c positions; False otherwise (including any shorter or divergent prefix,
bool

which must trigger a cache rebuild).
Source code in steerability/algorithms/output_control/common/kv_cache.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def extends_prefix(cached_ids: torch.Tensor | None, ids: torch.Tensor) -> bool:
    """Whether `ids` extends `cached_ids` row-for-row (a shared, no-shorter prefix).

    Args:
        cached_ids: The previously cached prefix ids `[B, T_c]`, or None when nothing is cached.
        ids: The candidate new prefix ids `[B, T]`.

    Returns:
        True when `cached_ids` is not None and `ids` is at least as long and matches `cached_ids`
        over its first `T_c` positions; False otherwise (including any shorter or divergent prefix,
        which must trigger a cache rebuild).
    """
    if cached_ids is None or ids.size(1) < cached_ids.size(1):
        return False
    return bool(torch.equal(ids[:, : cached_ids.size(1)], cached_ids.to(ids.device)))

full_prefix_mask(prefix_ids, attention_mask)

The provided mask right-extended with ones to the prefix length.

The mask must span the full prefix; generated tokens are always real, so right-extension with ones is exact.

Parameters:

Name Type Description Default
prefix_ids Tensor

The prefix ids [B, T].

 required
attention_mask Tensor | None

The prefix mask [B, T'] with T' <= T, or None for an all-ones mask.

 required

Returns:

Type Description
Tensor

An attention mask [B, T].

Raises:

Type Description
ValueError

If attention_mask is longer than prefix_ids.
Source code in steerability/algorithms/output_control/common/kv_cache.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def full_prefix_mask(prefix_ids: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
    """The provided mask right-extended with ones to the prefix length.

    The mask must span the full prefix; generated tokens are always real, so right-extension with
    ones is exact.

    Args:
        prefix_ids: The prefix ids `[B, T]`.
        attention_mask: The prefix mask `[B, T']` with `T' <= T`, or None for an all-ones mask.

    Returns:
        An attention mask `[B, T]`.

    Raises:
        ValueError: If `attention_mask` is longer than `prefix_ids`.
    """
    if attention_mask is None:
        return torch.ones_like(prefix_ids)
    pad = prefix_ids.size(1) - attention_mask.size(1)
    if pad < 0:
        raise ValueError("attention_mask is longer than prefix_ids.")
    if pad == 0:
        return attention_mask
    ones = torch.ones(attention_mask.size(0), pad, device=attention_mask.device, dtype=attention_mask.dtype)
    return torch.cat([attention_mask, ones], dim=1)

repeat_cache(cache, n, *, preserve_input=False)

Repeat every cache entry n times along the batch dimension.

Parameters:

Name Type Description Default
cache Cache | tuple

A KV cache (a layer-based Cache such as DynamicCache, a key/value-list style object, or a raw per-layer tuple).

 required
n int

Number of repeats per entry.

 required
preserve_input bool

When True, never mutate cache; take the copying paths (a fresh DynamicCache built from the input's per-layer tensors, the tensor-building path for raw tuples) and raise TypeError for cache types that offer only in-place repetition. The returned cache shares no batch-repeated storage with the input.

 False

Returns:

Type Description
Cache | tuple

The repeated cache (may alias the input unless preserve_input=True; see the module
Cache | tuple

mutation contract).

Raises:

Type Description
TypeError

If the cache type is unsupported, or if preserve_input=True and the cache offers only in-place repetition.
Source code in steerability/algorithms/output_control/common/kv_cache.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def repeat_cache(cache: Cache | tuple, n: int, *, preserve_input: bool = False) -> Cache | tuple:
    """Repeat every cache entry `n` times along the batch dimension.

    Args:
        cache: A KV cache (a layer-based `Cache` such as `DynamicCache`, a key/value-list style
            object, or a raw per-layer tuple).
        n: Number of repeats per entry.
        preserve_input: When True, never mutate `cache`; take the copying paths (a fresh
            `DynamicCache` built from the input's per-layer tensors, the tensor-building path for
            raw tuples) and raise `TypeError` for cache types that offer only in-place repetition.
            The returned cache shares no batch-repeated storage with the input.

    Returns:
        The repeated cache (may alias the input unless `preserve_input=True`; see the module
        mutation contract).

    Raises:
        TypeError: If the cache type is unsupported, or if `preserve_input=True` and the cache offers
            only in-place repetition.
    """
    if hasattr(cache, "batch_repeat_interleave") and not preserve_input:
        cache.batch_repeat_interleave(n)
        return cache

    layer_tensors = _layer_tensors(cache)
    if layer_tensors is not None:
        fresh = DynamicCache()
        for layer_idx, (keys, values) in enumerate(layer_tensors):
            if keys is None or values is None:
                raise TypeError(
                    f"{type(cache).__name__} has an unmaterialized layer {layer_idx}; cannot repeat."
                )
            fresh.update(
                keys.repeat_interleave(n, dim=0),
                values.repeat_interleave(n, dim=0),
                layer_idx,
            )
        return fresh

    if hasattr(cache, "key_cache") and hasattr(cache, "value_cache"):
        if preserve_input:
            raise TypeError(
                f"{type(cache).__name__} offers only in-place repetition; cannot preserve_input."
            )
        for i in range(len(cache.key_cache)):
            cache.key_cache[i] = cache.key_cache[i].repeat_interleave(n, dim=0)
            cache.value_cache[i] = cache.value_cache[i].repeat_interleave(n, dim=0)
        return cache

    if isinstance(cache, tuple):
        return tuple(
            tuple(t.repeat_interleave(n, dim=0) for t in layer)
            for layer in cache
        )

    raise TypeError(f"Unsupported cache type: {type(cache).__name__}")

select_cache(cache, idx)

Select cache entries along the batch dimension by index.

Parameters:

Name Type Description Default
cache Cache | tuple

A KV cache (a layer-based Cache such as DynamicCache, a key/value-list style object, or a raw per-layer tuple).

 required
idx Tensor

1-D index tensor of rows to keep.

 required

Returns:

Type Description
Cache | tuple

The selected cache (may alias the input; see the module mutation contract).

Raises:

Type Description
ValueError

If idx is not 1-D.
TypeError

If the cache type is unsupported.
Source code in steerability/algorithms/output_control/common/kv_cache.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
def select_cache(cache: Cache | tuple, idx: torch.Tensor) -> Cache | tuple:
    """Select cache entries along the batch dimension by index.

    Args:
        cache: A KV cache (a layer-based `Cache` such as `DynamicCache`, a key/value-list style
            object, or a raw per-layer tuple).
        idx: 1-D index tensor of rows to keep.

    Returns:
        The selected cache (may alias the input; see the module mutation contract).

    Raises:
        ValueError: If `idx` is not 1-D.
        TypeError: If the cache type is unsupported.
    """
    if not torch.is_tensor(idx):
        idx = torch.as_tensor(idx)
    if idx.dtype != torch.long:
        idx = idx.long()
    if idx.dim() != 1:
        raise ValueError(f"idx must be 1D, got shape {tuple(idx.shape)}")

    if hasattr(cache, "batch_select"):
        cache.batch_select(idx)
        return cache

    if hasattr(cache, "batch_gather"):
        cache.batch_gather(idx)
        return cache

    layers = getattr(cache, "layers", None)
    if layers is not None and _layer_tensors(cache) is not None:
        for layer in layers:
            if layer.keys is not None:
                layer.keys = layer.keys.index_select(dim=0, index=idx.to(layer.keys.device))
            if layer.values is not None:
                layer.values = layer.values.index_select(dim=0, index=idx.to(layer.values.device))
        return cache

    if hasattr(cache, "key_cache") and hasattr(cache, "value_cache"):
        for i in range(len(cache.key_cache)):
            if cache.key_cache[i] is not None:
                cache.key_cache[i] = cache.key_cache[i].index_select(dim=0, index=idx.to(cache.key_cache[i].device))
            if cache.value_cache[i] is not None:
                cache.value_cache[i] = cache.value_cache[i].index_select(
                    dim=0, index=idx.to(cache.value_cache[i].device)
                )
        return cache

    if isinstance(cache, tuple):
        return tuple(
            tuple(t.index_select(dim=0, index=idx.to(t.device)) for t in layer)
            for layer in cache
        )

    raise TypeError(f"Unsupported cache type: {type(cache).__name__}")

loading

The one auxiliary sequence-classifier loader.

Loads and configures an AutoModelForSequenceClassification reward model / attribute classifier the way the output category needs it: eval mode, pad-token fallback to EOS, right padding, and a clamp of the sentinel model_max_length some tokenizers carry onto a plain max_length attribute that RewardModelValue reads. Used by the reward_model / classifier value loaders, the reward_model scorer loader, and RAD's HF-classifier path. Granite checkpoints, for which transformers ships no head, resolve through the toolkit heads selected by sequence_classifier_class.

logger = logging.getLogger(__name__) module-attribute

load_sequence_classifier(model_id, *, device, hf_model_kwargs=None, max_length_clamp=512)

Load an HF sequence classifier and its tokenizer, configured for candidate/continuation scoring.

Parameters:

Name Type Description Default
model_id str

HF hub id or local path for an AutoModelForSequenceClassification.

 required
device str | device

Device to move the loaded model onto.

 required
hf_model_kwargs dict | None

Extra kwargs forwarded to from_pretrained.

 None
max_length_clamp int

Value to use for tokenizer.max_length when the tokenizer's model_max_length is unset or an absurd sentinel (> 100_000).

 512

Returns:

Type Description
tuple

A tuple (model, tokenizer); model is in eval mode on device, and tokenizer has
tuple

pad_token set (falling back to eos_token), padding_side="right", and a max_length
tuple

attribute.
Source code in steerability/algorithms/output_control/common/loading.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def load_sequence_classifier(
    model_id: str,
    *,
    device: str | torch.device,
    hf_model_kwargs: dict | None = None,
    max_length_clamp: int = 512,
) -> tuple:
    """Load an HF sequence classifier and its tokenizer, configured for candidate/continuation scoring.

    Args:
        model_id: HF hub id or local path for an `AutoModelForSequenceClassification`.
        device: Device to move the loaded model onto.
        hf_model_kwargs: Extra kwargs forwarded to `from_pretrained`.
        max_length_clamp: Value to use for `tokenizer.max_length` when the tokenizer's
            `model_max_length` is unset or an absurd sentinel (`> 100_000`).

    Returns:
        A tuple `(model, tokenizer)`; `model` is in eval mode on `device`, and `tokenizer` has
        `pad_token` set (falling back to `eos_token`), `padding_side="right"`, and a `max_length`
        attribute.
    """
    kwargs = dict(hf_model_kwargs or {})
    model = sequence_classifier_class(model_id, **kwargs).from_pretrained(model_id, **kwargs)
    model = model.to(device)
    model.eval()

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    max_len = getattr(tokenizer, "model_max_length", None)
    if max_len is None or max_len > 100_000:
        max_len = max_length_clamp
    tokenizer.max_length = max_len

    logger.debug("Loaded sequence classifier from %s", model_id)
    return model, tokenizer

logit_sources

Full-vocabulary logit sources for contrastive mixing.

A logit source produces next-token log-probs [B, V] for the current prefix from some distribution, either an auxiliary LM (expert / anti-expert / amateur) or the pipeline's own model on a transformed prompt (classifier-free guidance's unconditional branch). Sources own their model handles and any prefix-keyed KV caching.

AuxModelSource

Bases: BaseLogitSource

A separate causal LM (expert, anti-expert, amateur, class-conditional guide).

Loads its own model in prepare() and maintains its own tokenizer. Optionally applies a prompt transform (e.g. prepend a control code) before its tokenization. The auxiliary vocabulary must match the base model's; pass shared_vocab=True to assert this, since the contrast family in the literature assumes shared vocabularies.

Parameters:

Name Type Description Default
name_or_path str

HF hub id or local path for the auxiliary LM.

 required
base_tokenizer PreTrainedTokenizerBase | None

The base model's tokenizer (candidate ids are the base vocabulary).

 None
prompt_transform Callable[[str], str] | None

Optional str -> str transform applied before aux tokenization.

 None
shared_vocab bool

When True (default), validate that the auxiliary vocabulary matches the base vocabulary and raise on mismatch. The mixture assumes shared vocabularies; setting False only skips the check.

 True
hf_model_kwargs dict | None

Extra kwargs for from_pretrained.

 None
Source code in steerability/algorithms/output_control/common/logit_sources.py 
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
class AuxModelSource(BaseLogitSource):
    """A separate causal LM (expert, anti-expert, amateur, class-conditional guide).

    Loads its own model in `prepare()` and maintains its own tokenizer. Optionally applies a prompt
    transform (e.g. prepend a control code) before its tokenization. The auxiliary vocabulary must
    match the base model's; pass `shared_vocab=True` to assert this, since the contrast family in the
    literature assumes shared vocabularies.

    Args:
        name_or_path: HF hub id or local path for the auxiliary LM.
        base_tokenizer: The base model's tokenizer (candidate ids are the base vocabulary).
        prompt_transform: Optional `str -> str` transform applied before aux tokenization.
        shared_vocab: When True (default), validate that the auxiliary vocabulary matches the base
            vocabulary and raise on mismatch. The mixture assumes shared vocabularies; setting False
            only skips the check.
        hf_model_kwargs: Extra kwargs for `from_pretrained`.
    """

    def __init__(
        self,
        name_or_path: str,
        base_tokenizer: PreTrainedTokenizerBase | None = None,
        prompt_transform: Callable[[str], str] | None = None,
        shared_vocab: bool = True,
        hf_model_kwargs: dict | None = None,
    ):
        self.name_or_path = name_or_path
        self.base_tokenizer = base_tokenizer
        self.prompt_transform = prompt_transform
        self.shared_vocab = shared_vocab
        self.hf_model_kwargs = hf_model_kwargs or {}
        self.model = None
        self.tokenizer = None
        self._device = None
        self._base_pad_id = None

    def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
        if self.base_tokenizer is None:
            self.base_tokenizer = tokenizer
        self.model = AutoModelForCausalLM.from_pretrained(self.name_or_path, **self.hf_model_kwargs)
        if model is not None:
            self.model = self.model.to(next(model.parameters()).device)
        self.model.eval()
        self._device = next(self.model.parameters()).device
        self.tokenizer = AutoTokenizer.from_pretrained(self.name_or_path)
        self.tokenizer.padding_side = "left"
        self._base_pad_id = getattr(self.base_tokenizer, "pad_token_id", None)

        if self.shared_vocab:
            aux_vocab = getattr(self.model.config, "vocab_size", None)
            # the base vocabulary is the model's output dimension (what the mixed logits must align
            # with); fall back to the tokenizer length only when no base model is available
            if model is not None:
                base_vocab = getattr(model.config, "vocab_size", None)
            elif self.base_tokenizer is not None:
                base_vocab = len(self.base_tokenizer)
            else:
                base_vocab = None
            if aux_vocab is not None and base_vocab is not None and aux_vocab != base_vocab:
                raise ValueError(
                    f"AuxModelSource vocab mismatch: aux={aux_vocab} vs base={base_vocab}. The "
                    "contrast family requires shared vocabularies; provide a token-id mapping or a "
                    "shared-vocab auxiliary model."
                )

    def set_model(self, model, tokenizer) -> None:
        """Bind an already-loaded auxiliary model/tokenizer directly (bypasses `prepare` loading)."""
        self.model = model
        self.tokenizer = tokenizer or self.tokenizer
        self.model.eval()
        self._device = next(self.model.parameters()).device
        if self.tokenizer is not None:
            self.tokenizer.padding_side = "left"
        self._base_pad_id = getattr(self.base_tokenizer, "pad_token_id", None)

    @torch.no_grad()
    def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
        """Next-token log-probs `[B, V]`. Mask-correct for left-padded batches; ragged-batch
        equivalence is exact only up to the aux model's own padding-invariance."""
        if self.model is None:
            raise RuntimeError("AuxModelSource is not prepared; call prepare() from steer().")
        if self.prompt_transform is not None and self.base_tokenizer is not None:
            texts = self.base_tokenizer.decode(prefix_ids, skip_special_tokens=True)
            texts = [self.prompt_transform(t) for t in texts]
            enc = self.tokenizer(texts, return_tensors="pt", padding=True).to(self._device)
            ids = enc["input_ids"]
            mask = enc["attention_mask"]
        else:
            ids = prefix_ids.to(self._device)
            mask = infer_attention_mask_from_ids(ids, self._base_pad_id).to(self._device)
        logits = self.model(input_ids=ids, attention_mask=mask).logits[:, -1, :]
        return torch.log_softmax(logits, dim=-1)

    def cleanup(self) -> None:
        """Drop the auxiliary model / tokenizer references and reclaim memory."""
        self.model = None
        self.tokenizer = None
        self._device = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
_base_pad_id = None instance-attribute
_device = None instance-attribute
base_tokenizer = base_tokenizer instance-attribute
hf_model_kwargs = hf_model_kwargs or {} instance-attribute
model = None instance-attribute
name_or_path = name_or_path instance-attribute
prompt_transform = prompt_transform instance-attribute
same_model_forwards = False class-attribute instance-attribute
shared_vocab = shared_vocab instance-attribute
tokenizer = None instance-attribute
__init__(name_or_path, base_tokenizer=None, prompt_transform=None, shared_vocab=True, hf_model_kwargs=None)
Source code in steerability/algorithms/output_control/common/logit_sources.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def __init__(
    self,
    name_or_path: str,
    base_tokenizer: PreTrainedTokenizerBase | None = None,
    prompt_transform: Callable[[str], str] | None = None,
    shared_vocab: bool = True,
    hf_model_kwargs: dict | None = None,
):
    self.name_or_path = name_or_path
    self.base_tokenizer = base_tokenizer
    self.prompt_transform = prompt_transform
    self.shared_vocab = shared_vocab
    self.hf_model_kwargs = hf_model_kwargs or {}
    self.model = None
    self.tokenizer = None
    self._device = None
    self._base_pad_id = None
cleanup()

Drop the auxiliary model / tokenizer references and reclaim memory.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
143
144
145
146
147
148
149
150
def cleanup(self) -> None:
    """Drop the auxiliary model / tokenizer references and reclaim memory."""
    self.model = None
    self.tokenizer = None
    self._device = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
logprobs(prefix_ids)

Next-token log-probs [B, V]. Mask-correct for left-padded batches; ragged-batch equivalence is exact only up to the aux model's own padding-invariance.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
@torch.no_grad()
def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
    """Next-token log-probs `[B, V]`. Mask-correct for left-padded batches; ragged-batch
    equivalence is exact only up to the aux model's own padding-invariance."""
    if self.model is None:
        raise RuntimeError("AuxModelSource is not prepared; call prepare() from steer().")
    if self.prompt_transform is not None and self.base_tokenizer is not None:
        texts = self.base_tokenizer.decode(prefix_ids, skip_special_tokens=True)
        texts = [self.prompt_transform(t) for t in texts]
        enc = self.tokenizer(texts, return_tensors="pt", padding=True).to(self._device)
        ids = enc["input_ids"]
        mask = enc["attention_mask"]
    else:
        ids = prefix_ids.to(self._device)
        mask = infer_attention_mask_from_ids(ids, self._base_pad_id).to(self._device)
    logits = self.model(input_ids=ids, attention_mask=mask).logits[:, -1, :]
    return torch.log_softmax(logits, dim=-1)
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's steer(). Default no-op (mirrors the input-control selector prepare() convention).

Source code in steerability/algorithms/output_control/common/logit_sources.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    if self.base_tokenizer is None:
        self.base_tokenizer = tokenizer
    self.model = AutoModelForCausalLM.from_pretrained(self.name_or_path, **self.hf_model_kwargs)
    if model is not None:
        self.model = self.model.to(next(model.parameters()).device)
    self.model.eval()
    self._device = next(self.model.parameters()).device
    self.tokenizer = AutoTokenizer.from_pretrained(self.name_or_path)
    self.tokenizer.padding_side = "left"
    self._base_pad_id = getattr(self.base_tokenizer, "pad_token_id", None)

    if self.shared_vocab:
        aux_vocab = getattr(self.model.config, "vocab_size", None)
        # the base vocabulary is the model's output dimension (what the mixed logits must align
        # with); fall back to the tokenizer length only when no base model is available
        if model is not None:
            base_vocab = getattr(model.config, "vocab_size", None)
        elif self.base_tokenizer is not None:
            base_vocab = len(self.base_tokenizer)
        else:
            base_vocab = None
        if aux_vocab is not None and base_vocab is not None and aux_vocab != base_vocab:
            raise ValueError(
                f"AuxModelSource vocab mismatch: aux={aux_vocab} vs base={base_vocab}. The "
                "contrast family requires shared vocabularies; provide a token-id mapping or a "
                "shared-vocab auxiliary model."
            )
set_model(model, tokenizer)

Bind an already-loaded auxiliary model/tokenizer directly (bypasses prepare loading).

Source code in steerability/algorithms/output_control/common/logit_sources.py 
115
116
117
118
119
120
121
122
123
def set_model(self, model, tokenizer) -> None:
    """Bind an already-loaded auxiliary model/tokenizer directly (bypasses `prepare` loading)."""
    self.model = model
    self.tokenizer = tokenizer or self.tokenizer
    self.model.eval()
    self._device = next(self.model.parameters()).device
    if self.tokenizer is not None:
        self.tokenizer.padding_side = "left"
    self._base_pad_id = getattr(self.base_tokenizer, "pad_token_id", None)

BaseLogitSource

Bases: ABC

Produce next-token log-probs [B, V] for the current prefix from another distribution.

Class attributes

same_model_forwards: Whether this source issues additional forward passes through the pipeline's own model during decoding. Such passes must be wrapped in auxiliary_pass() (see steerability.algorithms.core.utils.auxiliary_pass), which keeps them out of state-control condition scoring, gate updates, and fallback position counting. Defaults to False; the flag is declarative metadata and is not read by the pipeline.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class BaseLogitSource(ABC):
    """Produce next-token log-probs `[B, V]` for the current prefix from another distribution.

    Class attributes:
        same_model_forwards: Whether this source issues additional forward passes through the
            pipeline's own model during decoding. Such passes must be wrapped in
            `auxiliary_pass()` (see `steerability.algorithms.core.utils.auxiliary_pass`), which
            keeps them out of state-control condition scoring, gate updates, and fallback
            position counting. Defaults to False; the flag is declarative metadata and is not
            read by the pipeline.
    """

    same_model_forwards: bool = False

    @abstractmethod
    def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
        """Return next-token log-probabilities `[B, V]` for `prefix_ids`."""
        ...

    def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
        """Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's
        `steer()`. Default no-op (mirrors the input-control selector `prepare()` convention)."""
        pass

    def cleanup(self) -> None:
        """Release any resources allocated during `prepare()`."""
        pass
same_model_forwards = False class-attribute instance-attribute
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/logit_sources.py 
45
46
47
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
logprobs(prefix_ids) abstractmethod

Return next-token log-probabilities [B, V] for prefix_ids.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
35
36
37
38
@abstractmethod
def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
    """Return next-token log-probabilities `[B, V]` for `prefix_ids`."""
    ...
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's steer(). Default no-op (mirrors the input-control selector prepare() convention).

Source code in steerability/algorithms/output_control/common/logit_sources.py 
40
41
42
43
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's
    `steer()`. Default no-op (mirrors the input-control selector `prepare()` convention)."""
    pass

CallableSource

Bases: BaseLogitSource

Adapt a (prefix_ids) -> Tensor[B, V] log-prob callable into a BaseLogitSource.

The escape hatch of the source family: any function returning next-token log-probs for the current prefix becomes a source without a subclass. The wrapped function must return a [B, V] tensor aligned with the base model's vocabulary.

same_model_forwards keeps the BaseLogitSource default of False, since it cannot be inferred from a bare callable. A callable that forwards the pipeline's model should be wrapped in a component that declares same_model_forwards.

Parameters:

Name Type Description Default
fn Callable[[Tensor], Tensor]

A callable mapping prefix_ids [B, T] to next-token log-probs [B, V].

 required
Source code in steerability/algorithms/output_control/common/logit_sources.py 
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class CallableSource(BaseLogitSource):
    """Adapt a `(prefix_ids) -> Tensor[B, V]` log-prob callable into a `BaseLogitSource`.

    The escape hatch of the source family: any function returning next-token log-probs for the
    current prefix becomes a source without a subclass. The wrapped function must return a
    `[B, V]` tensor aligned with the base model's vocabulary.

    `same_model_forwards` keeps the `BaseLogitSource` default of False, since it cannot be inferred
    from a bare callable. A callable that forwards the pipeline's model should be wrapped in a
    component that declares `same_model_forwards`.

    Args:
        fn: A callable mapping `prefix_ids [B, T]` to next-token log-probs `[B, V]`.
    """

    def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
        if not callable(fn):
            raise TypeError("CallableSource requires a callable (prefix_ids) -> Tensor[B, V].")
        self.fn = fn

    def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
        return torch.as_tensor(self.fn(prefix_ids))
fn = fn instance-attribute
same_model_forwards = False class-attribute instance-attribute
__init__(fn)
Source code in steerability/algorithms/output_control/common/logit_sources.py 
222
223
224
225
def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
    if not callable(fn):
        raise TypeError("CallableSource requires a callable (prefix_ids) -> Tensor[B, V].")
    self.fn = fn
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/logit_sources.py 
45
46
47
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
logprobs(prefix_ids)

Return next-token log-probabilities [B, V] for prefix_ids.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
227
228
def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
    return torch.as_tensor(self.fn(prefix_ids))
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's steer(). Default no-op (mirrors the input-control selector prepare() convention).

Source code in steerability/algorithms/output_control/common/logit_sources.py 
40
41
42
43
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's
    `steer()`. Default no-op (mirrors the input-control selector `prepare()` convention)."""
    pass

PromptVariantSource

Bases: BaseLogitSource

The pipeline's own model on a transformed prompt (CFG's unconditional branch, CAD's context-free branch, self-debias's biased branch).

This source forwards the pipeline's own model, so it declares same_model_forwards=True and marks its forwards via auxiliary_pass(aligned=False). The variant prompt is a different sequence whose positions do not map onto the generation, so state-control transforms, condition scoring, and gate updates all skip these passes, and the variant branch reflects the unsteered model.

Parameters:

Name Type Description Default
prompt_transform Callable[[str], str]

str -> str transform producing the variant prompt.

 required
base_tokenizer PreTrainedTokenizerBase | None

The base model's tokenizer.

 None
Source code in steerability/algorithms/output_control/common/logit_sources.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class PromptVariantSource(BaseLogitSource):
    """The pipeline's own model on a transformed prompt (CFG's unconditional branch, CAD's
    context-free branch, self-debias's biased branch).

    This source forwards the pipeline's own model, so it declares `same_model_forwards=True` and
    marks its forwards via `auxiliary_pass(aligned=False)`. The variant prompt is a different
    sequence whose positions do not map onto the generation, so state-control transforms, condition
    scoring, and gate updates all skip these passes, and the variant branch reflects the unsteered
    model.

    Args:
        prompt_transform: `str -> str` transform producing the variant prompt.
        base_tokenizer: The base model's tokenizer.
    """

    same_model_forwards: bool = True

    def __init__(self, prompt_transform: Callable[[str], str], base_tokenizer: PreTrainedTokenizerBase | None = None):
        self.prompt_transform = prompt_transform
        self.base_tokenizer = base_tokenizer
        self.model = None
        self._device = None

    def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
        self.model = model
        if self.base_tokenizer is None:
            self.base_tokenizer = tokenizer
        if self.base_tokenizer is not None:
            self.base_tokenizer.padding_side = "left"
        self._device = next(model.parameters()).device

    @torch.no_grad()
    def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
        """Next-token log-probs `[B, V]`. Mask-correct for left-padded batches; ragged-batch
        equivalence is exact only up to the model's own padding-invariance."""
        if self.model is None:
            raise RuntimeError("PromptVariantSource is not prepared; call prepare() from steer().")
        texts = self.base_tokenizer.decode(prefix_ids, skip_special_tokens=True)
        texts = [self.prompt_transform(t) for t in texts]
        enc = self.base_tokenizer(texts, return_tensors="pt", padding=True).to(self._device)
        with auxiliary_pass(aligned=False):
            logits = self.model(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).logits[:, -1, :]
        return torch.log_softmax(logits, dim=-1)

    def cleanup(self) -> None:
        """Drop the (pipeline-owned) model reference; the pipeline owns the model's lifecycle."""
        self.model = None
        self.base_tokenizer = None
        self._device = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
_device = None instance-attribute
base_tokenizer = base_tokenizer instance-attribute
model = None instance-attribute
prompt_transform = prompt_transform instance-attribute
same_model_forwards = True class-attribute instance-attribute
__init__(prompt_transform, base_tokenizer=None)
Source code in steerability/algorithms/output_control/common/logit_sources.py 
170
171
172
173
174
def __init__(self, prompt_transform: Callable[[str], str], base_tokenizer: PreTrainedTokenizerBase | None = None):
    self.prompt_transform = prompt_transform
    self.base_tokenizer = base_tokenizer
    self.model = None
    self._device = None
cleanup()

Drop the (pipeline-owned) model reference; the pipeline owns the model's lifecycle.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
197
198
199
200
201
202
203
204
def cleanup(self) -> None:
    """Drop the (pipeline-owned) model reference; the pipeline owns the model's lifecycle."""
    self.model = None
    self.base_tokenizer = None
    self._device = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
logprobs(prefix_ids)

Next-token log-probs [B, V]. Mask-correct for left-padded batches; ragged-batch equivalence is exact only up to the model's own padding-invariance.

Source code in steerability/algorithms/output_control/common/logit_sources.py 
184
185
186
187
188
189
190
191
192
193
194
195
@torch.no_grad()
def logprobs(self, prefix_ids: torch.Tensor) -> torch.Tensor:
    """Next-token log-probs `[B, V]`. Mask-correct for left-padded batches; ragged-batch
    equivalence is exact only up to the model's own padding-invariance."""
    if self.model is None:
        raise RuntimeError("PromptVariantSource is not prepared; call prepare() from steer().")
    texts = self.base_tokenizer.decode(prefix_ids, skip_special_tokens=True)
    texts = [self.prompt_transform(t) for t in texts]
    enc = self.base_tokenizer(texts, return_tensors="pt", padding=True).to(self._device)
    with auxiliary_pass(aligned=False):
        logits = self.model(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).logits[:, -1, :]
    return torch.log_softmax(logits, dim=-1)
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup (e.g. load the auxiliary model), invoked from the owning control's steer(). Default no-op (mirrors the input-control selector prepare() convention).

Source code in steerability/algorithms/output_control/common/logit_sources.py 
176
177
178
179
180
181
182
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    self.model = model
    if self.base_tokenizer is None:
        self.base_tokenizer = tokenizer
    if self.base_tokenizer is not None:
        self.base_tokenizer.padding_side = "left"
    self._device = next(model.parameters()).device

processors

Output-control logits processors (composed into the decoding stack).

base

Base class for stateful logits processors.

A processor must behave as a function of (prefix_ids, scores); subclasses key any internal state on the prefix and re-derive it on a prefix mismatch, so drivers may restart, rewind, or reorder sequences and scoring may replay prefixes teacher-forced.

PrefixKeyedProcessor

Bases: LogitsProcessor

A LogitsProcessor whose internal state is memoization keyed on the prefix.

Drivers restart, rewind, and reorder sequences (segment search re-enters from a shorter frontier; beam reordering permutes rows; compute_logprobs replays prefixes teacher-forced). A processor holding hidden per-call state silently corrupts under any of these. This base makes the purity contract of OutputControl.get_logits_processors structural: process() runs only after the base has verified the observed input_ids extend the last-seen prefix row-for-row; otherwise reset_state() re-derives from scratch.

Subclasses implement:

- `reset_state(input_ids)`: rebuild all internal state for this (possibly new) prefix.
- `process(input_ids, scores) -> scores`: the actual logits edit.

Stateless subclasses implement process and inherit the no-op reset_state.

Source code in steerability/algorithms/output_control/common/processors/base.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
class PrefixKeyedProcessor(LogitsProcessor):
    """A `LogitsProcessor` whose internal state is memoization keyed on the prefix.

    Drivers restart, rewind, and reorder sequences (segment search re-enters from a shorter
    frontier; beam reordering permutes rows; `compute_logprobs` replays prefixes teacher-forced). A
    processor holding hidden per-call state silently corrupts under any of these. This base makes
    the purity contract of `OutputControl.get_logits_processors` structural: `process()` runs only
    after the base has verified the observed `input_ids` extend the last-seen prefix row-for-row;
    otherwise `reset_state()` re-derives from scratch.

    Subclasses implement:

        - `reset_state(input_ids)`: rebuild all internal state for this (possibly new) prefix.
        - `process(input_ids, scores) -> scores`: the actual logits edit.

    Stateless subclasses implement `process` and inherit the no-op `reset_state`.
    """

    def __init__(self):
        self._last_ids: torch.Tensor | None = None

    def _extends(self, ids: torch.Tensor) -> bool:
        last = self._last_ids
        if last is None or ids.size(0) != last.size(0) or ids.size(1) < last.size(1):
            return False
        return bool(torch.equal(ids[:, : last.size(1)], last.to(ids.device)))

    def reset_state(self, input_ids: torch.Tensor) -> None:  # default: stateless
        """Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass)."""
        pass

    def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
        """Apply the logits edit. Subclasses must override."""
        raise NotImplementedError

    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
        if not self._extends(input_ids):
            self.reset_state(input_ids)
        out = self.process(input_ids, scores)
        self._last_ids = input_ids.detach()
        return out
_last_ids = None instance-attribute
__call__(input_ids, scores)
Source code in steerability/algorithms/output_control/common/processors/base.py 
48
49
50
51
52
53
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    if not self._extends(input_ids):
        self.reset_state(input_ids)
    out = self.process(input_ids, scores)
    self._last_ids = input_ids.detach()
    return out
__init__()
Source code in steerability/algorithms/output_control/common/processors/base.py 
31
32
def __init__(self):
    self._last_ids: torch.Tensor | None = None
_extends(ids)
Source code in steerability/algorithms/output_control/common/processors/base.py 
34
35
36
37
38
def _extends(self, ids: torch.Tensor) -> bool:
    last = self._last_ids
    if last is None or ids.size(0) != last.size(0) or ids.size(1) < last.size(1):
        return False
    return bool(torch.equal(ids[:, : last.size(1)], last.to(ids.device)))
process(input_ids, scores)

Apply the logits edit. Subclasses must override.

Source code in steerability/algorithms/output_control/common/processors/base.py 
44
45
46
def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    """Apply the logits edit. Subclasses must override."""
    raise NotImplementedError
reset_state(input_ids)

Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass).

Source code in steerability/algorithms/output_control/common/processors/base.py 
40
41
42
def reset_state(self, input_ids: torch.Tensor) -> None:  # default: stateless
    """Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass)."""
    pass

constraint

ConstraintProcessor, the integration point for constrained decoding.

Takes an automaton object exposing reset(prefix_ids) and allowed(prefix_ids) -> LongTensor of permitted token ids, and masks everything else to -inf. Grammar, JSON-schema, or regex backends plug in here by supplying the automaton (outlines-core / xgrammar adapted to the automaton protocol); no FSM engine is built in-repo.

ConstraintAutomaton

Bases: Protocol

Minimal automaton protocol a ConstraintProcessor drives.

Source code in steerability/algorithms/output_control/common/processors/constraint.py 
17
18
19
20
21
22
23
24
25
26
class ConstraintAutomaton(Protocol):
    """Minimal automaton protocol a `ConstraintProcessor` drives."""

    def reset(self, prefix_ids: torch.Tensor) -> None:
        """Reset automaton state for a new/rewound prefix."""
        ...

    def allowed(self, prefix_ids: torch.Tensor) -> torch.Tensor:
        """Return a 1-D `LongTensor` of token ids permitted at the current step."""
        ...
allowed(prefix_ids)

Return a 1-D LongTensor of token ids permitted at the current step.

Source code in steerability/algorithms/output_control/common/processors/constraint.py 
24
25
26
def allowed(self, prefix_ids: torch.Tensor) -> torch.Tensor:
    """Return a 1-D `LongTensor` of token ids permitted at the current step."""
    ...
reset(prefix_ids)

Reset automaton state for a new/rewound prefix.

Source code in steerability/algorithms/output_control/common/processors/constraint.py 
20
21
22
def reset(self, prefix_ids: torch.Tensor) -> None:
    """Reset automaton state for a new/rewound prefix."""
    ...
ConstraintProcessor

Bases: PrefixKeyedProcessor

Mask all logits except the automaton's currently permitted token ids.

Parameters:

Name Type Description Default
automaton ConstraintAutomaton

An object implementing the ConstraintAutomaton protocol.

 required
Source code in steerability/algorithms/output_control/common/processors/constraint.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class ConstraintProcessor(PrefixKeyedProcessor):
    """Mask all logits except the automaton's currently permitted token ids.

    Args:
        automaton: An object implementing the `ConstraintAutomaton` protocol.
    """

    def __init__(self, automaton: ConstraintAutomaton):
        super().__init__()
        self.automaton = automaton

    def reset_state(self, input_ids: torch.Tensor) -> None:
        self.automaton.reset(input_ids)

    def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
        allowed = self.automaton.allowed(input_ids).to(scores.device)
        out = torch.full_like(scores, float("-inf"))
        out[:, allowed] = scores[:, allowed]
        return out
_last_ids = None instance-attribute
automaton = automaton instance-attribute
__call__(input_ids, scores)
Source code in steerability/algorithms/output_control/common/processors/base.py 
48
49
50
51
52
53
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    if not self._extends(input_ids):
        self.reset_state(input_ids)
    out = self.process(input_ids, scores)
    self._last_ids = input_ids.detach()
    return out
__init__(automaton)
Source code in steerability/algorithms/output_control/common/processors/constraint.py 
36
37
38
def __init__(self, automaton: ConstraintAutomaton):
    super().__init__()
    self.automaton = automaton
_extends(ids)
Source code in steerability/algorithms/output_control/common/processors/base.py 
34
35
36
37
38
def _extends(self, ids: torch.Tensor) -> bool:
    last = self._last_ids
    if last is None or ids.size(0) != last.size(0) or ids.size(1) < last.size(1):
        return False
    return bool(torch.equal(ids[:, : last.size(1)], last.to(ids.device)))
process(input_ids, scores)

Apply the logits edit. Subclasses must override.

Source code in steerability/algorithms/output_control/common/processors/constraint.py 
43
44
45
46
47
def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    allowed = self.automaton.allowed(input_ids).to(scores.device)
    out = torch.full_like(scores, float("-inf"))
    out[:, allowed] = scores[:, allowed]
    return out
reset_state(input_ids)

Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass).

Source code in steerability/algorithms/output_control/common/processors/constraint.py 
40
41
def reset_state(self, input_ids: torch.Tensor) -> None:
    self.automaton.reset(input_ids)

contrastive_mixture

Mix full-vocabulary log-prob sources into the base distribution.

scores' = w0 * log p_base + sum_i w_i * log p_source_i, with an optional plausibility mask keep = {t : p_base(t) >= alpha * max_t p_base(t)} applied before mixing (contrastive decoding's alpha-filter; alpha=None disables; masked-out tokens -> -inf). Mixing is in log-prob space; the return is unnormalized scores, which HF's loop accepts.

ContrastiveMixtureProcessor

Bases: PrefixKeyedProcessor

Mix the base log-probs with weighted auxiliary log-prob sources.

Parameters:

Name Type Description Default
sources list[tuple[BaseLogitSource, float]]

A list of (BaseLogitSource, weight) pairs.

 required
base_weight float

Weight w0 on the base log-probs.

 1.0
alpha float | None

Plausibility-mask threshold in [0, 1]; keep tokens with p_base(t) >= alpha * max_t p_base(t). None disables the mask.

 None
Source code in steerability/algorithms/output_control/common/processors/contrastive_mixture.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class ContrastiveMixtureProcessor(PrefixKeyedProcessor):
    """Mix the base log-probs with weighted auxiliary log-prob sources.

    Args:
        sources: A list of `(BaseLogitSource, weight)` pairs.
        base_weight: Weight `w0` on the base log-probs.
        alpha: Plausibility-mask threshold in `[0, 1]`; keep tokens with
            `p_base(t) >= alpha * max_t p_base(t)`. `None` disables the mask.
    """

    def __init__(
        self,
        sources: list[tuple[BaseLogitSource, float]],
        *,
        base_weight: float = 1.0,
        alpha: float | None = None,
    ):
        super().__init__()
        self.sources = sources
        self.base_weight = base_weight
        self.alpha = alpha

    def reset_state(self, input_ids: torch.Tensor) -> None:
        # sources recompute per call; no per-generation state to reset here,
        # but a source that caches can override its own reset via this hook path.
        pass

    def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
        base_logprobs = torch.log_softmax(scores, dim=-1)
        mixed = self.base_weight * base_logprobs
        for source, weight in self.sources:
            mixed = mixed + weight * source.logprobs(input_ids).to(mixed.dtype)

        if self.alpha is not None:
            base_probs = base_logprobs.exp()
            threshold = self.alpha * base_probs.max(dim=-1, keepdim=True).values
            keep = base_probs >= threshold
            mixed = mixed.masked_fill(~keep, float("-inf"))
        return mixed
_last_ids = None instance-attribute
alpha = alpha instance-attribute
base_weight = base_weight instance-attribute
sources = sources instance-attribute
__call__(input_ids, scores)
Source code in steerability/algorithms/output_control/common/processors/base.py 
48
49
50
51
52
53
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    if not self._extends(input_ids):
        self.reset_state(input_ids)
    out = self.process(input_ids, scores)
    self._last_ids = input_ids.detach()
    return out
__init__(sources, *, base_weight=1.0, alpha=None)
Source code in steerability/algorithms/output_control/common/processors/contrastive_mixture.py 
26
27
28
29
30
31
32
33
34
35
36
def __init__(
    self,
    sources: list[tuple[BaseLogitSource, float]],
    *,
    base_weight: float = 1.0,
    alpha: float | None = None,
):
    super().__init__()
    self.sources = sources
    self.base_weight = base_weight
    self.alpha = alpha
_extends(ids)
Source code in steerability/algorithms/output_control/common/processors/base.py 
34
35
36
37
38
def _extends(self, ids: torch.Tensor) -> bool:
    last = self._last_ids
    if last is None or ids.size(0) != last.size(0) or ids.size(1) < last.size(1):
        return False
    return bool(torch.equal(ids[:, : last.size(1)], last.to(ids.device)))
process(input_ids, scores)

Apply the logits edit. Subclasses must override.

Source code in steerability/algorithms/output_control/common/processors/contrastive_mixture.py 
43
44
45
46
47
48
49
50
51
52
53
54
def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    base_logprobs = torch.log_softmax(scores, dim=-1)
    mixed = self.base_weight * base_logprobs
    for source, weight in self.sources:
        mixed = mixed + weight * source.logprobs(input_ids).to(mixed.dtype)

    if self.alpha is not None:
        base_probs = base_logprobs.exp()
        threshold = self.alpha * base_probs.max(dim=-1, keepdim=True).values
        keep = base_probs >= threshold
        mixed = mixed.masked_fill(~keep, float("-inf"))
    return mixed
reset_state(input_ids)

Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass).

Source code in steerability/algorithms/output_control/common/processors/contrastive_mixture.py 
38
39
40
41
def reset_state(self, input_ids: torch.Tensor) -> None:
    # sources recompute per call; no per-generation state to reset here,
    # but a source that caches can override its own reset via this hook path.
    pass

value_guided

Shift candidate-token logits by a per-candidate value (select candidates, score, combine).

One processor covers many methods. RAD = (RewardModelValue, top_k, clamp, mask=True). SASA = (SubspaceMarginValue, surviving, softmax, mask=False). FUDGE = (ClassifierValue, top_k, none, beta=1). ARGS = (RewardModelValue, top_k, none).

LARGE_CANDIDATE_SET_WARN_THRESHOLD = 1024 module-attribute
Normalize = Literal['none', 'minmax', 'softmax', 'clamp'] module-attribute
ValueGuidedProcessor

Bases: PrefixKeyedProcessor

scores[cand] += beta * normalize(value(prefix, cand)); optionally -inf elsewhere.

Parameters:

Name Type Description Default
value BaseCandidateValue

A BaseCandidateValue.

 required
policy str

Candidate policy ("top_k", "top_p", "surviving"; see select_candidates).

 'top_k'
k int | None

Candidate count for top_k.

 20
p float | None

Nucleus threshold for top_p.

 None
beta float

Shift scale.

 1.0
normalize Normalize

"minmax" (in-set; degenerate set -> 0.5), "softmax" (over the set), "clamp" (element-wise to [0, 1], absolute rather than in-set), or "none". Applied per row.

 'none'
invert bool

Post-normalization v <- 1 - v (steer away from the scored attribute).

 False
mask_non_candidates bool

Set non-candidate logits to -inf (RAD semantics). Forced False when policy="surviving" (everything finite is already a candidate).

 True
max_candidates int | None

Optional clamp on the candidate-set size. After the policy selects candidates, keep only the top max_candidates by current score. Applies to every policy (it is a clamp, not a policy); the main use is bounding the unbounded surviving set under a model-forward value. None (default) disables the clamp.

 None
lm_tokenizer PreTrainedTokenizerBase | None

The language-model tokenizer, forwarded into StepContext.

 None
model PreTrainedModel | None

The pipeline's model (for same-model values), forwarded into StepContext.

 None
attention_mask Tensor | None

The prefix attention mask, forwarded into StepContext.

 None
trace list | None

Optional caller-owned list that receives one ValueStepRecord per process call. The record is a write-only sink, so the processor stays a function of (prefix_ids, scores); compute_logprobs replays also append when the owning control's include_in_scoring is True. None (default) records nothing.

 None
Source code in steerability/algorithms/output_control/common/processors/value_guided.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
class ValueGuidedProcessor(PrefixKeyedProcessor):
    """`scores[cand] += beta * normalize(value(prefix, cand))`; optionally `-inf` elsewhere.

    Args:
        value: A `BaseCandidateValue`.
        policy: Candidate policy (`"top_k"`, `"top_p"`, `"surviving"`; see `select_candidates`).
        k: Candidate count for `top_k`.
        p: Nucleus threshold for `top_p`.
        beta: Shift scale.
        normalize: `"minmax"` (in-set; degenerate set -> 0.5), `"softmax"` (over the set), `"clamp"`
            (element-wise to `[0, 1]`, absolute rather than in-set), or `"none"`. Applied per row.
        invert: Post-normalization `v <- 1 - v` (steer away from the scored attribute).
        mask_non_candidates: Set non-candidate logits to `-inf` (RAD semantics). Forced False when
            `policy="surviving"` (everything finite is already a candidate).
        max_candidates: Optional clamp on the candidate-set size. After the policy selects candidates,
            keep only the top `max_candidates` by current score. Applies to every policy (it is a
            clamp, not a policy); the main use is bounding the unbounded `surviving` set under a
            model-forward value. `None` (default) disables the clamp.
        lm_tokenizer: The language-model tokenizer, forwarded into `StepContext`.
        model: The pipeline's model (for same-model values), forwarded into `StepContext`.
        attention_mask: The prefix attention mask, forwarded into `StepContext`.
        trace: Optional caller-owned list that receives one `ValueStepRecord` per `process` call.
            The record is a write-only sink, so the processor stays a function of
            `(prefix_ids, scores)`; `compute_logprobs` replays also append when the owning control's
            `include_in_scoring` is True. `None` (default) records nothing.
    """

    def __init__(
        self,
        value: BaseCandidateValue,
        *,
        policy: str = "top_k",
        k: int | None = 20,
        p: float | None = None,
        beta: float = 1.0,
        normalize: Normalize = "none",
        invert: bool = False,
        mask_non_candidates: bool = True,
        max_candidates: int | None = None,
        lm_tokenizer: PreTrainedTokenizerBase | None = None,
        model: PreTrainedModel | None = None,
        attention_mask: torch.Tensor | None = None,
        trace: list | None = None,
    ):
        super().__init__()
        self.value = value
        self.policy = policy
        self.k = k
        self.p = p
        self.beta = beta
        self.normalize = normalize
        self.invert = invert
        self.mask_non_candidates = mask_non_candidates and policy != "surviving"
        self.max_candidates = max_candidates
        self.lm_tokenizer = lm_tokenizer
        self.model = model
        self.attention_mask = attention_mask
        self.trace = trace
        self._warned_large_set = False

    def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
        cand_ids, cand_scores = select_candidates(scores, self.policy, k=self.k, p=self.p)

        if self.max_candidates is not None and cand_ids.size(1) > self.max_candidates:
            keep = torch.topk(cand_scores, self.max_candidates, dim=-1).indices
            cand_ids = cand_ids.gather(1, keep)

        num_candidates = cand_ids.size(1)
        if (
            not self._warned_large_set
            and num_candidates > LARGE_CANDIDATE_SET_WARN_THRESHOLD
            and getattr(self.value, "scoring_cost", None) == "model_forward"
        ):
            warnings.warn(
                f"ValueGuidedProcessor is evaluating {num_candidates} candidates with a model-forward "
                "value; pass top_k/top_p in gen_kwargs or set max_candidates to bound the per-step cost.",
                UserWarning,
            )
            self._warned_large_set = True

        ctx = StepContext(
            prefix_ids=input_ids,
            candidate_ids=cand_ids,
            lm_tokenizer=self.lm_tokenizer,
            model=self.model,
            attention_mask=self.attention_mask,
        )
        raw = self.value.score(ctx).to(scores.dtype)  # [B, K]
        v = _normalize(raw, self.normalize, self.invert)

        if self.trace is not None:
            self.trace.append(
                ValueStepRecord(
                    prefix_length=int(input_ids.size(1)),
                    candidate_ids=cand_ids.detach().cpu(),
                    candidate_scores=scores.gather(1, cand_ids).detach().cpu(),
                    values=raw.detach().cpu(),
                    normalized=v.detach().cpu(),
                )
            )

        if self.mask_non_candidates:
            out = torch.full_like(scores, float("-inf"))
            out.scatter_(1, cand_ids, scores.gather(1, cand_ids))
        else:
            out = scores.clone()
        out.scatter_add_(1, cand_ids, self.beta * v)
        return out
_last_ids = None instance-attribute
_warned_large_set = False instance-attribute
attention_mask = attention_mask instance-attribute
beta = beta instance-attribute
invert = invert instance-attribute
k = k instance-attribute
lm_tokenizer = lm_tokenizer instance-attribute
mask_non_candidates = mask_non_candidates and policy != 'surviving' instance-attribute
max_candidates = max_candidates instance-attribute
model = model instance-attribute
normalize = normalize instance-attribute
p = p instance-attribute
policy = policy instance-attribute
trace = trace instance-attribute
value = value instance-attribute
__call__(input_ids, scores)
Source code in steerability/algorithms/output_control/common/processors/base.py 
48
49
50
51
52
53
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    if not self._extends(input_ids):
        self.reset_state(input_ids)
    out = self.process(input_ids, scores)
    self._last_ids = input_ids.detach()
    return out
__init__(value, *, policy='top_k', k=20, p=None, beta=1.0, normalize='none', invert=False, mask_non_candidates=True, max_candidates=None, lm_tokenizer=None, model=None, attention_mask=None, trace=None)
Source code in steerability/algorithms/output_control/common/processors/value_guided.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def __init__(
    self,
    value: BaseCandidateValue,
    *,
    policy: str = "top_k",
    k: int | None = 20,
    p: float | None = None,
    beta: float = 1.0,
    normalize: Normalize = "none",
    invert: bool = False,
    mask_non_candidates: bool = True,
    max_candidates: int | None = None,
    lm_tokenizer: PreTrainedTokenizerBase | None = None,
    model: PreTrainedModel | None = None,
    attention_mask: torch.Tensor | None = None,
    trace: list | None = None,
):
    super().__init__()
    self.value = value
    self.policy = policy
    self.k = k
    self.p = p
    self.beta = beta
    self.normalize = normalize
    self.invert = invert
    self.mask_non_candidates = mask_non_candidates and policy != "surviving"
    self.max_candidates = max_candidates
    self.lm_tokenizer = lm_tokenizer
    self.model = model
    self.attention_mask = attention_mask
    self.trace = trace
    self._warned_large_set = False
_extends(ids)
Source code in steerability/algorithms/output_control/common/processors/base.py 
34
35
36
37
38
def _extends(self, ids: torch.Tensor) -> bool:
    last = self._last_ids
    if last is None or ids.size(0) != last.size(0) or ids.size(1) < last.size(1):
        return False
    return bool(torch.equal(ids[:, : last.size(1)], last.to(ids.device)))
process(input_ids, scores)

Apply the logits edit. Subclasses must override.

Source code in steerability/algorithms/output_control/common/processors/value_guided.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def process(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
    cand_ids, cand_scores = select_candidates(scores, self.policy, k=self.k, p=self.p)

    if self.max_candidates is not None and cand_ids.size(1) > self.max_candidates:
        keep = torch.topk(cand_scores, self.max_candidates, dim=-1).indices
        cand_ids = cand_ids.gather(1, keep)

    num_candidates = cand_ids.size(1)
    if (
        not self._warned_large_set
        and num_candidates > LARGE_CANDIDATE_SET_WARN_THRESHOLD
        and getattr(self.value, "scoring_cost", None) == "model_forward"
    ):
        warnings.warn(
            f"ValueGuidedProcessor is evaluating {num_candidates} candidates with a model-forward "
            "value; pass top_k/top_p in gen_kwargs or set max_candidates to bound the per-step cost.",
            UserWarning,
        )
        self._warned_large_set = True

    ctx = StepContext(
        prefix_ids=input_ids,
        candidate_ids=cand_ids,
        lm_tokenizer=self.lm_tokenizer,
        model=self.model,
        attention_mask=self.attention_mask,
    )
    raw = self.value.score(ctx).to(scores.dtype)  # [B, K]
    v = _normalize(raw, self.normalize, self.invert)

    if self.trace is not None:
        self.trace.append(
            ValueStepRecord(
                prefix_length=int(input_ids.size(1)),
                candidate_ids=cand_ids.detach().cpu(),
                candidate_scores=scores.gather(1, cand_ids).detach().cpu(),
                values=raw.detach().cpu(),
                normalized=v.detach().cpu(),
            )
        )

    if self.mask_non_candidates:
        out = torch.full_like(scores, float("-inf"))
        out.scatter_(1, cand_ids, scores.gather(1, cand_ids))
    else:
        out = scores.clone()
    out.scatter_add_(1, cand_ids, self.beta * v)
    return out
reset_state(input_ids)

Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass).

Source code in steerability/algorithms/output_control/common/processors/base.py 
40
41
42
def reset_state(self, input_ids: torch.Tensor) -> None:  # default: stateless
    """Rebuild internal state for a new/rewound prefix. Default no-op (stateless subclass)."""
    pass
ValueStepRecord dataclass

One ValueGuidedProcessor.process step, recorded for the caller-owned trace.

All tensors are detached and on CPU. candidate_scores are the processor's input scores for the candidates before the shift; normalized is the value after normalize and invert, i.e. the quantity the processor multiplies by beta.

Attributes:

Name Type Description
prefix_length int

The prefix length at this step (input_ids.size(1)).
candidate_ids Tensor

Selected candidate token ids [B, K].
candidate_scores Tensor

Input scores of the candidates before the shift [B, K].
values Tensor

Raw per-candidate value output [B, K].
normalized Tensor

Values after normalize and invert [B, K].
Source code in steerability/algorithms/output_control/common/processors/value_guided.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
@dataclass(frozen=True, slots=True)
class ValueStepRecord:
    """One `ValueGuidedProcessor.process` step, recorded for the caller-owned trace.

    All tensors are detached and on CPU. `candidate_scores` are the processor's input scores for
    the candidates before the shift; `normalized` is the value after `normalize` and `invert`, i.e.
    the quantity the processor multiplies by `beta`.

    Attributes:
        prefix_length: The prefix length at this step (`input_ids.size(1)`).
        candidate_ids: Selected candidate token ids `[B, K]`.
        candidate_scores: Input scores of the candidates before the shift `[B, K]`.
        values: Raw per-candidate value output `[B, K]`.
        normalized: Values after `normalize` and `invert` `[B, K]`.
    """

    prefix_length: int
    candidate_ids: torch.Tensor
    candidate_scores: torch.Tensor
    values: torch.Tensor
    normalized: torch.Tensor
candidate_ids instance-attribute
candidate_scores instance-attribute
normalized instance-attribute
prefix_length instance-attribute
values instance-attribute
__init__(prefix_length, candidate_ids, candidate_scores, values, normalized)
_normalize(v, mode, invert)

Normalize per-candidate values row-wise, then optionally invert.

Parameters:

Name Type Description Default
v Tensor

Values [B, K].

 required
mode Normalize

"minmax" (per-row min-max, relative to the set; degenerate row -> 0.5), "softmax" (per-row softmax), "clamp" (element-wise clamp to [0, 1], absolute rather than relative to the set), or "none".

 required
invert bool

When True, v <- 1 - v after normalization (steer away from the scored attribute).

 required

Returns:

Type Description
Tensor

Normalized values [B, K].
Source code in steerability/algorithms/output_control/common/processors/value_guided.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def _normalize(v: torch.Tensor, mode: Normalize, invert: bool) -> torch.Tensor:
    """Normalize per-candidate values row-wise, then optionally invert.

    Args:
        v: Values `[B, K]`.
        mode: `"minmax"` (per-row min-max, relative to the set; degenerate row -> 0.5), `"softmax"`
            (per-row softmax), `"clamp"` (element-wise clamp to `[0, 1]`, absolute rather than
            relative to the set), or `"none"`.
        invert: When True, `v <- 1 - v` after normalization (steer away from the scored attribute).

    Returns:
        Normalized values `[B, K]`.
    """
    if mode == "clamp":
        normalized = v.clamp(0.0, 1.0)
    elif mode == "minmax":
        r_min = v.min(dim=-1, keepdim=True).values
        r_max = v.max(dim=-1, keepdim=True).values
        span = r_max - r_min
        normalized = torch.where(span > 1e-8, (v - r_min) / span.clamp_min(1e-8), 0.5)
    elif mode == "softmax":
        normalized = torch.softmax(v, dim=-1)
    elif mode == "none":
        normalized = v
    else:
        raise ValueError(f"Unknown normalize mode: {mode!r}.")

    if invert:
        normalized = 1 - normalized
    return normalized

resolve

Component-spec resolution for the generic output controls.

The generics expose the common slots (values, sources, scorers) through flat Args, and resolve a spec into a live component at steer() time, when the base model and tokenizer are available. A spec is one of:

- an instance of the target protocol (its `prepare(model=, tokenizer=)` is called, then it is
  returned);
- a callable (wrapped: values -> `CallableValue`, sources -> `CallableSource`, scorers are
  already `SequenceScorer`s and are returned as-is);
- a str (sources only: shorthand for an auxiliary model);
- a dict with a `"kind"` key (dispatched to a loader; see the tables below).

Sweepability rule: ControlSpec sweeps flat constructor kwargs, so scalar knobs (beta, alpha, k, policy, weights, budgets) are top-level Args fields and dict specs are treated as identities; you sweep over configs by listing whole dicts in a spec's vars, not by reaching inside them.

SampleSequenceScorer is served by passing an instance (it wraps a per-row SampleScorer callable, which has no meaningful string form); no dict kind is added for it.

_require(spec, key, kind)

Read a required key from a dict spec, raising a ValueError that names the offending kind.

Source code in steerability/algorithms/output_control/common/resolve.py 
46
47
48
49
50
def _require(spec: dict, key: str, kind: str):
    """Read a required key from a dict spec, raising a `ValueError` that names the offending kind."""
    if key not in spec or spec[key] is None:
        raise ValueError(f"{kind!r} spec requires a {key!r} key.")
    return spec[key]

resolve_scorer(spec, *, device)

Resolve a scorer spec into a SequenceScorer.

Accepted forms: any SequenceScorer callable (returned as-is, such as a plain function, a MajorityVoteScorer, or a SampleSequenceScorer); or a dict with a "kind" key:

- `"reward_model"`: `model_id` (required); `score_index=0`, `batch_size=8`,
  `hf_model_kwargs`. Loads a classifier and wraps it in a `RewardModelScorer`.
- `"majority_vote"`: `answer_extractor` (optional). Builds a `MajorityVoteScorer`.

Parameters:

Name Type Description Default
spec Any

The scorer spec.

 required
device device | str

Device to load an auxiliary reward model onto.

 required

Returns:

Type Description
SequenceScorer

A SequenceScorer (a callable (prompt, continuations, params) -> list[float]).

Raises:

Type Description
ValueError

On an unknown kind or a malformed dict spec.
Source code in steerability/algorithms/output_control/common/resolve.py 
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
def resolve_scorer(spec: Any, *, device: torch.device | str) -> SequenceScorer:
    """Resolve a scorer spec into a `SequenceScorer`.

    Accepted forms: any `SequenceScorer` callable (returned as-is, such as a plain function, a
    `MajorityVoteScorer`, or a `SampleSequenceScorer`); or a dict with a `"kind"` key:

        - `"reward_model"`: `model_id` (required); `score_index=0`, `batch_size=8`,
          `hf_model_kwargs`. Loads a classifier and wraps it in a `RewardModelScorer`.
        - `"majority_vote"`: `answer_extractor` (optional). Builds a `MajorityVoteScorer`.

    Args:
        spec: The scorer spec.
        device: Device to load an auxiliary reward model onto.

    Returns:
        A `SequenceScorer` (a callable `(prompt, continuations, params) -> list[float]`).

    Raises:
        ValueError: On an unknown kind or a malformed dict spec.
    """
    if isinstance(spec, dict):
        kind = spec.get("kind")
        if kind == "reward_model":
            model_id = _require(spec, "model_id", "reward_model")
            rm, rm_tokenizer = load_sequence_classifier(
                model_id, device=device, hf_model_kwargs=spec.get("hf_model_kwargs"),
            )
            return RewardModelScorer(
                rm,
                rm_tokenizer,
                score_index=spec.get("score_index", 0),
                batch_size=spec.get("batch_size", 8),
            )
        if kind == "majority_vote":
            return MajorityVoteScorer(answer_extractor=spec.get("answer_extractor"))
        raise ValueError(
            f"Unknown scorer kind {kind!r}; accepted kinds are 'reward_model', 'majority_vote' "
            "(or pass a SequenceScorer callable / SampleSequenceScorer instance)."
        )

    if callable(spec):
        return spec

    raise ValueError(
        "A scorer spec must be a SequenceScorer callable (e.g. a function, MajorityVoteScorer, or "
        "SampleSequenceScorer) or a dict with a 'kind' key."
    )

resolve_source(spec, *, model, tokenizer)

Resolve a source spec into a prepared BaseLogitSource.

Accepted forms: a BaseLogitSource instance; a (prefix_ids) -> Tensor[B, V] callable (wrapped in CallableSource); a str (shorthand for {"kind": "aux_model", "name_or_path": <str>}); or a dict with a "kind" key:

- `"aux_model"`: `name_or_path` (required); `prompt_transform`, `hf_model_kwargs`,
  `shared_vocab=True`. Builds an `AuxModelSource` (its `prepare` loads and vocab-checks).
- `"prompt_variant"`: `prompt_transform` (required). Builds a `PromptVariantSource`
  (see its compatibility note on same-model forwards).

Parameters:

Name Type Description Default
spec Any

The source spec.

 required
model PreTrainedModel

The pipeline's base model.

 required
tokenizer PreTrainedTokenizerBase

The base tokenizer.

 required

Returns:

Type Description
BaseLogitSource

A prepared BaseLogitSource.

Raises:

Type Description
ValueError

On an unknown kind or a malformed dict spec.
Source code in steerability/algorithms/output_control/common/resolve.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def resolve_source(spec: Any, *, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase) -> BaseLogitSource:
    """Resolve a source spec into a prepared `BaseLogitSource`.

    Accepted forms: a `BaseLogitSource` instance; a `(prefix_ids) -> Tensor[B, V]` callable (wrapped
    in `CallableSource`); a `str` (shorthand for `{"kind": "aux_model", "name_or_path": <str>}`); or
    a dict with a `"kind"` key:

        - `"aux_model"`: `name_or_path` (required); `prompt_transform`, `hf_model_kwargs`,
          `shared_vocab=True`. Builds an `AuxModelSource` (its `prepare` loads and vocab-checks).
        - `"prompt_variant"`: `prompt_transform` (required). Builds a `PromptVariantSource`
          (see its compatibility note on same-model forwards).

    Args:
        spec: The source spec.
        model: The pipeline's base model.
        tokenizer: The base tokenizer.

    Returns:
        A prepared `BaseLogitSource`.

    Raises:
        ValueError: On an unknown kind or a malformed dict spec.
    """
    if isinstance(spec, BaseLogitSource):
        spec.prepare(model=model, tokenizer=tokenizer)
        return spec

    if isinstance(spec, str):
        spec = {"kind": "aux_model", "name_or_path": spec}

    if callable(spec) and not isinstance(spec, dict):
        source = CallableSource(spec)
        source.prepare(model=model, tokenizer=tokenizer)
        return source

    if isinstance(spec, dict):
        kind = spec.get("kind")
        if kind == "aux_model":
            name_or_path = _require(spec, "name_or_path", "aux_model")
            source = AuxModelSource(
                name_or_path,
                base_tokenizer=tokenizer,
                prompt_transform=spec.get("prompt_transform"),
                shared_vocab=spec.get("shared_vocab", True),
                hf_model_kwargs=spec.get("hf_model_kwargs"),
            )
            source.prepare(model=model, tokenizer=tokenizer)
            return source
        if kind == "prompt_variant":
            prompt_transform = _require(spec, "prompt_transform", "prompt_variant")
            source = PromptVariantSource(prompt_transform, base_tokenizer=tokenizer)
            source.prepare(model=model, tokenizer=tokenizer)
            return source
        raise ValueError(
            f"Unknown source kind {kind!r}; accepted kinds are 'aux_model', 'prompt_variant' "
            "(or pass a BaseLogitSource instance, a callable, or an aux-model name/path string)."
        )

    raise ValueError(
        "A source spec must be a BaseLogitSource instance, a callable, a str, or a dict with a "
        "'kind' key."
    )

resolve_value(spec, *, model, tokenizer, device)

Resolve a value spec into a BaseCandidateValue.

Accepted forms: a BaseCandidateValue instance; a (StepContext) -> Tensor[B, K] callable (wrapped in CallableValue); or a dict with a "kind" key:

- `"reward_model"`: `model_id` (required); `score_index=0`, `score_transform="none"`,
  `hf_model_kwargs`. Loads an `AutoModelForSequenceClassification` and wraps it in a
  `RewardModelValue` (the text-round-trip path; `shared_vocab` defaults False).
- `"classifier"`: `model_id` or `fn` (required); `label_index=1`, `hf_model_kwargs`.
  Wraps a loaded classifier (or a `list[str] -> Tensor` callable) in a `ClassifierValue`.
- `"subspace_margin"`: `probe_path` or `data` (required); `batch_size=4`,
  `max_length=1024`, `save_path`. Loads a probe from a directory artifact (`Probe.load`)
  or a single-file checkpoint (`.probe` JSON or legacy `{'wv', 'mu_mu'}` tensor), or fits
  one on the base model via `fit_probe` (fisher direction over last-token features at the
  raw final-layer boundary, midpoint calibration); `save_path` writes the fitted probe's
  directory artifact.
- `"callable"`: `fn` (required); `supports_batching`, `scoring_cost`. Wraps `fn` in a
  `CallableValue` with explicit flags.

Parameters:

Name Type Description Default
spec Any

The value spec.

 required
model PreTrainedModel

The pipeline's base model.

 required
tokenizer PreTrainedTokenizerBase

The base tokenizer.

 required
device device | str

Device to load auxiliary models / align probes onto.

 required

Returns:

Type Description
BaseCandidateValue

A prepared BaseCandidateValue.

Raises:

Type Description
ValueError

On an unknown kind or a malformed dict spec.
Source code in steerability/algorithms/output_control/common/resolve.py 
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def resolve_value(
    spec: Any,
    *,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    device: torch.device | str,
) -> BaseCandidateValue:
    """Resolve a value spec into a `BaseCandidateValue`.

    Accepted forms: a `BaseCandidateValue` instance; a `(StepContext) -> Tensor[B, K]` callable
    (wrapped in `CallableValue`); or a dict with a `"kind"` key:

        - `"reward_model"`: `model_id` (required); `score_index=0`, `score_transform="none"`,
          `hf_model_kwargs`. Loads an `AutoModelForSequenceClassification` and wraps it in a
          `RewardModelValue` (the text-round-trip path; `shared_vocab` defaults False).
        - `"classifier"`: `model_id` or `fn` (required); `label_index=1`, `hf_model_kwargs`.
          Wraps a loaded classifier (or a `list[str] -> Tensor` callable) in a `ClassifierValue`.
        - `"subspace_margin"`: `probe_path` or `data` (required); `batch_size=4`,
          `max_length=1024`, `save_path`. Loads a probe from a directory artifact (`Probe.load`)
          or a single-file checkpoint (`.probe` JSON or legacy `{'wv', 'mu_mu'}` tensor), or fits
          one on the base model via `fit_probe` (fisher direction over last-token features at the
          raw final-layer boundary, midpoint calibration); `save_path` writes the fitted probe's
          directory artifact.
        - `"callable"`: `fn` (required); `supports_batching`, `scoring_cost`. Wraps `fn` in a
          `CallableValue` with explicit flags.

    Args:
        spec: The value spec.
        model: The pipeline's base model.
        tokenizer: The base tokenizer.
        device: Device to load auxiliary models / align probes onto.

    Returns:
        A prepared `BaseCandidateValue`.

    Raises:
        ValueError: On an unknown kind or a malformed dict spec.
    """
    if isinstance(spec, BaseCandidateValue):
        spec.prepare(model=model, tokenizer=tokenizer)
        return spec

    if callable(spec) and not isinstance(spec, dict):
        return CallableValue(spec)

    if isinstance(spec, dict):
        kind = spec.get("kind")
        if kind == "reward_model":
            model_id = _require(spec, "model_id", "reward_model")
            rm, rm_tokenizer = load_sequence_classifier(
                model_id, device=device, hf_model_kwargs=spec.get("hf_model_kwargs"),
            )
            return RewardModelValue(
                reward_model=rm,
                rm_tokenizer=rm_tokenizer,
                score_index=spec.get("score_index", 0),
                score_transform=spec.get("score_transform", "none"),
            )
        if kind == "classifier":
            label_index = spec.get("label_index", 1)
            if spec.get("fn") is not None:
                return ClassifierValue(spec["fn"], label_index=label_index)
            model_id = _require(spec, "model_id", "classifier")
            clf, clf_tokenizer = load_sequence_classifier(
                model_id, device=device, hf_model_kwargs=spec.get("hf_model_kwargs"),
            )
            return ClassifierValue(clf, classifier_tokenizer=clf_tokenizer, label_index=label_index)
        if kind == "subspace_margin":
            # imported lazily to keep probe-fitting imports out of unrelated resolves
            import os

            from steerability.algorithms.core.internals.data import as_labeled_examples
            from steerability.algorithms.core.internals.model_layout import resolve_model_layout
            from steerability.algorithms.core.internals.probes.fitting import ProbeFitSpec, fit_probe
            from steerability.algorithms.core.internals.probes.probe import Probe
            from steerability.algorithms.output_control.common.values.subspace_margin import load_single_file_probe

            final_layer = resolve_model_layout(model).num_layers - 1
            if spec.get("probe_path") is not None:
                path = spec["probe_path"]
                if os.path.isdir(path):
                    probe = Probe.load(path)
                else:
                    probe = load_single_file_probe(path, layer_id=final_layer)
            elif spec.get("data") is not None:
                fit_spec = ProbeFitSpec(
                    method="fisher",
                    pooling="last",
                    location="layer_output",
                    prompt_format="raw",
                    candidate_layers=[final_layer],
                    calibration="midpoint",
                )
                probe = fit_probe(
                    model,
                    tokenizer,
                    data=as_labeled_examples(spec["data"]),
                    spec=fit_spec,
                    batch_size=spec.get("batch_size", 4),
                    max_length=spec.get("max_length", 1024),
                )
                if spec.get("save_path") is not None:
                    probe.save(spec["save_path"])
            else:
                raise ValueError("'subspace_margin' spec requires a 'probe_path' or 'data' key.")
            return SubspaceMarginValue(probe)
        if kind == "callable":
            fn = _require(spec, "fn", "callable")
            return CallableValue(
                fn,
                supports_batching=spec.get("supports_batching", False),
                scoring_cost=spec.get("scoring_cost", "cheap"),
            )
        raise ValueError(
            f"Unknown value kind {kind!r}; accepted kinds are 'reward_model', 'classifier', "
            "'subspace_margin', 'callable' (or pass a BaseCandidateValue instance or a callable)."
        )

    raise ValueError(
        "A value spec must be a BaseCandidateValue instance, a callable, or a dict with a 'kind' key."
    )

scorers

Sequence scorers (score whole continuations, per-sequence floats).

base

SequenceScorer — the sequence-level scoring protocol.

DeAL's reward_func signature is adopted verbatim as the library convention, so every existing user-supplied callable is already a SequenceScorer.

SequenceScorer

Bases: Protocol

Score continuations of a prompt; higher = more desired.

A callable (prompt, continuations, params) -> list[float], one score per continuation.

Source code in steerability/algorithms/output_control/common/scorers/base.py 
11
12
13
14
15
16
17
18
19
@runtime_checkable
class SequenceScorer(Protocol):
    """Score continuations of a prompt; higher = more desired.

    A callable `(prompt, continuations, params) -> list[float]`, one score per continuation.
    """

    def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
        ...
__call__(prompt, continuations, params)
Source code in steerability/algorithms/output_control/common/scorers/base.py 
18
19
def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
    ...

majority_vote

MajorityVoteScorer — self-consistency scoring by answer agreement.

MajorityVoteScorer

Score each continuation by how many others share its extracted answer.

Used under BestOfN, this is self-consistency (Wang et al. 2022): the continuation whose answer the plurality agrees with wins.

Parameters:

Name Type Description Default
answer_extractor Callable[[str], str] | None

str -> str mapping a continuation to its answer. Defaults to final-line extraction.

 None
Source code in steerability/algorithms/output_control/common/scorers/majority_vote.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class MajorityVoteScorer:
    """Score each continuation by how many others share its extracted answer.

    Used under `BestOfN`, this is self-consistency (Wang et al. 2022): the continuation whose answer
    the plurality agrees with wins.

    Args:
        answer_extractor: `str -> str` mapping a continuation to its answer. Defaults to final-line
            extraction.
    """

    def __init__(self, answer_extractor: Callable[[str], str] | None = None):
        self.answer_extractor = answer_extractor or _final_line

    def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
        answers = [self.answer_extractor(c) for c in continuations]
        counts = Counter(answers)
        # score = agreement count with the others (exclude self)
        return [float(counts[a] - 1) for a in answers]
answer_extractor = answer_extractor or _final_line instance-attribute
__call__(prompt, continuations, params)
Source code in steerability/algorithms/output_control/common/scorers/majority_vote.py 
28
29
30
31
32
def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
    answers = [self.answer_extractor(c) for c in continuations]
    counts = Counter(answers)
    # score = agreement count with the others (exclude self)
    return [float(counts[a] - 1) for a in answers]
__init__(answer_extractor=None)
Source code in steerability/algorithms/output_control/common/scorers/majority_vote.py 
25
26
def __init__(self, answer_extractor: Callable[[str], str] | None = None):
    self.answer_extractor = answer_extractor or _final_line
_final_line(text)

Default answer extractor: the last non-empty line, stripped.

Source code in steerability/algorithms/output_control/common/scorers/majority_vote.py 
 8
 9
10
11
def _final_line(text: str) -> str:
    """Default answer extractor: the last non-empty line, stripped."""
    lines = [line.strip() for line in text.strip().splitlines() if line.strip()]
    return lines[-1] if lines else ""

reward_model

RewardModelScorer — score prompt + continuation with an HF sequence classifier.

RewardModelScorer

Score each continuation with an HF sequence classifier over prompt + continuation.

Parameters:

Name Type Description Default
model PreTrainedModel

The reward model (HF sequence classifier, in eval mode).

 required
tokenizer PreTrainedTokenizerBase

Tokenizer for the reward model.

 required
score_index int

Logit index to read as the scalar reward.

 0
batch_size int

Batch size for scoring.

 8
Source code in steerability/algorithms/output_control/common/scorers/reward_model.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class RewardModelScorer:
    """Score each continuation with an HF sequence classifier over `prompt + continuation`.

    Args:
        model: The reward model (HF sequence classifier, in eval mode).
        tokenizer: Tokenizer for the reward model.
        score_index: Logit index to read as the scalar reward.
        batch_size: Batch size for scoring.
    """

    def __init__(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        score_index: int = 0,
        batch_size: int = 8,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.score_index = score_index
        self.batch_size = batch_size
        self._device = next(model.parameters()).device

    @torch.inference_mode()
    def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
        texts = [f"{prompt}{continuation}" for continuation in continuations]
        scores: list[float] = []
        for start in range(0, len(texts), self.batch_size):
            batch = texts[start:start + self.batch_size]
            inputs = self.tokenizer(
                batch, return_tensors="pt", padding=True, truncation=True
            ).to(self._device)
            logits = self.model(**inputs).logits
            scores.extend(logits[:, self.score_index].detach().cpu().tolist())
        return scores
_device = next(model.parameters()).device instance-attribute
batch_size = batch_size instance-attribute
model = model instance-attribute
score_index = score_index instance-attribute
tokenizer = tokenizer instance-attribute
__call__(prompt, continuations, params)
Source code in steerability/algorithms/output_control/common/scorers/reward_model.py 
31
32
33
34
35
36
37
38
39
40
41
42
@torch.inference_mode()
def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
    texts = [f"{prompt}{continuation}" for continuation in continuations]
    scores: list[float] = []
    for start in range(0, len(texts), self.batch_size):
        batch = texts[start:start + self.batch_size]
        inputs = self.tokenizer(
            batch, return_tensors="pt", padding=True, truncation=True
        ).to(self._device)
        logits = self.model(**inputs).logits
        scores.extend(logits[:, self.score_index].detach().cpu().tolist())
    return scores
__init__(model, tokenizer, score_index=0, batch_size=8)
Source code in steerability/algorithms/output_control/common/scorers/reward_model.py 
18
19
20
21
22
23
24
25
26
27
28
29
def __init__(
    self,
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    score_index: int = 0,
    batch_size: int = 8,
):
    self.model = model
    self.tokenizer = tokenizer
    self.score_index = score_index
    self.batch_size = batch_size
    self._device = next(model.parameters()).device

sample

SampleSequenceScorer: score continuations with a per-row SampleScorer.

SampleSequenceScorer

Score each continuation with a per-row SampleScorer.

Each continuation is scored individually against a row built from the prompt and the call's scoring params, as row_scorer(continuation, {"input": prompt, **params}). The row always carries the prompt as "input".

Parameters:

Name Type Description Default
row_scorer SampleScorer

SampleScorer scoring one (response, row) pair; higher is better.

 required
Source code in steerability/algorithms/output_control/common/scorers/sample.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class SampleSequenceScorer:
    """Score each continuation with a per-row `SampleScorer`.

    Each continuation is scored individually against a row built from the prompt and the call's
    scoring params, as `row_scorer(continuation, {"input": prompt, **params})`. The row always
    carries the prompt as `"input"`.

    Args:
        row_scorer: `SampleScorer` scoring one `(response, row)` pair; higher is better.
    """

    def __init__(self, row_scorer: SampleScorer):
        self.row_scorer = row_scorer

    def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
        row = {"input": prompt, **(params or {})}
        return [float(self.row_scorer(continuation, row)) for continuation in continuations]
row_scorer = row_scorer instance-attribute
__call__(prompt, continuations, params)
Source code in steerability/algorithms/output_control/common/scorers/sample.py 
21
22
23
def __call__(self, prompt: str, continuations: list[str], params: dict) -> list[float]:
    row = {"input": prompt, **(params or {})}
    return [float(self.row_scorer(continuation, row)) for continuation in continuations]
__init__(row_scorer)
Source code in steerability/algorithms/output_control/common/scorers/sample.py 
18
19
def __init__(self, row_scorer: SampleScorer):
    self.row_scorer = row_scorer

values

Per-candidate value functions (score candidate continuations of a prefix).

base

StepContext and BaseCandidateValue, the candidate-value interface.

values/ scores candidates (per-token scalars for one decode step); the sibling scorers/ scores continuations (per-sequence floats). Values are consumed by ValueGuidedProcessor.

BaseCandidateValue

Bases: ABC

Score candidate continuations of a prefix; higher = more desired.

score() returns a float tensor [B, K] aligned with ctx.candidate_ids. Values must be pure in the PrefixKeyedProcessor sense; any caching keys on the prefix.

Class attributes

supports_batching: Whether the value handles B > 1 candidate rows. scoring_cost: Rough cost tier, surfaced so controls can warn when scoring with expensive values is enabled. "cheap" (no model forward), "aux_forward" (an auxiliary-model forward), or "model_forward" (a forward of the pipeline's own model). same_model_forwards: Whether this value issues additional forward passes through the pipeline's own model during decoding. Such passes must be wrapped in auxiliary_pass() (see steerability.algorithms.core.utils.auxiliary_pass), which keeps them out of state-control condition scoring, gate updates, and fallback position counting. Defaults to False; the flag is declarative metadata and is not read by the pipeline.

Source code in steerability/algorithms/output_control/common/values/base.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class BaseCandidateValue(ABC):
    """Score candidate continuations of a prefix; higher = more desired.

    `score()` returns a float tensor `[B, K]` aligned with `ctx.candidate_ids`. Values must be pure
    in the `PrefixKeyedProcessor` sense; any caching keys on the prefix.

    Class attributes:
        supports_batching: Whether the value handles `B > 1` candidate rows.
        scoring_cost: Rough cost tier, surfaced so controls can warn when scoring with expensive
            values is enabled. `"cheap"` (no model forward), `"aux_forward"` (an auxiliary-model
            forward), or `"model_forward"` (a forward of the pipeline's own model).
        same_model_forwards: Whether this value issues additional forward passes through the
            pipeline's own model during decoding. Such passes must be wrapped in
            `auxiliary_pass()` (see `steerability.algorithms.core.utils.auxiliary_pass`), which
            keeps them out of state-control condition scoring, gate updates, and fallback
            position counting. Defaults to False; the flag is declarative metadata and is not
            read by the pipeline.
    """

    supports_batching: bool = False
    scoring_cost: Literal["cheap", "aux_forward", "model_forward"] = "cheap"
    same_model_forwards: bool = False

    @abstractmethod
    def score(self, ctx: StepContext) -> torch.Tensor:
        """Return per-candidate values `[B, K]` aligned with `ctx.candidate_ids`."""
        ...

    def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
        """Optional offline setup, invoked from the owning control's `steer()`.

        Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
        control calls during its own `steer()` when the value needs preparation.
        """
        pass

    def cleanup(self) -> None:
        """Release any resources allocated during `prepare()`."""
        pass
same_model_forwards = False class-attribute instance-attribute
scoring_cost = 'cheap' class-attribute instance-attribute
supports_batching = False class-attribute instance-attribute
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/values/base.py 
71
72
73
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup, invoked from the owning control's steer().

Mirrors the input-control selector prepare() convention: a default no-op that the parent control calls during its own steer() when the value needs preparation.

Source code in steerability/algorithms/output_control/common/values/base.py 
63
64
65
66
67
68
69
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup, invoked from the owning control's `steer()`.

    Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
    control calls during its own `steer()` when the value needs preparation.
    """
    pass
score(ctx) abstractmethod

Return per-candidate values [B, K] aligned with ctx.candidate_ids.

Source code in steerability/algorithms/output_control/common/values/base.py 
58
59
60
61
@abstractmethod
def score(self, ctx: StepContext) -> torch.Tensor:
    """Return per-candidate values `[B, K]` aligned with `ctx.candidate_ids`."""
    ...
StepContext dataclass

Everything a candidate value may consult at one decoding step.

Attributes:

Name Type Description
prefix_ids Tensor

The current prefix token ids [B, T].
candidate_ids Tensor

The candidate next-token ids [B, K].
lm_tokenizer PreTrainedTokenizerBase

The language model's tokenizer.
model PreTrainedModel | None

The pipeline's model (only used by same-model values; see CandidateForward).
attention_mask Tensor | None

The prefix attention mask [B, T], when available.
Source code in steerability/algorithms/output_control/common/values/base.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
@dataclass(slots=True)
class StepContext:
    """Everything a candidate value may consult at one decoding step.

    Attributes:
        prefix_ids: The current prefix token ids `[B, T]`.
        candidate_ids: The candidate next-token ids `[B, K]`.
        lm_tokenizer: The language model's tokenizer.
        model: The pipeline's model (only used by same-model values; see `CandidateForward`).
        attention_mask: The prefix attention mask `[B, T]`, when available.
    """

    prefix_ids: torch.Tensor
    candidate_ids: torch.Tensor
    lm_tokenizer: PreTrainedTokenizerBase
    model: PreTrainedModel | None = None
    attention_mask: torch.Tensor | None = None
attention_mask = None class-attribute instance-attribute
candidate_ids instance-attribute
lm_tokenizer instance-attribute
model = None class-attribute instance-attribute
prefix_ids instance-attribute
__init__(prefix_ids, candidate_ids, lm_tokenizer, model=None, attention_mask=None)

callable

Wrap a user function as a candidate value.

The escape hatch of the value family: any (StepContext) -> Tensor[B, K] callable becomes a BaseCandidateValue without a subclass. The wrapped function is responsible for returning per-candidate scores aligned with ctx.candidate_ids; the batching and cost flags are supplied explicitly at construction (they cannot be inferred from a bare callable).

CallableValue

Bases: BaseCandidateValue

Adapt a (StepContext) -> Tensor[B, K] callable into a BaseCandidateValue.

same_model_forwards keeps the BaseCandidateValue default of False, since it cannot be inferred from a bare callable. A callable that forwards the pipeline's model should be wrapped in a component that declares same_model_forwards.

Parameters:

Name Type Description Default
fn Callable[[StepContext], Tensor]

A callable mapping a StepContext to per-candidate values [B, K] aligned with ctx.candidate_ids; higher = more desired.

 required
supports_batching bool

Whether fn handles B > 1 candidate rows. Defaults to False.

 False
scoring_cost Literal['cheap', 'aux_forward', 'model_forward']

Cost tier surfaced to owning controls ("cheap", "aux_forward", or "model_forward"). Defaults to "cheap".

 'cheap'
Source code in steerability/algorithms/output_control/common/values/callable.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class CallableValue(BaseCandidateValue):
    """Adapt a `(StepContext) -> Tensor[B, K]` callable into a `BaseCandidateValue`.

    `same_model_forwards` keeps the `BaseCandidateValue` default of False, since it cannot be
    inferred from a bare callable. A callable that forwards the pipeline's model should be wrapped in
    a component that declares `same_model_forwards`.

    Args:
        fn: A callable mapping a `StepContext` to per-candidate values `[B, K]` aligned with
            `ctx.candidate_ids`; higher = more desired.
        supports_batching: Whether `fn` handles `B > 1` candidate rows. Defaults to False.
        scoring_cost: Cost tier surfaced to owning controls (`"cheap"`, `"aux_forward"`, or
            `"model_forward"`). Defaults to `"cheap"`.
    """

    def __init__(
        self,
        fn: Callable[[StepContext], torch.Tensor],
        *,
        supports_batching: bool = False,
        scoring_cost: Literal["cheap", "aux_forward", "model_forward"] = "cheap",
    ):
        if not callable(fn):
            raise TypeError("CallableValue requires a callable (StepContext) -> Tensor[B, K].")
        self.fn = fn
        self.supports_batching = supports_batching
        self.scoring_cost = scoring_cost

    def score(self, ctx: StepContext) -> torch.Tensor:
        return torch.as_tensor(self.fn(ctx))
fn = fn instance-attribute
same_model_forwards = False class-attribute instance-attribute
scoring_cost = scoring_cost instance-attribute
supports_batching = supports_batching instance-attribute
__init__(fn, *, supports_batching=False, scoring_cost='cheap')
Source code in steerability/algorithms/output_control/common/values/callable.py 
32
33
34
35
36
37
38
39
40
41
42
43
def __init__(
    self,
    fn: Callable[[StepContext], torch.Tensor],
    *,
    supports_batching: bool = False,
    scoring_cost: Literal["cheap", "aux_forward", "model_forward"] = "cheap",
):
    if not callable(fn):
        raise TypeError("CallableValue requires a callable (StepContext) -> Tensor[B, K].")
    self.fn = fn
    self.supports_batching = supports_batching
    self.scoring_cost = scoring_cost
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/values/base.py 
71
72
73
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup, invoked from the owning control's steer().

Mirrors the input-control selector prepare() convention: a default no-op that the parent control calls during its own steer() when the value needs preparation.

Source code in steerability/algorithms/output_control/common/values/base.py 
63
64
65
66
67
68
69
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup, invoked from the owning control's `steer()`.

    Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
    control calls during its own `steer()` when the value needs preparation.
    """
    pass
score(ctx)

Return per-candidate values [B, K] aligned with ctx.candidate_ids.

Source code in steerability/algorithms/output_control/common/values/callable.py 
45
46
def score(self, ctx: StepContext) -> torch.Tensor:
    return torch.as_tensor(self.fn(ctx))

classifier

ClassifierValue (FUDGE-family) — candidate value from an attribute classifier.

Returns log p(attribute | prefix + candidate) per candidate. Accepts either an HF sequence classifier (with a label_index selecting the attribute logit) or any callable mapping a list of texts to a [batch] tensor of attribute log-probabilities.

ClassifierValue

Bases: BaseCandidateValue

Score prefix + candidate with an attribute classifier (FUDGE-family).

Parameters:

Name Type Description Default
classifier Callable[[list[str]], Tensor] | PreTrainedModel

Either a callable list[str] -> Tensor[batch] returning attribute log-probs, or an HF sequence classifier (used with classifier_tokenizer + label_index).

 required
classifier_tokenizer PreTrainedTokenizerBase | None

Tokenizer for an HF classifier (required when classifier is a model).

 None
label_index int

Index of the attribute logit for an HF classifier.

 1
Note

scoring_cost="aux_forward", supports_batching=True.

Source code in steerability/algorithms/output_control/common/values/classifier.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class ClassifierValue(BaseCandidateValue):
    """Score `prefix + candidate` with an attribute classifier (FUDGE-family).

    Args:
        classifier: Either a callable `list[str] -> Tensor[batch]` returning attribute log-probs, or
            an HF sequence classifier (used with `classifier_tokenizer` + `label_index`).
        classifier_tokenizer: Tokenizer for an HF classifier (required when `classifier` is a model).
        label_index: Index of the attribute logit for an HF classifier.

    Note:
        `scoring_cost="aux_forward"`, `supports_batching=True`.
    """

    supports_batching: bool = True
    scoring_cost = "aux_forward"

    def __init__(
        self,
        classifier: Callable[[list[str]], torch.Tensor] | PreTrainedModel,
        classifier_tokenizer: PreTrainedTokenizerBase | None = None,
        label_index: int = 1,
    ):
        self.classifier = classifier
        self.classifier_tokenizer = classifier_tokenizer
        self.label_index = label_index
        self._callable = not hasattr(classifier, "parameters")
        self._device = None if self._callable else next(classifier.parameters()).device

    @torch.inference_mode()
    def score(self, ctx: StepContext) -> torch.Tensor:
        batch_size = ctx.prefix_ids.size(0)
        num_candidates = ctx.candidate_ids.size(1)

        prefix = ctx.prefix_ids.unsqueeze(1).expand(-1, num_candidates, -1)
        combined = torch.cat([prefix, ctx.candidate_ids.unsqueeze(-1)], dim=-1)
        flat = combined.reshape(batch_size * num_candidates, -1)
        texts = ctx.lm_tokenizer.decode(flat, skip_special_tokens=True)

        if self._callable:
            logp = self.classifier(texts)
            logp = torch.as_tensor(logp, dtype=torch.float32)
        else:
            inputs = self.classifier_tokenizer(
                texts, return_tensors="pt", padding=True, truncation=True
            ).to(self._device)
            logits = self.classifier(**inputs).logits
            logp = torch.log_softmax(logits, dim=-1)[:, self.label_index]

        return logp.reshape(batch_size, num_candidates)
_callable = not hasattr(classifier, 'parameters') instance-attribute
_device = None if self._callable else next(classifier.parameters()).device instance-attribute
classifier = classifier instance-attribute
classifier_tokenizer = classifier_tokenizer instance-attribute
label_index = label_index instance-attribute
same_model_forwards = False class-attribute instance-attribute
scoring_cost = 'aux_forward' class-attribute instance-attribute
supports_batching = True class-attribute instance-attribute
__init__(classifier, classifier_tokenizer=None, label_index=1)
Source code in steerability/algorithms/output_control/common/values/classifier.py 
33
34
35
36
37
38
39
40
41
42
43
def __init__(
    self,
    classifier: Callable[[list[str]], torch.Tensor] | PreTrainedModel,
    classifier_tokenizer: PreTrainedTokenizerBase | None = None,
    label_index: int = 1,
):
    self.classifier = classifier
    self.classifier_tokenizer = classifier_tokenizer
    self.label_index = label_index
    self._callable = not hasattr(classifier, "parameters")
    self._device = None if self._callable else next(classifier.parameters()).device
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/values/base.py 
71
72
73
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup, invoked from the owning control's steer().

Mirrors the input-control selector prepare() convention: a default no-op that the parent control calls during its own steer() when the value needs preparation.

Source code in steerability/algorithms/output_control/common/values/base.py 
63
64
65
66
67
68
69
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup, invoked from the owning control's `steer()`.

    Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
    control calls during its own `steer()` when the value needs preparation.
    """
    pass
score(ctx)

Return per-candidate values [B, K] aligned with ctx.candidate_ids.

Source code in steerability/algorithms/output_control/common/values/classifier.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
@torch.inference_mode()
def score(self, ctx: StepContext) -> torch.Tensor:
    batch_size = ctx.prefix_ids.size(0)
    num_candidates = ctx.candidate_ids.size(1)

    prefix = ctx.prefix_ids.unsqueeze(1).expand(-1, num_candidates, -1)
    combined = torch.cat([prefix, ctx.candidate_ids.unsqueeze(-1)], dim=-1)
    flat = combined.reshape(batch_size * num_candidates, -1)
    texts = ctx.lm_tokenizer.decode(flat, skip_special_tokens=True)

    if self._callable:
        logp = self.classifier(texts)
        logp = torch.as_tensor(logp, dtype=torch.float32)
    else:
        inputs = self.classifier_tokenizer(
            texts, return_tensors="pt", padding=True, truncation=True
        ).to(self._device)
        logits = self.classifier(**inputs).logits
        logp = torch.log_softmax(logits, dim=-1)[:, self.label_index]

    return logp.reshape(batch_size, num_candidates)

reward_model

Candidate values from an auxiliary sequence classifier.

For each row-candidate, scores prefix + candidate with an auxiliary reward model. Loading and configuration of the reward model is done by the owning control (RAD), which constructs the value with the loaded model and tokenizer.

Two values live here. RewardModelValue scores each step statelessly. When the reward model shares the language model's vocabulary (shared_vocab=True) it feeds the raw prefix + candidate ids to the reward model; otherwise it decodes to text and re-encodes with the reward-model tokenizer (the mismatched-vocabulary fallback). CachedRewardModelValue produces the same scores as RewardModelValue(shared_vocab=True) at lower cost by memoizing the reward model's prefix past_key_values across decode steps.

ScoreTransform = Literal['none', 'sigmoid', 'softmax'] module-attribute
CachedRewardModelValue

Bases: BaseCandidateValue

Score prefix + candidate with a unidirectional reward model, caching prefix activations.

Produces the same scores as RewardModelValue(shared_vocab=True) under the same score_index and score_transform, differing only in cost. The reward model's prefix past_key_values are memoized keyed on the prefix ids: each step forwards only the delta tokens to extend the cache, then evaluates the K candidates as single-token forwards over a repeated copy of the cache. Any non-extension of the cached prefix (rewind, reorder, restart, teacher-forced replay) rebuilds the cache from scratch, which reproduces a fresh full forward.

The reward model must be decoder-only (accept past_key_values and cache_position) and share the language model's vocabulary; the owning control checks both preconditions and a smoke forward at steer time and falls back to RewardModelValue(shared_vocab=True) on failure. Batch size 1 only.

Parameters:

Name Type Description Default
reward_model PreTrainedModel

The loaded unidirectional reward model (in eval mode).

 required
rm_tokenizer PreTrainedTokenizerBase

Tokenizer for the reward model. Its max_length attribute (if set) bounds the reward-model input length; over-length prefixes are scored without the cache over the left-truncated tail.

 required
score_index int

Column of the reward model's output read as the score.

 0
score_transform ScoreTransform

Map the reward model's output to a score before selecting score_index ("none", "sigmoid", or "softmax").

 'none'
Note

scoring_cost="aux_forward" and supports_batching=False.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
class CachedRewardModelValue(BaseCandidateValue):
    """Score `prefix + candidate` with a unidirectional reward model, caching prefix activations.

    Produces the same scores as `RewardModelValue(shared_vocab=True)` under the same `score_index`
    and `score_transform`, differing only in cost. The reward model's prefix `past_key_values` are
    memoized keyed on the prefix ids: each step forwards only the delta tokens to extend the cache,
    then evaluates the K candidates as single-token forwards over a repeated copy of the cache. Any
    non-extension of the cached prefix (rewind, reorder, restart, teacher-forced replay) rebuilds the
    cache from scratch, which reproduces a fresh full forward.

    The reward model must be decoder-only (accept `past_key_values` and `cache_position`) and share
    the language model's vocabulary; the owning control checks both preconditions and a smoke forward
    at steer time and falls back to `RewardModelValue(shared_vocab=True)` on failure. Batch size 1
    only.

    Args:
        reward_model: The loaded unidirectional reward model (in eval mode).
        rm_tokenizer: Tokenizer for the reward model. Its `max_length` attribute (if set) bounds the
            reward-model input length; over-length prefixes are scored without the cache over the
            left-truncated tail.
        score_index: Column of the reward model's output read as the score.
        score_transform: Map the reward model's output to a score before selecting `score_index`
            (`"none"`, `"sigmoid"`, or `"softmax"`).

    Note:
        `scoring_cost="aux_forward"` and `supports_batching=False`.
    """

    scoring_cost = "aux_forward"

    def __init__(
        self,
        reward_model: PreTrainedModel,
        rm_tokenizer: PreTrainedTokenizerBase,
        score_index: int = 0,
        score_transform: ScoreTransform = "none",
    ):
        self.reward_model = reward_model
        self.rm_tokenizer = rm_tokenizer
        self.score_index = score_index
        self.score_transform = score_transform
        self.supports_batching = False
        self._device = next(reward_model.parameters()).device
        self._cached_ids: torch.Tensor | None = None  # [1, T_c]
        self._cache = None  # past_key_values covering _cached_ids

    def _sync_cache(self, prefix_ids: torch.Tensor, full_mask: torch.Tensor) -> None:
        """Bring the internal cache up to `prefix_ids` (extend by the delta, or rebuild)."""
        if not extends_prefix(self._cached_ids, prefix_ids):
            out = self.reward_model(
                input_ids=prefix_ids, attention_mask=full_mask, use_cache=True, return_dict=True,
                cache_position=torch.arange(prefix_ids.size(1), device=prefix_ids.device),
            )
            self._cache = out.past_key_values
        else:
            cached_len = self._cached_ids.size(1)
            if prefix_ids.size(1) > cached_len:
                delta = prefix_ids[:, cached_len:]
                positions = torch.arange(cached_len, prefix_ids.size(1), device=prefix_ids.device)
                out = self.reward_model(
                    input_ids=delta, attention_mask=full_mask, past_key_values=self._cache,
                    use_cache=True, cache_position=positions, return_dict=True,
                )
                self._cache = out.past_key_values
        self._cached_ids = prefix_ids.detach()

    @torch.inference_mode()
    def score(self, ctx: StepContext) -> torch.Tensor:
        if ctx.prefix_ids.size(0) != 1:
            raise ValueError("CachedRewardModelValue supports batch size 1 only.")

        num_candidates = ctx.candidate_ids.size(1)
        prefix_ids = ctx.prefix_ids.to(self._device)
        candidate_ids = ctx.candidate_ids.to(self._device)
        full_mask = full_prefix_mask(prefix_ids, ctx.attention_mask.to(self._device) if ctx.attention_mask is not None else None)

        max_length = getattr(self.rm_tokenizer, "max_length", None)
        if max_length is not None and prefix_ids.size(1) + 1 > max_length:
            return self._score_no_cache(prefix_ids, candidate_ids, full_mask, max_length)

        self._sync_cache(prefix_ids, full_mask)

        prefix_len = prefix_ids.size(1)
        repeated = repeat_cache(self._cache, num_candidates, preserve_input=True)
        cand_tokens = candidate_ids.reshape(num_candidates, 1)
        cand_mask = torch.cat(
            [full_mask.repeat(num_candidates, 1),
             torch.ones(num_candidates, 1, device=self._device, dtype=full_mask.dtype)],
            dim=1,
        )
        positions = torch.arange(prefix_len, prefix_len + 1, device=self._device)
        output = self.reward_model(
            input_ids=cand_tokens,
            attention_mask=cand_mask,
            past_key_values=repeated,
            use_cache=True,
            cache_position=positions,
            return_dict=True,
        )
        rewards = extract_score(output, self.score_index, self.score_transform)  # [K]
        return rewards.reshape(1, num_candidates)

    def _score_no_cache(
        self,
        prefix_ids: torch.Tensor,
        candidate_ids: torch.Tensor,
        full_mask: torch.Tensor,
        max_length: int,
    ) -> torch.Tensor:
        """Full forward over the left-truncated `prefix + candidate` for an over-length prefix."""
        num_candidates = candidate_ids.size(1)
        prefix = prefix_ids.expand(num_candidates, -1)
        combined = torch.cat([prefix, candidate_ids.reshape(num_candidates, 1)], dim=-1)  # [K, T+1]
        mask = torch.cat(
            [full_mask.repeat(num_candidates, 1),
             torch.ones(num_candidates, 1, device=self._device, dtype=full_mask.dtype)],
            dim=1,
        )
        combined = combined[:, -max_length:]
        mask = mask[:, -max_length:]
        output = self.reward_model(input_ids=combined, attention_mask=mask, return_dict=True)
        rewards = extract_score(output, self.score_index, self.score_transform)  # [K]
        return rewards.reshape(1, num_candidates)

    def cleanup(self) -> None:
        """Release the reward model and cached activations."""
        self.reward_model = None
        self._cache = None
        self._cached_ids = None
_cache = None instance-attribute
_cached_ids = None instance-attribute
_device = next(reward_model.parameters()).device instance-attribute
reward_model = reward_model instance-attribute
rm_tokenizer = rm_tokenizer instance-attribute
same_model_forwards = False class-attribute instance-attribute
score_index = score_index instance-attribute
score_transform = score_transform instance-attribute
scoring_cost = 'aux_forward' class-attribute instance-attribute
supports_batching = False instance-attribute
__init__(reward_model, rm_tokenizer, score_index=0, score_transform='none')
Source code in steerability/algorithms/output_control/common/values/reward_model.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
def __init__(
    self,
    reward_model: PreTrainedModel,
    rm_tokenizer: PreTrainedTokenizerBase,
    score_index: int = 0,
    score_transform: ScoreTransform = "none",
):
    self.reward_model = reward_model
    self.rm_tokenizer = rm_tokenizer
    self.score_index = score_index
    self.score_transform = score_transform
    self.supports_batching = False
    self._device = next(reward_model.parameters()).device
    self._cached_ids: torch.Tensor | None = None  # [1, T_c]
    self._cache = None  # past_key_values covering _cached_ids
_score_no_cache(prefix_ids, candidate_ids, full_mask, max_length)

Full forward over the left-truncated prefix + candidate for an over-length prefix.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def _score_no_cache(
    self,
    prefix_ids: torch.Tensor,
    candidate_ids: torch.Tensor,
    full_mask: torch.Tensor,
    max_length: int,
) -> torch.Tensor:
    """Full forward over the left-truncated `prefix + candidate` for an over-length prefix."""
    num_candidates = candidate_ids.size(1)
    prefix = prefix_ids.expand(num_candidates, -1)
    combined = torch.cat([prefix, candidate_ids.reshape(num_candidates, 1)], dim=-1)  # [K, T+1]
    mask = torch.cat(
        [full_mask.repeat(num_candidates, 1),
         torch.ones(num_candidates, 1, device=self._device, dtype=full_mask.dtype)],
        dim=1,
    )
    combined = combined[:, -max_length:]
    mask = mask[:, -max_length:]
    output = self.reward_model(input_ids=combined, attention_mask=mask, return_dict=True)
    rewards = extract_score(output, self.score_index, self.score_transform)  # [K]
    return rewards.reshape(1, num_candidates)
_sync_cache(prefix_ids, full_mask)

Bring the internal cache up to prefix_ids (extend by the delta, or rebuild).

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def _sync_cache(self, prefix_ids: torch.Tensor, full_mask: torch.Tensor) -> None:
    """Bring the internal cache up to `prefix_ids` (extend by the delta, or rebuild)."""
    if not extends_prefix(self._cached_ids, prefix_ids):
        out = self.reward_model(
            input_ids=prefix_ids, attention_mask=full_mask, use_cache=True, return_dict=True,
            cache_position=torch.arange(prefix_ids.size(1), device=prefix_ids.device),
        )
        self._cache = out.past_key_values
    else:
        cached_len = self._cached_ids.size(1)
        if prefix_ids.size(1) > cached_len:
            delta = prefix_ids[:, cached_len:]
            positions = torch.arange(cached_len, prefix_ids.size(1), device=prefix_ids.device)
            out = self.reward_model(
                input_ids=delta, attention_mask=full_mask, past_key_values=self._cache,
                use_cache=True, cache_position=positions, return_dict=True,
            )
            self._cache = out.past_key_values
    self._cached_ids = prefix_ids.detach()
cleanup()

Release the reward model and cached activations.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
290
291
292
293
294
def cleanup(self) -> None:
    """Release the reward model and cached activations."""
    self.reward_model = None
    self._cache = None
    self._cached_ids = None
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup, invoked from the owning control's steer().

Mirrors the input-control selector prepare() convention: a default no-op that the parent control calls during its own steer() when the value needs preparation.

Source code in steerability/algorithms/output_control/common/values/base.py 
63
64
65
66
67
68
69
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup, invoked from the owning control's `steer()`.

    Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
    control calls during its own `steer()` when the value needs preparation.
    """
    pass
score(ctx)

Return per-candidate values [B, K] aligned with ctx.candidate_ids.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
@torch.inference_mode()
def score(self, ctx: StepContext) -> torch.Tensor:
    if ctx.prefix_ids.size(0) != 1:
        raise ValueError("CachedRewardModelValue supports batch size 1 only.")

    num_candidates = ctx.candidate_ids.size(1)
    prefix_ids = ctx.prefix_ids.to(self._device)
    candidate_ids = ctx.candidate_ids.to(self._device)
    full_mask = full_prefix_mask(prefix_ids, ctx.attention_mask.to(self._device) if ctx.attention_mask is not None else None)

    max_length = getattr(self.rm_tokenizer, "max_length", None)
    if max_length is not None and prefix_ids.size(1) + 1 > max_length:
        return self._score_no_cache(prefix_ids, candidate_ids, full_mask, max_length)

    self._sync_cache(prefix_ids, full_mask)

    prefix_len = prefix_ids.size(1)
    repeated = repeat_cache(self._cache, num_candidates, preserve_input=True)
    cand_tokens = candidate_ids.reshape(num_candidates, 1)
    cand_mask = torch.cat(
        [full_mask.repeat(num_candidates, 1),
         torch.ones(num_candidates, 1, device=self._device, dtype=full_mask.dtype)],
        dim=1,
    )
    positions = torch.arange(prefix_len, prefix_len + 1, device=self._device)
    output = self.reward_model(
        input_ids=cand_tokens,
        attention_mask=cand_mask,
        past_key_values=repeated,
        use_cache=True,
        cache_position=positions,
        return_dict=True,
    )
    rewards = extract_score(output, self.score_index, self.score_transform)  # [K]
    return rewards.reshape(1, num_candidates)
RewardModelValue

Bases: BaseCandidateValue

Score prefix + candidate with an auxiliary reward model (RAD).

When shared_vocab is True, the prefix and candidate ids are fed to the reward model directly (no text round-trip), so the reward model scores exactly the sequence the language model generated. When False, the ids are decoded to text and re-encoded with the reward-model tokenizer, which is the correct behavior for a reward model whose vocabulary differs from the language model's. The two paths are not numerically equivalent (the text round-trip may drop special tokens and re-segment the boundary token), so a reward model that shares the vocabulary should use the id path. On the id path a reward model with a score head is read at the candidate position explicitly (the last position of every row), so a candidate equal to the reward model's pad id is scored as a candidate rather than pooled onto the prefix.

Parameters:

Name Type Description Default
reward_model PreTrainedModel

The loaded auxiliary reward model (in eval mode).

 required
rm_tokenizer PreTrainedTokenizerBase

Tokenizer for the reward model. Its max_length attribute (if set) bounds the reward-model input length. On the id path (shared_vocab=True), over-length prefixes are left-truncated (the tail is kept so the candidate token survives), unlike the text path's inherited right-truncation.

 required
score_index int

Column of the reward model's output read as the score.

 0
score_transform ScoreTransform

Map the reward model's output to a score before selecting score_index ("none", "sigmoid", or "softmax").

 'none'
shared_vocab bool

When True, score via the raw-id path (requires the reward-model vocabulary to equal the language model's). When False (default), score via the text round-trip.

 False
Note

scoring_cost="aux_forward" and supports_batching=True; candidate scoring batches B * K rows through the auxiliary model.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
class RewardModelValue(BaseCandidateValue):
    """Score `prefix + candidate` with an auxiliary reward model (RAD).

    When `shared_vocab` is True, the prefix and candidate ids are fed to the reward model directly
    (no text round-trip), so the reward model scores exactly the sequence the language model
    generated. When False, the ids are decoded to text and re-encoded with the reward-model
    tokenizer, which is the correct behavior for a reward model whose vocabulary differs from the
    language model's. The two paths are not numerically equivalent (the text round-trip may drop
    special tokens and re-segment the boundary token), so a reward model that shares the vocabulary
    should use the id path. On the id path a reward model with a `score` head is read at the
    candidate position explicitly (the last position of every row), so a candidate equal to the
    reward model's pad id is scored as a candidate rather than pooled onto the prefix.

    Args:
        reward_model: The loaded auxiliary reward model (in eval mode).
        rm_tokenizer: Tokenizer for the reward model. Its `max_length` attribute (if set) bounds the
            reward-model input length. On the id path (`shared_vocab=True`), over-length prefixes are
            left-truncated (the tail is kept so the candidate token survives), unlike the text path's
            inherited right-truncation.
        score_index: Column of the reward model's output read as the score.
        score_transform: Map the reward model's output to a score before selecting `score_index`
            (`"none"`, `"sigmoid"`, or `"softmax"`).
        shared_vocab: When True, score via the raw-id path (requires the reward-model vocabulary to
            equal the language model's). When False (default), score via the text round-trip.

    Note:
        `scoring_cost="aux_forward"` and `supports_batching=True`; candidate scoring batches `B * K`
        rows through the auxiliary model.
    """

    scoring_cost = "aux_forward"

    def __init__(
        self,
        reward_model: PreTrainedModel,
        rm_tokenizer: PreTrainedTokenizerBase,
        score_index: int = 0,
        score_transform: ScoreTransform = "none",
        shared_vocab: bool = False,
    ):
        self.reward_model = reward_model
        self.rm_tokenizer = rm_tokenizer
        self.score_index = score_index
        self.score_transform = score_transform
        self.shared_vocab = shared_vocab
        self.supports_batching = True
        self._device = next(reward_model.parameters()).device

    @torch.inference_mode()
    def score(self, ctx: StepContext) -> torch.Tensor:
        if self.shared_vocab:
            return self._score_ids(ctx)
        return self._score_text(ctx)

    def _score_text(self, ctx: StepContext) -> torch.Tensor:
        """Score by decoding `prefix + candidate` to text and re-encoding (mismatched-vocab path).

        The ids are decoded with `skip_special_tokens=True`, so a special-token candidate such as
        eos contributes nothing to the scored text.
        """
        batch_size = ctx.prefix_ids.size(0)
        num_candidates = ctx.candidate_ids.size(1)

        prefix = ctx.prefix_ids.unsqueeze(1).expand(-1, num_candidates, -1)  # [B, K, T]
        combined = torch.cat([prefix, ctx.candidate_ids.unsqueeze(-1)], dim=-1)  # [B, K, T+1]
        flat = combined.reshape(batch_size * num_candidates, -1)  # [B*K, T+1]
        texts = ctx.lm_tokenizer.decode(flat, skip_special_tokens=True)

        max_length = getattr(self.rm_tokenizer, "max_length", None)
        inputs = self.rm_tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=max_length,
        ).to(self._device)
        output = self.reward_model(**inputs)
        rewards = extract_score(output, self.score_index, self.score_transform)  # [B*K]
        return rewards.reshape(batch_size, num_candidates)

    def _score_ids(self, ctx: StepContext) -> torch.Tensor:
        """Score by feeding raw `prefix + candidate` ids to the reward model (shared-vocab path)."""
        batch_size = ctx.prefix_ids.size(0)
        num_candidates = ctx.candidate_ids.size(1)

        prefix_mask = full_prefix_mask(ctx.prefix_ids, ctx.attention_mask)  # [B, T]
        prefix = ctx.prefix_ids.unsqueeze(1).expand(-1, num_candidates, -1)  # [B, K, T]
        combined = torch.cat([prefix, ctx.candidate_ids.unsqueeze(-1)], dim=-1)  # [B, K, T+1]
        mask = torch.cat(
            [prefix_mask.unsqueeze(1).expand(-1, num_candidates, -1),
             torch.ones(batch_size, num_candidates, 1, device=prefix_mask.device, dtype=prefix_mask.dtype)],
            dim=-1,
        )  # [B, K, T+1]
        ids = combined.reshape(batch_size * num_candidates, -1).to(self._device)  # [B*K, T+1]
        mask = mask.reshape(batch_size * num_candidates, -1).to(self._device)

        max_length = getattr(self.rm_tokenizer, "max_length", None)
        if max_length is not None and ids.size(1) > max_length:
            ids = ids[:, -max_length:]
            mask = mask[:, -max_length:]

        if hasattr(self.reward_model, "score"):
            # the head is applied at the last position, which is the candidate for every row; the
            # classifier's own forward pools at the last id not equal to the configured pad id
            backbone = getattr(self.reward_model, self.reward_model.base_model_prefix)
            hidden = backbone(input_ids=ids, attention_mask=mask, return_dict=True).last_hidden_state
            output = self.reward_model.score(hidden[:, -1, :])
        else:
            output = self.reward_model(input_ids=ids, attention_mask=mask)
        rewards = extract_score(output, self.score_index, self.score_transform)  # [B*K]
        return rewards.reshape(batch_size, num_candidates)
_device = next(reward_model.parameters()).device instance-attribute
reward_model = reward_model instance-attribute
rm_tokenizer = rm_tokenizer instance-attribute
same_model_forwards = False class-attribute instance-attribute
score_index = score_index instance-attribute
score_transform = score_transform instance-attribute
scoring_cost = 'aux_forward' class-attribute instance-attribute
shared_vocab = shared_vocab instance-attribute
supports_batching = True instance-attribute
__init__(reward_model, rm_tokenizer, score_index=0, score_transform='none', shared_vocab=False)
Source code in steerability/algorithms/output_control/common/values/reward_model.py 
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __init__(
    self,
    reward_model: PreTrainedModel,
    rm_tokenizer: PreTrainedTokenizerBase,
    score_index: int = 0,
    score_transform: ScoreTransform = "none",
    shared_vocab: bool = False,
):
    self.reward_model = reward_model
    self.rm_tokenizer = rm_tokenizer
    self.score_index = score_index
    self.score_transform = score_transform
    self.shared_vocab = shared_vocab
    self.supports_batching = True
    self._device = next(reward_model.parameters()).device
_score_ids(ctx)

Score by feeding raw prefix + candidate ids to the reward model (shared-vocab path).

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
def _score_ids(self, ctx: StepContext) -> torch.Tensor:
    """Score by feeding raw `prefix + candidate` ids to the reward model (shared-vocab path)."""
    batch_size = ctx.prefix_ids.size(0)
    num_candidates = ctx.candidate_ids.size(1)

    prefix_mask = full_prefix_mask(ctx.prefix_ids, ctx.attention_mask)  # [B, T]
    prefix = ctx.prefix_ids.unsqueeze(1).expand(-1, num_candidates, -1)  # [B, K, T]
    combined = torch.cat([prefix, ctx.candidate_ids.unsqueeze(-1)], dim=-1)  # [B, K, T+1]
    mask = torch.cat(
        [prefix_mask.unsqueeze(1).expand(-1, num_candidates, -1),
         torch.ones(batch_size, num_candidates, 1, device=prefix_mask.device, dtype=prefix_mask.dtype)],
        dim=-1,
    )  # [B, K, T+1]
    ids = combined.reshape(batch_size * num_candidates, -1).to(self._device)  # [B*K, T+1]
    mask = mask.reshape(batch_size * num_candidates, -1).to(self._device)

    max_length = getattr(self.rm_tokenizer, "max_length", None)
    if max_length is not None and ids.size(1) > max_length:
        ids = ids[:, -max_length:]
        mask = mask[:, -max_length:]

    if hasattr(self.reward_model, "score"):
        # the head is applied at the last position, which is the candidate for every row; the
        # classifier's own forward pools at the last id not equal to the configured pad id
        backbone = getattr(self.reward_model, self.reward_model.base_model_prefix)
        hidden = backbone(input_ids=ids, attention_mask=mask, return_dict=True).last_hidden_state
        output = self.reward_model.score(hidden[:, -1, :])
    else:
        output = self.reward_model(input_ids=ids, attention_mask=mask)
    rewards = extract_score(output, self.score_index, self.score_transform)  # [B*K]
    return rewards.reshape(batch_size, num_candidates)
_score_text(ctx)

Score by decoding prefix + candidate to text and re-encoding (mismatched-vocab path).

The ids are decoded with skip_special_tokens=True, so a special-token candidate such as eos contributes nothing to the scored text.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def _score_text(self, ctx: StepContext) -> torch.Tensor:
    """Score by decoding `prefix + candidate` to text and re-encoding (mismatched-vocab path).

    The ids are decoded with `skip_special_tokens=True`, so a special-token candidate such as
    eos contributes nothing to the scored text.
    """
    batch_size = ctx.prefix_ids.size(0)
    num_candidates = ctx.candidate_ids.size(1)

    prefix = ctx.prefix_ids.unsqueeze(1).expand(-1, num_candidates, -1)  # [B, K, T]
    combined = torch.cat([prefix, ctx.candidate_ids.unsqueeze(-1)], dim=-1)  # [B, K, T+1]
    flat = combined.reshape(batch_size * num_candidates, -1)  # [B*K, T+1]
    texts = ctx.lm_tokenizer.decode(flat, skip_special_tokens=True)

    max_length = getattr(self.rm_tokenizer, "max_length", None)
    inputs = self.rm_tokenizer(
        texts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
    ).to(self._device)
    output = self.reward_model(**inputs)
    rewards = extract_score(output, self.score_index, self.score_transform)  # [B*K]
    return rewards.reshape(batch_size, num_candidates)
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/values/base.py 
71
72
73
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup, invoked from the owning control's steer().

Mirrors the input-control selector prepare() convention: a default no-op that the parent control calls during its own steer() when the value needs preparation.

Source code in steerability/algorithms/output_control/common/values/base.py 
63
64
65
66
67
68
69
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup, invoked from the owning control's `steer()`.

    Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
    control calls during its own `steer()` when the value needs preparation.
    """
    pass
score(ctx)

Return per-candidate values [B, K] aligned with ctx.candidate_ids.

Source code in steerability/algorithms/output_control/common/values/reward_model.py 
101
102
103
104
105
@torch.inference_mode()
def score(self, ctx: StepContext) -> torch.Tensor:
    if self.shared_vocab:
        return self._score_ids(ctx)
    return self._score_text(ctx)
extract_score(output, score_index, score_transform)

Read a per-row scalar score from a reward-model output.

Takes output.logits when present, else the raw output tensor [N, C]. Applies the transform over the class dimension, then selects column score_index:

- `"none"`: select the column directly (a raw, possibly unbounded logit).
- `"sigmoid"`: element-wise sigmoid, then select.
- `"softmax"`: softmax over the class dimension, then select.

Parameters:

Name Type Description Default
output Any

The reward model's output (a classifier output with logits, or a raw tensor).

 required
score_index int

Column of the class dimension read as the score.

 required
score_transform ScoreTransform

One of "none", "sigmoid", "softmax".

 required

Returns:

Type Description
Tensor

A tensor [N] of per-row scores.
Source code in steerability/algorithms/output_control/common/values/reward_model.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def extract_score(output: Any, score_index: int, score_transform: ScoreTransform) -> torch.Tensor:
    """Read a per-row scalar score from a reward-model output.

    Takes `output.logits` when present, else the raw output tensor `[N, C]`. Applies the transform
    over the class dimension, then selects column `score_index`:

        - `"none"`: select the column directly (a raw, possibly unbounded logit).
        - `"sigmoid"`: element-wise sigmoid, then select.
        - `"softmax"`: softmax over the class dimension, then select.

    Args:
        output: The reward model's output (a classifier output with `logits`, or a raw tensor).
        score_index: Column of the class dimension read as the score.
        score_transform: One of `"none"`, `"sigmoid"`, `"softmax"`.

    Returns:
        A tensor `[N]` of per-row scores.
    """
    logits = output.logits if hasattr(output, "logits") else output
    if score_transform == "sigmoid":
        logits = torch.sigmoid(logits)
    elif score_transform == "softmax":
        logits = torch.softmax(logits, dim=-1)
    return logits[:, score_index]

subspace_margin

Candidate value from a linear-probe margin, plus the single-file probe checkpoint loader.

For each candidate token, SubspaceMarginValue forwards prefix + candidate through the pipeline's own model via CandidateForward and reads the final-layer hidden state h at the raw output boundary; the value is the probe score h @ w + bias. This is the value underlying SASA.

SubspaceMarginValue

Bases: BaseCandidateValue

Per-candidate margin against a fitted single-layer Probe.

The margin of a candidate hidden state h is h @ w + bias, with w the probe's weight vector at its single layer and bias its calibrated offset (the documented Probe score). Hidden states are read by CandidateForward at the raw output boundary of the final decoder layer, so the probe is expected to be fitted at that boundary (location="layer_output", pooling="last", final decoder layer).

This value forwards the pipeline's own model per candidate, so it declares same_model_forwards=True; the forwards run inside auxiliary_pass(aligned=True) (via CandidateForward), so state-control transforms apply to them at the candidates' true positions while condition scoring and gates ignore them.

Parameters:

Name Type Description Default
probe Probe

The fitted single-layer Probe.

 required
Note

scoring_cost="model_forward" and supports_batching is False (the prefix cache tracks a single row).

Source code in steerability/algorithms/output_control/common/values/subspace_margin.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class SubspaceMarginValue(BaseCandidateValue):
    """Per-candidate margin against a fitted single-layer `Probe`.

    The margin of a candidate hidden state `h` is `h @ w + bias`, with `w` the probe's weight
    vector at its single layer and `bias` its calibrated offset (the documented `Probe` score).
    Hidden states are read by `CandidateForward` at the raw output boundary of the final decoder
    layer, so the probe is expected to be fitted at that boundary (`location="layer_output"`,
    `pooling="last"`, final decoder layer).

    This value forwards the pipeline's own model per candidate, so it declares
    `same_model_forwards=True`; the forwards run inside `auxiliary_pass(aligned=True)` (via
    `CandidateForward`), so state-control transforms apply to them at the candidates' true
    positions while condition scoring and gates ignore them.

    Args:
        probe: The fitted single-layer `Probe`.

    Note:
        `scoring_cost="model_forward"` and `supports_batching` is False (the prefix cache tracks a
        single row).
    """

    supports_batching: bool = False
    scoring_cost = "model_forward"
    same_model_forwards: bool = True

    def __init__(self, probe: Probe):
        self.probe = probe
        self._forward: CandidateForward | None = None
        self._aligned: tuple[torch.device, torch.dtype] | None = None
        self._weights: torch.Tensor | None = None

    def _align_probe(self, device: torch.device, dtype: torch.dtype) -> None:
        """Cache the probe weights aligned to (device, dtype); the probe is fixed per generation."""
        if self._aligned != (device, dtype):
            self._weights = self.probe.weights[self.probe.layer_ids[0]].to(device, dtype)
            self._aligned = (device, dtype)

    def score(self, ctx: StepContext) -> torch.Tensor:
        if ctx.model is None:
            raise RuntimeError("SubspaceMarginValue requires the pipeline model in StepContext.")
        if self._forward is None or self._forward.model is not ctx.model:
            self._forward = CandidateForward(ctx.model)
        hidden = self._forward.last_hidden_states(
            ctx.prefix_ids, ctx.candidate_ids, ctx.attention_mask
        )  # [K, H]
        self._align_probe(hidden.device, hidden.dtype)
        margins = hidden @ self._weights + self.probe.bias  # [K]
        return margins.unsqueeze(0)  # [1, K]
_aligned = None instance-attribute
_forward = None instance-attribute
_weights = None instance-attribute
probe = probe instance-attribute
same_model_forwards = True class-attribute instance-attribute
scoring_cost = 'model_forward' class-attribute instance-attribute
supports_batching = False class-attribute instance-attribute
__init__(probe)
Source code in steerability/algorithms/output_control/common/values/subspace_margin.py 
 98
 99
100
101
102
def __init__(self, probe: Probe):
    self.probe = probe
    self._forward: CandidateForward | None = None
    self._aligned: tuple[torch.device, torch.dtype] | None = None
    self._weights: torch.Tensor | None = None
_align_probe(device, dtype)

Cache the probe weights aligned to (device, dtype); the probe is fixed per generation.

Source code in steerability/algorithms/output_control/common/values/subspace_margin.py 
104
105
106
107
108
def _align_probe(self, device: torch.device, dtype: torch.dtype) -> None:
    """Cache the probe weights aligned to (device, dtype); the probe is fixed per generation."""
    if self._aligned != (device, dtype):
        self._weights = self.probe.weights[self.probe.layer_ids[0]].to(device, dtype)
        self._aligned = (device, dtype)
cleanup()

Release any resources allocated during prepare().

Source code in steerability/algorithms/output_control/common/values/base.py 
71
72
73
def cleanup(self) -> None:
    """Release any resources allocated during `prepare()`."""
    pass
prepare(model=None, tokenizer=None, **kwargs)

Optional offline setup, invoked from the owning control's steer().

Mirrors the input-control selector prepare() convention: a default no-op that the parent control calls during its own steer() when the value needs preparation.

Source code in steerability/algorithms/output_control/common/values/base.py 
63
64
65
66
67
68
69
def prepare(self, model=None, tokenizer=None, **kwargs) -> None:
    """Optional offline setup, invoked from the owning control's `steer()`.

    Mirrors the input-control selector `prepare()` convention: a default no-op that the parent
    control calls during its own `steer()` when the value needs preparation.
    """
    pass
score(ctx)

Return per-candidate values [B, K] aligned with ctx.candidate_ids.

Source code in steerability/algorithms/output_control/common/values/subspace_margin.py 
110
111
112
113
114
115
116
117
118
119
120
def score(self, ctx: StepContext) -> torch.Tensor:
    if ctx.model is None:
        raise RuntimeError("SubspaceMarginValue requires the pipeline model in StepContext.")
    if self._forward is None or self._forward.model is not ctx.model:
        self._forward = CandidateForward(ctx.model)
    hidden = self._forward.last_hidden_states(
        ctx.prefix_ids, ctx.candidate_ids, ctx.attention_mask
    )  # [K, H]
    self._align_probe(hidden.device, hidden.dtype)
    margins = hidden @ self._weights + self.probe.bias  # [K]
    return margins.unsqueeze(0)  # [1, K]
load_single_file_probe(file_path, layer_id)

Load a single-file probe checkpoint into a Probe.

Accepts two forms: a .probe JSON file with direction and midpoint lists, and a legacy {'wv', 'mu_mu'} torch tensor checkpoint. Both record a direction and a class-mean midpoint and no space metadata; they are assumed fitted at the raw output boundary of the final decoder layer over last-token features. The returned probe carries the direction as the weight vector of layer_id, bias = -direction . midpoint (so its score equals the margin direction . (h - midpoint)), location="layer_output", pooling="last", model_type="unknown", and empty meta (which disarms fingerprint checks).

Parameters:

Name Type Description Default
file_path str

Path to the checkpoint.

 required
layer_id int

Layer id the direction is registered under (the model's final decoder layer).

 required

Returns:

Type Description
Probe

The adapted Probe.

Raises:

Type Description
ValueError

If the checkpoint is not one of the supported forms.
Source code in steerability/algorithms/output_control/common/values/subspace_margin.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def load_single_file_probe(file_path: str, layer_id: int) -> Probe:
    """Load a single-file probe checkpoint into a `Probe`.

    Accepts two forms: a `.probe` JSON file with `direction` and `midpoint` lists, and a legacy
    `{'wv', 'mu_mu'}` torch tensor checkpoint. Both record a direction and a class-mean midpoint
    and no space metadata; they are assumed fitted at the raw output boundary of the final decoder
    layer over last-token features. The returned probe carries the direction as the weight vector
    of `layer_id`, `bias = -direction . midpoint` (so its score equals the margin
    `direction . (h - midpoint)`), `location="layer_output"`, `pooling="last"`,
    `model_type="unknown"`, and empty `meta` (which disarms fingerprint checks).

    Args:
        file_path: Path to the checkpoint.
        layer_id: Layer id the direction is registered under (the model's final decoder layer).

    Returns:
        The adapted `Probe`.

    Raises:
        ValueError: If the checkpoint is not one of the supported forms.
    """
    if str(file_path).endswith(".probe"):
        with open(file_path) as f:
            payload = json.load(f)
        direction = torch.tensor(payload["direction"], dtype=torch.float32)
        midpoint = torch.tensor(payload["midpoint"], dtype=torch.float32)
    else:
        try:
            loaded = torch.load(file_path, map_location="cpu", weights_only=True)
        except Exception as err:
            raise ValueError(
                f"Unrecognized probe checkpoint at {file_path!r}; expected a .probe JSON file or "
                "a legacy {'wv', 'mu_mu'} torch checkpoint."
            ) from err
        if not (isinstance(loaded, dict) and "wv" in loaded and "mu_mu" in loaded):
            raise ValueError(
                f"Unrecognized probe checkpoint at {file_path!r}; expected a .probe JSON file or "
                "a legacy {'wv', 'mu_mu'} torch checkpoint."
            )
        direction = loaded["wv"].float()
        midpoint = loaded["mu_mu"].float()

    direction = direction.reshape(-1)
    bias = -float(torch.dot(direction, midpoint.reshape(-1)))
    return Probe(
        model_type="unknown",
        location="layer_output",
        pooling="last",
        layer_ids=[layer_id],
        weights={layer_id: direction},
        bias=bias,
    )