Skip to content

SASA

steerability.algorithms.output_control.sasa

args

control

logger = logging.getLogger(__name__) module-attribute

SASA

Bases: OutputControl

Implementation of SASA (Self-disciplined autoregressive sampling) from Ko et al., 2024.

SASA steers generation toward a target attribute defined by labeled examples. It works in two phases:

  1. Subspace learning: From labeled positive (desired) and negative (undesired) examples, it fits a linear classifier in the model's own final-layer space. The classifier's weight vector defines a subspace separating the two attribute classes. Data is passed via gen_wv_data as unpaired classes (a {'pos', 'neg'} dict or a LabeledExamples) or as paired prompt/response data (a dict carrying a 'prompts' key or a ContrastivePairs), and prompt_format selects how each example is rendered before capture. prompt_format='chat_completion' renders each pair as a user turn plus the response and requires the paired form (the paper's {prompt, response, annotation} format); 'raw' and 'chat_prompt' accept unpaired classes.

  2. Controlled decoding: At every decoding step the candidate-token logits are shifted by beta * margin, where margin is the classifier distance of the updated context from the undesired side of the subspace. Sampling from the softmax of the adjusted logits nudges generation toward the desired attribute while staying close to the original distribution.

Any binary-labeled attribute works, since the attribute is defined by the labels on gen_wv_data. A previously fitted probe is loaded via wv_path.

SASA is a step-level control. steer() fits (or loads) a Probe, and get_logits_processors() returns a ValueGuidedProcessor whose per-candidate value is the subspace margin, obtained via a single same-model forward per step. candidate_policy selects which tokens are scored: 'surviving' (every token earlier processors left finite), 'top_p' (the nucleus of the raw logits, the paper's setting), or 'top_k'. The margins are softmax-normalized over the candidate set and added (scaled by beta) to the candidate logits with no non-candidate masking. max_candidates clamps the set on top of any policy to bound the per-step forward. As a step-level control, SASA composes with other output controls and with a decoding driver.

A value_trace list passed via runtime_kwargs receives one ValueStepRecord per scored step (the candidate ids, their pre-shift scores, the raw margins, and the softmax-normalized shift), for inspecting the per-step redistribution.

include_in_scoring defaults to False, since scoring under SASA costs a K-candidate model forward per reference position; opt in explicitly if needed.

Reference:

  • "Large Language Models can Become Strong Self-Detoxifiers" Ching-Yun Ko, Pin-Yu Chen, Payel Das, Youssef Mroueh, Soham Dan, Georgios Kollias, Subhajit Chaudhury, Tejaswini Pedapati, Luca Daniel https://arxiv.org/abs/2410.03818
Source code in steerability/algorithms/output_control/sasa/control.py 
 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
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
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
class SASA(OutputControl):
    """Implementation of SASA (Self-disciplined autoregressive sampling) from Ko et al., 2024.

    SASA steers generation toward a target attribute defined by labeled examples. It works in two phases:

    1. **Subspace learning**: From labeled positive (desired) and negative (undesired) examples, it fits a linear
    classifier in the model's own final-layer space. The classifier's weight vector defines a subspace separating
    the two attribute classes. Data is passed via `gen_wv_data` as unpaired classes (a `{'pos', 'neg'}` dict or a
    `LabeledExamples`) or as paired prompt/response data (a dict carrying a `'prompts'` key or a `ContrastivePairs`),
    and `prompt_format` selects how each example is rendered before capture. `prompt_format='chat_completion'`
    renders each pair as a user turn plus the response and requires the paired form (the paper's
    `{prompt, response, annotation}` format); `'raw'` and `'chat_prompt'` accept unpaired classes.

    2. **Controlled decoding**: At every decoding step the candidate-token logits are shifted by `beta * margin`,
    where `margin` is the classifier distance of the updated context from the undesired side of the subspace.
    Sampling from the softmax of the adjusted logits nudges generation toward the desired attribute while staying
    close to the original distribution.

    Any binary-labeled attribute works, since the attribute is defined by the labels on `gen_wv_data`. A previously
    fitted probe is loaded via `wv_path`.

    SASA is a step-level control. `steer()` fits (or loads) a `Probe`, and `get_logits_processors()` returns a
    `ValueGuidedProcessor` whose per-candidate value is the subspace margin, obtained via a single same-model
    forward per step. `candidate_policy` selects which tokens are scored: `'surviving'` (every token earlier
    processors left finite), `'top_p'` (the nucleus of the raw logits, the paper's setting), or `'top_k'`. The
    margins are softmax-normalized over the candidate set and added (scaled by `beta`) to the candidate logits with
    no non-candidate masking. `max_candidates` clamps the set on top of any policy to bound the per-step forward. As
    a step-level control, SASA composes with other output controls and with a decoding driver.

    A `value_trace` list passed via `runtime_kwargs` receives one `ValueStepRecord` per scored step (the candidate
    ids, their pre-shift scores, the raw margins, and the softmax-normalized shift), for inspecting the per-step
    redistribution.

    `include_in_scoring` defaults to False, since scoring under SASA costs a K-candidate model forward per
    reference position; opt in explicitly if needed.

    Reference:

    - "Large Language Models can Become Strong Self-Detoxifiers"
      Ching-Yun Ko, Pin-Yu Chen, Payel Das, Youssef Mroueh, Soham Dan, Georgios Kollias, Subhajit Chaudhury,
      Tejaswini Pedapati, Luca Daniel
      [https://arxiv.org/abs/2410.03818](https://arxiv.org/abs/2410.03818)
    """
    Args = SASAArgs

    RUNTIME_KWARGS_SCHEMA: list[dict] = [
        {
            "name": "value_trace",
            "type": "list",
            "scope": "call",
            "help": "A caller-owned list that receives one ValueStepRecord per scored step of this call.",
        },
    ]

    include_in_scoring: bool = False
    same_model_forwards: bool = True

    # placeholders (filled by steer)
    model: PreTrainedModel | None = None
    tokenizer: PreTrainedTokenizerBase | None = None
    probe: Probe | None = None

    beta: float

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.MODULE`; the probe fits on the live model, which is retained for the
        per-step value forwards (the generate phase is in-process)."""
        return ModelAccess.MODULE

    def export_state(self) -> dict:
        """The fitted value-subspace probe under the `"probe"` key (after `steer()`)."""
        return {"probe": self.probe} if self.probe is not None else {}

    def frozen_form(self, state: dict) -> tuple[str, dict]:
        """A same-class frozen form: `wv_path` points at the exported probe directory and the
        fit-only `gen_wv_*` fields are dropped. The decode-time settings (`candidate_policy`,
        `top_p`, `top_k`, `max_candidates`) are carried forward."""
        from steerability.spipe.codec import AsPath

        return "output_control/sasa", {
            "beta": self.beta,
            "wv_path": AsPath(state["probe"]),
            "gen_wv_data": None,
            "candidate_policy": self.candidate_policy,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "max_candidates": self.max_candidates,
        }

    def fit_identity(self):
        """The probe-fitting inputs (`gen_wv_*` fields and `prompt_format`), or None when a saved
        probe is loaded."""
        if getattr(self, "wv_path", None):
            return None
        return {
            "gen_wv_data": self.gen_wv_data,
            "prompt_format": self.prompt_format,
            "gen_wv_length": self.gen_wv_length,
            "gen_wv_batch_size": self.gen_wv_batch_size,
        }

    def steer(
            self,
            model: PreTrainedModel,
            tokenizer: PreTrainedTokenizerBase | None = None,
            **__,
    ) -> PreTrainedModel:
        """Load or fit the linear probe defining the attribute subspace.

        A `wv_path` naming a directory loads a saved `Probe` artifact; a `.probe` JSON file or a
        legacy `{'wv', 'mu_mu'}` tensor checkpoint is adapted into a `Probe` over the final
        decoder layer. Without `wv_path`, a probe is fitted on `gen_wv_data` (fisher direction over
        last-token features at the raw final-layer boundary, rendered per `prompt_format`, midpoint
        calibration). The probe's recorded space is validated against the boundary the margins are
        evaluated at.

        Args:
            model (PreTrainedModel): The base language model to be steered.
            tokenizer (PreTrainedTokenizerBase | None): Tokenizer for the base model.

        Returns:
            PreTrainedModel: The input model (unchanged).

        Raises:
            ValueError: If neither `wv_path` nor `gen_wv_data` is set, or a loaded probe's
                `location`, `pooling`, or layer ids do not match last-token features at the raw
                output boundary of the final decoder layer.
        """
        self.model = model
        self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
        if self.tokenizer.pad_token_id is None:
            if self.tokenizer.eos_token_id is not None:
                self.tokenizer = ensure_pad_token(self.tokenizer)
            else:
                self.tokenizer.add_special_tokens({"pad_token": "<pad>"})

        final_layer = resolve_model_layout(model).num_layers - 1
        if getattr(self, "wv_path", None):
            logger.info("Loading SASA probe.")
            if os.path.isdir(self.wv_path):
                self.probe = Probe.load(self.wv_path)
            else:
                self.probe = load_single_file_probe(self.wv_path, layer_id=final_layer)
        else:
            if self.gen_wv_data is None:
                raise ValueError("SASA.steer() requires either gen_wv_data (to fit a probe) or wv_path (to load one).")
            logger.info("Fitting SASA probe.")
            data = self._resolve_fit_data()
            spec = ProbeFitSpec(
                method="fisher",
                pooling="last",
                location="layer_output",
                prompt_format=self.prompt_format,
                candidate_layers=[final_layer],
                calibration="midpoint",
            )
            self.probe = fit_probe(
                model,
                self.tokenizer,
                data=data,
                spec=spec,
                batch_size=self.gen_wv_batch_size,
                max_length=1024,
            )
        _validate_probe_space(self.probe, final_layer)
        return model

    def _resolve_fit_data(self) -> LabeledExamples | ContrastivePairs:
        """Normalize `gen_wv_data` into a fit-data container and truncate each class to `gen_wv_length`.

        A `{'pos', 'neg'}` dict becomes `LabeledExamples`; a dict carrying a `'prompts'` key becomes
        `ContrastivePairs`; a `LabeledExamples` or `ContrastivePairs` instance passes through. When
        `0 < gen_wv_length`, each class (and the aligned prompts for paired data) is truncated to at
        most `gen_wv_length` entries.
        """
        data = self.gen_wv_data
        if isinstance(data, dict):
            if "prompts" in data:
                data = ContrastivePairs(
                    positives=data["pos"], negatives=data["neg"], prompts=data["prompts"]
                )
            else:
                data = LabeledExamples(positives=data["pos"], negatives=data["neg"])

        limit = self.gen_wv_length
        if limit is None or limit <= 0:
            return data

        if isinstance(data, ContrastivePairs):
            prompts = None if data.prompts is None else list(data.prompts[:limit])
            return ContrastivePairs(
                positives=list(data.positives[:limit]),
                negatives=list(data.negatives[:limit]),
                prompts=prompts,
            )
        return LabeledExamples(
            positives=list(data.positives[:limit]), negatives=list(data.negatives[:limit])
        )

    def get_logits_processors(self, input_ids, runtime_kwargs, attention_mask=None, **kwargs) -> list:
        """Return a fresh `ValueGuidedProcessor` implementing SASA's margin-based shift.

        The candidate set follows `candidate_policy`: `surviving` (every token left finite by earlier
        processors), `top_p` (the nucleus of the raw logits, the paper's setting), or `top_k`. Margins
        are softmax-normalized over the candidate set and added (scaled by `beta`) with no
        non-candidate masking. `max_candidates` clamps the set on top of the policy to bound the
        per-step model forward. A `value_trace` list in `runtime_kwargs` receives one `ValueStepRecord`
        per step.
        """
        if self.probe is None:
            raise RuntimeError("SASA.steer() must run before generation (probe not fitted/loaded).")
        return [
            ValueGuidedProcessor(
                SubspaceMarginValue(self.probe),
                policy=self.candidate_policy,
                p=self.top_p,
                k=self.top_k,
                beta=self.beta,
                normalize="softmax",
                mask_non_candidates=False,
                max_candidates=self.max_candidates,
                lm_tokenizer=self.tokenizer,
                model=self.model,
                attention_mask=attention_mask,
                trace=(runtime_kwargs or {}).get("value_trace"),
            )
        ]

    def cleanup(self) -> None:
        """Release the probe and model references."""
        self.probe = None
        self.model = None
        self.tokenizer = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        logger.debug("SASA cleanup completed")
RUNTIME_KWARGS_SCHEMA = [{'name': 'value_trace', 'type': 'list', 'scope': 'call', 'help': 'A caller-owned list that receives one ValueStepRecord per scored step of this call.'}] class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
beta instance-attribute
enabled = True class-attribute instance-attribute
include_in_scoring = False class-attribute instance-attribute
model = None class-attribute instance-attribute
probe = None class-attribute instance-attribute
same_model_forwards = True class-attribute instance-attribute
supports_batching = False class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
cleanup()

Release the probe and model references.

Source code in steerability/algorithms/output_control/sasa/control.py 
279
280
281
282
283
284
285
286
287
def cleanup(self) -> None:
    """Release the probe and model references."""
    self.probe = None
    self.model = None
    self.tokenizer = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    logger.debug("SASA cleanup completed")
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
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()

The fitted value-subspace probe under the "probe" key (after steer()).

Source code in steerability/algorithms/output_control/sasa/control.py 
120
121
122
def export_state(self) -> dict:
    """The fitted value-subspace probe under the `"probe"` key (after `steer()`)."""
    return {"probe": self.probe} if self.probe is not None else {}
fit_identity()

The probe-fitting inputs (gen_wv_* fields and prompt_format), or None when a saved probe is loaded.

Source code in steerability/algorithms/output_control/sasa/control.py 
140
141
142
143
144
145
146
147
148
149
150
def fit_identity(self):
    """The probe-fitting inputs (`gen_wv_*` fields and `prompt_format`), or None when a saved
    probe is loaded."""
    if getattr(self, "wv_path", None):
        return None
    return {
        "gen_wv_data": self.gen_wv_data,
        "prompt_format": self.prompt_format,
        "gen_wv_length": self.gen_wv_length,
        "gen_wv_batch_size": self.gen_wv_batch_size,
    }
frozen_form(state)

A same-class frozen form: wv_path points at the exported probe directory and the fit-only gen_wv_* fields are dropped. The decode-time settings (candidate_policy, top_p, top_k, max_candidates) are carried forward.

Source code in steerability/algorithms/output_control/sasa/control.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class frozen form: `wv_path` points at the exported probe directory and the
    fit-only `gen_wv_*` fields are dropped. The decode-time settings (`candidate_policy`,
    `top_p`, `top_k`, `max_candidates`) are carried forward."""
    from steerability.spipe.codec import AsPath

    return "output_control/sasa", {
        "beta": self.beta,
        "wv_path": AsPath(state["probe"]),
        "gen_wv_data": None,
        "candidate_policy": self.candidate_policy,
        "top_p": self.top_p,
        "top_k": self.top_k,
        "max_candidates": self.max_candidates,
    }
get_logits_processors(input_ids, runtime_kwargs, attention_mask=None, **kwargs)

Return a fresh ValueGuidedProcessor implementing SASA's margin-based shift.

The candidate set follows candidate_policy: surviving (every token left finite by earlier processors), top_p (the nucleus of the raw logits, the paper's setting), or top_k. Margins are softmax-normalized over the candidate set and added (scaled by beta) with no non-candidate masking. max_candidates clamps the set on top of the policy to bound the per-step model forward. A value_trace list in runtime_kwargs receives one ValueStepRecord per step.

Source code in steerability/algorithms/output_control/sasa/control.py 
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
def get_logits_processors(self, input_ids, runtime_kwargs, attention_mask=None, **kwargs) -> list:
    """Return a fresh `ValueGuidedProcessor` implementing SASA's margin-based shift.

    The candidate set follows `candidate_policy`: `surviving` (every token left finite by earlier
    processors), `top_p` (the nucleus of the raw logits, the paper's setting), or `top_k`. Margins
    are softmax-normalized over the candidate set and added (scaled by `beta`) with no
    non-candidate masking. `max_candidates` clamps the set on top of the policy to bound the
    per-step model forward. A `value_trace` list in `runtime_kwargs` receives one `ValueStepRecord`
    per step.
    """
    if self.probe is None:
        raise RuntimeError("SASA.steer() must run before generation (probe not fitted/loaded).")
    return [
        ValueGuidedProcessor(
            SubspaceMarginValue(self.probe),
            policy=self.candidate_policy,
            p=self.top_p,
            k=self.top_k,
            beta=self.beta,
            normalize="softmax",
            mask_non_candidates=False,
            max_candidates=self.max_candidates,
            lm_tokenizer=self.tokenizer,
            model=self.model,
            attention_mask=attention_mask,
            trace=(runtime_kwargs or {}).get("value_trace"),
        )
    ]
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 []
requirements()

Backend requirements computed from this instance's configuration, per phase.

The default requires Capability.IN_PROCESS_TORCH at generate and, when include_in_scoring is True, at score as well, since remote prompt-logprob computation applies neither live processors nor engine-registered sampling processors to prefill logits. Setting include_in_scoring=False removes the score-phase requirement.

Returns:

Type Description
Requirements

The control's phase-keyed requirements.
Source code in steerability/algorithms/output_control/base.py 
214
215
216
217
218
219
220
221
222
223
224
225
226
def requirements(self) -> Requirements:
    """Backend requirements computed from this instance's configuration, per phase.

    The default requires `Capability.IN_PROCESS_TORCH` at generate and, when
    `include_in_scoring` is True, at score as well, since remote prompt-logprob computation
    applies neither live processors nor engine-registered sampling processors to prefill
    logits. Setting `include_in_scoring=False` removes the score-phase requirement.

    Returns:
        The control's phase-keyed requirements.
    """
    score = needs(Capability.IN_PROCESS_TORCH) if self.include_in_scoring else ()
    return Requirements(generate=needs(Capability.IN_PROCESS_TORCH), score=score)
steer(model, tokenizer=None, **__)

Load or fit the linear probe defining the attribute subspace.

A wv_path naming a directory loads a saved Probe artifact; a .probe JSON file or a legacy {'wv', 'mu_mu'} tensor checkpoint is adapted into a Probe over the final decoder layer. Without wv_path, a probe is fitted on gen_wv_data (fisher direction over last-token features at the raw final-layer boundary, rendered per prompt_format, midpoint calibration). The probe's recorded space is validated against the boundary the margins are evaluated at.

Parameters:

Name Type Description Default
model PreTrainedModel

The base language model to be steered.

 required
tokenizer PreTrainedTokenizerBase | None

Tokenizer for the base model.

 None

Returns:

Name Type Description
PreTrainedModel PreTrainedModel

The input model (unchanged).

Raises:

Type Description
ValueError

If neither wv_path nor gen_wv_data is set, or a loaded probe's location, pooling, or layer ids do not match last-token features at the raw output boundary of the final decoder layer.
Source code in steerability/algorithms/output_control/sasa/control.py 
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
def steer(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase | None = None,
        **__,
) -> PreTrainedModel:
    """Load or fit the linear probe defining the attribute subspace.

    A `wv_path` naming a directory loads a saved `Probe` artifact; a `.probe` JSON file or a
    legacy `{'wv', 'mu_mu'}` tensor checkpoint is adapted into a `Probe` over the final
    decoder layer. Without `wv_path`, a probe is fitted on `gen_wv_data` (fisher direction over
    last-token features at the raw final-layer boundary, rendered per `prompt_format`, midpoint
    calibration). The probe's recorded space is validated against the boundary the margins are
    evaluated at.

    Args:
        model (PreTrainedModel): The base language model to be steered.
        tokenizer (PreTrainedTokenizerBase | None): Tokenizer for the base model.

    Returns:
        PreTrainedModel: The input model (unchanged).

    Raises:
        ValueError: If neither `wv_path` nor `gen_wv_data` is set, or a loaded probe's
            `location`, `pooling`, or layer ids do not match last-token features at the raw
            output boundary of the final decoder layer.
    """
    self.model = model
    self.tokenizer = tokenizer or getattr(model, "tokenizer", None)
    if self.tokenizer.pad_token_id is None:
        if self.tokenizer.eos_token_id is not None:
            self.tokenizer = ensure_pad_token(self.tokenizer)
        else:
            self.tokenizer.add_special_tokens({"pad_token": "<pad>"})

    final_layer = resolve_model_layout(model).num_layers - 1
    if getattr(self, "wv_path", None):
        logger.info("Loading SASA probe.")
        if os.path.isdir(self.wv_path):
            self.probe = Probe.load(self.wv_path)
        else:
            self.probe = load_single_file_probe(self.wv_path, layer_id=final_layer)
    else:
        if self.gen_wv_data is None:
            raise ValueError("SASA.steer() requires either gen_wv_data (to fit a probe) or wv_path (to load one).")
        logger.info("Fitting SASA probe.")
        data = self._resolve_fit_data()
        spec = ProbeFitSpec(
            method="fisher",
            pooling="last",
            location="layer_output",
            prompt_format=self.prompt_format,
            candidate_layers=[final_layer],
            calibration="midpoint",
        )
        self.probe = fit_probe(
            model,
            self.tokenizer,
            data=data,
            spec=spec,
            batch_size=self.gen_wv_batch_size,
            max_length=1024,
        )
    _validate_probe_space(self.probe, final_layer)
    return model
steer_access()

ModelAccess.MODULE; the probe fits on the live model, which is retained for the per-step value forwards (the generate phase is in-process).

Source code in steerability/algorithms/output_control/sasa/control.py 
115
116
117
118
def steer_access(self) -> ModelAccess:
    """`ModelAccess.MODULE`; the probe fits on the live model, which is retained for the
    per-step value forwards (the generate phase is in-process)."""
    return ModelAccess.MODULE
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 ()