Skip to content

Utils

steerability.utils

answers

Extracting a canonical numeric answer from a decoded model response.

extract_numeric_answer(text)

Extract the final numeric answer from a response and canonicalize it.

The last value anchored on an Answer: label or a \boxed{...} wrapper wins. Anchored values may be integers, decimals, plain fractions (2/3), or LaTeX fractions (\frac{2}{3}, \dfrac{2}{3}), optionally wrapped in $ or * markers. When no anchored value is present, the last number-like token in the text is used instead. Thousands separators are ignored.

The extracted value is canonicalized through fractions.Fraction, so equivalent forms map to one string (4/6, \frac{2}{3}, and 2/3 all yield "2/3"; exact decimals merge with their fraction, e.g. 0.5 yields "1/2"). This makes the result usable as an agreement key, e.g. as the answer_extractor of MajorityVoteScorer.

Parameters:

Name Type Description Default
text str

The decoded response to extract from.

 required

Returns:

Type Description
str

The canonical answer string, or the empty string when no value is found or the anchored
str

value does not parse (e.g. a zero denominator).
Source code in steerability/utils/answers.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
def extract_numeric_answer(text: str) -> str:
    """Extract the final numeric answer from a response and canonicalize it.

    The last value anchored on an `Answer:` label or a `\\boxed{...}` wrapper wins. Anchored values
    may be integers, decimals, plain fractions (`2/3`), or LaTeX fractions (`\\frac{2}{3}`,
    `\\dfrac{2}{3}`), optionally wrapped in `$` or `*` markers. When no anchored value is present,
    the last number-like token in the text is used instead. Thousands separators are ignored.

    The extracted value is canonicalized through `fractions.Fraction`, so equivalent forms map to
    one string (`4/6`, `\\frac{2}{3}`, and `2/3` all yield `"2/3"`; exact decimals merge with their
    fraction, e.g. `0.5` yields `"1/2"`). This makes the result usable as an agreement key, e.g. as
    the `answer_extractor` of `MajorityVoteScorer`.

    Args:
        text: The decoded response to extract from.

    Returns:
        The canonical answer string, or the empty string when no value is found or the anchored
        value does not parse (e.g. a zero denominator).
    """
    text = text.replace(",", "")

    matches = list(_ANCHORED_VALUE.finditer(text))
    if matches:
        frac_num, frac_den, slash_num, slash_den, plain = matches[-1].groups()
        try:
            if frac_num is not None:
                return str(Fraction(int(frac_num), int(frac_den)))
            if slash_num is not None:
                return str(Fraction(int(slash_num), int(slash_den)))
            return str(Fraction(plain))
        except (ZeroDivisionError, ValueError):
            return ""

    tokens = _BARE_VALUE.findall(text)
    if not tokens:
        return ""
    try:
        return str(Fraction(tokens[-1].replace(" ", "")))
    except (ZeroDivisionError, ValueError):
        return ""

model_utils

find_project_root(current_path)

Finds root dir by looking for pyproject.toml

Source code in steerability/utils/model_utils.py 
 4
 5
 6
 7
 8
 9
10
def find_project_root(current_path: Path) -> Path:
    """Finds root dir by looking for pyproject.toml"""
    while current_path.parent != current_path:
        if (current_path / 'pyproject.toml').exists():
            return current_path
        current_path = current_path.parent
    raise FileNotFoundError("no pyproject.toml found")

is_valid_model(config, model_id, service)

Source code in steerability/utils/model_utils.py 
13
14
15
16
17
18
def is_valid_model(config, model_id, service):
    model_config = config['model-config']
    return (
            model_id in model_config and
            service in model_config[model_id]['access']
    )

optional

Optional-dependency declarations and import guard.

Single source of truth mapping optional third-party modules to the extra that provides them. The registry consults it for skip messages; feature modules call require() at their import boundary to fail with an actionable hint.

OPTIONAL_MODULE_EXTRAS = {'mergekit': 'merging', 'inspect_ai': 'eval', 'inspect_evals': 'eval', 'vllm': 'vllm', 'vllm_hook_plugins': 'vllm', 'matplotlib': 'eval', 'seaborn': 'eval'} module-attribute

require(module_name)

Import module_name or raise ModuleNotFoundError naming the extra that provides it.

Parameters:

Name Type Description Default
module_name str

The importable module name (e.g., "mergekit" or "inspect_ai").

 required

Returns:

Type Description
ModuleType

The imported module.

Raises:

Type Description
ModuleNotFoundError

If the module is not installed. The message names the steerability[<extra>] extra when the top-level module is a declared optional dependency, otherwise it names the missing package directly. The raised error preserves the missing module's name, so the registry can classify it, and is an ImportError subclass, so callers guarding on ImportError keep working.
Source code in steerability/utils/optional.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
48
49
50
51
def require(module_name: str) -> ModuleType:
    """Import `module_name` or raise `ModuleNotFoundError` naming the extra that provides it.

    Args:
        module_name: The importable module name (e.g., ``"mergekit"`` or ``"inspect_ai"``).

    Returns:
        The imported module.

    Raises:
        ModuleNotFoundError: If the module is not installed. The message names the
            ``steerability[<extra>]`` extra when the top-level module is a declared
            optional dependency, otherwise it names the missing package directly. The
            raised error preserves the missing module's `name`, so the registry can
            classify it, and is an `ImportError` subclass, so callers guarding on
            `ImportError` keep working.
    """
    try:
        return import_module(module_name)
    except ModuleNotFoundError as exc:
        top_level = module_name.split(".")[0]
        extra = OPTIONAL_MODULE_EXTRAS.get(top_level)
        hint = (
            f'Install it via: pip install "steerability[{extra}]"'
            if extra
            else f"Install the '{top_level}' package to use this feature."
        )
        raise ModuleNotFoundError(
            f"'{module_name}' is required for this feature but is not installed. {hint}",
            name=exc.name,
        ) from exc

rendering

General-purpose, tokenizer-level rendering.

PromptFormat = Literal['raw', 'chat_completion', 'chat_prompt'] module-attribute

logger = logging.getLogger(__name__) module-attribute

encode_for_model(tokenizer, *, prompt=None, completion=None, messages=None, mode='chat_prompt', **tokenizer_kwargs)

Render then tokenize with the correct add_special_tokens.

Convenience for single-prompt call sites (e.g. judges).

A template is applied iff has_chat_template(tokenizer) and (for the prompt path) mode != "raw"; in that case the rendered string already contains the special tokens, so it is tokenized with add_special_tokens=False, otherwise True.

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

Tokenizer to render and tokenize with.

 required
prompt str | None

User-turn content for the render_for_model path.

 None
completion str | None

Continuation for the render_for_model path.

 None
messages list[dict[str, str]] | None

Chat messages for the render_messages path. Takes precedence over prompt/completion when provided.

 None
mode PromptFormat

Rendering policy for the render_for_model path (ignored when messages is given).

 'chat_prompt'

Returns:

Type Description
BatchEncoding

The tokenizer output (a BatchEncoding).
Source code in steerability/utils/rendering.py 
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
def encode_for_model(
    tokenizer: PreTrainedTokenizerBase,
    *,
    prompt: str | None = None,
    completion: str | None = None,
    messages: list[dict[str, str]] | None = None,
    mode: PromptFormat = "chat_prompt",
    **tokenizer_kwargs,
) -> BatchEncoding:
    """Render then tokenize with the correct `add_special_tokens`.

    Convenience for single-prompt call sites (e.g. judges).

    A template is applied iff `has_chat_template(tokenizer)` and
    (for the `prompt` path) `mode != "raw"`; in that case the
    rendered string already contains the special tokens, so it is
    tokenized with `add_special_tokens=False`, otherwise `True`.

    Args:
        tokenizer: Tokenizer to render and tokenize with.
        prompt: User-turn content for the `render_for_model` path.
        completion: Continuation for the `render_for_model` path.
        messages: Chat messages for the `render_messages` path. Takes precedence
            over `prompt`/`completion` when provided.
        mode: Rendering policy for the `render_for_model` path (ignored when
            `messages` is given).

    Returns:
        The tokenizer output (a `BatchEncoding`).
    """
    if messages is not None:
        text = render_messages(tokenizer, messages)
        template_applied = has_chat_template(tokenizer)
    else:
        text = render_for_model(tokenizer, prompt=prompt, completion=completion, mode=mode)
        template_applied = has_chat_template(tokenizer) and mode != "raw"
    return tokenizer(text, add_special_tokens=not template_applied, **tokenizer_kwargs)

has_chat_template(tokenizer)

Single source of truth for chat-template detection.

Source code in steerability/utils/rendering.py 
15
16
17
def has_chat_template(tokenizer: PreTrainedTokenizerBase) -> bool:
    """Single source of truth for chat-template detection."""
    return getattr(tokenizer, "chat_template", None) is not None

render_for_model(tokenizer, *, prompt=None, completion=None, mode='chat_completion')

Render a single example into the exact string the model will process.

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

Tokenizer whose chat template defines the rendering.

 required
prompt str | None

The user-turn content (or, for raw, the left part of the concatenation).

 None
completion str | None

The continuation appended after the generation prompt (or, for raw, the right part of the concatenation).

 None
mode PromptFormat

Rendering policy.

  • "raw": concatenate prompt + completion verbatim (no template). For base-model methods (ActAdd) and standalone statements (ITI).
  • "chat_completion": render prompt as a user turn (with generation prompt) and append completion as the continuation. The last content token is the completion's final token (correct for accumulate="last_token"). For prompt+answer contrastive pairs (CAA).
  • "chat_prompt": render prompt as a user turn with generation prompt and no completion, producing the prompt as inference renders it. For standalone-prompt contrasts (CAST condition) and for the inference path itself.
 'chat_completion'

Returns:

Type Description
str

The rendered string. Falls back to raw (with a warning) when chat
str

formatting is requested but the tokenizer has no chat template.
Note

The returned string already contains the template's special tokens (e.g. BOS). Callers MUST tokenize it with add_special_tokens=False.

Source code in steerability/utils/rendering.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
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def render_for_model(
    tokenizer: PreTrainedTokenizerBase,
    *,
    prompt: str | None = None,
    completion: str | None = None,
    mode: PromptFormat = "chat_completion",
) -> str:
    """Render a single example into the exact string the model will process.

    Args:
        tokenizer: Tokenizer whose chat template defines the rendering.
        prompt: The user-turn content (or, for `raw`, the left part of the
            concatenation).
        completion: The continuation appended after the generation prompt (or,
            for `raw`, the right part of the concatenation).
        mode: Rendering policy.

            - `"raw"`: concatenate `prompt + completion` verbatim (no template).
              For base-model methods (ActAdd) and standalone statements (ITI).
            - `"chat_completion"`: render `prompt` as a user turn (with generation
              prompt) and append `completion` as the continuation. The last
              content token is the completion's final token (correct for
              `accumulate="last_token"`). For prompt+answer contrastive pairs (CAA).
            - `"chat_prompt"`: render `prompt` as a user turn with generation
              prompt and no completion, producing the prompt as inference renders
              it. For standalone-prompt contrasts (CAST condition) and for the
              inference path itself.

    Returns:
        The rendered string. Falls back to `raw` (with a warning) when chat
        formatting is requested but the tokenizer has no chat template.

    Note:
        The returned string already contains the template's special tokens
        (e.g. BOS). Callers MUST tokenize it with `add_special_tokens=False`.
    """
    text = prompt or ""
    comp = completion or ""

    if mode == "raw":
        return text + comp

    if not has_chat_template(tokenizer):
        logger.warning(
            "render_for_model: mode=%r requested but tokenizer has no chat_template; "
            "falling back to raw concatenation.",
            mode,
        )
        return text + comp

    head = render_messages(tokenizer, [{"role": "user", "content": text}], add_generation_prompt=True)
    return head if mode == "chat_prompt" else head + comp

render_messages(tokenizer, messages, *, add_generation_prompt=True)

Render chat messages into a string via the tokenizer's chat template.

Wrapper around tokenizer.apply_chat_template(tokenize=False, ...).

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

Tokenizer whose chat template defines the rendering.

 required
messages list[dict[str, str]]

Chat messages as {"role": ..., "content": ...} dicts.

 required
add_generation_prompt bool

Whether to append the assistant generation prompt.

 True

Returns:

Type Description
str

The rendered string. When the tokenizer has no chat template, falls back
str

(with a warning) to joining message contents with blank lines.
Note

The returned string already contains the template's special tokens (e.g. BOS). Callers MUST tokenize it with add_special_tokens=False (see encode_for_model).

Source code in steerability/utils/rendering.py 
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
def render_messages(
    tokenizer: PreTrainedTokenizerBase,
    messages: list[dict[str, str]],
    *,
    add_generation_prompt: bool = True,
) -> str:
    """Render chat messages into a string via the tokenizer's chat template.

    Wrapper around `tokenizer.apply_chat_template(tokenize=False, ...)`.

    Args:
        tokenizer: Tokenizer whose chat template defines the rendering.
        messages: Chat messages as `{"role": ..., "content": ...}` dicts.
        add_generation_prompt: Whether to append the assistant generation prompt.

    Returns:
        The rendered string. When the tokenizer has no chat template, falls back
        (with a warning) to joining message contents with blank lines.

    Note:
        The returned string already contains the template's special tokens (e.g.
        BOS). Callers MUST tokenize it with `add_special_tokens=False` (see
        `encode_for_model`).
    """
    if not has_chat_template(tokenizer):
        logger.warning("render_messages: tokenizer has no chat_template; concatenating contents.")
        return "\n\n".join(m["content"] for m in messages)
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=add_generation_prompt,
    )

thinking

Splitting continuations from reasoning models into thinking and answer segments.

The split runs in one of two modes. Text mode is a substring split on the decoded continuation. Token mode is an id-level split on the continuation ids, decoding each segment afterwards; it exists because some tokenizers encode the reasoning delimiters as special tokens that skip_special_tokens=True strips before a substring split could see them. resolve_split_mode picks between the two for a given tokenizer and tag pair.

DEFAULT_THINK_TAGS = ('<think>', '</think>') module-attribute

ThinkingSplit

Bases: NamedTuple

The two segments of one continuation.

Attributes:

Name Type Description
thinking str | None

The reasoning segment, or None when no think tag is present.
answer str

The answer segment; the full text when no think tag is present, and the empty string when an opened thinking segment never closes.
Source code in steerability/utils/thinking.py 
16
17
18
19
20
21
22
23
24
25
26
class ThinkingSplit(NamedTuple):
    """The two segments of one continuation.

    Attributes:
        thinking: The reasoning segment, or None when no think tag is present.
        answer: The answer segment; the full text when no think tag is present, and the empty
            string when an opened thinking segment never closes.
    """

    thinking: str | None
    answer: str
answer instance-attribute
thinking instance-attribute

find_subsequence(haystack, needle, start=0)

Return the first index at or after start where needle occurs in haystack, or -1.

An empty needle never matches (returns -1).

Source code in steerability/utils/thinking.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def find_subsequence(haystack: Sequence[int], needle: Sequence[int], start: int = 0) -> int:
    """Return the first index at or after `start` where `needle` occurs in `haystack`, or -1.

    An empty `needle` never matches (returns -1).
    """
    needle = list(needle)
    span = len(needle)
    if span == 0:
        return -1
    last = len(haystack) - span
    for index in range(start, last + 1):
        if list(haystack[index:index + span]) == needle:
            return index
    return -1

resolve_split_mode(tokenizer, tags=DEFAULT_THINK_TAGS)

Resolve "auto" to "text" or "tokens" for a tokenizer and tag pair.

Each tag is encoded with add_special_tokens=False and the result decoded with skip_special_tokens=True. When both tags round-trip to their exact original strings, a substring split on the decoded continuation sees the delimiters, so the mode is "text". Otherwise at least one delimiter is stripped by skip_special_tokens=True (it is a special token) and the split must run on the ids, so the mode is "tokens".

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

The pipeline tokenizer whose vocabulary decides the delimiter treatment.

 required
tags tuple[str, str]

The (open_tag, close_tag) pair. Both entries must be non-empty strings.

 DEFAULT_THINK_TAGS

Returns:

Type Description
Literal['text', 'tokens']

"text" when both tags survive skip_special_tokens=True, "tokens" otherwise.

Raises:

Type Description
ValueError

If either tag entry is not a non-empty string, or a tag encodes to an empty id sequence (checked so "auto" never routes an unmatchable tag into token mode).
Source code in steerability/utils/thinking.py 
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
def resolve_split_mode(
    tokenizer: PreTrainedTokenizerBase,
    tags: tuple[str, str] = DEFAULT_THINK_TAGS,
) -> Literal["text", "tokens"]:
    """Resolve `"auto"` to `"text"` or `"tokens"` for a tokenizer and tag pair.

    Each tag is encoded with `add_special_tokens=False` and the result decoded with
    `skip_special_tokens=True`. When both tags round-trip to their exact original strings, a
    substring split on the decoded continuation sees the delimiters, so the mode is `"text"`.
    Otherwise at least one delimiter is stripped by `skip_special_tokens=True` (it is a special
    token) and the split must run on the ids, so the mode is `"tokens"`.

    Args:
        tokenizer: The pipeline tokenizer whose vocabulary decides the delimiter treatment.
        tags: The `(open_tag, close_tag)` pair. Both entries must be non-empty strings.

    Returns:
        `"text"` when both tags survive `skip_special_tokens=True`, `"tokens"` otherwise.

    Raises:
        ValueError: If either tag entry is not a non-empty string, or a tag encodes to an empty id
            sequence (checked so `"auto"` never routes an unmatchable tag into token mode).
    """
    open_tag, close_tag = tags
    if not (isinstance(open_tag, str) and open_tag) or not (isinstance(close_tag, str) and close_tag):
        raise ValueError("tags must be a pair of non-empty strings.")
    open_ids = _encode_tag(tokenizer, open_tag)
    close_ids = _encode_tag(tokenizer, close_tag)
    open_survives = tokenizer.decode(open_ids, skip_special_tokens=True) == open_tag
    close_survives = tokenizer.decode(close_ids, skip_special_tokens=True) == close_tag
    return "text" if open_survives and close_survives else "tokens"

split_thinking(text, tags=DEFAULT_THINK_TAGS, *, opened_at_start=False)

Split a decoded continuation into its thinking and answer segments by substring matching.

Matching is plain substring, case-sensitive. Let open_tag, close_tag = tags. The result depends on which tags are present and on opened_at_start:

  • Close tag present (open tag optional): the split is at the last occurrence of close_tag. thinking is everything before it, with one leading open_tag removed when the thinking segment starts with open_tag after leading whitespace. answer is everything after, left-stripped. The open tag is optional because thinking-mode chat templates commonly end the generation prompt with the open tag, so the continuation carries only the closing tag.
  • Open tag present, close tag absent (thinking truncated): thinking is everything after the first open_tag and answer is the empty string, since an unclosed thinking segment means no final answer was produced.
  • Neither tag present, opened_at_start=False: thinking is None and answer is the full text, so the split is a no-op for non-reasoning models.
  • Neither tag present, opened_at_start=True: the generation prompt opened the channel and it never closed, so thinking is the full text and answer is the empty string.

An empty thinking segment yields thinking == "" (not None), since a tag was present or the channel was opened by the prompt.

The split is at the last occurrence of close_tag, so any earlier close_tag occurrences fold into thinking and any open_tag/close_tag occurrences after the split stay in answer verbatim.

Parameters:

Name Type Description Default
text str

The decoded continuation to split.

 required
tags tuple[str, str]

The (open_tag, close_tag) pair. Both entries must be non-empty strings.

 DEFAULT_THINK_TAGS
opened_at_start bool

Whether the generation prompt already opened the reasoning channel, so a continuation carrying neither tag is treated as unclosed reasoning rather than a plain answer.

 False

Returns:

Type Description
ThinkingSplit

A ThinkingSplit with the thinking and answer segments.

Raises:

Type Description
ValueError

If either tag entry is not a non-empty string.
Source code in steerability/utils/thinking.py 
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
def split_thinking(
    text: str,
    tags: tuple[str, str] = DEFAULT_THINK_TAGS,
    *,
    opened_at_start: bool = False,
) -> ThinkingSplit:
    """Split a decoded continuation into its thinking and answer segments by substring matching.

    Matching is plain substring, case-sensitive. Let `open_tag, close_tag = tags`. The result
    depends on which tags are present and on `opened_at_start`:

    - Close tag present (open tag optional): the split is at the last occurrence of `close_tag`.
      `thinking` is everything before it, with one leading `open_tag` removed when the thinking
      segment starts with `open_tag` after leading whitespace. `answer` is everything after,
      left-stripped. The open tag is optional because thinking-mode chat templates commonly end
      the generation prompt with the open tag, so the continuation carries only the closing tag.
    - Open tag present, close tag absent (thinking truncated): `thinking` is everything after the
      first `open_tag` and `answer` is the empty string, since an unclosed thinking segment means
      no final answer was produced.
    - Neither tag present, `opened_at_start=False`: `thinking` is None and `answer` is the full
      text, so the split is a no-op for non-reasoning models.
    - Neither tag present, `opened_at_start=True`: the generation prompt opened the channel and it
      never closed, so `thinking` is the full text and `answer` is the empty string.

    An empty thinking segment yields `thinking == ""` (not None), since a tag was present or the
    channel was opened by the prompt.

    The split is at the last occurrence of `close_tag`, so any earlier `close_tag` occurrences fold
    into `thinking` and any `open_tag`/`close_tag` occurrences after the split stay in `answer`
    verbatim.

    Args:
        text: The decoded continuation to split.
        tags: The `(open_tag, close_tag)` pair. Both entries must be non-empty strings.
        opened_at_start: Whether the generation prompt already opened the reasoning channel, so a
            continuation carrying neither tag is treated as unclosed reasoning rather than a plain
            answer.

    Returns:
        A `ThinkingSplit` with the `thinking` and `answer` segments.

    Raises:
        ValueError: If either tag entry is not a non-empty string.
    """
    open_tag, close_tag = tags
    if not (isinstance(open_tag, str) and open_tag) or not (isinstance(close_tag, str) and close_tag):
        raise ValueError("tags must be a pair of non-empty strings.")

    if close_tag in text:
        before, _, after = text.rpartition(close_tag)
        thinking = before
        if thinking.lstrip().startswith(open_tag):
            thinking = thinking.lstrip()[len(open_tag):]
        return ThinkingSplit(thinking=thinking, answer=after.lstrip())

    if open_tag in text:
        thinking = text.split(open_tag, 1)[1]
        return ThinkingSplit(thinking=thinking, answer="")

    if opened_at_start:
        return ThinkingSplit(thinking=text, answer="")

    return ThinkingSplit(thinking=None, answer=text)

split_thinking_ids(output_ids, tokenizer, tags=DEFAULT_THINK_TAGS, *, opened_at_start=False)

Split a continuation into thinking and answer segments at the token-id level.

Each tag is the id sequence obtained by encoding it with add_special_tokens=False, and matching is first-occurrence subsequence search over output_ids. The delimiter ids belong to neither segment; each resulting segment is decoded with skip_special_tokens=True. Let open_ids, close_ids = tags encoded this way:

  • opened_at_start=False, open subsequence found: the reasoning starts after the open subsequence. Any ids preceding the open subsequence are prepended to the answer, since output emitted before the channel opened is not reasoning.
  • opened_at_start=False, no open subsequence but a close subsequence: the reasoning starts at position 0. The open tag is optional because thinking-mode chat templates commonly end the generation prompt with it, so the continuation carries only the close.
  • opened_at_start=True: the continuation begins inside the channel, so the reasoning starts at position 0 and no open subsequence is searched.
  • Once inside the channel, the split is at the first close subsequence: thinking is the ids up to it and answer is the ids after it. Close subsequences after the first stay in the answer verbatim.
  • Inside the channel with no close subsequence (truncation, or a caller stop string halted generation before the close was emitted): thinking is the rest of the ids and answer is the empty string.
  • opened_at_start=False and neither subsequence: thinking is None and answer is the full decoded continuation.

An empty thinking segment yields thinking == "" (not None), since the channel was entered.

A multi-token ordinary-text portion of a composed tag (for example the thought\n following a special open token) can tokenize context-dependently, so its id sequence in isolation need not be the id sequence it forms inside the continuation. When opened_at_start=True the open tag is never searched, so this affects the open tag only when the channel is opened within the continuation.

Parameters:

Name Type Description Default
output_ids Sequence[int]

The continuation token ids (one row, the prompt excluded). Trailing padding or eos ids may be present; delimiter matching is unaffected since those ids do not match a delimiter, and each segment is decoded with skip_special_tokens=True.

 required
tokenizer PreTrainedTokenizerBase

The tokenizer used to encode the tags and decode the segments.

 required
tags tuple[str, str]

The (open_tag, close_tag) pair. Both entries must be non-empty strings.

 DEFAULT_THINK_TAGS
opened_at_start bool

Whether the generation prompt already opened the reasoning channel. The flag matters only for a continuation carrying neither tag, which is classified as unclosed reasoning when set and as a plain answer otherwise.

 False

Returns:

Type Description
ThinkingSplit

A ThinkingSplit with the decoded thinking and answer segments.

Raises:

Type Description
ValueError

If either tag entry is not a non-empty string, or a tag encodes to an empty id sequence under this tokenizer.
Source code in steerability/utils/thinking.py 
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
def split_thinking_ids(
    output_ids: Sequence[int],
    tokenizer: PreTrainedTokenizerBase,
    tags: tuple[str, str] = DEFAULT_THINK_TAGS,
    *,
    opened_at_start: bool = False,
) -> ThinkingSplit:
    """Split a continuation into thinking and answer segments at the token-id level.

    Each tag is the id sequence obtained by encoding it with `add_special_tokens=False`, and
    matching is first-occurrence subsequence search over `output_ids`. The delimiter ids belong to
    neither segment; each resulting segment is decoded with `skip_special_tokens=True`. Let
    `open_ids, close_ids = tags` encoded this way:

    - `opened_at_start=False`, open subsequence found: the reasoning starts after the open
      subsequence. Any ids preceding the open subsequence are prepended to the answer, since output
      emitted before the channel opened is not reasoning.
    - `opened_at_start=False`, no open subsequence but a close subsequence: the reasoning starts
      at position 0. The open tag is optional because thinking-mode chat templates commonly end
      the generation prompt with it, so the continuation carries only the close.
    - `opened_at_start=True`: the continuation begins inside the channel, so the reasoning starts at
      position 0 and no open subsequence is searched.
    - Once inside the channel, the split is at the first close subsequence: `thinking` is the ids
      up to it and `answer` is the ids after it. Close subsequences after the first stay in the
      answer verbatim.
    - Inside the channel with no close subsequence (truncation, or a caller stop string halted
      generation before the close was emitted): `thinking` is the rest of the ids and `answer` is
      the empty string.
    - `opened_at_start=False` and neither subsequence: `thinking` is None and `answer` is the full
      decoded continuation.

    An empty thinking segment yields `thinking == ""` (not None), since the channel was entered.

    A multi-token ordinary-text portion of a composed tag (for example the `thought\\n` following a
    special open token) can tokenize context-dependently, so its id sequence in isolation need not
    be the id sequence it forms inside the continuation. When `opened_at_start=True` the open tag is
    never searched, so this affects the open tag only when the channel is opened within the
    continuation.

    Args:
        output_ids: The continuation token ids (one row, the prompt excluded). Trailing padding or
            eos ids may be present; delimiter matching is unaffected since those ids do not match a
            delimiter, and each segment is decoded with `skip_special_tokens=True`.
        tokenizer: The tokenizer used to encode the tags and decode the segments.
        tags: The `(open_tag, close_tag)` pair. Both entries must be non-empty strings.
        opened_at_start: Whether the generation prompt already opened the reasoning channel. The
            flag matters only for a continuation carrying neither tag, which is classified as
            unclosed reasoning when set and as a plain answer otherwise.

    Returns:
        A `ThinkingSplit` with the decoded `thinking` and `answer` segments.

    Raises:
        ValueError: If either tag entry is not a non-empty string, or a tag encodes to an empty id
            sequence under this tokenizer.
    """
    open_tag, close_tag = tags
    if not (isinstance(open_tag, str) and open_tag) or not (isinstance(close_tag, str) and close_tag):
        raise ValueError("tags must be a pair of non-empty strings.")

    ids = list(output_ids)
    close_ids = _encode_tag(tokenizer, close_tag)

    def decode(segment: list[int]) -> str:
        return tokenizer.decode(segment, skip_special_tokens=True)

    answer_prefix: list[int] = []
    if opened_at_start:
        reasoning_start = 0
    else:
        open_ids = _encode_tag(tokenizer, open_tag)
        open_at = find_subsequence(ids, open_ids)
        if open_at == -1:
            if find_subsequence(ids, close_ids) == -1:
                return ThinkingSplit(thinking=None, answer=decode(ids))
            # the open tag is optional: a generation prompt that opens the channel leaves only the
            # close subsequence in the continuation
            reasoning_start = 0
        else:
            answer_prefix = ids[:open_at]
            reasoning_start = open_at + len(open_ids)

    close_at = find_subsequence(ids, close_ids, start=reasoning_start)
    if close_at == -1:
        return ThinkingSplit(thinking=decode(ids[reasoning_start:]), answer="")

    thinking = decode(ids[reasoning_start:close_at])
    answer = decode(answer_prefix + ids[close_at + len(close_ids):])
    return ThinkingSplit(thinking=thinking, answer=answer)

tokenization

Tokenizer- and attention-mask-level hygiene: pad-token setup, token counting, mask inference, batch re-padding, and BOS sanity checks.

logger = logging.getLogger(__name__) module-attribute

count_tokens(tokenizer, text)

Count the tokens in text, without special tokens.

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

HuggingFace tokenizer instance.

 required
text str

The text to tokenize.

 required

Returns:

Type Description
int

The token count.
Source code in steerability/utils/tokenization.py 
28
29
30
31
32
33
34
35
36
37
38
def count_tokens(tokenizer: PreTrainedTokenizerBase, text: str) -> int:
    """Count the tokens in `text`, without special tokens.

    Args:
        tokenizer: HuggingFace tokenizer instance.
        text: The text to tokenize.

    Returns:
        The token count.
    """
    return len(tokenizer(text, add_special_tokens=False)["input_ids"])

ensure_pad_token(tokenizer)

Set pad token to eos token if not already defined.

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

HuggingFace tokenizer instance

 required

Returns:

Type Description
PreTrainedTokenizerBase

The same tokenizer with pad_token configured
Source code in steerability/utils/tokenization.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
def ensure_pad_token(tokenizer: PreTrainedTokenizerBase) -> PreTrainedTokenizerBase:
    """Set pad token to eos token if not already defined.

    Args:
       tokenizer: HuggingFace tokenizer instance

    Returns:
       The same tokenizer with pad_token configured
    """
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    return tokenizer

infer_attention_mask_from_ids(input_ids, pad_token_id)

Best-effort attention mask when the caller supplied none.

Marks only leading and trailing runs of pad_token_id as padding; interior occurrences are kept as real tokens. This matters when pad == eos (set by ensure_pad_token for tokenizers without a pad token): chat templates legitimately contain the eos mid-prompt (e.g. ChatML <|im_end|>), and a token-identity mask (ids != pad_id) would wrongly mask those positions.

Rules:

  • pad_token_id is None -> all ones.
  • a row that is entirely pad -> all ones (degenerate; safest).
  • otherwise ones on [first_nonpad, last_nonpad], zeros outside.

Parameters:

Name Type Description Default
input_ids Tensor

Token ids of shape [batch, seq_len].

 required
pad_token_id int | None

The tokenizer's pad token id, or None.

 required

Returns:

Type Description
Tensor

A long tensor of shape [batch, seq_len] on input_ids.device.
Source code in steerability/utils/tokenization.py 
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
def infer_attention_mask_from_ids(
        input_ids: torch.Tensor,
        pad_token_id: int | None,
) -> torch.Tensor:
    """Best-effort attention mask when the caller supplied none.

    Marks only leading and trailing runs of `pad_token_id` as padding; interior occurrences are kept
    as real tokens. This matters when pad == eos (set by `ensure_pad_token` for tokenizers without a
    pad token): chat templates legitimately contain the eos mid-prompt (e.g. ChatML `<|im_end|>`), and
    a token-identity mask (`ids != pad_id`) would wrongly mask those positions.

    Rules:

    - `pad_token_id` is None -> all ones.
    - a row that is entirely pad -> all ones (degenerate; safest).
    - otherwise ones on `[first_nonpad, last_nonpad]`, zeros outside.

    Args:
        input_ids: Token ids of shape [batch, seq_len].
        pad_token_id: The tokenizer's pad token id, or None.

    Returns:
        A long tensor of shape [batch, seq_len] on `input_ids.device`.
    """
    if pad_token_id is None:
        return torch.ones_like(input_ids, dtype=torch.long)

    _, seq_len = input_ids.shape
    nonpad = input_ids != pad_token_id
    positions = torch.arange(seq_len, device=input_ids.device)
    first = torch.where(nonpad, positions, seq_len).min(dim=1).values  # [batch]
    last = torch.where(nonpad, positions, -1).max(dim=1).values  # [batch]
    mask = (positions.unsqueeze(0) >= first.unsqueeze(1)) & (positions.unsqueeze(0) <= last.unsqueeze(1))
    mask[first == seq_len] = True  # all-pad rows -> all ones
    return mask.long()

to_left_pad(input_ids, attention_mask)

Rearrange a batch to left-padded layout for correct batched causal scoring.

Left-padding ensures that all sequences in the batch end at the same position, allowing a uniform logit slice after concatenation with reference tokens. Works correctly regardless of whether the input is right-padded, left-padded, or unpadded.

Parameters:

Name Type Description Default
input_ids Tensor

Input token IDs [batch, seq_len]

 required
attention_mask Tensor

Corresponding attention mask [batch, seq_len]

 required

Returns:

Type Description
tuple[Tensor, Tensor]

tuple[torch.Tensor, torch.Tensor]: Left-padded (input_ids, attention_mask)
Source code in steerability/utils/tokenization.py 
 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
def to_left_pad(
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Rearrange a batch to left-padded layout for correct batched causal scoring.

    Left-padding ensures that all sequences in the batch end at the same position, allowing a uniform logit slice
    after concatenation with reference tokens. Works correctly regardless of whether the input is right-padded,
    left-padded, or unpadded.

    Args:
        input_ids: Input token IDs [batch, seq_len]
        attention_mask: Corresponding attention mask [batch, seq_len]

    Returns:
        tuple[torch.Tensor, torch.Tensor]: Left-padded (input_ids, attention_mask)
    """
    batch_size, max_len = input_ids.shape
    seq_lens = attention_mask.sum(dim=1)

    left_ids = input_ids.clone()
    left_mask = torch.zeros_like(attention_mask)

    for i in range(batch_size):
        length = seq_lens[i]
        pad = max_len - length
        if pad > 0:
            real_tokens = input_ids[i][attention_mask[i].bool()]
            pad_tokens = input_ids[i][~attention_mask[i].bool()]
            left_ids[i] = torch.cat([pad_tokens, real_tokens])
        left_mask[i, max_len - length:] = 1

    return left_ids, left_mask

warn_if_duplicate_bos(input_ids, attention_mask, tokenizer, already_warned)

Warn once if any row starts with two BOS tokens; returns the updated warned-state.

The dominant cause is rendering a chat template to a string and re-tokenizing it with add_special_tokens=True; the template already contains the BOS.

Returns already_warned unchanged (no re-check) when already_warned is True, when tokenizer is None, when tokenizer.bos_token_id is None, or when seq_len < 2. Returns True after emitting the warning.

Parameters:

Name Type Description Default
input_ids Tensor

The (2D) prompt token ids about to be generated from.

 required
attention_mask Tensor

The matching attention mask (used to locate the first real token, so the check is left-padding safe).

 required
tokenizer PreTrainedTokenizerBase | None

The pipeline tokenizer (or None).

 required
already_warned bool

Whether the duplicate-BOS warning has already fired for this pipeline.

 required

Returns:

Type Description
bool

The updated warned-state.
Source code in steerability/utils/tokenization.py 
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
def warn_if_duplicate_bos(
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        tokenizer: PreTrainedTokenizerBase | None,
        already_warned: bool,
) -> bool:
    """Warn once if any row starts with two BOS tokens; returns the updated warned-state.

    The dominant cause is rendering a chat template to a string and re-tokenizing it with
    `add_special_tokens=True`; the template already contains the BOS.

    Returns `already_warned` unchanged (no re-check) when `already_warned` is True, when `tokenizer`
    is None, when `tokenizer.bos_token_id` is None, or when seq_len < 2. Returns True after emitting
    the warning.

    Args:
        input_ids: The (2D) prompt token ids about to be generated from.
        attention_mask: The matching attention mask (used to locate the first real token, so the
            check is left-padding safe).
        tokenizer: The pipeline tokenizer (or None).
        already_warned: Whether the duplicate-BOS warning has already fired for this pipeline.

    Returns:
        The updated warned-state.
    """
    if already_warned or tokenizer is None:
        return already_warned
    bos = getattr(tokenizer, "bos_token_id", None)
    if bos is None or input_ids.size(1) < 2:
        return already_warned
    first = attention_mask.long().argmax(dim=1)  # first real token (left-pad safe)
    rows = torch.arange(input_ids.size(0), device=input_ids.device)
    second = (first + 1).clamp(max=input_ids.size(1) - 1)
    dup = (
        (input_ids[rows, first] == bos)
        & (input_ids[rows, second] == bos)
        & (first + 1 < input_ids.size(1))
    )
    if bool(dup.any()):
        logger.warning(
            "Duplicate BOS detected at the start of the prompt (token id %d twice). "
            "Likely cause: chat-templated text re-tokenized with add_special_tokens=True. "
            "Tokenize with add_special_tokens=False, use "
            "steerability.utils.rendering.encode_for_model, or pass chat messages "
            "directly to generate(). Steering methods calibrated on single-BOS "
            "inputs will misbehave on double-BOS inputs.", bos,
        )
        return True
    return already_warned

verbosity

Opt-in verbosity controls for the steerability package logger.

The package attaches a logging.NullHandler to its root logger at import, so toolkit logging is silent by default and never emits "no handlers could be found" warnings. This module exposes the supported way to turn that logging up or down and to quiet known-noisy third-party libraries. None of these functions run at import time, and the library never calls them on its own; the package mutates no global logging state unless a caller asks.

The named module is verbosity rather than logging so it does not shadow the standard library logging module from inside the package.

get_verbosity()

Return the effective level of the steerability logger.

Applies the STEERABILITY_VERBOSITY environment default once (on first library use) before reading, so an environment-configured level is reflected without an explicit set_verbosity call.

Returns:

Type Description
int

The effective integer level, resolved up the logger hierarchy when the package logger has
int

no level of its own.
Source code in steerability/utils/verbosity.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def get_verbosity() -> int:
    """Return the effective level of the `steerability` logger.

    Applies the `STEERABILITY_VERBOSITY` environment default once (on first library use) before
    reading, so an environment-configured level is reflected without an explicit `set_verbosity`
    call.

    Returns:
        The effective integer level, resolved up the logger hierarchy when the package logger has
        no level of its own.
    """
    _apply_env_default()
    return logging.getLogger(_ROOT_LOGGER_NAME).getEffectiveLevel()

quiet_third_party()

Reduce output from known-noisy third-party libraries.

Sets the transformers logging facility to ERROR, disables Hugging Face Hub progress bars, and (when importable) sets the datasets logging facility to ERROR. Each step is guarded so a missing optional dependency is a no-op. This is opt-in and the library never calls it. It does not add any warnings filter and does not set TQDM_DISABLE, so user-owned progress bars and Python warnings are left alone.

Source code in steerability/utils/verbosity.py 
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
def quiet_third_party() -> None:
    """Reduce output from known-noisy third-party libraries.

    Sets the `transformers` logging facility to ERROR, disables Hugging Face Hub progress bars, and
    (when importable) sets the `datasets` logging facility to ERROR. Each step is guarded so a
    missing optional dependency is a no-op. This is opt-in and the library never calls it. It does
    not add any `warnings` filter and does not set `TQDM_DISABLE`, so user-owned progress bars and
    Python warnings are left alone.
    """
    try:
        from transformers.utils import logging as hf_logging

        hf_logging.set_verbosity_error()
    except ImportError:
        pass

    try:
        from huggingface_hub.utils import logging as hub_logging  # noqa: F401

        os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
        from huggingface_hub.utils import disable_progress_bars

        disable_progress_bars()
    except ImportError:
        pass

    try:
        from datasets.utils import logging as datasets_logging

        datasets_logging.set_verbosity_error()
    except ImportError:
        pass

set_verbosity(level)

Set the level of the steerability logger and attach one stream handler if none is attached.

Sets the level of the package root logger (steerability), so every module logger under it is affected. If the package root logger has no handler other than the import-time NullHandler, a single StreamHandler with a plain %(levelname)s %(name)s: %(message)s format is attached so records reach the console. The function is idempotent: a second call updates the level and does not attach a second handler.

Parameters:

Name Type Description Default
level int | str

One of "debug", "info", "warning", "error" (case-insensitive) or a logging level constant.

 required

Raises:

Type Description
ValueError

If a string level name is not recognized.
TypeError

If level is neither a string nor an integer.
Source code in steerability/utils/verbosity.py 
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
def set_verbosity(level: int | str) -> None:
    """Set the level of the `steerability` logger and attach one stream handler if none is attached.

    Sets the level of the package root logger (`steerability`), so every module logger under it is
    affected. If the package root logger has no handler other than the import-time `NullHandler`, a
    single `StreamHandler` with a plain `%(levelname)s %(name)s: %(message)s` format is attached so
    records reach the console. The function is idempotent: a second call updates the level and does
    not attach a second handler.

    Args:
        level: One of `"debug"`, `"info"`, `"warning"`, `"error"` (case-insensitive) or a
            `logging` level constant.

    Raises:
        ValueError: If a string level name is not recognized.
        TypeError: If `level` is neither a string nor an integer.
    """
    global _env_default_applied
    _env_default_applied = True  # an explicit call supersedes the env default
    logger = logging.getLogger(_ROOT_LOGGER_NAME)
    logger.setLevel(_coerce_level(level))
    if not _has_real_handler(logger):
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter(_HANDLER_FORMAT))
        logger.addHandler(handler)