Skip to content

Common library

steerability.algorithms.input_control.common

Reusable building blocks for input controls.

This package holds components whose interface is self-explanatory and shared across unrelated methods. Method-specific procedures stay in each method's own utils/ directory.

__all__ = ['RolloutBudget', 'ParetoFrontier', 'generate_with_system_prompt'] module-attribute

ParetoFrontier

Pareto-frontier utility over a [n_candidates, n_instances] score matrix.

Maximizes scores by default; pass minimize=True for the opposite. Scores are handled internally as a maximization problem.

Source code in steerability/algorithms/input_control/common/pareto.py 
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
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 ParetoFrontier:
    """Pareto-frontier utility over a `[n_candidates, n_instances]` score matrix.

    Maximizes scores by default; pass `minimize=True` for the opposite. Scores are handled internally as a
    maximization problem.
    """

    def __init__(self, scores: np.ndarray, minimize: bool = False) -> None:
        scores = np.asarray(scores, dtype=float)
        if scores.ndim != 2:
            raise ValueError(f"scores must be 2-D [n_candidates, n_instances]; got shape {scores.shape}.")
        if scores.size == 0:
            raise ValueError("scores must be non-empty.")
        self._scores = -scores if minimize else scores
        self._minimize = minimize

    @property
    def per_instance_best(self) -> np.ndarray:
        """Per-instance maximum score across candidates. Shape `[n_instances]`. Returned in the original
        sign convention (i.e. minimum, when `minimize=True`)."""
        best = self._scores.max(axis=0)
        return -best if self._minimize else best

    @property
    def per_instance_winners(self) -> list[set[int]]:
        """For each instance, the set of candidate indices that achieve the per-instance best.

        Multiple candidates may tie.
        """
        best = self._scores.max(axis=0, keepdims=True)
        wins = self._scores == best  # [n_candidates, n_instances]
        return [set(np.where(wins[:, j])[0].tolist()) for j in range(wins.shape[1])]

    @property
    def non_dominated(self) -> list[int]:
        """Candidate indices that are not strictly dominated by any other candidate.

        Candidate `a` is dominated by `b` iff `b >= a` on every instance and `b > a` on at least one.
        """
        n = self._scores.shape[0]
        keep: list[int] = []
        for i in range(n):
            dominated = False
            for j in range(n):
                if i == j:
                    continue
                if np.all(self._scores[j] >= self._scores[i]) and np.any(self._scores[j] > self._scores[i]):
                    dominated = True
                    break
            if not dominated:
                keep.append(i)
        return keep

    def coverage(self, candidate: int) -> int:
        """Number of instances on which `candidate` is a per-instance winner (ties count)."""
        winners = self.per_instance_winners
        return sum(1 for s in winners if candidate in s)

    @property
    def illumination_set(self) -> list[tuple[int, int]]:
        """Algorithm 2's pruned candidate set Ĉ with win-frequencies f[·].

        Returns `(candidate_index, coverage)` for every candidate that (a) is a per-instance winner on
        at least one instance and (b) is not strictly dominated by any other candidate. Equivalent to
        `[(k, self.coverage(k)) for k in self.non_dominated if self.coverage(k) >= 1]`.

        Empty only in degenerate cases; callers should fall back to uniform sampling then.
        """
        result: list[tuple[int, int]] = []
        for k in self.non_dominated:
            cov = self.coverage(k)
            if cov >= 1:
                result.append((k, cov))
        return result

_minimize = minimize instance-attribute

_scores = -scores if minimize else scores instance-attribute

illumination_set property

Algorithm 2's pruned candidate set Ĉ with win-frequencies f[·].

Returns (candidate_index, coverage) for every candidate that (a) is a per-instance winner on at least one instance and (b) is not strictly dominated by any other candidate. Equivalent to [(k, self.coverage(k)) for k in self.non_dominated if self.coverage(k) >= 1].

Empty only in degenerate cases; callers should fall back to uniform sampling then.

non_dominated property

Candidate indices that are not strictly dominated by any other candidate.

Candidate a is dominated by b iff b >= a on every instance and b > a on at least one.

per_instance_best property

Per-instance maximum score across candidates. Shape [n_instances]. Returned in the original sign convention (i.e. minimum, when minimize=True).

per_instance_winners property

For each instance, the set of candidate indices that achieve the per-instance best.

Multiple candidates may tie.

__init__(scores, minimize=False)

Source code in steerability/algorithms/input_control/common/pareto.py 
17
18
19
20
21
22
23
24
def __init__(self, scores: np.ndarray, minimize: bool = False) -> None:
    scores = np.asarray(scores, dtype=float)
    if scores.ndim != 2:
        raise ValueError(f"scores must be 2-D [n_candidates, n_instances]; got shape {scores.shape}.")
    if scores.size == 0:
        raise ValueError("scores must be non-empty.")
    self._scores = -scores if minimize else scores
    self._minimize = minimize

coverage(candidate)

Number of instances on which candidate is a per-instance winner (ties count).

Source code in steerability/algorithms/input_control/common/pareto.py 
63
64
65
66
def coverage(self, candidate: int) -> int:
    """Number of instances on which `candidate` is a per-instance winner (ties count)."""
    winners = self.per_instance_winners
    return sum(1 for s in winners if candidate in s)

RolloutBudget

Counter for capping the number of rollouts an optimizer may consume.

Used by methods that bound offline search (GEPA, CPO).

Example

budget = RolloutBudget(1000) while budget: budget.charge(minibatch_size) ...

Source code in steerability/algorithms/input_control/common/budget.py 
 5
 6
 7
 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
class RolloutBudget:
    """Counter for capping the number of rollouts an optimizer may consume.

    Used by methods that bound offline search (GEPA, CPO).

    Example:
        budget = RolloutBudget(1000)
        while budget:
            budget.charge(minibatch_size)
            ...
    """

    def __init__(self, max_rollouts: int) -> None:
        if max_rollouts < 0:
            raise ValueError(f"max_rollouts must be non-negative; got {max_rollouts}.")
        self.max = max_rollouts
        self.spent = 0

    @property
    def remaining(self) -> int:
        return max(self.max - self.spent, 0)

    def __bool__(self) -> bool:
        return self.spent < self.max

    def charge(self, n: int = 1) -> None:
        if n < 0:
            raise ValueError(f"Cannot charge a negative amount; got {n}.")
        self.spent += n

max = max_rollouts instance-attribute

remaining property

spent = 0 instance-attribute

__bool__()

Source code in steerability/algorithms/input_control/common/budget.py 
27
28
def __bool__(self) -> bool:
    return self.spent < self.max

__init__(max_rollouts)

Source code in steerability/algorithms/input_control/common/budget.py 
17
18
19
20
21
def __init__(self, max_rollouts: int) -> None:
    if max_rollouts < 0:
        raise ValueError(f"max_rollouts must be non-negative; got {max_rollouts}.")
    self.max = max_rollouts
    self.spent = 0

charge(n=1)

Source code in steerability/algorithms/input_control/common/budget.py 
30
31
32
33
def charge(self, n: int = 1) -> None:
    if n < 0:
        raise ValueError(f"Cannot charge a negative amount; got {n}.")
    self.spent += n

generate_with_system_prompt(task_lm, tokenizer, system_prompt, queries, gen_kwargs=None)

Generate one continuation per query under a shared system prompt.

Parameters:

Name Type Description Default
task_lm PreTrainedModel

Causal language model.

 required
tokenizer PreTrainedTokenizerBase

Tokenizer paired with task_lm.

 required
system_prompt str

System prompt prepended to every query.

 required
queries list[str]

User-facing query texts (one continuation generated per query).

 required
gen_kwargs dict | None

Forwarded to task_lm.generate. Defaults to greedy, max_new_tokens=32.

 None

Returns:

Type Description
list[str]

Decoded continuations, one per query, in order.
Source code in steerability/algorithms/input_control/common/generation.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
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 generate_with_system_prompt(
    task_lm: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    system_prompt: str,
    queries: list[str],
    gen_kwargs: dict | None = None,
) -> list[str]:
    """Generate one continuation per query under a shared system prompt.

    Args:
        task_lm: Causal language model.
        tokenizer: Tokenizer paired with `task_lm`.
        system_prompt: System prompt prepended to every query.
        queries: User-facing query texts (one continuation generated per query).
        gen_kwargs: Forwarded to `task_lm.generate`. Defaults to greedy, `max_new_tokens=32`.

    Returns:
        Decoded continuations, one per query, in order.
    """
    if not queries:
        return []
    gen_kwargs = dict(gen_kwargs or {"max_new_tokens": 32, "do_sample": False})

    used_template = bool(getattr(tokenizer, "chat_template", None))
    if used_template:
        prompts_text = [
            tokenizer.apply_chat_template(
                [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": query},
                ],
                tokenize=False,
                add_generation_prompt=True,
            )
            for query in queries
        ]
    else:
        prompts_text = [f"{system_prompt}\n\n{query}" for query in queries]

    original_padding_side = tokenizer.padding_side
    tokenizer.padding_side = "left"
    try:
        encoded = tokenizer(
            prompts_text,
            return_tensors="pt",
            padding=True,
            add_special_tokens=not used_template,
        ).to(task_lm.device)

        pad_id = tokenizer.pad_token_id
        if pad_id is None:
            pad_id = tokenizer.eos_token_id
        if pad_id is not None:
            gen_kwargs.setdefault("pad_token_id", pad_id)

        with torch.no_grad():
            output_ids = task_lm.generate(**encoded, **gen_kwargs)

        prompt_len = encoded["input_ids"].size(1)
        new_ids = output_ids[:, prompt_len:]
        decoded = tokenizer.decode(new_ids, skip_special_tokens=True)
    finally:
        tokenizer.padding_side = original_padding_side

    return decoded

budget

Counter for capping the number of rollouts an optimizer may consume.

RolloutBudget

Counter for capping the number of rollouts an optimizer may consume.

Used by methods that bound offline search (GEPA, CPO).

Example

budget = RolloutBudget(1000) while budget: budget.charge(minibatch_size) ...

Source code in steerability/algorithms/input_control/common/budget.py 
 5
 6
 7
 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
class RolloutBudget:
    """Counter for capping the number of rollouts an optimizer may consume.

    Used by methods that bound offline search (GEPA, CPO).

    Example:
        budget = RolloutBudget(1000)
        while budget:
            budget.charge(minibatch_size)
            ...
    """

    def __init__(self, max_rollouts: int) -> None:
        if max_rollouts < 0:
            raise ValueError(f"max_rollouts must be non-negative; got {max_rollouts}.")
        self.max = max_rollouts
        self.spent = 0

    @property
    def remaining(self) -> int:
        return max(self.max - self.spent, 0)

    def __bool__(self) -> bool:
        return self.spent < self.max

    def charge(self, n: int = 1) -> None:
        if n < 0:
            raise ValueError(f"Cannot charge a negative amount; got {n}.")
        self.spent += n
max = max_rollouts instance-attribute
remaining property
spent = 0 instance-attribute
__bool__()
Source code in steerability/algorithms/input_control/common/budget.py 
27
28
def __bool__(self) -> bool:
    return self.spent < self.max
__init__(max_rollouts)
Source code in steerability/algorithms/input_control/common/budget.py 
17
18
19
20
21
def __init__(self, max_rollouts: int) -> None:
    if max_rollouts < 0:
        raise ValueError(f"max_rollouts must be non-negative; got {max_rollouts}.")
    self.max = max_rollouts
    self.spent = 0
charge(n=1)
Source code in steerability/algorithms/input_control/common/budget.py 
30
31
32
33
def charge(self, n: int = 1) -> None:
    if n < 0:
        raise ValueError(f"Cannot charge a negative amount; got {n}.")
    self.spent += n

formatters

Formatters render Memory content into adapted prompts (token-level or message-level).

__all__ = ['BaseFormatter', 'ChatTemplateSlotFormatter', 'FewShotBlockFormatter', 'PrependTextFormatter', 'SystemPromptFormatter'] module-attribute

BaseFormatter

Bases: ABC

Renders memory content into an adapted prompt.

Two entry points: token-level (apply_to_ids) and message-level (apply_to_messages). Subclasses implement one or both; the default of the unused one raises NotImplementedError. Methods using adapt_messages should prefer apply_to_messages since it operates before chat-template tokenization.

Source code in steerability/algorithms/input_control/common/formatters/base.py 
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
class BaseFormatter(ABC):
    """Renders memory content into an adapted prompt.

    Two entry points: token-level (`apply_to_ids`) and message-level (`apply_to_messages`). Subclasses
    implement one or both; the default of the unused one raises NotImplementedError. Methods using
    `adapt_messages` should prefer `apply_to_messages` since it operates before chat-template tokenization.
    """

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError(
            f"{type(self).__name__} does not implement `apply_to_ids`. Use `apply_to_messages` instead."
        )

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        raise NotImplementedError(
            f"{type(self).__name__} does not implement `apply_to_messages`. Use `apply_to_ids` instead."
        )
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/base.py 
20
21
22
23
24
25
26
27
28
29
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    raise NotImplementedError(
        f"{type(self).__name__} does not implement `apply_to_ids`. Use `apply_to_messages` instead."
    )
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/base.py 
31
32
33
34
35
36
37
38
39
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    raise NotImplementedError(
        f"{type(self).__name__} does not implement `apply_to_messages`. Use `apply_to_ids` instead."
    )

ChatTemplateSlotFormatter

Bases: BaseFormatter

Replaces {{slot}} placeholders inside message content with the corresponding memory slot.

Slot values may be str or list[str] (lists are joined as bullet points). Unmatched placeholders are left as-is.

Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
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
43
44
45
46
47
48
49
50
51
class ChatTemplateSlotFormatter(BaseFormatter):
    """Replaces `{{slot}}` placeholders inside message content with the corresponding memory slot.

    Slot values may be `str` or `list[str]` (lists are joined as bullet points). Unmatched placeholders
    are left as-is.
    """

    _PATTERN = re.compile(r"\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}")

    def __init__(self, list_separator: str = "\n- ", list_prefix: str = "- ") -> None:
        self.list_separator = list_separator
        self.list_prefix = list_prefix

    def _render_value(self, value: str | list[str]) -> str:
        if isinstance(value, list):
            return self.list_prefix + self.list_separator.join(str(v) for v in value)
        return str(value)

    def _substitute(self, content: str, memory: Memory) -> str:
        def _repl(m: re.Match) -> str:
            key = m.group(1)
            if key in memory:
                return self._render_value(memory[key])
            return m.group(0)
        return self._PATTERN.sub(_repl, content)

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        out: list[list[dict]] = []
        for chat in messages:
            new_chat: list[dict] = []
            for msg in chat:
                content = msg.get("content", "")
                if isinstance(content, str):
                    msg = {**msg, "content": self._substitute(content, memory)}
                new_chat.append(msg)
            out.append(new_chat)
        return out
_PATTERN = re.compile('\\{\\{\\s*([a-zA-Z_][a-zA-Z0-9_]*)\\s*\\}\\}') class-attribute instance-attribute
list_prefix = list_prefix instance-attribute
list_separator = list_separator instance-attribute
__init__(list_separator='\n- ', list_prefix='- ')
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
19
20
21
def __init__(self, list_separator: str = "\n- ", list_prefix: str = "- ") -> None:
    self.list_separator = list_separator
    self.list_prefix = list_prefix
_render_value(value)
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
23
24
25
26
def _render_value(self, value: str | list[str]) -> str:
    if isinstance(value, list):
        return self.list_prefix + self.list_separator.join(str(v) for v in value)
    return str(value)
_substitute(content, memory)
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
28
29
30
31
32
33
34
def _substitute(self, content: str, memory: Memory) -> str:
    def _repl(m: re.Match) -> str:
        key = m.group(1)
        if key in memory:
            return self._render_value(memory[key])
        return m.group(0)
    return self._PATTERN.sub(_repl, content)
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/base.py 
20
21
22
23
24
25
26
27
28
29
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    raise NotImplementedError(
        f"{type(self).__name__} does not implement `apply_to_ids`. Use `apply_to_messages` instead."
    )
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    out: list[list[dict]] = []
    for chat in messages:
        new_chat: list[dict] = []
        for msg in chat:
            content = msg.get("content", "")
            if isinstance(content, str):
                msg = {**msg, "content": self._substitute(content, memory)}
            new_chat.append(msg)
        out.append(new_chat)
    return out

FewShotBlockFormatter

Bases: BaseFormatter

Renders memory['directive'] and memory['examples'] as a single system message.

Each example dict's non-private fields (those not prefixed with _) are emitted as Title-Cased Key: value lines under a header determined by the example's _polarity. The _polarity key is set internally by FewShot and is the only required field.

Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
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
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
class FewShotBlockFormatter(BaseFormatter):
    """Renders memory['directive'] and memory['examples'] as a single system message.

    Each example dict's non-private fields (those not prefixed with `_`) are emitted as
    `Title-Cased Key: value` lines under a header determined by the example's `_polarity`.
    The `_polarity` key is set internally by `FewShot` and is the only required field.
    """

    DEFAULT_POSITIVE_HEADER = "### Positive example (behavior to follow)"
    DEFAULT_NEGATIVE_HEADER = "### Negative example (behavior to avoid)"

    def __init__(
        self,
        positive_header: str = DEFAULT_POSITIVE_HEADER,
        negative_header: str = DEFAULT_NEGATIVE_HEADER,
    ) -> None:
        self.positive_header = positive_header
        self.negative_header = negative_header

    def _render_example(self, example: dict) -> str:
        header = (
            self.negative_header if example.get("_polarity") == "negative"
            else self.positive_header
        )
        body = "\n".join(
            f"{key.replace('_', ' ').title()}: {value}"
            for key, value in example.items()
            if not key.startswith("_")
        )
        return f"{header}\n{body}"

    def _render_block(self, memory: Memory) -> str:
        examples = memory.get("examples") or []
        directive = memory.get("directive") or ""
        parts: list[str] = []
        if directive:
            parts.append(directive)
        parts.extend(self._render_example(ex) for ex in examples)
        return "\n\n".join(parts)

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        block = self._render_block(memory)
        if not block:
            return [list(chat) for chat in messages]

        out: list[list[dict]] = []
        for chat in messages:
            chat = [dict(m) for m in chat]
            insert_at = 1 if chat and chat[0].get("role") == "system" else 0
            out.append(chat[:insert_at] + [{"role": "system", "content": block}] + chat[insert_at:])
        return out

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        block = self._render_block(memory)
        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        if not block:
            return input_ids
        prefix_ids = tokenizer.encode(block + "\n\n", add_special_tokens=False)
        prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
        return torch.cat([prefix.unsqueeze(0).expand(input_ids.size(0), -1), input_ids], dim=1)
DEFAULT_NEGATIVE_HEADER = '### Negative example (behavior to avoid)' class-attribute instance-attribute
DEFAULT_POSITIVE_HEADER = '### Positive example (behavior to follow)' class-attribute instance-attribute
negative_header = negative_header instance-attribute
positive_header = positive_header instance-attribute
__init__(positive_header=DEFAULT_POSITIVE_HEADER, negative_header=DEFAULT_NEGATIVE_HEADER)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
22
23
24
25
26
27
28
def __init__(
    self,
    positive_header: str = DEFAULT_POSITIVE_HEADER,
    negative_header: str = DEFAULT_NEGATIVE_HEADER,
) -> None:
    self.positive_header = positive_header
    self.negative_header = negative_header
_render_block(memory)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
42
43
44
45
46
47
48
49
def _render_block(self, memory: Memory) -> str:
    examples = memory.get("examples") or []
    directive = memory.get("directive") or ""
    parts: list[str] = []
    if directive:
        parts.append(directive)
    parts.extend(self._render_example(ex) for ex in examples)
    return "\n\n".join(parts)
_render_example(example)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
30
31
32
33
34
35
36
37
38
39
40
def _render_example(self, example: dict) -> str:
    header = (
        self.negative_header if example.get("_polarity") == "negative"
        else self.positive_header
    )
    body = "\n".join(
        f"{key.replace('_', ' ').title()}: {value}"
        for key, value in example.items()
        if not key.startswith("_")
    )
    return f"{header}\n{body}"
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    block = self._render_block(memory)
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    if not block:
        return input_ids
    prefix_ids = tokenizer.encode(block + "\n\n", add_special_tokens=False)
    prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
    return torch.cat([prefix.unsqueeze(0).expand(input_ids.size(0), -1), input_ids], dim=1)
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    block = self._render_block(memory)
    if not block:
        return [list(chat) for chat in messages]

    out: list[list[dict]] = []
    for chat in messages:
        chat = [dict(m) for m in chat]
        insert_at = 1 if chat and chat[0].get("role") == "system" else 0
        out.append(chat[:insert_at] + [{"role": "system", "content": block}] + chat[insert_at:])
    return out

PrependTextFormatter

Bases: BaseFormatter

Prepends memory["text"] to a user message or to the raw token ids.

When operating on messages, the prepended text is concatenated to the content of the user-role message(s) selected by target, joined with separator. target is one of:

- `"first_user"`: the first user-role message.
- `"last_user"`: the last user-role message.
- `"all_user"`: every user-role message.

When a chat contains no user-role message, all targets append a new {"role": "user", "content": text} turn. Message dicts are copied, so the caller's structures are not mutated. When operating on token ids, the encoded text is prefixed to each sequence in the batch; token streams carry no turn structure, so target does not apply on the apply_to_ids path.

Parameters:

Name Type Description Default
separator str

String inserted between the prepended text and the existing message content.

 '\n\n'
target str

Which user-role message(s) to prepend to ("first_user", "last_user", or "all_user").

 'first_user'

Raises:

Type Description
ValueError

If target is not one of the supported values.
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.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
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
class PrependTextFormatter(BaseFormatter):
    """Prepends `memory["text"]` to a user message or to the raw token ids.

    When operating on messages, the prepended text is concatenated to the content of the user-role
    message(s) selected by `target`, joined with `separator`. `target` is one of:

        - `"first_user"`: the first user-role message.
        - `"last_user"`: the last user-role message.
        - `"all_user"`: every user-role message.

    When a chat contains no user-role message, all targets append a new `{"role": "user", "content": text}`
    turn. Message dicts are copied, so the caller's structures are not mutated. When operating on token ids,
    the encoded text is prefixed to each sequence in the batch; token streams carry no turn structure, so
    `target` does not apply on the `apply_to_ids` path.

    Args:
        separator: String inserted between the prepended text and the existing message content.
        target: Which user-role message(s) to prepend to (`"first_user"`, `"last_user"`, or `"all_user"`).

    Raises:
        ValueError: If `target` is not one of the supported values.
    """

    def __init__(self, separator: str = "\n\n", target: str = "first_user") -> None:
        if target not in _TARGETS:
            raise ValueError(
                f"PrependTextFormatter target must be one of {sorted(_TARGETS)}; got {target!r}."
            )
        self.separator = separator
        self.target = target

    def _resolve_text(self, memory: Memory) -> str:
        text = memory["text"] if "text" in memory else memory.get("text")
        if not isinstance(text, str):
            raise TypeError(
                f"PrependTextFormatter expects memory['text'] to be a str; got {type(text).__name__}."
            )
        return text

    def _target_indices(self, chat: list[dict]) -> list[int]:
        user_indices = [i for i, m in enumerate(chat) if m.get("role") == "user"]
        if not user_indices:
            return []
        if self.target == "first_user":
            return [user_indices[0]]
        if self.target == "last_user":
            return [user_indices[-1]]
        return user_indices

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        text = self._resolve_text(memory)
        out: list[list[dict]] = []
        for chat in messages:
            chat = [dict(m) for m in chat]
            indices = self._target_indices(chat)
            if not indices:
                chat.append({"role": "user", "content": text})
            else:
                for idx in indices:
                    chat[idx] = {
                        "role": "user",
                        "content": text + self.separator + chat[idx].get("content", ""),
                    }
            out.append(chat)
        return out

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        text = self._resolve_text(memory)
        prefix_ids = tokenizer.encode(text + self.separator, add_special_tokens=False)
        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
        prefix = prefix.unsqueeze(0).expand(input_ids.size(0), -1)
        return torch.cat([prefix, input_ids], dim=1)
separator = separator instance-attribute
target = target instance-attribute
__init__(separator='\n\n', target='first_user')
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
36
37
38
39
40
41
42
def __init__(self, separator: str = "\n\n", target: str = "first_user") -> None:
    if target not in _TARGETS:
        raise ValueError(
            f"PrependTextFormatter target must be one of {sorted(_TARGETS)}; got {target!r}."
        )
    self.separator = separator
    self.target = target
_resolve_text(memory)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
44
45
46
47
48
49
50
def _resolve_text(self, memory: Memory) -> str:
    text = memory["text"] if "text" in memory else memory.get("text")
    if not isinstance(text, str):
        raise TypeError(
            f"PrependTextFormatter expects memory['text'] to be a str; got {type(text).__name__}."
        )
    return text
_target_indices(chat)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
52
53
54
55
56
57
58
59
60
def _target_indices(self, chat: list[dict]) -> list[int]:
    user_indices = [i for i, m in enumerate(chat) if m.get("role") == "user"]
    if not user_indices:
        return []
    if self.target == "first_user":
        return [user_indices[0]]
    if self.target == "last_user":
        return [user_indices[-1]]
    return user_indices
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    text = self._resolve_text(memory)
    prefix_ids = tokenizer.encode(text + self.separator, add_special_tokens=False)
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
    prefix = prefix.unsqueeze(0).expand(input_ids.size(0), -1)
    return torch.cat([prefix, input_ids], dim=1)
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    text = self._resolve_text(memory)
    out: list[list[dict]] = []
    for chat in messages:
        chat = [dict(m) for m in chat]
        indices = self._target_indices(chat)
        if not indices:
            chat.append({"role": "user", "content": text})
        else:
            for idx in indices:
                chat[idx] = {
                    "role": "user",
                    "content": text + self.separator + chat[idx].get("content", ""),
                }
        out.append(chat)
    return out

SystemPromptFormatter

Bases: BaseFormatter

Sets or merges the leading system message in each chat with memory["instruction"].

Only the leading system message (chat[0] when its role is "system") participates; system messages elsewhere in the chat are left untouched. When a leading system message is present, mode controls how the instruction combines with its content:

- `"replace"`: the instruction becomes the content.
- `"prepend"`: `instruction + separator + existing`.
- `"append"`: `existing + separator + instruction`.

When no leading system message is present, all three modes insert a single {"role": "system", "content": instruction} at position 0, so separator has no effect. Message dicts are copied, so the caller's structures are not mutated. Every input shape yields exactly one leading system message.

The token path (apply_to_ids) decodes, re-templates, and re-encodes, so it cannot merge with a system prompt that decoding has flattened into text; mode does not apply there, and every mode behaves as "replace". Prefer message-level entry for chat input.

Parameters:

Name Type Description Default
mode str

How the instruction combines with an existing leading system message ("replace" (default), "prepend", or "append"). Ignored when no leading system message is present.

 'replace'
separator str

String inserted between the instruction and the existing content for "prepend" and "append". Empty string allowed.

 '\n\n'

Raises:

Type Description
ValueError

If mode is not one of the supported values.
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.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
 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
class SystemPromptFormatter(BaseFormatter):
    """Sets or merges the leading system message in each chat with `memory["instruction"]`.

    Only the leading system message (`chat[0]` when its role is `"system"`) participates; system messages
    elsewhere in the chat are left untouched. When a leading system message is present, `mode` controls how
    the instruction combines with its content:

        - `"replace"`: the instruction becomes the content.
        - `"prepend"`: `instruction + separator + existing`.
        - `"append"`: `existing + separator + instruction`.

    When no leading system message is present, all three modes insert a single
    `{"role": "system", "content": instruction}` at position 0, so `separator` has no effect. Message dicts are
    copied, so the caller's structures are not mutated. Every input shape yields exactly one leading system
    message.

    The token path (`apply_to_ids`) decodes, re-templates, and re-encodes, so it cannot merge with a system
    prompt that decoding has flattened into text; `mode` does not apply there, and every mode behaves as
    `"replace"`. Prefer message-level entry for chat input.

    Args:
        mode: How the instruction combines with an existing leading system message (`"replace"` (default),
            `"prepend"`, or `"append"`). Ignored when no leading system message is present.
        separator: String inserted between the instruction and the existing content for `"prepend"` and
            `"append"`. Empty string allowed.

    Raises:
        ValueError: If `mode` is not one of the supported values.
    """

    def __init__(self, mode: str = "replace", separator: str = "\n\n") -> None:
        if mode not in _MODES:
            raise ValueError(f"SystemPromptFormatter mode must be one of {sorted(_MODES)}; got {mode!r}.")
        self.mode = mode
        self.separator = separator

    def _resolve_instruction(self, memory: Memory) -> str:
        instruction = memory["instruction"] if "instruction" in memory else memory.get("instruction")
        if not isinstance(instruction, str):
            raise TypeError(
                f"SystemPromptFormatter expects memory['instruction'] to be a str; got {type(instruction).__name__}."
            )
        return instruction

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        instruction = self._resolve_instruction(memory)

        out: list[list[dict]] = []
        for chat in messages:
            chat = list(chat)
            if chat and chat[0].get("role") == "system":
                existing = chat[0].get("content", "")
                if self.mode == "prepend":
                    content = instruction + self.separator + existing
                elif self.mode == "append":
                    content = existing + self.separator + instruction
                else:
                    content = instruction
                chat[0] = {"role": "system", "content": content}
            else:
                chat.insert(0, {"role": "system", "content": instruction})
            out.append(chat)
        return out

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        warnings.warn(
            "SystemPromptFormatter.apply_to_ids decodes → edits → re-tokenizes; prefer message-level entry "
            "(pass chat input to the pipeline so `apply_to_messages` runs).",
            UserWarning,
        )
        instruction = self._resolve_instruction(memory)

        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        decoded = tokenizer.batch_decode(input_ids, skip_special_tokens=True)

        rebuilt: list[list[int]] = []
        for user_text in decoded:
            if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
                templated = tokenizer.apply_chat_template(
                    [
                        {"role": "system", "content": instruction},
                        {"role": "user", "content": user_text},
                    ],
                    tokenize=False,
                    add_generation_prompt=True,
                )
                rebuilt.append(tokenizer.encode(templated, add_special_tokens=False))
            else:
                rebuilt.append(tokenizer.encode(instruction + "\n\n" + user_text, add_special_tokens=False))

        max_len = max(len(seq) for seq in rebuilt)
        pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
        padded = [seq + [pad_id] * (max_len - len(seq)) for seq in rebuilt]
        return torch.tensor(padded, dtype=input_ids.dtype, device=input_ids.device)
mode = mode instance-attribute
separator = separator instance-attribute
__init__(mode='replace', separator='\n\n')
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.py 
45
46
47
48
49
def __init__(self, mode: str = "replace", separator: str = "\n\n") -> None:
    if mode not in _MODES:
        raise ValueError(f"SystemPromptFormatter mode must be one of {sorted(_MODES)}; got {mode!r}.")
    self.mode = mode
    self.separator = separator
_resolve_instruction(memory)
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.py 
51
52
53
54
55
56
57
def _resolve_instruction(self, memory: Memory) -> str:
    instruction = memory["instruction"] if "instruction" in memory else memory.get("instruction")
    if not isinstance(instruction, str):
        raise TypeError(
            f"SystemPromptFormatter expects memory['instruction'] to be a str; got {type(instruction).__name__}."
        )
    return instruction
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.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
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    warnings.warn(
        "SystemPromptFormatter.apply_to_ids decodes → edits → re-tokenizes; prefer message-level entry "
        "(pass chat input to the pipeline so `apply_to_messages` runs).",
        UserWarning,
    )
    instruction = self._resolve_instruction(memory)

    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    decoded = tokenizer.batch_decode(input_ids, skip_special_tokens=True)

    rebuilt: list[list[int]] = []
    for user_text in decoded:
        if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
            templated = tokenizer.apply_chat_template(
                [
                    {"role": "system", "content": instruction},
                    {"role": "user", "content": user_text},
                ],
                tokenize=False,
                add_generation_prompt=True,
            )
            rebuilt.append(tokenizer.encode(templated, add_special_tokens=False))
        else:
            rebuilt.append(tokenizer.encode(instruction + "\n\n" + user_text, add_special_tokens=False))

    max_len = max(len(seq) for seq in rebuilt)
    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
    padded = [seq + [pad_id] * (max_len - len(seq)) for seq in rebuilt]
    return torch.tensor(padded, dtype=input_ids.dtype, device=input_ids.device)
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    instruction = self._resolve_instruction(memory)

    out: list[list[dict]] = []
    for chat in messages:
        chat = list(chat)
        if chat and chat[0].get("role") == "system":
            existing = chat[0].get("content", "")
            if self.mode == "prepend":
                content = instruction + self.separator + existing
            elif self.mode == "append":
                content = existing + self.separator + instruction
            else:
                content = instruction
            chat[0] = {"role": "system", "content": content}
        else:
            chat.insert(0, {"role": "system", "content": instruction})
        out.append(chat)
    return out

base

Token-level and message-level rendering of Memory content.

BaseFormatter

Bases: ABC

Renders memory content into an adapted prompt.

Two entry points: token-level (apply_to_ids) and message-level (apply_to_messages). Subclasses implement one or both; the default of the unused one raises NotImplementedError. Methods using adapt_messages should prefer apply_to_messages since it operates before chat-template tokenization.

Source code in steerability/algorithms/input_control/common/formatters/base.py 
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
class BaseFormatter(ABC):
    """Renders memory content into an adapted prompt.

    Two entry points: token-level (`apply_to_ids`) and message-level (`apply_to_messages`). Subclasses
    implement one or both; the default of the unused one raises NotImplementedError. Methods using
    `adapt_messages` should prefer `apply_to_messages` since it operates before chat-template tokenization.
    """

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError(
            f"{type(self).__name__} does not implement `apply_to_ids`. Use `apply_to_messages` instead."
        )

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        raise NotImplementedError(
            f"{type(self).__name__} does not implement `apply_to_messages`. Use `apply_to_ids` instead."
        )
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/base.py 
20
21
22
23
24
25
26
27
28
29
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    raise NotImplementedError(
        f"{type(self).__name__} does not implement `apply_to_ids`. Use `apply_to_messages` instead."
    )
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/base.py 
31
32
33
34
35
36
37
38
39
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    raise NotImplementedError(
        f"{type(self).__name__} does not implement `apply_to_messages`. Use `apply_to_ids` instead."
    )

chat_template_slot

Replace {{slot}} placeholders in messages with memory content.

ChatTemplateSlotFormatter

Bases: BaseFormatter

Replaces {{slot}} placeholders inside message content with the corresponding memory slot.

Slot values may be str or list[str] (lists are joined as bullet points). Unmatched placeholders are left as-is.

Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
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
43
44
45
46
47
48
49
50
51
class ChatTemplateSlotFormatter(BaseFormatter):
    """Replaces `{{slot}}` placeholders inside message content with the corresponding memory slot.

    Slot values may be `str` or `list[str]` (lists are joined as bullet points). Unmatched placeholders
    are left as-is.
    """

    _PATTERN = re.compile(r"\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}")

    def __init__(self, list_separator: str = "\n- ", list_prefix: str = "- ") -> None:
        self.list_separator = list_separator
        self.list_prefix = list_prefix

    def _render_value(self, value: str | list[str]) -> str:
        if isinstance(value, list):
            return self.list_prefix + self.list_separator.join(str(v) for v in value)
        return str(value)

    def _substitute(self, content: str, memory: Memory) -> str:
        def _repl(m: re.Match) -> str:
            key = m.group(1)
            if key in memory:
                return self._render_value(memory[key])
            return m.group(0)
        return self._PATTERN.sub(_repl, content)

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        out: list[list[dict]] = []
        for chat in messages:
            new_chat: list[dict] = []
            for msg in chat:
                content = msg.get("content", "")
                if isinstance(content, str):
                    msg = {**msg, "content": self._substitute(content, memory)}
                new_chat.append(msg)
            out.append(new_chat)
        return out
_PATTERN = re.compile('\\{\\{\\s*([a-zA-Z_][a-zA-Z0-9_]*)\\s*\\}\\}') class-attribute instance-attribute
list_prefix = list_prefix instance-attribute
list_separator = list_separator instance-attribute
__init__(list_separator='\n- ', list_prefix='- ')
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
19
20
21
def __init__(self, list_separator: str = "\n- ", list_prefix: str = "- ") -> None:
    self.list_separator = list_separator
    self.list_prefix = list_prefix
_render_value(value)
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
23
24
25
26
def _render_value(self, value: str | list[str]) -> str:
    if isinstance(value, list):
        return self.list_prefix + self.list_separator.join(str(v) for v in value)
    return str(value)
_substitute(content, memory)
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
28
29
30
31
32
33
34
def _substitute(self, content: str, memory: Memory) -> str:
    def _repl(m: re.Match) -> str:
        key = m.group(1)
        if key in memory:
            return self._render_value(memory[key])
        return m.group(0)
    return self._PATTERN.sub(_repl, content)
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/base.py 
20
21
22
23
24
25
26
27
28
29
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    raise NotImplementedError(
        f"{type(self).__name__} does not implement `apply_to_ids`. Use `apply_to_messages` instead."
    )
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/chat_template_slot.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    out: list[list[dict]] = []
    for chat in messages:
        new_chat: list[dict] = []
        for msg in chat:
            content = msg.get("content", "")
            if isinstance(content, str):
                msg = {**msg, "content": self._substitute(content, memory)}
            new_chat.append(msg)
        out.append(new_chat)
    return out

few_shot_block

Render memory['examples'] as a labeled system-message block.

FewShotBlockFormatter

Bases: BaseFormatter

Renders memory['directive'] and memory['examples'] as a single system message.

Each example dict's non-private fields (those not prefixed with _) are emitted as Title-Cased Key: value lines under a header determined by the example's _polarity. The _polarity key is set internally by FewShot and is the only required field.

Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
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
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
class FewShotBlockFormatter(BaseFormatter):
    """Renders memory['directive'] and memory['examples'] as a single system message.

    Each example dict's non-private fields (those not prefixed with `_`) are emitted as
    `Title-Cased Key: value` lines under a header determined by the example's `_polarity`.
    The `_polarity` key is set internally by `FewShot` and is the only required field.
    """

    DEFAULT_POSITIVE_HEADER = "### Positive example (behavior to follow)"
    DEFAULT_NEGATIVE_HEADER = "### Negative example (behavior to avoid)"

    def __init__(
        self,
        positive_header: str = DEFAULT_POSITIVE_HEADER,
        negative_header: str = DEFAULT_NEGATIVE_HEADER,
    ) -> None:
        self.positive_header = positive_header
        self.negative_header = negative_header

    def _render_example(self, example: dict) -> str:
        header = (
            self.negative_header if example.get("_polarity") == "negative"
            else self.positive_header
        )
        body = "\n".join(
            f"{key.replace('_', ' ').title()}: {value}"
            for key, value in example.items()
            if not key.startswith("_")
        )
        return f"{header}\n{body}"

    def _render_block(self, memory: Memory) -> str:
        examples = memory.get("examples") or []
        directive = memory.get("directive") or ""
        parts: list[str] = []
        if directive:
            parts.append(directive)
        parts.extend(self._render_example(ex) for ex in examples)
        return "\n\n".join(parts)

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        block = self._render_block(memory)
        if not block:
            return [list(chat) for chat in messages]

        out: list[list[dict]] = []
        for chat in messages:
            chat = [dict(m) for m in chat]
            insert_at = 1 if chat and chat[0].get("role") == "system" else 0
            out.append(chat[:insert_at] + [{"role": "system", "content": block}] + chat[insert_at:])
        return out

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        block = self._render_block(memory)
        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        if not block:
            return input_ids
        prefix_ids = tokenizer.encode(block + "\n\n", add_special_tokens=False)
        prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
        return torch.cat([prefix.unsqueeze(0).expand(input_ids.size(0), -1), input_ids], dim=1)
DEFAULT_NEGATIVE_HEADER = '### Negative example (behavior to avoid)' class-attribute instance-attribute
DEFAULT_POSITIVE_HEADER = '### Positive example (behavior to follow)' class-attribute instance-attribute
negative_header = negative_header instance-attribute
positive_header = positive_header instance-attribute
__init__(positive_header=DEFAULT_POSITIVE_HEADER, negative_header=DEFAULT_NEGATIVE_HEADER)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
22
23
24
25
26
27
28
def __init__(
    self,
    positive_header: str = DEFAULT_POSITIVE_HEADER,
    negative_header: str = DEFAULT_NEGATIVE_HEADER,
) -> None:
    self.positive_header = positive_header
    self.negative_header = negative_header
_render_block(memory)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
42
43
44
45
46
47
48
49
def _render_block(self, memory: Memory) -> str:
    examples = memory.get("examples") or []
    directive = memory.get("directive") or ""
    parts: list[str] = []
    if directive:
        parts.append(directive)
    parts.extend(self._render_example(ex) for ex in examples)
    return "\n\n".join(parts)
_render_example(example)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
30
31
32
33
34
35
36
37
38
39
40
def _render_example(self, example: dict) -> str:
    header = (
        self.negative_header if example.get("_polarity") == "negative"
        else self.positive_header
    )
    body = "\n".join(
        f"{key.replace('_', ' ').title()}: {value}"
        for key, value in example.items()
        if not key.startswith("_")
    )
    return f"{header}\n{body}"
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    block = self._render_block(memory)
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    if not block:
        return input_ids
    prefix_ids = tokenizer.encode(block + "\n\n", add_special_tokens=False)
    prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
    return torch.cat([prefix.unsqueeze(0).expand(input_ids.size(0), -1), input_ids], dim=1)
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/few_shot_block.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    block = self._render_block(memory)
    if not block:
        return [list(chat) for chat in messages]

    out: list[list[dict]] = []
    for chat in messages:
        chat = [dict(m) for m in chat]
        insert_at = 1 if chat and chat[0].get("role") == "system" else 0
        out.append(chat[:insert_at] + [{"role": "system", "content": block}] + chat[insert_at:])
    return out

prepend_text

Prepend a raw text block to a user turn or input_ids.

_TARGETS = frozenset({'first_user', 'last_user', 'all_user'}) module-attribute
PrependTextFormatter

Bases: BaseFormatter

Prepends memory["text"] to a user message or to the raw token ids.

When operating on messages, the prepended text is concatenated to the content of the user-role message(s) selected by target, joined with separator. target is one of:

- `"first_user"`: the first user-role message.
- `"last_user"`: the last user-role message.
- `"all_user"`: every user-role message.

When a chat contains no user-role message, all targets append a new {"role": "user", "content": text} turn. Message dicts are copied, so the caller's structures are not mutated. When operating on token ids, the encoded text is prefixed to each sequence in the batch; token streams carry no turn structure, so target does not apply on the apply_to_ids path.

Parameters:

Name Type Description Default
separator str

String inserted between the prepended text and the existing message content.

 '\n\n'
target str

Which user-role message(s) to prepend to ("first_user", "last_user", or "all_user").

 'first_user'

Raises:

Type Description
ValueError

If target is not one of the supported values.
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.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
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
class PrependTextFormatter(BaseFormatter):
    """Prepends `memory["text"]` to a user message or to the raw token ids.

    When operating on messages, the prepended text is concatenated to the content of the user-role
    message(s) selected by `target`, joined with `separator`. `target` is one of:

        - `"first_user"`: the first user-role message.
        - `"last_user"`: the last user-role message.
        - `"all_user"`: every user-role message.

    When a chat contains no user-role message, all targets append a new `{"role": "user", "content": text}`
    turn. Message dicts are copied, so the caller's structures are not mutated. When operating on token ids,
    the encoded text is prefixed to each sequence in the batch; token streams carry no turn structure, so
    `target` does not apply on the `apply_to_ids` path.

    Args:
        separator: String inserted between the prepended text and the existing message content.
        target: Which user-role message(s) to prepend to (`"first_user"`, `"last_user"`, or `"all_user"`).

    Raises:
        ValueError: If `target` is not one of the supported values.
    """

    def __init__(self, separator: str = "\n\n", target: str = "first_user") -> None:
        if target not in _TARGETS:
            raise ValueError(
                f"PrependTextFormatter target must be one of {sorted(_TARGETS)}; got {target!r}."
            )
        self.separator = separator
        self.target = target

    def _resolve_text(self, memory: Memory) -> str:
        text = memory["text"] if "text" in memory else memory.get("text")
        if not isinstance(text, str):
            raise TypeError(
                f"PrependTextFormatter expects memory['text'] to be a str; got {type(text).__name__}."
            )
        return text

    def _target_indices(self, chat: list[dict]) -> list[int]:
        user_indices = [i for i, m in enumerate(chat) if m.get("role") == "user"]
        if not user_indices:
            return []
        if self.target == "first_user":
            return [user_indices[0]]
        if self.target == "last_user":
            return [user_indices[-1]]
        return user_indices

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        text = self._resolve_text(memory)
        out: list[list[dict]] = []
        for chat in messages:
            chat = [dict(m) for m in chat]
            indices = self._target_indices(chat)
            if not indices:
                chat.append({"role": "user", "content": text})
            else:
                for idx in indices:
                    chat[idx] = {
                        "role": "user",
                        "content": text + self.separator + chat[idx].get("content", ""),
                    }
            out.append(chat)
        return out

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        text = self._resolve_text(memory)
        prefix_ids = tokenizer.encode(text + self.separator, add_special_tokens=False)
        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
        prefix = prefix.unsqueeze(0).expand(input_ids.size(0), -1)
        return torch.cat([prefix, input_ids], dim=1)
separator = separator instance-attribute
target = target instance-attribute
__init__(separator='\n\n', target='first_user')
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
36
37
38
39
40
41
42
def __init__(self, separator: str = "\n\n", target: str = "first_user") -> None:
    if target not in _TARGETS:
        raise ValueError(
            f"PrependTextFormatter target must be one of {sorted(_TARGETS)}; got {target!r}."
        )
    self.separator = separator
    self.target = target
_resolve_text(memory)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
44
45
46
47
48
49
50
def _resolve_text(self, memory: Memory) -> str:
    text = memory["text"] if "text" in memory else memory.get("text")
    if not isinstance(text, str):
        raise TypeError(
            f"PrependTextFormatter expects memory['text'] to be a str; got {type(text).__name__}."
        )
    return text
_target_indices(chat)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
52
53
54
55
56
57
58
59
60
def _target_indices(self, chat: list[dict]) -> list[int]:
    user_indices = [i for i, m in enumerate(chat) if m.get("role") == "user"]
    if not user_indices:
        return []
    if self.target == "first_user":
        return [user_indices[0]]
    if self.target == "last_user":
        return [user_indices[-1]]
    return user_indices
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    text = self._resolve_text(memory)
    prefix_ids = tokenizer.encode(text + self.separator, add_special_tokens=False)
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    prefix = torch.tensor(prefix_ids, dtype=input_ids.dtype, device=input_ids.device)
    prefix = prefix.unsqueeze(0).expand(input_ids.size(0), -1)
    return torch.cat([prefix, input_ids], dim=1)
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/prepend_text.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    text = self._resolve_text(memory)
    out: list[list[dict]] = []
    for chat in messages:
        chat = [dict(m) for m in chat]
        indices = self._target_indices(chat)
        if not indices:
            chat.append({"role": "user", "content": text})
        else:
            for idx in indices:
                chat[idx] = {
                    "role": "user",
                    "content": text + self.separator + chat[idx].get("content", ""),
                }
        out.append(chat)
    return out

system_prompt

Set or merge the leading system message from memory["instruction"].

_MODES = frozenset({'replace', 'prepend', 'append'}) module-attribute
SystemPromptFormatter

Bases: BaseFormatter

Sets or merges the leading system message in each chat with memory["instruction"].

Only the leading system message (chat[0] when its role is "system") participates; system messages elsewhere in the chat are left untouched. When a leading system message is present, mode controls how the instruction combines with its content:

- `"replace"`: the instruction becomes the content.
- `"prepend"`: `instruction + separator + existing`.
- `"append"`: `existing + separator + instruction`.

When no leading system message is present, all three modes insert a single {"role": "system", "content": instruction} at position 0, so separator has no effect. Message dicts are copied, so the caller's structures are not mutated. Every input shape yields exactly one leading system message.

The token path (apply_to_ids) decodes, re-templates, and re-encodes, so it cannot merge with a system prompt that decoding has flattened into text; mode does not apply there, and every mode behaves as "replace". Prefer message-level entry for chat input.

Parameters:

Name Type Description Default
mode str

How the instruction combines with an existing leading system message ("replace" (default), "prepend", or "append"). Ignored when no leading system message is present.

 'replace'
separator str

String inserted between the instruction and the existing content for "prepend" and "append". Empty string allowed.

 '\n\n'

Raises:

Type Description
ValueError

If mode is not one of the supported values.
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.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
 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
class SystemPromptFormatter(BaseFormatter):
    """Sets or merges the leading system message in each chat with `memory["instruction"]`.

    Only the leading system message (`chat[0]` when its role is `"system"`) participates; system messages
    elsewhere in the chat are left untouched. When a leading system message is present, `mode` controls how
    the instruction combines with its content:

        - `"replace"`: the instruction becomes the content.
        - `"prepend"`: `instruction + separator + existing`.
        - `"append"`: `existing + separator + instruction`.

    When no leading system message is present, all three modes insert a single
    `{"role": "system", "content": instruction}` at position 0, so `separator` has no effect. Message dicts are
    copied, so the caller's structures are not mutated. Every input shape yields exactly one leading system
    message.

    The token path (`apply_to_ids`) decodes, re-templates, and re-encodes, so it cannot merge with a system
    prompt that decoding has flattened into text; `mode` does not apply there, and every mode behaves as
    `"replace"`. Prefer message-level entry for chat input.

    Args:
        mode: How the instruction combines with an existing leading system message (`"replace"` (default),
            `"prepend"`, or `"append"`). Ignored when no leading system message is present.
        separator: String inserted between the instruction and the existing content for `"prepend"` and
            `"append"`. Empty string allowed.

    Raises:
        ValueError: If `mode` is not one of the supported values.
    """

    def __init__(self, mode: str = "replace", separator: str = "\n\n") -> None:
        if mode not in _MODES:
            raise ValueError(f"SystemPromptFormatter mode must be one of {sorted(_MODES)}; got {mode!r}.")
        self.mode = mode
        self.separator = separator

    def _resolve_instruction(self, memory: Memory) -> str:
        instruction = memory["instruction"] if "instruction" in memory else memory.get("instruction")
        if not isinstance(instruction, str):
            raise TypeError(
                f"SystemPromptFormatter expects memory['instruction'] to be a str; got {type(instruction).__name__}."
            )
        return instruction

    def apply_to_messages(
        self,
        messages: list[list[dict]],
        memory: Memory,
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]]:
        instruction = self._resolve_instruction(memory)

        out: list[list[dict]] = []
        for chat in messages:
            chat = list(chat)
            if chat and chat[0].get("role") == "system":
                existing = chat[0].get("content", "")
                if self.mode == "prepend":
                    content = instruction + self.separator + existing
                elif self.mode == "append":
                    content = existing + self.separator + instruction
                else:
                    content = instruction
                chat[0] = {"role": "system", "content": content}
            else:
                chat.insert(0, {"role": "system", "content": instruction})
            out.append(chat)
        return out

    def apply_to_ids(
        self,
        input_ids: torch.Tensor,
        memory: Memory,
        tokenizer: PreTrainedTokenizerBase,
        runtime_kwargs: dict | None = None,
    ) -> torch.Tensor:
        warnings.warn(
            "SystemPromptFormatter.apply_to_ids decodes → edits → re-tokenizes; prefer message-level entry "
            "(pass chat input to the pipeline so `apply_to_messages` runs).",
            UserWarning,
        )
        instruction = self._resolve_instruction(memory)

        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        decoded = tokenizer.batch_decode(input_ids, skip_special_tokens=True)

        rebuilt: list[list[int]] = []
        for user_text in decoded:
            if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
                templated = tokenizer.apply_chat_template(
                    [
                        {"role": "system", "content": instruction},
                        {"role": "user", "content": user_text},
                    ],
                    tokenize=False,
                    add_generation_prompt=True,
                )
                rebuilt.append(tokenizer.encode(templated, add_special_tokens=False))
            else:
                rebuilt.append(tokenizer.encode(instruction + "\n\n" + user_text, add_special_tokens=False))

        max_len = max(len(seq) for seq in rebuilt)
        pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
        padded = [seq + [pad_id] * (max_len - len(seq)) for seq in rebuilt]
        return torch.tensor(padded, dtype=input_ids.dtype, device=input_ids.device)
mode = mode instance-attribute
separator = separator instance-attribute
__init__(mode='replace', separator='\n\n')
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.py 
45
46
47
48
49
def __init__(self, mode: str = "replace", separator: str = "\n\n") -> None:
    if mode not in _MODES:
        raise ValueError(f"SystemPromptFormatter mode must be one of {sorted(_MODES)}; got {mode!r}.")
    self.mode = mode
    self.separator = separator
_resolve_instruction(memory)
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.py 
51
52
53
54
55
56
57
def _resolve_instruction(self, memory: Memory) -> str:
    instruction = memory["instruction"] if "instruction" in memory else memory.get("instruction")
    if not isinstance(instruction, str):
        raise TypeError(
            f"SystemPromptFormatter expects memory['instruction'] to be a str; got {type(instruction).__name__}."
        )
    return instruction
apply_to_ids(input_ids, memory, tokenizer, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.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
def apply_to_ids(
    self,
    input_ids: torch.Tensor,
    memory: Memory,
    tokenizer: PreTrainedTokenizerBase,
    runtime_kwargs: dict | None = None,
) -> torch.Tensor:
    warnings.warn(
        "SystemPromptFormatter.apply_to_ids decodes → edits → re-tokenizes; prefer message-level entry "
        "(pass chat input to the pipeline so `apply_to_messages` runs).",
        UserWarning,
    )
    instruction = self._resolve_instruction(memory)

    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    decoded = tokenizer.batch_decode(input_ids, skip_special_tokens=True)

    rebuilt: list[list[int]] = []
    for user_text in decoded:
        if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
            templated = tokenizer.apply_chat_template(
                [
                    {"role": "system", "content": instruction},
                    {"role": "user", "content": user_text},
                ],
                tokenize=False,
                add_generation_prompt=True,
            )
            rebuilt.append(tokenizer.encode(templated, add_special_tokens=False))
        else:
            rebuilt.append(tokenizer.encode(instruction + "\n\n" + user_text, add_special_tokens=False))

    max_len = max(len(seq) for seq in rebuilt)
    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
    padded = [seq + [pad_id] * (max_len - len(seq)) for seq in rebuilt]
    return torch.tensor(padded, dtype=input_ids.dtype, device=input_ids.device)
apply_to_messages(messages, memory, runtime_kwargs=None)
Source code in steerability/algorithms/input_control/common/formatters/system_prompt.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def apply_to_messages(
    self,
    messages: list[list[dict]],
    memory: Memory,
    runtime_kwargs: dict | None = None,
) -> list[list[dict]]:
    instruction = self._resolve_instruction(memory)

    out: list[list[dict]] = []
    for chat in messages:
        chat = list(chat)
        if chat and chat[0].get("role") == "system":
            existing = chat[0].get("content", "")
            if self.mode == "prepend":
                content = instruction + self.separator + existing
            elif self.mode == "append":
                content = existing + self.separator + instruction
            else:
                content = instruction
            chat[0] = {"role": "system", "content": content}
        else:
            chat.insert(0, {"role": "system", "content": instruction})
        out.append(chat)
    return out

generation

Shared batched generation helper for single-system-prompt input controls.

generate_with_system_prompt runs the task LM over a batch of queries under one system prompt. It applies the tokenizer's chat template when available (otherwise a plain system\n\nquery join), left-pads for the causal-LM batch, restores the tokenizer's padding side, and decodes only the continuation.

generate_with_system_prompt(task_lm, tokenizer, system_prompt, queries, gen_kwargs=None)

Generate one continuation per query under a shared system prompt.

Parameters:

Name Type Description Default
task_lm PreTrainedModel

Causal language model.

 required
tokenizer PreTrainedTokenizerBase

Tokenizer paired with task_lm.

 required
system_prompt str

System prompt prepended to every query.

 required
queries list[str]

User-facing query texts (one continuation generated per query).

 required
gen_kwargs dict | None

Forwarded to task_lm.generate. Defaults to greedy, max_new_tokens=32.

 None

Returns:

Type Description
list[str]

Decoded continuations, one per query, in order.
Source code in steerability/algorithms/input_control/common/generation.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
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 generate_with_system_prompt(
    task_lm: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    system_prompt: str,
    queries: list[str],
    gen_kwargs: dict | None = None,
) -> list[str]:
    """Generate one continuation per query under a shared system prompt.

    Args:
        task_lm: Causal language model.
        tokenizer: Tokenizer paired with `task_lm`.
        system_prompt: System prompt prepended to every query.
        queries: User-facing query texts (one continuation generated per query).
        gen_kwargs: Forwarded to `task_lm.generate`. Defaults to greedy, `max_new_tokens=32`.

    Returns:
        Decoded continuations, one per query, in order.
    """
    if not queries:
        return []
    gen_kwargs = dict(gen_kwargs or {"max_new_tokens": 32, "do_sample": False})

    used_template = bool(getattr(tokenizer, "chat_template", None))
    if used_template:
        prompts_text = [
            tokenizer.apply_chat_template(
                [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": query},
                ],
                tokenize=False,
                add_generation_prompt=True,
            )
            for query in queries
        ]
    else:
        prompts_text = [f"{system_prompt}\n\n{query}" for query in queries]

    original_padding_side = tokenizer.padding_side
    tokenizer.padding_side = "left"
    try:
        encoded = tokenizer(
            prompts_text,
            return_tensors="pt",
            padding=True,
            add_special_tokens=not used_template,
        ).to(task_lm.device)

        pad_id = tokenizer.pad_token_id
        if pad_id is None:
            pad_id = tokenizer.eos_token_id
        if pad_id is not None:
            gen_kwargs.setdefault("pad_token_id", pad_id)

        with torch.no_grad():
            output_ids = task_lm.generate(**encoded, **gen_kwargs)

        prompt_len = encoded["input_ids"].size(1)
        new_ids = output_ids[:, prompt_len:]
        decoded = tokenizer.decode(new_ids, skip_special_tokens=True)
    finally:
        tokenizer.padding_side = original_padding_side

    return decoded

memory

Method-owned, serializable state for input controls.

__all__ = ['Memory', 'PoolMemory', 'TextMemory'] module-attribute

Memory

Bases: Protocol

Method-owned, serializable state.

Methods place their persistent state behind this Protocol so that the framework can save/load it without knowing the concrete type. Subclass or define a new class that satisfies this Protocol; see TextMemory for the primary usage.

Source code in steerability/algorithms/input_control/common/memory/base.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
@runtime_checkable
class Memory(Protocol):
    """Method-owned, serializable state.

    Methods place their persistent state behind this Protocol so that the framework can save/load it without
    knowing the concrete type. Subclass or define a new class that satisfies this Protocol; see `TextMemory`
    for the primary usage.
    """

    def save(self, path: Path) -> None:
        ...

    @classmethod
    def load(cls, path: Path) -> "Memory":
        ...
load(path) classmethod
Source code in steerability/algorithms/input_control/common/memory/base.py 
20
21
22
@classmethod
def load(cls, path: Path) -> "Memory":
    ...
save(path)
Source code in steerability/algorithms/input_control/common/memory/base.py 
17
18
def save(self, path: Path) -> None:
    ...

PoolMemory dataclass

Bases: Generic[T]

Typed pool of items with optional parallel-indexed metadata.

items[i] has metadata[k][i] for each metadata key k. All metadata lists must have len() == len(items). None is used as a placeholder for missing per-item metadata.

Stored as pickle by default (T may be arbitrary). For text-only pools, methods can implement their own JSON path.

Source code in steerability/algorithms/input_control/common/memory/pool.py 
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
43
44
45
46
47
48
49
@dataclass
class PoolMemory(Generic[T]):
    """Typed pool of items with optional parallel-indexed metadata.

    `items[i]` has `metadata[k][i]` for each metadata key `k`. All metadata lists must have
    `len() == len(items)`. None is used as a placeholder for missing per-item metadata.

    Stored as pickle by default (T may be arbitrary). For text-only pools, methods can implement their
    own JSON path.
    """
    items: list[T] = field(default_factory=list)
    metadata: dict[str, list[Any]] = field(default_factory=dict)

    def __len__(self) -> int:
        return len(self.items)

    def add(self, item: T, **per_item_meta: Any) -> None:
        self.items.append(item)
        new_index = len(self.items) - 1
        # extend supplied metadata keys (back-filling earlier indices with None)
        for key, value in per_item_meta.items():
            if key not in self.metadata:
                self.metadata[key] = [None] * new_index
            self.metadata[key].append(value)
        # backfill any pre-existing keys not supplied this call
        for key, lst in self.metadata.items():
            if len(lst) < len(self.items):
                lst.append(None)

    def save(self, path: Path) -> None:
        with open(path, "wb") as f:
            pickle.dump({"items": self.items, "metadata": self.metadata}, f)

    @classmethod
    def load(cls, path: Path) -> "PoolMemory":
        with open(path, "rb") as f:
            data = pickle.load(f)
        return cls(items=data["items"], metadata=data["metadata"])
items = field(default_factory=list) class-attribute instance-attribute
metadata = field(default_factory=dict) class-attribute instance-attribute
__init__(items=list(), metadata=dict())
__len__()
Source code in steerability/algorithms/input_control/common/memory/pool.py 
25
26
def __len__(self) -> int:
    return len(self.items)
add(item, **per_item_meta)
Source code in steerability/algorithms/input_control/common/memory/pool.py 
28
29
30
31
32
33
34
35
36
37
38
39
def add(self, item: T, **per_item_meta: Any) -> None:
    self.items.append(item)
    new_index = len(self.items) - 1
    # extend supplied metadata keys (back-filling earlier indices with None)
    for key, value in per_item_meta.items():
        if key not in self.metadata:
            self.metadata[key] = [None] * new_index
        self.metadata[key].append(value)
    # backfill any pre-existing keys not supplied this call
    for key, lst in self.metadata.items():
        if len(lst) < len(self.items):
            lst.append(None)
load(path) classmethod
Source code in steerability/algorithms/input_control/common/memory/pool.py 
45
46
47
48
49
@classmethod
def load(cls, path: Path) -> "PoolMemory":
    with open(path, "rb") as f:
        data = pickle.load(f)
    return cls(items=data["items"], metadata=data["metadata"])
save(path)
Source code in steerability/algorithms/input_control/common/memory/pool.py 
41
42
43
def save(self, path: Path) -> None:
    with open(path, "wb") as f:
        pickle.dump({"items": self.items, "metadata": self.metadata}, f)

TextMemory dataclass

Named text slots.

Slot values are str or list[str]; nested structures should use a custom Memory implementation instead.

Example

memory = TextMemory(slots={"instruction": "...", "examples": [...]})

Source code in steerability/algorithms/input_control/common/memory/text.py 
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
@dataclass
class TextMemory:
    """Named text slots.

    Slot values are str or list[str]; nested structures should use a custom Memory implementation instead.

    Example:
        memory = TextMemory(slots={"instruction": "...", "examples": [...]})
    """
    slots: dict[str, str | list[str]] = field(default_factory=dict)

    def __getitem__(self, key: str) -> Any:
        return self.slots[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.slots[key] = value

    def __contains__(self, key: str) -> bool:
        return key in self.slots

    def get(self, key: str, default: Any = None) -> Any:
        return self.slots.get(key, default)

    def save(self, path: Path) -> None:
        Path(path).write_text(json.dumps(self.slots, ensure_ascii=False, indent=2))

    @classmethod
    def load(cls, path: Path) -> "TextMemory":
        return cls(slots=json.loads(Path(path).read_text()))
slots = field(default_factory=dict) class-attribute instance-attribute
__contains__(key)
Source code in steerability/algorithms/input_control/common/memory/text.py 
27
28
def __contains__(self, key: str) -> bool:
    return key in self.slots
__getitem__(key)
Source code in steerability/algorithms/input_control/common/memory/text.py 
21
22
def __getitem__(self, key: str) -> Any:
    return self.slots[key]
__init__(slots=dict())
__setitem__(key, value)
Source code in steerability/algorithms/input_control/common/memory/text.py 
24
25
def __setitem__(self, key: str, value: Any) -> None:
    self.slots[key] = value
get(key, default=None)
Source code in steerability/algorithms/input_control/common/memory/text.py 
30
31
def get(self, key: str, default: Any = None) -> Any:
    return self.slots.get(key, default)
load(path) classmethod
Source code in steerability/algorithms/input_control/common/memory/text.py 
36
37
38
@classmethod
def load(cls, path: Path) -> "TextMemory":
    return cls(slots=json.loads(Path(path).read_text()))
save(path)
Source code in steerability/algorithms/input_control/common/memory/text.py 
33
34
def save(self, path: Path) -> None:
    Path(path).write_text(json.dumps(self.slots, ensure_ascii=False, indent=2))

base

The structural contract for input-control state.

Memory

Bases: Protocol

Method-owned, serializable state.

Methods place their persistent state behind this Protocol so that the framework can save/load it without knowing the concrete type. Subclass or define a new class that satisfies this Protocol; see TextMemory for the primary usage.

Source code in steerability/algorithms/input_control/common/memory/base.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
@runtime_checkable
class Memory(Protocol):
    """Method-owned, serializable state.

    Methods place their persistent state behind this Protocol so that the framework can save/load it without
    knowing the concrete type. Subclass or define a new class that satisfies this Protocol; see `TextMemory`
    for the primary usage.
    """

    def save(self, path: Path) -> None:
        ...

    @classmethod
    def load(cls, path: Path) -> "Memory":
        ...
load(path) classmethod
Source code in steerability/algorithms/input_control/common/memory/base.py 
20
21
22
@classmethod
def load(cls, path: Path) -> "Memory":
    ...
save(path)
Source code in steerability/algorithms/input_control/common/memory/base.py 
17
18
def save(self, path: Path) -> None:
    ...

pool

Typed pool with parallel-indexed metadata.

T = TypeVar('T') module-attribute
PoolMemory dataclass

Bases: Generic[T]

Typed pool of items with optional parallel-indexed metadata.

items[i] has metadata[k][i] for each metadata key k. All metadata lists must have len() == len(items). None is used as a placeholder for missing per-item metadata.

Stored as pickle by default (T may be arbitrary). For text-only pools, methods can implement their own JSON path.

Source code in steerability/algorithms/input_control/common/memory/pool.py 
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
43
44
45
46
47
48
49
@dataclass
class PoolMemory(Generic[T]):
    """Typed pool of items with optional parallel-indexed metadata.

    `items[i]` has `metadata[k][i]` for each metadata key `k`. All metadata lists must have
    `len() == len(items)`. None is used as a placeholder for missing per-item metadata.

    Stored as pickle by default (T may be arbitrary). For text-only pools, methods can implement their
    own JSON path.
    """
    items: list[T] = field(default_factory=list)
    metadata: dict[str, list[Any]] = field(default_factory=dict)

    def __len__(self) -> int:
        return len(self.items)

    def add(self, item: T, **per_item_meta: Any) -> None:
        self.items.append(item)
        new_index = len(self.items) - 1
        # extend supplied metadata keys (back-filling earlier indices with None)
        for key, value in per_item_meta.items():
            if key not in self.metadata:
                self.metadata[key] = [None] * new_index
            self.metadata[key].append(value)
        # backfill any pre-existing keys not supplied this call
        for key, lst in self.metadata.items():
            if len(lst) < len(self.items):
                lst.append(None)

    def save(self, path: Path) -> None:
        with open(path, "wb") as f:
            pickle.dump({"items": self.items, "metadata": self.metadata}, f)

    @classmethod
    def load(cls, path: Path) -> "PoolMemory":
        with open(path, "rb") as f:
            data = pickle.load(f)
        return cls(items=data["items"], metadata=data["metadata"])
items = field(default_factory=list) class-attribute instance-attribute
metadata = field(default_factory=dict) class-attribute instance-attribute
__init__(items=list(), metadata=dict())
__len__()
Source code in steerability/algorithms/input_control/common/memory/pool.py 
25
26
def __len__(self) -> int:
    return len(self.items)
add(item, **per_item_meta)
Source code in steerability/algorithms/input_control/common/memory/pool.py 
28
29
30
31
32
33
34
35
36
37
38
39
def add(self, item: T, **per_item_meta: Any) -> None:
    self.items.append(item)
    new_index = len(self.items) - 1
    # extend supplied metadata keys (back-filling earlier indices with None)
    for key, value in per_item_meta.items():
        if key not in self.metadata:
            self.metadata[key] = [None] * new_index
        self.metadata[key].append(value)
    # backfill any pre-existing keys not supplied this call
    for key, lst in self.metadata.items():
        if len(lst) < len(self.items):
            lst.append(None)
load(path) classmethod
Source code in steerability/algorithms/input_control/common/memory/pool.py 
45
46
47
48
49
@classmethod
def load(cls, path: Path) -> "PoolMemory":
    with open(path, "rb") as f:
        data = pickle.load(f)
    return cls(items=data["items"], metadata=data["metadata"])
save(path)
Source code in steerability/algorithms/input_control/common/memory/pool.py 
41
42
43
def save(self, path: Path) -> None:
    with open(path, "wb") as f:
        pickle.dump({"items": self.items, "metadata": self.metadata}, f)

text

Named text slots, JSON-serializable.

TextMemory dataclass

Named text slots.

Slot values are str or list[str]; nested structures should use a custom Memory implementation instead.

Example

memory = TextMemory(slots={"instruction": "...", "examples": [...]})

Source code in steerability/algorithms/input_control/common/memory/text.py 
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
@dataclass
class TextMemory:
    """Named text slots.

    Slot values are str or list[str]; nested structures should use a custom Memory implementation instead.

    Example:
        memory = TextMemory(slots={"instruction": "...", "examples": [...]})
    """
    slots: dict[str, str | list[str]] = field(default_factory=dict)

    def __getitem__(self, key: str) -> Any:
        return self.slots[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.slots[key] = value

    def __contains__(self, key: str) -> bool:
        return key in self.slots

    def get(self, key: str, default: Any = None) -> Any:
        return self.slots.get(key, default)

    def save(self, path: Path) -> None:
        Path(path).write_text(json.dumps(self.slots, ensure_ascii=False, indent=2))

    @classmethod
    def load(cls, path: Path) -> "TextMemory":
        return cls(slots=json.loads(Path(path).read_text()))
slots = field(default_factory=dict) class-attribute instance-attribute
__contains__(key)
Source code in steerability/algorithms/input_control/common/memory/text.py 
27
28
def __contains__(self, key: str) -> bool:
    return key in self.slots
__getitem__(key)
Source code in steerability/algorithms/input_control/common/memory/text.py 
21
22
def __getitem__(self, key: str) -> Any:
    return self.slots[key]
__init__(slots=dict())
__setitem__(key, value)
Source code in steerability/algorithms/input_control/common/memory/text.py 
24
25
def __setitem__(self, key: str, value: Any) -> None:
    self.slots[key] = value
get(key, default=None)
Source code in steerability/algorithms/input_control/common/memory/text.py 
30
31
def get(self, key: str, default: Any = None) -> Any:
    return self.slots.get(key, default)
load(path) classmethod
Source code in steerability/algorithms/input_control/common/memory/text.py 
36
37
38
@classmethod
def load(cls, path: Path) -> "TextMemory":
    return cls(slots=json.loads(Path(path).read_text()))
save(path)
Source code in steerability/algorithms/input_control/common/memory/text.py 
33
34
def save(self, path: Path) -> None:
    Path(path).write_text(json.dumps(self.slots, ensure_ascii=False, indent=2))

pareto

Utility for Pareto analysis over a [n_candidates, n_instances] score matrix.

Used by methods that maintain instance-wise best-so-far records during search (GEPA).

ParetoFrontier

Pareto-frontier utility over a [n_candidates, n_instances] score matrix.

Maximizes scores by default; pass minimize=True for the opposite. Scores are handled internally as a maximization problem.

Source code in steerability/algorithms/input_control/common/pareto.py 
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
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 ParetoFrontier:
    """Pareto-frontier utility over a `[n_candidates, n_instances]` score matrix.

    Maximizes scores by default; pass `minimize=True` for the opposite. Scores are handled internally as a
    maximization problem.
    """

    def __init__(self, scores: np.ndarray, minimize: bool = False) -> None:
        scores = np.asarray(scores, dtype=float)
        if scores.ndim != 2:
            raise ValueError(f"scores must be 2-D [n_candidates, n_instances]; got shape {scores.shape}.")
        if scores.size == 0:
            raise ValueError("scores must be non-empty.")
        self._scores = -scores if minimize else scores
        self._minimize = minimize

    @property
    def per_instance_best(self) -> np.ndarray:
        """Per-instance maximum score across candidates. Shape `[n_instances]`. Returned in the original
        sign convention (i.e. minimum, when `minimize=True`)."""
        best = self._scores.max(axis=0)
        return -best if self._minimize else best

    @property
    def per_instance_winners(self) -> list[set[int]]:
        """For each instance, the set of candidate indices that achieve the per-instance best.

        Multiple candidates may tie.
        """
        best = self._scores.max(axis=0, keepdims=True)
        wins = self._scores == best  # [n_candidates, n_instances]
        return [set(np.where(wins[:, j])[0].tolist()) for j in range(wins.shape[1])]

    @property
    def non_dominated(self) -> list[int]:
        """Candidate indices that are not strictly dominated by any other candidate.

        Candidate `a` is dominated by `b` iff `b >= a` on every instance and `b > a` on at least one.
        """
        n = self._scores.shape[0]
        keep: list[int] = []
        for i in range(n):
            dominated = False
            for j in range(n):
                if i == j:
                    continue
                if np.all(self._scores[j] >= self._scores[i]) and np.any(self._scores[j] > self._scores[i]):
                    dominated = True
                    break
            if not dominated:
                keep.append(i)
        return keep

    def coverage(self, candidate: int) -> int:
        """Number of instances on which `candidate` is a per-instance winner (ties count)."""
        winners = self.per_instance_winners
        return sum(1 for s in winners if candidate in s)

    @property
    def illumination_set(self) -> list[tuple[int, int]]:
        """Algorithm 2's pruned candidate set Ĉ with win-frequencies f[·].

        Returns `(candidate_index, coverage)` for every candidate that (a) is a per-instance winner on
        at least one instance and (b) is not strictly dominated by any other candidate. Equivalent to
        `[(k, self.coverage(k)) for k in self.non_dominated if self.coverage(k) >= 1]`.

        Empty only in degenerate cases; callers should fall back to uniform sampling then.
        """
        result: list[tuple[int, int]] = []
        for k in self.non_dominated:
            cov = self.coverage(k)
            if cov >= 1:
                result.append((k, cov))
        return result
_minimize = minimize instance-attribute
_scores = -scores if minimize else scores instance-attribute
illumination_set property

Algorithm 2's pruned candidate set Ĉ with win-frequencies f[·].

Returns (candidate_index, coverage) for every candidate that (a) is a per-instance winner on at least one instance and (b) is not strictly dominated by any other candidate. Equivalent to [(k, self.coverage(k)) for k in self.non_dominated if self.coverage(k) >= 1].

Empty only in degenerate cases; callers should fall back to uniform sampling then.

non_dominated property

Candidate indices that are not strictly dominated by any other candidate.

Candidate a is dominated by b iff b >= a on every instance and b > a on at least one.

per_instance_best property

Per-instance maximum score across candidates. Shape [n_instances]. Returned in the original sign convention (i.e. minimum, when minimize=True).

per_instance_winners property

For each instance, the set of candidate indices that achieve the per-instance best.

Multiple candidates may tie.

__init__(scores, minimize=False)
Source code in steerability/algorithms/input_control/common/pareto.py 
17
18
19
20
21
22
23
24
def __init__(self, scores: np.ndarray, minimize: bool = False) -> None:
    scores = np.asarray(scores, dtype=float)
    if scores.ndim != 2:
        raise ValueError(f"scores must be 2-D [n_candidates, n_instances]; got shape {scores.shape}.")
    if scores.size == 0:
        raise ValueError("scores must be non-empty.")
    self._scores = -scores if minimize else scores
    self._minimize = minimize
coverage(candidate)

Number of instances on which candidate is a per-instance winner (ties count).

Source code in steerability/algorithms/input_control/common/pareto.py 
63
64
65
66
def coverage(self, candidate: int) -> int:
    """Number of instances on which `candidate` is a per-instance winner (ties count)."""
    winners = self.per_instance_winners
    return sum(1 for s in winners if candidate in s)

proposers

Proposers produce candidate items from a seed.

__all__ = ['BaseProposer', 'LLMMetaPromptProposer', 'RetrievalProposer', 'parse_whole', 'parse_fenced_or_whole', 'parse_concise_instruction'] module-attribute

BaseProposer

Bases: ABC

Produce n candidate items from a seed.

Source code in steerability/algorithms/input_control/common/proposers/base.py 
 8
 9
10
11
12
13
14
15
16
17
18
class BaseProposer(ABC):
    """Produce `n` candidate items from a seed."""

    @abstractmethod
    def propose(
        self,
        seed: Any,
        n: int = 1,
        context: dict | None = None,
    ) -> list[Any]:
        ...
propose(seed, n=1, context=None) abstractmethod
Source code in steerability/algorithms/input_control/common/proposers/base.py 
11
12
13
14
15
16
17
18
@abstractmethod
def propose(
    self,
    seed: Any,
    n: int = 1,
    context: dict | None = None,
) -> list[Any]:
    ...

LLMMetaPromptProposer

Bases: BaseProposer

Format a meta-prompt with seed (and optional context), sample n responses from an LLM, and parse each response into one or more candidates.

Parameters:

Name Type Description Default
llm PreTrainedModel

A causal language model.

 required
tokenizer PreTrainedTokenizerBase

Tokenizer paired with llm.

 required
meta_prompt_template str

Python format string. Always receives seed; may also reference any keys present in the context dict passed to propose().

 required
parse_fn Callable[[str], list[Any]]

Maps a single decoded LLM response to a list of candidates (dropping empties). Defaults to parse_whole ("treat the whole stripped response as one candidate"). Pass parse_fenced_or_whole for long instructions (GEPA) or parse_concise_instruction for short system prompts (CPO / PRewrite). See proposers.utils.parsing.

 parse_whole
gen_kwargs dict | None

Forwarded to llm.generate. The default uses sampling with temperature 0.9 so that requesting n > 1 does not yield identical candidates.

 None
use_chat_template bool | None

Controls whether the rendered meta-prompt is wrapped as a single user turn via the tokenizer's chat template before sampling. None (default) wraps iff the tokenizer has a chat template; True wraps iff a template exists (silently raw otherwise); False never wraps. Base models without a template are unaffected.

 None
max_attempts int

Maximum number of sampling rounds in propose.

 4
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.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
 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
class LLMMetaPromptProposer(BaseProposer):
    """Format a meta-prompt with `seed` (and optional `context`), sample `n` responses from an LLM, and
    parse each response into one or more candidates.

    Args:
        llm: A causal language model.
        tokenizer: Tokenizer paired with `llm`.
        meta_prompt_template: Python format string. Always receives `seed`; may also reference any keys
            present in the `context` dict passed to `propose()`.
        parse_fn: Maps a single decoded LLM response to a list of candidates (dropping empties).
            Defaults to `parse_whole` ("treat the whole stripped response as one candidate"). Pass
            `parse_fenced_or_whole` for long instructions (GEPA) or `parse_concise_instruction` for
            short system prompts (CPO / PRewrite). See `proposers.utils.parsing`.
        gen_kwargs: Forwarded to `llm.generate`. The default uses sampling with temperature 0.9 so that
            requesting `n > 1` does not yield identical candidates.
        use_chat_template: Controls whether the rendered meta-prompt is wrapped as a single user turn
            via the tokenizer's chat template before sampling. `None` (default) wraps iff the tokenizer
            has a chat template; `True` wraps iff a template exists (silently raw otherwise); `False`
            never wraps. Base models without a template are unaffected.
        max_attempts: Maximum number of sampling rounds in `propose`.
    """

    def __init__(
        self,
        llm: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        meta_prompt_template: str,
        parse_fn: Callable[[str], list[Any]] = parse_whole,
        gen_kwargs: dict | None = None,
        use_chat_template: bool | None = None,
        max_attempts: int = 4,
    ) -> None:
        self.llm = llm
        self.tokenizer = tokenizer
        self.meta_prompt_template = meta_prompt_template
        self.parse_fn = parse_fn
        self.gen_kwargs = gen_kwargs or {
            "max_new_tokens": 128,
            "do_sample": True,
            "temperature": 0.9,
            "top_p": 0.95,
        }
        self.use_chat_template = use_chat_template
        self.max_attempts = max_attempts

    def _render_meta_prompt(self, seed: Any, context: dict | None) -> str:
        kwargs: dict[str, Any] = {"seed": seed}
        if context:
            kwargs.update(context)
        return self.meta_prompt_template.format(**kwargs)

    def _chat_mode(self) -> bool:
        template_available = bool(getattr(self.tokenizer, "chat_template", None))
        if self.use_chat_template is None:
            return template_available
        return bool(self.use_chat_template) and template_available

    def _sample_responses(self, rendered: str, n: int) -> list[str]:
        chat_mode = self._chat_mode()
        if chat_mode:
            text = self.tokenizer.apply_chat_template(
                [{"role": "user", "content": rendered}],
                tokenize=False,
                add_generation_prompt=True,
            )
        else:
            text = rendered
        encoded = self.tokenizer(
            text, return_tensors="pt", add_special_tokens=not chat_mode
        ).to(self.llm.device)
        input_ids = encoded["input_ids"]
        attention_mask = encoded.get("attention_mask")

        gen_kwargs = dict(self.gen_kwargs)
        pad_id = self.tokenizer.pad_token_id
        if pad_id is None:
            pad_id = self.tokenizer.eos_token_id
        if pad_id is not None:
            gen_kwargs.setdefault("pad_token_id", pad_id)

        # greedy / beam-1 decoding does not support `num_return_sequences > 1`. In those cases call
        # `generate` sequentially. (Greedy with n > 1 is mostly useful for tests; in production callers
        # pass `do_sample=True` so the parallel path applies.)
        do_sample = gen_kwargs.get("do_sample", False)
        num_beams = gen_kwargs.get("num_beams", 1)
        parallel_ok = do_sample or num_beams > 1

        prompt_len = input_ids.size(1)
        decoded: list[str] = []

        if parallel_ok and n > 1:
            gen_kwargs.setdefault("num_return_sequences", n)
            with torch.no_grad():
                output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **gen_kwargs)
            decoded.extend(
                self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
            )

        while len(decoded) < n:
            single_kwargs = dict(gen_kwargs)
            single_kwargs["num_return_sequences"] = 1
            with torch.no_grad():
                output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **single_kwargs)
            decoded.extend(
                self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
            )

        return decoded[:n]

    def propose(
        self,
        seed: Any,
        n: int = 1,
        context: dict | None = None,
    ) -> list[Any]:
        rendered = self._render_meta_prompt(seed, context)

        # resample until `n` survive or `max_attempts`; return whatever was collected
        candidates: list[Any] = []
        for _ in range(max(self.max_attempts, 1)):
            if len(candidates) >= n:
                break
            for response in self._sample_responses(rendered, n):
                candidates.extend(self.parse_fn(response))
        return candidates[:n] if len(candidates) >= n else candidates
gen_kwargs = gen_kwargs or {'max_new_tokens': 128, 'do_sample': True, 'temperature': 0.9, 'top_p': 0.95} instance-attribute
llm = llm instance-attribute
max_attempts = max_attempts instance-attribute
meta_prompt_template = meta_prompt_template instance-attribute
parse_fn = parse_fn instance-attribute
tokenizer = tokenizer instance-attribute
use_chat_template = use_chat_template instance-attribute
__init__(llm, tokenizer, meta_prompt_template, parse_fn=parse_whole, gen_kwargs=None, use_chat_template=None, max_attempts=4)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __init__(
    self,
    llm: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    meta_prompt_template: str,
    parse_fn: Callable[[str], list[Any]] = parse_whole,
    gen_kwargs: dict | None = None,
    use_chat_template: bool | None = None,
    max_attempts: int = 4,
) -> None:
    self.llm = llm
    self.tokenizer = tokenizer
    self.meta_prompt_template = meta_prompt_template
    self.parse_fn = parse_fn
    self.gen_kwargs = gen_kwargs or {
        "max_new_tokens": 128,
        "do_sample": True,
        "temperature": 0.9,
        "top_p": 0.95,
    }
    self.use_chat_template = use_chat_template
    self.max_attempts = max_attempts
_chat_mode()
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
64
65
66
67
68
def _chat_mode(self) -> bool:
    template_available = bool(getattr(self.tokenizer, "chat_template", None))
    if self.use_chat_template is None:
        return template_available
    return bool(self.use_chat_template) and template_available
_render_meta_prompt(seed, context)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
58
59
60
61
62
def _render_meta_prompt(self, seed: Any, context: dict | None) -> str:
    kwargs: dict[str, Any] = {"seed": seed}
    if context:
        kwargs.update(context)
    return self.meta_prompt_template.format(**kwargs)
_sample_responses(rendered, n)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
 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
def _sample_responses(self, rendered: str, n: int) -> list[str]:
    chat_mode = self._chat_mode()
    if chat_mode:
        text = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": rendered}],
            tokenize=False,
            add_generation_prompt=True,
        )
    else:
        text = rendered
    encoded = self.tokenizer(
        text, return_tensors="pt", add_special_tokens=not chat_mode
    ).to(self.llm.device)
    input_ids = encoded["input_ids"]
    attention_mask = encoded.get("attention_mask")

    gen_kwargs = dict(self.gen_kwargs)
    pad_id = self.tokenizer.pad_token_id
    if pad_id is None:
        pad_id = self.tokenizer.eos_token_id
    if pad_id is not None:
        gen_kwargs.setdefault("pad_token_id", pad_id)

    # greedy / beam-1 decoding does not support `num_return_sequences > 1`. In those cases call
    # `generate` sequentially. (Greedy with n > 1 is mostly useful for tests; in production callers
    # pass `do_sample=True` so the parallel path applies.)
    do_sample = gen_kwargs.get("do_sample", False)
    num_beams = gen_kwargs.get("num_beams", 1)
    parallel_ok = do_sample or num_beams > 1

    prompt_len = input_ids.size(1)
    decoded: list[str] = []

    if parallel_ok and n > 1:
        gen_kwargs.setdefault("num_return_sequences", n)
        with torch.no_grad():
            output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **gen_kwargs)
        decoded.extend(
            self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
        )

    while len(decoded) < n:
        single_kwargs = dict(gen_kwargs)
        single_kwargs["num_return_sequences"] = 1
        with torch.no_grad():
            output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **single_kwargs)
        decoded.extend(
            self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
        )

    return decoded[:n]
propose(seed, n=1, context=None)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def propose(
    self,
    seed: Any,
    n: int = 1,
    context: dict | None = None,
) -> list[Any]:
    rendered = self._render_meta_prompt(seed, context)

    # resample until `n` survive or `max_attempts`; return whatever was collected
    candidates: list[Any] = []
    for _ in range(max(self.max_attempts, 1)):
        if len(candidates) >= n:
            break
        for response in self._sample_responses(rendered, n):
            candidates.extend(self.parse_fn(response))
    return candidates[:n] if len(candidates) >= n else candidates

RetrievalProposer

Bases: BaseProposer

Encode the seed (treated as a query) and retrieve top-n items from a corpus index.

Parameters:

Name Type Description Default
encoder _Encoder

Object exposing encode(text) -> embedding.

 required
corpus_index _CorpusIndex

Object exposing query(embedding, top_k) -> list[dict]. Each returned dict should include at least "text" and "score"; methods may attach a "meta" field.

 required
k_max int

Hard upper bound on n accepted by propose.

 100
Source code in steerability/algorithms/input_control/common/proposers/retrieval.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
class RetrievalProposer(BaseProposer):
    """Encode the seed (treated as a query) and retrieve top-n items from a corpus index.

    Args:
        encoder: Object exposing `encode(text) -> embedding`.
        corpus_index: Object exposing `query(embedding, top_k) -> list[dict]`. Each returned dict should
            include at least `"text"` and `"score"`; methods may attach a `"meta"` field.
        k_max: Hard upper bound on `n` accepted by `propose`.
    """

    def __init__(self, encoder: _Encoder, corpus_index: _CorpusIndex, k_max: int = 100) -> None:
        self.encoder = encoder
        self.corpus_index = corpus_index
        self.k_max = k_max

    def propose(
        self,
        seed: Any,
        n: int = 1,
        context: dict | None = None,
    ) -> list[dict]:
        if n > self.k_max:
            raise ValueError(f"Requested n={n} exceeds k_max={self.k_max}.")
        query_text = seed if isinstance(seed, str) else str(seed)
        embedding = self.encoder.encode(query_text)
        return list(self.corpus_index.query(embedding, top_k=n))
corpus_index = corpus_index instance-attribute
encoder = encoder instance-attribute
k_max = k_max instance-attribute
__init__(encoder, corpus_index, k_max=100)
Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
33
34
35
36
def __init__(self, encoder: _Encoder, corpus_index: _CorpusIndex, k_max: int = 100) -> None:
    self.encoder = encoder
    self.corpus_index = corpus_index
    self.k_max = k_max
propose(seed, n=1, context=None)
Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
38
39
40
41
42
43
44
45
46
47
48
def propose(
    self,
    seed: Any,
    n: int = 1,
    context: dict | None = None,
) -> list[dict]:
    if n > self.k_max:
        raise ValueError(f"Requested n={n} exceeds k_max={self.k_max}.")
    query_text = seed if isinstance(seed, str) else str(seed)
    embedding = self.encoder.encode(query_text)
    return list(self.corpus_index.query(embedding, top_k=n))

parse_concise_instruction(text)

Short-system-prompt parser (CPO / PRewrite).

Reduce a free-text LLM response to a single short instruction:

  1. Split into paragraphs; skip leading preamble paragraphs (single-line cue/lead-in/ acknowledgement) as long as a later paragraph exists; keep the first remaining paragraph.
  2. Strip a leading enumerator/bullet ("1. ", "2) ", "- ") from the kept paragraph.
  3. Drop a leading label line ("Refined prompt:") inside the kept paragraph.
  4. If >= 2 sentences and the last lacks terminal punctuation, drop the trailing fragment.
  5. Balance quotes AFTER fragment-dropping (matching pair, else orphaned edge quote).
  6. Reject degenerate results: preamble-only, fewer than 2 words, or an unterminated candidate longer than _MAX_UNTERMINATED_WORDS (max_new_tokens truncation junk).

Returns [instruction] for accepted content, else [] (so callers drop the candidate).

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
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
def parse_concise_instruction(text: str) -> list[str]:
    """Short-system-prompt parser (CPO / PRewrite).

    Reduce a free-text LLM response to a single short instruction:

    1. Split into paragraphs; skip leading preamble paragraphs (single-line cue/lead-in/
       acknowledgement) as long as a later paragraph exists; keep the first remaining paragraph.
    2. Strip a leading enumerator/bullet ("1. ", "2) ", "- ") from the kept paragraph.
    3. Drop a leading label line ("Refined prompt:") inside the kept paragraph.
    4. If >= 2 sentences and the last lacks terminal punctuation, drop the trailing fragment.
    5. Balance quotes AFTER fragment-dropping (matching pair, else orphaned edge quote).
    6. Reject degenerate results: preamble-only, fewer than 2 words, or an unterminated
       candidate longer than `_MAX_UNTERMINATED_WORDS` (max_new_tokens truncation junk).

    Returns `[instruction]` for accepted content, else `[]` (so callers drop the candidate).
    """
    raw = text.strip()
    if not raw:
        return []
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", raw) if p.strip()]
    idx = 0
    while idx < len(paragraphs) - 1 and _is_preamble(paragraphs[idx]):
        idx += 1
    para = paragraphs[idx]
    if _is_preamble(para):
        return []
    para = _ENUMERATOR_RE.sub("", para, count=1)
    lines = para.splitlines()
    if len(lines) > 1 and _LABEL_LINE_RE.match(lines[0].strip()):
        para = "\n".join(lines[1:]).strip()
    parts = _SENTENCE_SPLIT_RE.split(para)
    if len(parts) >= 2 and not _TERMINAL_RE.search(parts[-1].strip()):
        para = " ".join(parts[:-1]).strip()
    para = _strip_quotes(para)
    if not para or _is_preamble(para):
        return []
    if len(para.split()) < 2:
        return []
    if not _TERMINAL_RE.search(para) and len(para.split()) > _MAX_UNTERMINATED_WORDS:
        return []
    return [para]

parse_fenced_or_whole(text)

Long-instruction parser (GEPA).

  1. If one or more ``` fenced blocks exist, take the FIRST block's inner content. A fence that opens but never closes is treated as running to end-of-text.
  2. Else take the whole text, stripping only a single leading meta lead-in line (e.g. "Here is the revised prompt:", "Sure, here is ...:").
  3. Preserve all internal structure: paragraphs, blank lines, bullet lists, length. Do not collapse to one paragraph, drop trailing sentences, or cap length.

Returns [candidate] if non-empty, else [].

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
 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
def parse_fenced_or_whole(text: str) -> list[str]:
    """Long-instruction parser (GEPA).

    1. If one or more ``` fenced blocks exist, take the FIRST block's inner content. A fence that
       opens but never closes is treated as running to end-of-text.
    2. Else take the whole text, stripping only a single leading meta lead-in line
       (e.g. "Here is the revised prompt:", "Sure, here is ...:").
    3. Preserve all internal structure: paragraphs, blank lines, bullet lists, length. Do not collapse
       to one paragraph, drop trailing sentences, or cap length.

    Returns `[candidate]` if non-empty, else `[]`.
    """
    if not text.strip():
        return []

    first_fence = _FENCE_RE.search(text)
    if first_fence is not None:
        inner_start = first_fence.end()
        closing_fence = _FENCE_RE.search(text, inner_start)
        if closing_fence is not None:
            inner = text[inner_start:closing_fence.start()]
        else:
            inner = text[inner_start:]
        inner = inner.lstrip("\n").rstrip()
        # drop an optional language tag immediately after the opening fence (e.g. ```text)
        lines = inner.split("\n", 1)
        if lines and re.fullmatch(r"[A-Za-z0-9_+-]{1,20}", lines[0].strip()):
            inner = lines[1] if len(lines) > 1 else ""
        candidate = inner.strip()
        return [candidate] if candidate else []

    paragraphs = re.split(r"(\n\s*\n)", text.strip())
    body = text.strip()
    first_para = paragraphs[0].strip() if paragraphs else ""
    if _is_preamble(first_para):
        # drop the single leading lead-in paragraph, preserve the rest verbatim
        remainder = "".join(paragraphs[1:]).lstrip("\n")
        body = remainder.strip()
    candidate = body.strip()
    return [candidate] if candidate else []

parse_whole(text)

Default parser. Return the whole stripped response as a single candidate.

Returns [] when the response is empty or whitespace-only.

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
65
66
67
68
69
70
71
def parse_whole(text: str) -> list[str]:
    """Default parser. Return the whole stripped response as a single candidate.

    Returns `[]` when the response is empty or whitespace-only.
    """
    stripped = text.strip()
    return [stripped] if stripped else []

base

Produce n candidate items from a seed.

BaseProposer

Bases: ABC

Produce n candidate items from a seed.

Source code in steerability/algorithms/input_control/common/proposers/base.py 
 8
 9
10
11
12
13
14
15
16
17
18
class BaseProposer(ABC):
    """Produce `n` candidate items from a seed."""

    @abstractmethod
    def propose(
        self,
        seed: Any,
        n: int = 1,
        context: dict | None = None,
    ) -> list[Any]:
        ...
propose(seed, n=1, context=None) abstractmethod
Source code in steerability/algorithms/input_control/common/proposers/base.py 
11
12
13
14
15
16
17
18
@abstractmethod
def propose(
    self,
    seed: Any,
    n: int = 1,
    context: dict | None = None,
) -> list[Any]:
    ...

llm_meta_prompt

Sample candidates from an LLM via a templated meta-prompt.

LLMMetaPromptProposer

Bases: BaseProposer

Format a meta-prompt with seed (and optional context), sample n responses from an LLM, and parse each response into one or more candidates.

Parameters:

Name Type Description Default
llm PreTrainedModel

A causal language model.

 required
tokenizer PreTrainedTokenizerBase

Tokenizer paired with llm.

 required
meta_prompt_template str

Python format string. Always receives seed; may also reference any keys present in the context dict passed to propose().

 required
parse_fn Callable[[str], list[Any]]

Maps a single decoded LLM response to a list of candidates (dropping empties). Defaults to parse_whole ("treat the whole stripped response as one candidate"). Pass parse_fenced_or_whole for long instructions (GEPA) or parse_concise_instruction for short system prompts (CPO / PRewrite). See proposers.utils.parsing.

 parse_whole
gen_kwargs dict | None

Forwarded to llm.generate. The default uses sampling with temperature 0.9 so that requesting n > 1 does not yield identical candidates.

 None
use_chat_template bool | None

Controls whether the rendered meta-prompt is wrapped as a single user turn via the tokenizer's chat template before sampling. None (default) wraps iff the tokenizer has a chat template; True wraps iff a template exists (silently raw otherwise); False never wraps. Base models without a template are unaffected.

 None
max_attempts int

Maximum number of sampling rounds in propose.

 4
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.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
 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
class LLMMetaPromptProposer(BaseProposer):
    """Format a meta-prompt with `seed` (and optional `context`), sample `n` responses from an LLM, and
    parse each response into one or more candidates.

    Args:
        llm: A causal language model.
        tokenizer: Tokenizer paired with `llm`.
        meta_prompt_template: Python format string. Always receives `seed`; may also reference any keys
            present in the `context` dict passed to `propose()`.
        parse_fn: Maps a single decoded LLM response to a list of candidates (dropping empties).
            Defaults to `parse_whole` ("treat the whole stripped response as one candidate"). Pass
            `parse_fenced_or_whole` for long instructions (GEPA) or `parse_concise_instruction` for
            short system prompts (CPO / PRewrite). See `proposers.utils.parsing`.
        gen_kwargs: Forwarded to `llm.generate`. The default uses sampling with temperature 0.9 so that
            requesting `n > 1` does not yield identical candidates.
        use_chat_template: Controls whether the rendered meta-prompt is wrapped as a single user turn
            via the tokenizer's chat template before sampling. `None` (default) wraps iff the tokenizer
            has a chat template; `True` wraps iff a template exists (silently raw otherwise); `False`
            never wraps. Base models without a template are unaffected.
        max_attempts: Maximum number of sampling rounds in `propose`.
    """

    def __init__(
        self,
        llm: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        meta_prompt_template: str,
        parse_fn: Callable[[str], list[Any]] = parse_whole,
        gen_kwargs: dict | None = None,
        use_chat_template: bool | None = None,
        max_attempts: int = 4,
    ) -> None:
        self.llm = llm
        self.tokenizer = tokenizer
        self.meta_prompt_template = meta_prompt_template
        self.parse_fn = parse_fn
        self.gen_kwargs = gen_kwargs or {
            "max_new_tokens": 128,
            "do_sample": True,
            "temperature": 0.9,
            "top_p": 0.95,
        }
        self.use_chat_template = use_chat_template
        self.max_attempts = max_attempts

    def _render_meta_prompt(self, seed: Any, context: dict | None) -> str:
        kwargs: dict[str, Any] = {"seed": seed}
        if context:
            kwargs.update(context)
        return self.meta_prompt_template.format(**kwargs)

    def _chat_mode(self) -> bool:
        template_available = bool(getattr(self.tokenizer, "chat_template", None))
        if self.use_chat_template is None:
            return template_available
        return bool(self.use_chat_template) and template_available

    def _sample_responses(self, rendered: str, n: int) -> list[str]:
        chat_mode = self._chat_mode()
        if chat_mode:
            text = self.tokenizer.apply_chat_template(
                [{"role": "user", "content": rendered}],
                tokenize=False,
                add_generation_prompt=True,
            )
        else:
            text = rendered
        encoded = self.tokenizer(
            text, return_tensors="pt", add_special_tokens=not chat_mode
        ).to(self.llm.device)
        input_ids = encoded["input_ids"]
        attention_mask = encoded.get("attention_mask")

        gen_kwargs = dict(self.gen_kwargs)
        pad_id = self.tokenizer.pad_token_id
        if pad_id is None:
            pad_id = self.tokenizer.eos_token_id
        if pad_id is not None:
            gen_kwargs.setdefault("pad_token_id", pad_id)

        # greedy / beam-1 decoding does not support `num_return_sequences > 1`. In those cases call
        # `generate` sequentially. (Greedy with n > 1 is mostly useful for tests; in production callers
        # pass `do_sample=True` so the parallel path applies.)
        do_sample = gen_kwargs.get("do_sample", False)
        num_beams = gen_kwargs.get("num_beams", 1)
        parallel_ok = do_sample or num_beams > 1

        prompt_len = input_ids.size(1)
        decoded: list[str] = []

        if parallel_ok and n > 1:
            gen_kwargs.setdefault("num_return_sequences", n)
            with torch.no_grad():
                output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **gen_kwargs)
            decoded.extend(
                self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
            )

        while len(decoded) < n:
            single_kwargs = dict(gen_kwargs)
            single_kwargs["num_return_sequences"] = 1
            with torch.no_grad():
                output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **single_kwargs)
            decoded.extend(
                self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
            )

        return decoded[:n]

    def propose(
        self,
        seed: Any,
        n: int = 1,
        context: dict | None = None,
    ) -> list[Any]:
        rendered = self._render_meta_prompt(seed, context)

        # resample until `n` survive or `max_attempts`; return whatever was collected
        candidates: list[Any] = []
        for _ in range(max(self.max_attempts, 1)):
            if len(candidates) >= n:
                break
            for response in self._sample_responses(rendered, n):
                candidates.extend(self.parse_fn(response))
        return candidates[:n] if len(candidates) >= n else candidates
gen_kwargs = gen_kwargs or {'max_new_tokens': 128, 'do_sample': True, 'temperature': 0.9, 'top_p': 0.95} instance-attribute
llm = llm instance-attribute
max_attempts = max_attempts instance-attribute
meta_prompt_template = meta_prompt_template instance-attribute
parse_fn = parse_fn instance-attribute
tokenizer = tokenizer instance-attribute
use_chat_template = use_chat_template instance-attribute
__init__(llm, tokenizer, meta_prompt_template, parse_fn=parse_whole, gen_kwargs=None, use_chat_template=None, max_attempts=4)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __init__(
    self,
    llm: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    meta_prompt_template: str,
    parse_fn: Callable[[str], list[Any]] = parse_whole,
    gen_kwargs: dict | None = None,
    use_chat_template: bool | None = None,
    max_attempts: int = 4,
) -> None:
    self.llm = llm
    self.tokenizer = tokenizer
    self.meta_prompt_template = meta_prompt_template
    self.parse_fn = parse_fn
    self.gen_kwargs = gen_kwargs or {
        "max_new_tokens": 128,
        "do_sample": True,
        "temperature": 0.9,
        "top_p": 0.95,
    }
    self.use_chat_template = use_chat_template
    self.max_attempts = max_attempts
_chat_mode()
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
64
65
66
67
68
def _chat_mode(self) -> bool:
    template_available = bool(getattr(self.tokenizer, "chat_template", None))
    if self.use_chat_template is None:
        return template_available
    return bool(self.use_chat_template) and template_available
_render_meta_prompt(seed, context)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
58
59
60
61
62
def _render_meta_prompt(self, seed: Any, context: dict | None) -> str:
    kwargs: dict[str, Any] = {"seed": seed}
    if context:
        kwargs.update(context)
    return self.meta_prompt_template.format(**kwargs)
_sample_responses(rendered, n)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
 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
def _sample_responses(self, rendered: str, n: int) -> list[str]:
    chat_mode = self._chat_mode()
    if chat_mode:
        text = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": rendered}],
            tokenize=False,
            add_generation_prompt=True,
        )
    else:
        text = rendered
    encoded = self.tokenizer(
        text, return_tensors="pt", add_special_tokens=not chat_mode
    ).to(self.llm.device)
    input_ids = encoded["input_ids"]
    attention_mask = encoded.get("attention_mask")

    gen_kwargs = dict(self.gen_kwargs)
    pad_id = self.tokenizer.pad_token_id
    if pad_id is None:
        pad_id = self.tokenizer.eos_token_id
    if pad_id is not None:
        gen_kwargs.setdefault("pad_token_id", pad_id)

    # greedy / beam-1 decoding does not support `num_return_sequences > 1`. In those cases call
    # `generate` sequentially. (Greedy with n > 1 is mostly useful for tests; in production callers
    # pass `do_sample=True` so the parallel path applies.)
    do_sample = gen_kwargs.get("do_sample", False)
    num_beams = gen_kwargs.get("num_beams", 1)
    parallel_ok = do_sample or num_beams > 1

    prompt_len = input_ids.size(1)
    decoded: list[str] = []

    if parallel_ok and n > 1:
        gen_kwargs.setdefault("num_return_sequences", n)
        with torch.no_grad():
            output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **gen_kwargs)
        decoded.extend(
            self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
        )

    while len(decoded) < n:
        single_kwargs = dict(gen_kwargs)
        single_kwargs["num_return_sequences"] = 1
        with torch.no_grad():
            output_ids = self.llm.generate(input_ids, attention_mask=attention_mask, **single_kwargs)
        decoded.extend(
            self.tokenizer.decode(output_ids[:, prompt_len:], skip_special_tokens=True)
        )

    return decoded[:n]
propose(seed, n=1, context=None)
Source code in steerability/algorithms/input_control/common/proposers/llm_meta_prompt.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def propose(
    self,
    seed: Any,
    n: int = 1,
    context: dict | None = None,
) -> list[Any]:
    rendered = self._render_meta_prompt(seed, context)

    # resample until `n` survive or `max_attempts`; return whatever was collected
    candidates: list[Any] = []
    for _ in range(max(self.max_attempts, 1)):
        if len(candidates) >= n:
            break
        for response in self._sample_responses(rendered, n):
            candidates.extend(self.parse_fn(response))
    return candidates[:n] if len(candidates) >= n else candidates

retrieval

Retrieve top-n items from a corpus index given a query.

RetrievalProposer

Bases: BaseProposer

Encode the seed (treated as a query) and retrieve top-n items from a corpus index.

Parameters:

Name Type Description Default
encoder _Encoder

Object exposing encode(text) -> embedding.

 required
corpus_index _CorpusIndex

Object exposing query(embedding, top_k) -> list[dict]. Each returned dict should include at least "text" and "score"; methods may attach a "meta" field.

 required
k_max int

Hard upper bound on n accepted by propose.

 100
Source code in steerability/algorithms/input_control/common/proposers/retrieval.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
class RetrievalProposer(BaseProposer):
    """Encode the seed (treated as a query) and retrieve top-n items from a corpus index.

    Args:
        encoder: Object exposing `encode(text) -> embedding`.
        corpus_index: Object exposing `query(embedding, top_k) -> list[dict]`. Each returned dict should
            include at least `"text"` and `"score"`; methods may attach a `"meta"` field.
        k_max: Hard upper bound on `n` accepted by `propose`.
    """

    def __init__(self, encoder: _Encoder, corpus_index: _CorpusIndex, k_max: int = 100) -> None:
        self.encoder = encoder
        self.corpus_index = corpus_index
        self.k_max = k_max

    def propose(
        self,
        seed: Any,
        n: int = 1,
        context: dict | None = None,
    ) -> list[dict]:
        if n > self.k_max:
            raise ValueError(f"Requested n={n} exceeds k_max={self.k_max}.")
        query_text = seed if isinstance(seed, str) else str(seed)
        embedding = self.encoder.encode(query_text)
        return list(self.corpus_index.query(embedding, top_k=n))
corpus_index = corpus_index instance-attribute
encoder = encoder instance-attribute
k_max = k_max instance-attribute
__init__(encoder, corpus_index, k_max=100)
Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
33
34
35
36
def __init__(self, encoder: _Encoder, corpus_index: _CorpusIndex, k_max: int = 100) -> None:
    self.encoder = encoder
    self.corpus_index = corpus_index
    self.k_max = k_max
propose(seed, n=1, context=None)
Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
38
39
40
41
42
43
44
45
46
47
48
def propose(
    self,
    seed: Any,
    n: int = 1,
    context: dict | None = None,
) -> list[dict]:
    if n > self.k_max:
        raise ValueError(f"Requested n={n} exceeds k_max={self.k_max}.")
    query_text = seed if isinstance(seed, str) else str(seed)
    embedding = self.encoder.encode(query_text)
    return list(self.corpus_index.query(embedding, top_k=n))
_CorpusIndex

Bases: Protocol

Minimal duck-typed contract for the corpus index used by RetrievalProposer.

Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
 9
10
11
12
13
14
@runtime_checkable
class _CorpusIndex(Protocol):
    """Minimal duck-typed contract for the corpus index used by `RetrievalProposer`."""

    def query(self, query_embedding, top_k: int) -> list[dict]:
        ...
query(query_embedding, top_k)
Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
13
14
def query(self, query_embedding, top_k: int) -> list[dict]:
    ...
_Encoder

Bases: Protocol

Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
17
18
19
20
@runtime_checkable
class _Encoder(Protocol):
    def encode(self, text: str) -> Any:
        ...
encode(text)
Source code in steerability/algorithms/input_control/common/proposers/retrieval.py 
19
20
def encode(self, text: str) -> Any:
    ...

utils

Helpers used by proposer implementations (not proposers themselves).

parsing

Parser strategies for LLMMetaPromptProposer.

A parser maps a single decoded LLM response to a list of candidate strings (dropping empties). Methods opt into a parser explicitly; the proposer default is parse_whole. Three strategies are provided:

  • parse_whole: the whole stripped response as a single candidate (default; minimal processing).
  • parse_fenced_or_whole: long-instruction parser (GEPA). Extracts the first fenced block when present, else the whole text minus a single meta lead-in line; preserves all internal structure.
  • parse_concise_instruction: short-system-prompt parser (CPO / PRewrite). Reduces an LLM response to a single short instruction paragraph, stripping preamble/enumerators and balancing quotes.
_ACK_RE = re.compile('^(sure|certainly|of course|okay|absolutely|great)[!.…]?$', re.IGNORECASE) module-attribute
_ENUMERATOR_RE = re.compile('^\\s*(\\d{1,2}[.)]|[-*•])\\s+') module-attribute
_FENCE_RE = re.compile('```') module-attribute
_LABEL_LINE_RE = re.compile('^[A-Za-z][A-Za-z ]{0,40}:$') module-attribute
_LEADIN_RE = re.compile("^(sure|certainly|of course|okay|absolutely|here('s| is)|below is|the following is|i('d| would) be happy)\\b", re.IGNORECASE) module-attribute
_MAX_UNTERMINATED_WORDS = 25 module-attribute
_METAWORD_RE = re.compile('\\b(prompt|version|refin\\w*|rewrit\\w*|revis\\w*|improv\\w*|attempt\\w*)\\b', re.IGNORECASE) module-attribute
_QUOTE_CLOSERS = {c: o for o, c in _QUOTE_PAIRS} module-attribute
_QUOTE_OPENERS = {o: c for o, c in _QUOTE_PAIRS} module-attribute
_QUOTE_PAIRS = {('"', '"'), ("'", "'"), ('“', '”'), ('`', '`')} module-attribute
_SENTENCE_SPLIT_RE = re.compile('(?<=[.!?])\\s+') module-attribute
_TERMINAL_RE = re.compile('[.!?:]$') module-attribute
_is_preamble(paragraph)

Single-line meta-discourse: a cue line, a chatty lead-in about the prompt, or a bare acknowledgement. Multi-line paragraphs are never preamble.

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def _is_preamble(paragraph: str) -> bool:
    """Single-line meta-discourse: a cue line, a chatty lead-in about the prompt, or a bare
    acknowledgement. Multi-line paragraphs are never preamble."""
    lines = paragraph.splitlines()
    if len(lines) != 1:
        return False
    line = lines[0].strip()
    if line.endswith(":"):  # "Refined prompt:" / "Here is ...:"
        return True
    if _LEADIN_RE.match(line) and _METAWORD_RE.search(line):  # "Sure, here's a refined version of the prompt."
        return True
    if len(line.split()) <= 4 and _ACK_RE.match(line):  # "Sure!" / "Certainly."
        return True
    return False
_strip_quotes(para)

Strip a matching surrounding pair; otherwise strip an unmatched edge quote whose counterpart does not appear anywhere else in the text (contraction-safe).

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
37
38
39
40
41
42
43
44
45
46
def _strip_quotes(para: str) -> str:
    """Strip a matching surrounding pair; otherwise strip an unmatched edge quote whose
    counterpart does not appear anywhere else in the text (contraction-safe)."""
    if len(para) >= 2 and (para[0], para[-1]) in _QUOTE_PAIRS:
        return para[1:-1].strip()
    if para and para[0] in _QUOTE_OPENERS and _QUOTE_OPENERS[para[0]] not in para[1:]:
        para = para[1:].strip()
    if para and para[-1] in _QUOTE_CLOSERS and _QUOTE_CLOSERS[para[-1]] not in para[:-1]:
        para = para[:-1].strip()
    return para
parse_concise_instruction(text)

Short-system-prompt parser (CPO / PRewrite).

Reduce a free-text LLM response to a single short instruction:

  1. Split into paragraphs; skip leading preamble paragraphs (single-line cue/lead-in/ acknowledgement) as long as a later paragraph exists; keep the first remaining paragraph.
  2. Strip a leading enumerator/bullet ("1. ", "2) ", "- ") from the kept paragraph.
  3. Drop a leading label line ("Refined prompt:") inside the kept paragraph.
  4. If >= 2 sentences and the last lacks terminal punctuation, drop the trailing fragment.
  5. Balance quotes AFTER fragment-dropping (matching pair, else orphaned edge quote).
  6. Reject degenerate results: preamble-only, fewer than 2 words, or an unterminated candidate longer than _MAX_UNTERMINATED_WORDS (max_new_tokens truncation junk).

Returns [instruction] for accepted content, else [] (so callers drop the candidate).

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
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
def parse_concise_instruction(text: str) -> list[str]:
    """Short-system-prompt parser (CPO / PRewrite).

    Reduce a free-text LLM response to a single short instruction:

    1. Split into paragraphs; skip leading preamble paragraphs (single-line cue/lead-in/
       acknowledgement) as long as a later paragraph exists; keep the first remaining paragraph.
    2. Strip a leading enumerator/bullet ("1. ", "2) ", "- ") from the kept paragraph.
    3. Drop a leading label line ("Refined prompt:") inside the kept paragraph.
    4. If >= 2 sentences and the last lacks terminal punctuation, drop the trailing fragment.
    5. Balance quotes AFTER fragment-dropping (matching pair, else orphaned edge quote).
    6. Reject degenerate results: preamble-only, fewer than 2 words, or an unterminated
       candidate longer than `_MAX_UNTERMINATED_WORDS` (max_new_tokens truncation junk).

    Returns `[instruction]` for accepted content, else `[]` (so callers drop the candidate).
    """
    raw = text.strip()
    if not raw:
        return []
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", raw) if p.strip()]
    idx = 0
    while idx < len(paragraphs) - 1 and _is_preamble(paragraphs[idx]):
        idx += 1
    para = paragraphs[idx]
    if _is_preamble(para):
        return []
    para = _ENUMERATOR_RE.sub("", para, count=1)
    lines = para.splitlines()
    if len(lines) > 1 and _LABEL_LINE_RE.match(lines[0].strip()):
        para = "\n".join(lines[1:]).strip()
    parts = _SENTENCE_SPLIT_RE.split(para)
    if len(parts) >= 2 and not _TERMINAL_RE.search(parts[-1].strip()):
        para = " ".join(parts[:-1]).strip()
    para = _strip_quotes(para)
    if not para or _is_preamble(para):
        return []
    if len(para.split()) < 2:
        return []
    if not _TERMINAL_RE.search(para) and len(para.split()) > _MAX_UNTERMINATED_WORDS:
        return []
    return [para]
parse_fenced_or_whole(text)

Long-instruction parser (GEPA).

  1. If one or more ``` fenced blocks exist, take the FIRST block's inner content. A fence that opens but never closes is treated as running to end-of-text.
  2. Else take the whole text, stripping only a single leading meta lead-in line (e.g. "Here is the revised prompt:", "Sure, here is ...:").
  3. Preserve all internal structure: paragraphs, blank lines, bullet lists, length. Do not collapse to one paragraph, drop trailing sentences, or cap length.

Returns [candidate] if non-empty, else [].

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
 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
def parse_fenced_or_whole(text: str) -> list[str]:
    """Long-instruction parser (GEPA).

    1. If one or more ``` fenced blocks exist, take the FIRST block's inner content. A fence that
       opens but never closes is treated as running to end-of-text.
    2. Else take the whole text, stripping only a single leading meta lead-in line
       (e.g. "Here is the revised prompt:", "Sure, here is ...:").
    3. Preserve all internal structure: paragraphs, blank lines, bullet lists, length. Do not collapse
       to one paragraph, drop trailing sentences, or cap length.

    Returns `[candidate]` if non-empty, else `[]`.
    """
    if not text.strip():
        return []

    first_fence = _FENCE_RE.search(text)
    if first_fence is not None:
        inner_start = first_fence.end()
        closing_fence = _FENCE_RE.search(text, inner_start)
        if closing_fence is not None:
            inner = text[inner_start:closing_fence.start()]
        else:
            inner = text[inner_start:]
        inner = inner.lstrip("\n").rstrip()
        # drop an optional language tag immediately after the opening fence (e.g. ```text)
        lines = inner.split("\n", 1)
        if lines and re.fullmatch(r"[A-Za-z0-9_+-]{1,20}", lines[0].strip()):
            inner = lines[1] if len(lines) > 1 else ""
        candidate = inner.strip()
        return [candidate] if candidate else []

    paragraphs = re.split(r"(\n\s*\n)", text.strip())
    body = text.strip()
    first_para = paragraphs[0].strip() if paragraphs else ""
    if _is_preamble(first_para):
        # drop the single leading lead-in paragraph, preserve the rest verbatim
        remainder = "".join(paragraphs[1:]).lstrip("\n")
        body = remainder.strip()
    candidate = body.strip()
    return [candidate] if candidate else []
parse_whole(text)

Default parser. Return the whole stripped response as a single candidate.

Returns [] when the response is empty or whitespace-only.

Source code in steerability/algorithms/input_control/common/proposers/utils/parsing.py 
65
66
67
68
69
70
71
def parse_whole(text: str) -> list[str]:
    """Default parser. Return the whole stripped response as a single candidate.

    Returns `[]` when the response is empty or whitespace-only.
    """
    stripped = text.strip()
    return [stripped] if stripped else []

scorers

Scorers assign a scalar score to one or more candidate prompts.

__all__ = ['BaseScorer', 'TaskEvaluationScorer'] module-attribute

BaseScorer

Bases: ABC

Score one or more prompts, optionally conditioned on a parallel sequence of queries.

Returns a list[float] of length len(prompts). Query-independent scorers accept queries=None. Concrete subclasses define what is scored (task evaluation, judge model, reward function). Selecting from those scores belongs in a BaseSelector.

Source code in steerability/algorithms/input_control/common/scorers/base.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class BaseScorer(ABC):
    """Score one or more prompts, optionally conditioned on a parallel sequence of queries.

    Returns a `list[float]` of length `len(prompts)`. Query-independent scorers accept `queries=None`.
    Concrete subclasses define what is scored (task evaluation, judge model, reward function). Selecting
    from those scores belongs in a `BaseSelector`.
    """

    @abstractmethod
    def score(
        self,
        prompts: Sequence[str],
        queries: Sequence[dict] | None = None,
    ) -> list[float]:
        ...
score(prompts, queries=None) abstractmethod
Source code in steerability/algorithms/input_control/common/scorers/base.py 
16
17
18
19
20
21
22
@abstractmethod
def score(
    self,
    prompts: Sequence[str],
    queries: Sequence[dict] | None = None,
) -> list[float]:
    ...

TaskEvaluationScorer

Bases: BaseScorer

For each candidate prompt, run the task LM over a dev set and aggregate per-row scores into a single scalar.

Each candidate prompt is applied as the system prompt, one response is generated per dev row, and the candidate's score is the mean of row_scorer(response, row) over the dev rows.

Parameters:

Name Type Description Default
task_lm PreTrainedModel

Causal language model used to generate responses on the dev set.

 required
tokenizer PreTrainedTokenizerBase

Tokenizer paired with task_lm.

 required
dev_set Sequence[dict]

Dataset rows; each row must contain at least the keys consumed by format_query (default: "input" for the user prompt) and any fields row_scorer reads (e.g. "reference"). Must be non-empty.

 required
row_scorer SampleScorer

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

 required
gen_kwargs dict | None

Forwarded to task_lm.generate.

 None
max_dev_size int | None

Optional cap on the number of dev rows used per scoring call.

 None
format_query Callable[[dict], str] | None

Callable that turns a dev row into the user-facing query text. Defaults to row -> row["input"].

 None

Raises:

Type Description
ValueError

If dev_set is empty.
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.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
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
class TaskEvaluationScorer(BaseScorer):
    """For each candidate prompt, run the task LM over a dev set and aggregate per-row scores into a
    single scalar.

    Each candidate prompt is applied as the system prompt, one response is generated per dev row,
    and the candidate's score is the mean of `row_scorer(response, row)` over the dev rows.

    Args:
        task_lm: Causal language model used to generate responses on the dev set.
        tokenizer: Tokenizer paired with `task_lm`.
        dev_set: Dataset rows; each row must contain at least the keys consumed by `format_query`
            (default: `"input"` for the user prompt) and any fields `row_scorer` reads (e.g.
            `"reference"`). Must be non-empty.
        row_scorer: `SampleScorer` scoring one `(response, row)` pair; higher is better.
        gen_kwargs: Forwarded to `task_lm.generate`.
        max_dev_size: Optional cap on the number of dev rows used per scoring call.
        format_query: Callable that turns a dev row into the user-facing query text. Defaults to
            `row -> row["input"]`.

    Raises:
        ValueError: If `dev_set` is empty.
    """

    def __init__(
        self,
        task_lm: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        dev_set: Sequence[dict],
        row_scorer: SampleScorer,
        gen_kwargs: dict | None = None,
        max_dev_size: int | None = None,
        format_query: Callable[[dict], str] | None = None,
    ) -> None:
        self.task_lm = task_lm
        self.tokenizer = tokenizer
        self.dev_set = list(dev_set)
        if not self.dev_set:
            raise ValueError("dev_set must be non-empty.")
        self.row_scorer = row_scorer
        self.gen_kwargs = gen_kwargs or {"max_new_tokens": 32, "do_sample": False}
        self.max_dev_size = max_dev_size
        self.format_query = format_query or (lambda row: row["input"])

    def _resolve_dev(self) -> list[dict]:
        if self.max_dev_size is not None:
            return self.dev_set[: self.max_dev_size]
        return self.dev_set

    def _generate_responses(self, prompt: str, dev_rows: list[dict]) -> list[str]:
        """Generate responses for all dev rows under a single candidate prompt."""
        queries = [self.format_query(row) for row in dev_rows]
        return generate_with_system_prompt(
            self.task_lm, self.tokenizer, prompt, queries, gen_kwargs=self.gen_kwargs
        )

    def score(
        self,
        prompts: Sequence[str],
        queries: Sequence[dict] | None = None,
    ) -> list[float]:
        if queries is not None:
            logger.debug("TaskEvaluationScorer ignores per-prompt `queries`; uses dev_set instead.")
        dev = self._resolve_dev()

        scores: list[float] = []
        for prompt in prompts:
            responses = self._generate_responses(prompt, dev)
            row_scores = [float(self.row_scorer(response, row)) for response, row in zip(responses, dev)]
            scores.append(sum(row_scores) / len(row_scores))
        return scores
dev_set = list(dev_set) instance-attribute
format_query = format_query or (lambda row: row['input']) instance-attribute
gen_kwargs = gen_kwargs or {'max_new_tokens': 32, 'do_sample': False} instance-attribute
max_dev_size = max_dev_size instance-attribute
row_scorer = row_scorer instance-attribute
task_lm = task_lm instance-attribute
tokenizer = tokenizer instance-attribute
__init__(task_lm, tokenizer, dev_set, row_scorer, gen_kwargs=None, max_dev_size=None, format_query=None)
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __init__(
    self,
    task_lm: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    dev_set: Sequence[dict],
    row_scorer: SampleScorer,
    gen_kwargs: dict | None = None,
    max_dev_size: int | None = None,
    format_query: Callable[[dict], str] | None = None,
) -> None:
    self.task_lm = task_lm
    self.tokenizer = tokenizer
    self.dev_set = list(dev_set)
    if not self.dev_set:
        raise ValueError("dev_set must be non-empty.")
    self.row_scorer = row_scorer
    self.gen_kwargs = gen_kwargs or {"max_new_tokens": 32, "do_sample": False}
    self.max_dev_size = max_dev_size
    self.format_query = format_query or (lambda row: row["input"])
_generate_responses(prompt, dev_rows)

Generate responses for all dev rows under a single candidate prompt.

Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
64
65
66
67
68
69
def _generate_responses(self, prompt: str, dev_rows: list[dict]) -> list[str]:
    """Generate responses for all dev rows under a single candidate prompt."""
    queries = [self.format_query(row) for row in dev_rows]
    return generate_with_system_prompt(
        self.task_lm, self.tokenizer, prompt, queries, gen_kwargs=self.gen_kwargs
    )
_resolve_dev()
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
59
60
61
62
def _resolve_dev(self) -> list[dict]:
    if self.max_dev_size is not None:
        return self.dev_set[: self.max_dev_size]
    return self.dev_set
score(prompts, queries=None)
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def score(
    self,
    prompts: Sequence[str],
    queries: Sequence[dict] | None = None,
) -> list[float]:
    if queries is not None:
        logger.debug("TaskEvaluationScorer ignores per-prompt `queries`; uses dev_set instead.")
    dev = self._resolve_dev()

    scores: list[float] = []
    for prompt in prompts:
        responses = self._generate_responses(prompt, dev)
        row_scores = [float(self.row_scorer(response, row)) for response, row in zip(responses, dev)]
        scores.append(sum(row_scores) / len(row_scores))
    return scores

base

Score one or more prompts, optionally per-query.

BaseScorer

Bases: ABC

Score one or more prompts, optionally conditioned on a parallel sequence of queries.

Returns a list[float] of length len(prompts). Query-independent scorers accept queries=None. Concrete subclasses define what is scored (task evaluation, judge model, reward function). Selecting from those scores belongs in a BaseSelector.

Source code in steerability/algorithms/input_control/common/scorers/base.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class BaseScorer(ABC):
    """Score one or more prompts, optionally conditioned on a parallel sequence of queries.

    Returns a `list[float]` of length `len(prompts)`. Query-independent scorers accept `queries=None`.
    Concrete subclasses define what is scored (task evaluation, judge model, reward function). Selecting
    from those scores belongs in a `BaseSelector`.
    """

    @abstractmethod
    def score(
        self,
        prompts: Sequence[str],
        queries: Sequence[dict] | None = None,
    ) -> list[float]:
        ...
score(prompts, queries=None) abstractmethod
Source code in steerability/algorithms/input_control/common/scorers/base.py 
16
17
18
19
20
21
22
@abstractmethod
def score(
    self,
    prompts: Sequence[str],
    queries: Sequence[dict] | None = None,
) -> list[float]:
    ...

task_evaluation

Run the task LM on a dev set under each prompt and aggregate via a per-row scorer.

logger = logging.getLogger(__name__) module-attribute
TaskEvaluationScorer

Bases: BaseScorer

For each candidate prompt, run the task LM over a dev set and aggregate per-row scores into a single scalar.

Each candidate prompt is applied as the system prompt, one response is generated per dev row, and the candidate's score is the mean of row_scorer(response, row) over the dev rows.

Parameters:

Name Type Description Default
task_lm PreTrainedModel

Causal language model used to generate responses on the dev set.

 required
tokenizer PreTrainedTokenizerBase

Tokenizer paired with task_lm.

 required
dev_set Sequence[dict]

Dataset rows; each row must contain at least the keys consumed by format_query (default: "input" for the user prompt) and any fields row_scorer reads (e.g. "reference"). Must be non-empty.

 required
row_scorer SampleScorer

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

 required
gen_kwargs dict | None

Forwarded to task_lm.generate.

 None
max_dev_size int | None

Optional cap on the number of dev rows used per scoring call.

 None
format_query Callable[[dict], str] | None

Callable that turns a dev row into the user-facing query text. Defaults to row -> row["input"].

 None

Raises:

Type Description
ValueError

If dev_set is empty.
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.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
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
class TaskEvaluationScorer(BaseScorer):
    """For each candidate prompt, run the task LM over a dev set and aggregate per-row scores into a
    single scalar.

    Each candidate prompt is applied as the system prompt, one response is generated per dev row,
    and the candidate's score is the mean of `row_scorer(response, row)` over the dev rows.

    Args:
        task_lm: Causal language model used to generate responses on the dev set.
        tokenizer: Tokenizer paired with `task_lm`.
        dev_set: Dataset rows; each row must contain at least the keys consumed by `format_query`
            (default: `"input"` for the user prompt) and any fields `row_scorer` reads (e.g.
            `"reference"`). Must be non-empty.
        row_scorer: `SampleScorer` scoring one `(response, row)` pair; higher is better.
        gen_kwargs: Forwarded to `task_lm.generate`.
        max_dev_size: Optional cap on the number of dev rows used per scoring call.
        format_query: Callable that turns a dev row into the user-facing query text. Defaults to
            `row -> row["input"]`.

    Raises:
        ValueError: If `dev_set` is empty.
    """

    def __init__(
        self,
        task_lm: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase,
        dev_set: Sequence[dict],
        row_scorer: SampleScorer,
        gen_kwargs: dict | None = None,
        max_dev_size: int | None = None,
        format_query: Callable[[dict], str] | None = None,
    ) -> None:
        self.task_lm = task_lm
        self.tokenizer = tokenizer
        self.dev_set = list(dev_set)
        if not self.dev_set:
            raise ValueError("dev_set must be non-empty.")
        self.row_scorer = row_scorer
        self.gen_kwargs = gen_kwargs or {"max_new_tokens": 32, "do_sample": False}
        self.max_dev_size = max_dev_size
        self.format_query = format_query or (lambda row: row["input"])

    def _resolve_dev(self) -> list[dict]:
        if self.max_dev_size is not None:
            return self.dev_set[: self.max_dev_size]
        return self.dev_set

    def _generate_responses(self, prompt: str, dev_rows: list[dict]) -> list[str]:
        """Generate responses for all dev rows under a single candidate prompt."""
        queries = [self.format_query(row) for row in dev_rows]
        return generate_with_system_prompt(
            self.task_lm, self.tokenizer, prompt, queries, gen_kwargs=self.gen_kwargs
        )

    def score(
        self,
        prompts: Sequence[str],
        queries: Sequence[dict] | None = None,
    ) -> list[float]:
        if queries is not None:
            logger.debug("TaskEvaluationScorer ignores per-prompt `queries`; uses dev_set instead.")
        dev = self._resolve_dev()

        scores: list[float] = []
        for prompt in prompts:
            responses = self._generate_responses(prompt, dev)
            row_scores = [float(self.row_scorer(response, row)) for response, row in zip(responses, dev)]
            scores.append(sum(row_scores) / len(row_scores))
        return scores
dev_set = list(dev_set) instance-attribute
format_query = format_query or (lambda row: row['input']) instance-attribute
gen_kwargs = gen_kwargs or {'max_new_tokens': 32, 'do_sample': False} instance-attribute
max_dev_size = max_dev_size instance-attribute
row_scorer = row_scorer instance-attribute
task_lm = task_lm instance-attribute
tokenizer = tokenizer instance-attribute
__init__(task_lm, tokenizer, dev_set, row_scorer, gen_kwargs=None, max_dev_size=None, format_query=None)
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __init__(
    self,
    task_lm: PreTrainedModel,
    tokenizer: PreTrainedTokenizerBase,
    dev_set: Sequence[dict],
    row_scorer: SampleScorer,
    gen_kwargs: dict | None = None,
    max_dev_size: int | None = None,
    format_query: Callable[[dict], str] | None = None,
) -> None:
    self.task_lm = task_lm
    self.tokenizer = tokenizer
    self.dev_set = list(dev_set)
    if not self.dev_set:
        raise ValueError("dev_set must be non-empty.")
    self.row_scorer = row_scorer
    self.gen_kwargs = gen_kwargs or {"max_new_tokens": 32, "do_sample": False}
    self.max_dev_size = max_dev_size
    self.format_query = format_query or (lambda row: row["input"])
_generate_responses(prompt, dev_rows)

Generate responses for all dev rows under a single candidate prompt.

Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
64
65
66
67
68
69
def _generate_responses(self, prompt: str, dev_rows: list[dict]) -> list[str]:
    """Generate responses for all dev rows under a single candidate prompt."""
    queries = [self.format_query(row) for row in dev_rows]
    return generate_with_system_prompt(
        self.task_lm, self.tokenizer, prompt, queries, gen_kwargs=self.gen_kwargs
    )
_resolve_dev()
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
59
60
61
62
def _resolve_dev(self) -> list[dict]:
    if self.max_dev_size is not None:
        return self.dev_set[: self.max_dev_size]
    return self.dev_set
score(prompts, queries=None)
Source code in steerability/algorithms/input_control/common/scorers/task_evaluation.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def score(
    self,
    prompts: Sequence[str],
    queries: Sequence[dict] | None = None,
) -> list[float]:
    if queries is not None:
        logger.debug("TaskEvaluationScorer ignores per-prompt `queries`; uses dev_set instead.")
    dev = self._resolve_dev()

    scores: list[float] = []
    for prompt in prompts:
        responses = self._generate_responses(prompt, dev)
        row_scores = [float(self.row_scorer(response, row)) for response, row in zip(responses, dev)]
        scores.append(sum(row_scores) / len(row_scores))
    return scores

selectors

Selectors pick k items from a pool, optionally query-conditioned.

__all__ = ['BaseSelector', 'DenseRetrievalSelector', 'MMRSelector', 'RandomSelector', 'TopKSelector'] module-attribute

BaseSelector

Bases: ABC, Generic[T]

Pick k items from a pool, optionally query-conditioned.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
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
class BaseSelector(ABC, Generic[T]):
    """Pick `k` items from a pool, optionally query-conditioned."""

    @abstractmethod
    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        ...

    def prepare(
        self,
        model=None,
        tokenizer=None,
        data=None,
        **kwargs,
    ) -> None:
        """Optional offline setup (e.g. train an encoder, precompute embeddings).

        Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
        when a selector implements it.
        """
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None) abstractmethod
Source code in steerability/algorithms/input_control/common/selectors/base.py 
13
14
15
16
17
18
19
20
21
@abstractmethod
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    ...

DenseRetrievalSelector

Bases: BaseSelector[T]

Encode a query and pick the nearest neighbors from items in embedding space.

By default, items are encoded on the fly via item_to_text. Subclasses (e.g. EPR's selector) can precompute embeddings during prepare() and store them on the items themselves; pass embedding_key to look up a precomputed embedding from a dict-like item instead of re-encoding.

Parameters:

Name Type Description Default
encoder _Encoder

Object exposing encode(text) -> embedding.

 required
similarity str

Either "cosine" (default) or "dot".

 'cosine'
item_to_text Callable[[T], str] | None

Mapping (item) -> str used when an item lacks a precomputed embedding.

 None
embedding_key str | None

If items are dicts, the key under which a precomputed embedding lives.

 None
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
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
77
78
79
80
81
82
83
84
85
86
87
88
89
class DenseRetrievalSelector(BaseSelector[T]):
    """Encode a query and pick the nearest neighbors from `items` in embedding space.

    By default, items are encoded on the fly via `item_to_text`. Subclasses (e.g. EPR's selector) can
    precompute embeddings during `prepare()` and store them on the items themselves; pass
    `embedding_key` to look up a precomputed embedding from a dict-like item instead of re-encoding.

    Args:
        encoder: Object exposing `encode(text) -> embedding`.
        similarity: Either `"cosine"` (default) or `"dot"`.
        item_to_text: Mapping `(item) -> str` used when an item lacks a precomputed embedding.
        embedding_key: If items are dicts, the key under which a precomputed embedding lives.
    """

    def __init__(
        self,
        encoder: _Encoder,
        similarity: str = "cosine",
        item_to_text: Callable[[T], str] | None = None,
        embedding_key: str | None = None,
    ) -> None:
        if similarity not in ("cosine", "dot"):
            raise ValueError(f"similarity must be 'cosine' or 'dot'; got {similarity}.")
        self.encoder = encoder
        self.similarity = similarity
        self.item_to_text = item_to_text or str
        self.embedding_key = embedding_key

    def _embed_item(self, item: T) -> np.ndarray:
        if self.embedding_key is not None and isinstance(item, dict) and self.embedding_key in item:
            return np.asarray(item[self.embedding_key])
        return np.asarray(self.encoder.encode(self.item_to_text(item)))

    def _embed_query(self, query: Any) -> np.ndarray:
        if isinstance(query, np.ndarray):
            return query
        if isinstance(query, list) and query and isinstance(query[0], (float, int)):
            return np.asarray(query)
        text = query if isinstance(query, str) else self.item_to_text(query)
        return np.asarray(self.encoder.encode(text))

    def _score(self, q: np.ndarray, e: np.ndarray) -> float:
        if self.similarity == "cosine":
            denom = (np.linalg.norm(q) * np.linalg.norm(e)) or 1.0
            return float(np.dot(q, e) / denom)
        return float(np.dot(q, e))

    def _top_k_by_similarity(
        self,
        items: Sequence[T],
        query_embedding: np.ndarray,
        k: int,
    ) -> list[T]:
        items_list = list(items)
        if not items_list:
            return []
        scores = [self._score(query_embedding, self._embed_item(item)) for item in items_list]
        ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
        return [item for item, _ in ranked[:k]]

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        if query is None:
            raise ValueError("DenseRetrievalSelector requires a `query`.")
        query_emb = self._embed_query(query)
        return self._top_k_by_similarity(items, query_emb, k)
embedding_key = embedding_key instance-attribute
encoder = encoder instance-attribute
item_to_text = item_to_text or str instance-attribute
similarity = similarity instance-attribute
__init__(encoder, similarity='cosine', item_to_text=None, embedding_key=None)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    encoder: _Encoder,
    similarity: str = "cosine",
    item_to_text: Callable[[T], str] | None = None,
    embedding_key: str | None = None,
) -> None:
    if similarity not in ("cosine", "dot"):
        raise ValueError(f"similarity must be 'cosine' or 'dot'; got {similarity}.")
    self.encoder = encoder
    self.similarity = similarity
    self.item_to_text = item_to_text or str
    self.embedding_key = embedding_key
_embed_item(item)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
47
48
49
50
def _embed_item(self, item: T) -> np.ndarray:
    if self.embedding_key is not None and isinstance(item, dict) and self.embedding_key in item:
        return np.asarray(item[self.embedding_key])
    return np.asarray(self.encoder.encode(self.item_to_text(item)))
_embed_query(query)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
52
53
54
55
56
57
58
def _embed_query(self, query: Any) -> np.ndarray:
    if isinstance(query, np.ndarray):
        return query
    if isinstance(query, list) and query and isinstance(query[0], (float, int)):
        return np.asarray(query)
    text = query if isinstance(query, str) else self.item_to_text(query)
    return np.asarray(self.encoder.encode(text))
_score(q, e)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
60
61
62
63
64
def _score(self, q: np.ndarray, e: np.ndarray) -> float:
    if self.similarity == "cosine":
        denom = (np.linalg.norm(q) * np.linalg.norm(e)) or 1.0
        return float(np.dot(q, e) / denom)
    return float(np.dot(q, e))
_top_k_by_similarity(items, query_embedding, k)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
66
67
68
69
70
71
72
73
74
75
76
77
def _top_k_by_similarity(
    self,
    items: Sequence[T],
    query_embedding: np.ndarray,
    k: int,
) -> list[T]:
    items_list = list(items)
    if not items_list:
        return []
    scores = [self._score(query_embedding, self._embed_item(item)) for item in items_list]
    ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
    return [item for item, _ in ranked[:k]]
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
79
80
81
82
83
84
85
86
87
88
89
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    if query is None:
        raise ValueError("DenseRetrievalSelector requires a `query`.")
    query_emb = self._embed_query(query)
    return self._top_k_by_similarity(items, query_emb, k)

MMRSelector

Bases: BaseSelector[T]

Maximum-marginal-relevance selection.

At each step, picks the item that maximizes λ · sim(item, query) − (1 − λ) · max_j sim(item, picked_j). Used when callers want a diverse subset of candidates instead of the top-k by raw similarity.

Parameters:

Name Type Description Default
encoder _Encoder

Object exposing encode(text) -> embedding (a 1-D array-like).

 required
lambda_param float

Trade-off in [0, 1]. 1.0 = pure relevance, 0.0 = pure diversity. Defaults to 0.5.

 0.5
item_to_text Callable[[T], str] | None

Mapping (item) -> str. Default is str().

 None
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
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
77
78
79
80
81
82
class MMRSelector(BaseSelector[T]):
    """Maximum-marginal-relevance selection.

    At each step, picks the item that maximizes `λ · sim(item, query) − (1 − λ) · max_j sim(item, picked_j)`.
    Used when callers want a diverse subset of candidates instead of the top-k by raw similarity.

    Args:
        encoder: Object exposing `encode(text) -> embedding` (a 1-D array-like).
        lambda_param: Trade-off in [0, 1]. 1.0 = pure relevance, 0.0 = pure diversity. Defaults to 0.5.
        item_to_text: Mapping `(item) -> str`. Default is `str()`.
    """

    def __init__(
        self,
        encoder: _Encoder,
        lambda_param: float = 0.5,
        item_to_text: Callable[[T], str] | None = None,
    ) -> None:
        if not 0.0 <= lambda_param <= 1.0:
            raise ValueError(f"lambda_param must be in [0, 1]; got {lambda_param}.")
        self.encoder = encoder
        self.lambda_param = lambda_param
        self.item_to_text = item_to_text or str

    @staticmethod
    def _cosine(a: np.ndarray, b: np.ndarray) -> float:
        denom = (np.linalg.norm(a) * np.linalg.norm(b)) or 1.0
        return float(np.dot(a, b) / denom)

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        items_list = list(items)
        if not items_list:
            return []

        item_embs = [np.asarray(self.encoder.encode(self.item_to_text(item))) for item in items_list]
        if query is None:
            query_emb = np.zeros_like(item_embs[0])
        else:
            query_text = query if isinstance(query, str) else self.item_to_text(query)
            query_emb = np.asarray(self.encoder.encode(query_text))

        relevance = np.array([self._cosine(query_emb, emb) for emb in item_embs])
        picked: list[int] = []
        remaining = set(range(len(items_list)))

        while remaining and len(picked) < k:
            if not picked:
                best = max(remaining, key=lambda i: relevance[i])
            else:
                def mmr_score(i: int) -> float:
                    diversity_penalty = max(self._cosine(item_embs[i], item_embs[j]) for j in picked)
                    return self.lambda_param * relevance[i] - (1 - self.lambda_param) * diversity_penalty

                best = max(remaining, key=mmr_score)
            picked.append(best)
            remaining.discard(best)

        return [items_list[i] for i in picked]
encoder = encoder instance-attribute
item_to_text = item_to_text or str instance-attribute
lambda_param = lambda_param instance-attribute
__init__(encoder, lambda_param=0.5, item_to_text=None)
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
31
32
33
34
35
36
37
38
39
40
41
def __init__(
    self,
    encoder: _Encoder,
    lambda_param: float = 0.5,
    item_to_text: Callable[[T], str] | None = None,
) -> None:
    if not 0.0 <= lambda_param <= 1.0:
        raise ValueError(f"lambda_param must be in [0, 1]; got {lambda_param}.")
    self.encoder = encoder
    self.lambda_param = lambda_param
    self.item_to_text = item_to_text or str
_cosine(a, b) staticmethod
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
43
44
45
46
@staticmethod
def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    denom = (np.linalg.norm(a) * np.linalg.norm(b)) or 1.0
    return float(np.dot(a, b) / denom)
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/mmr.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
78
79
80
81
82
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    items_list = list(items)
    if not items_list:
        return []

    item_embs = [np.asarray(self.encoder.encode(self.item_to_text(item))) for item in items_list]
    if query is None:
        query_emb = np.zeros_like(item_embs[0])
    else:
        query_text = query if isinstance(query, str) else self.item_to_text(query)
        query_emb = np.asarray(self.encoder.encode(query_text))

    relevance = np.array([self._cosine(query_emb, emb) for emb in item_embs])
    picked: list[int] = []
    remaining = set(range(len(items_list)))

    while remaining and len(picked) < k:
        if not picked:
            best = max(remaining, key=lambda i: relevance[i])
        else:
            def mmr_score(i: int) -> float:
                diversity_penalty = max(self._cosine(item_embs[i], item_embs[j]) for j in picked)
                return self.lambda_param * relevance[i] - (1 - self.lambda_param) * diversity_penalty

            best = max(remaining, key=mmr_score)
        picked.append(best)
        remaining.discard(best)

    return [items_list[i] for i in picked]

RandomSelector

Bases: BaseSelector[T]

Sample k items uniformly without replacement.

If k >= len(items), returns a copy of all items in random order. The selector owns its RNG, so reseed() gives per-call determinism without disturbing other consumers of the process-global generator.

Source code in steerability/algorithms/input_control/common/selectors/random.py 
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
class RandomSelector(BaseSelector[T]):
    """Sample `k` items uniformly without replacement.

    If `k >= len(items)`, returns a copy of all items in random order. The selector owns its
    RNG, so `reseed()` gives per-call determinism without disturbing other consumers of the
    process-global generator.
    """

    def __init__(self, seed: int | None = None) -> None:
        self._rng = random.Random(seed)

    def reseed(self, seed: int) -> None:
        """Re-seed the owned RNG for one call's deterministic selection."""
        self._rng.seed(seed)

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        items_list = list(items)
        if k >= len(items_list):
            shuffled = items_list[:]
            self._rng.shuffle(shuffled)
            return shuffled
        return self._rng.sample(items_list, k)
_rng = random.Random(seed) instance-attribute
__init__(seed=None)
Source code in steerability/algorithms/input_control/common/selectors/random.py 
20
21
def __init__(self, seed: int | None = None) -> None:
    self._rng = random.Random(seed)
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
reseed(seed)

Re-seed the owned RNG for one call's deterministic selection.

Source code in steerability/algorithms/input_control/common/selectors/random.py 
23
24
25
def reseed(self, seed: int) -> None:
    """Re-seed the owned RNG for one call's deterministic selection."""
    self._rng.seed(seed)
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/random.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    items_list = list(items)
    if k >= len(items_list):
        shuffled = items_list[:]
        self._rng.shuffle(shuffled)
        return shuffled
    return self._rng.sample(items_list, k)

TopKSelector

Bases: BaseSelector[T]

Pick the k items with the highest scores, where scores come from a BaseScorer.

The scorer's score() is called once with the full item list. Items must be str-castable so they can be passed as prompts to the scorer.

Parameters:

Name Type Description Default
scorer BaseScorer

A BaseScorer instance.

 required
item_to_prompt Callable[[T], str] | None

Optional mapping (item) -> str used when items are not strings (e.g. dicts). Defaults to str().

 None
Source code in steerability/algorithms/input_control/common/selectors/top_k.py 
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
43
44
class TopKSelector(BaseSelector[T]):
    """Pick the `k` items with the highest scores, where scores come from a `BaseScorer`.

    The scorer's `score()` is called once with the full item list. Items must be `str`-castable so they
    can be passed as `prompts` to the scorer.

    Args:
        scorer: A `BaseScorer` instance.
        item_to_prompt: Optional mapping `(item) -> str` used when items are not strings (e.g. dicts).
            Defaults to `str()`.
    """

    def __init__(self, scorer: BaseScorer, item_to_prompt: Callable[[T], str] | None = None) -> None:
        self.scorer = scorer
        self.item_to_prompt = item_to_prompt or str

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        items_list = list(items)
        if not items_list:
            return []
        prompts = [self.item_to_prompt(item) for item in items_list]
        queries = None
        if query is not None:
            queries = [query if isinstance(query, dict) else {"query": query}] * len(prompts)
        scores = self.scorer.score(prompts, queries=queries)
        ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
        return [item for item, _ in ranked[:k]]
item_to_prompt = item_to_prompt or str instance-attribute
scorer = scorer instance-attribute
__init__(scorer, item_to_prompt=None)
Source code in steerability/algorithms/input_control/common/selectors/top_k.py 
24
25
26
def __init__(self, scorer: BaseScorer, item_to_prompt: Callable[[T], str] | None = None) -> None:
    self.scorer = scorer
    self.item_to_prompt = item_to_prompt or str
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/top_k.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    items_list = list(items)
    if not items_list:
        return []
    prompts = [self.item_to_prompt(item) for item in items_list]
    queries = None
    if query is not None:
        queries = [query if isinstance(query, dict) else {"query": query}] * len(prompts)
    scores = self.scorer.score(prompts, queries=queries)
    ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
    return [item for item, _ in ranked[:k]]

base

Pick k items from a pool, optionally query-conditioned.

T = TypeVar('T') module-attribute
BaseSelector

Bases: ABC, Generic[T]

Pick k items from a pool, optionally query-conditioned.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
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
class BaseSelector(ABC, Generic[T]):
    """Pick `k` items from a pool, optionally query-conditioned."""

    @abstractmethod
    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        ...

    def prepare(
        self,
        model=None,
        tokenizer=None,
        data=None,
        **kwargs,
    ) -> None:
        """Optional offline setup (e.g. train an encoder, precompute embeddings).

        Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
        when a selector implements it.
        """
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None) abstractmethod
Source code in steerability/algorithms/input_control/common/selectors/base.py 
13
14
15
16
17
18
19
20
21
@abstractmethod
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    ...

dense_retrieval

Pick nearest neighbors of a query in a dense embedding space.

T = TypeVar('T') module-attribute
DenseRetrievalSelector

Bases: BaseSelector[T]

Encode a query and pick the nearest neighbors from items in embedding space.

By default, items are encoded on the fly via item_to_text. Subclasses (e.g. EPR's selector) can precompute embeddings during prepare() and store them on the items themselves; pass embedding_key to look up a precomputed embedding from a dict-like item instead of re-encoding.

Parameters:

Name Type Description Default
encoder _Encoder

Object exposing encode(text) -> embedding.

 required
similarity str

Either "cosine" (default) or "dot".

 'cosine'
item_to_text Callable[[T], str] | None

Mapping (item) -> str used when an item lacks a precomputed embedding.

 None
embedding_key str | None

If items are dicts, the key under which a precomputed embedding lives.

 None
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
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
77
78
79
80
81
82
83
84
85
86
87
88
89
class DenseRetrievalSelector(BaseSelector[T]):
    """Encode a query and pick the nearest neighbors from `items` in embedding space.

    By default, items are encoded on the fly via `item_to_text`. Subclasses (e.g. EPR's selector) can
    precompute embeddings during `prepare()` and store them on the items themselves; pass
    `embedding_key` to look up a precomputed embedding from a dict-like item instead of re-encoding.

    Args:
        encoder: Object exposing `encode(text) -> embedding`.
        similarity: Either `"cosine"` (default) or `"dot"`.
        item_to_text: Mapping `(item) -> str` used when an item lacks a precomputed embedding.
        embedding_key: If items are dicts, the key under which a precomputed embedding lives.
    """

    def __init__(
        self,
        encoder: _Encoder,
        similarity: str = "cosine",
        item_to_text: Callable[[T], str] | None = None,
        embedding_key: str | None = None,
    ) -> None:
        if similarity not in ("cosine", "dot"):
            raise ValueError(f"similarity must be 'cosine' or 'dot'; got {similarity}.")
        self.encoder = encoder
        self.similarity = similarity
        self.item_to_text = item_to_text or str
        self.embedding_key = embedding_key

    def _embed_item(self, item: T) -> np.ndarray:
        if self.embedding_key is not None and isinstance(item, dict) and self.embedding_key in item:
            return np.asarray(item[self.embedding_key])
        return np.asarray(self.encoder.encode(self.item_to_text(item)))

    def _embed_query(self, query: Any) -> np.ndarray:
        if isinstance(query, np.ndarray):
            return query
        if isinstance(query, list) and query and isinstance(query[0], (float, int)):
            return np.asarray(query)
        text = query if isinstance(query, str) else self.item_to_text(query)
        return np.asarray(self.encoder.encode(text))

    def _score(self, q: np.ndarray, e: np.ndarray) -> float:
        if self.similarity == "cosine":
            denom = (np.linalg.norm(q) * np.linalg.norm(e)) or 1.0
            return float(np.dot(q, e) / denom)
        return float(np.dot(q, e))

    def _top_k_by_similarity(
        self,
        items: Sequence[T],
        query_embedding: np.ndarray,
        k: int,
    ) -> list[T]:
        items_list = list(items)
        if not items_list:
            return []
        scores = [self._score(query_embedding, self._embed_item(item)) for item in items_list]
        ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
        return [item for item, _ in ranked[:k]]

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        if query is None:
            raise ValueError("DenseRetrievalSelector requires a `query`.")
        query_emb = self._embed_query(query)
        return self._top_k_by_similarity(items, query_emb, k)
embedding_key = embedding_key instance-attribute
encoder = encoder instance-attribute
item_to_text = item_to_text or str instance-attribute
similarity = similarity instance-attribute
__init__(encoder, similarity='cosine', item_to_text=None, embedding_key=None)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    encoder: _Encoder,
    similarity: str = "cosine",
    item_to_text: Callable[[T], str] | None = None,
    embedding_key: str | None = None,
) -> None:
    if similarity not in ("cosine", "dot"):
        raise ValueError(f"similarity must be 'cosine' or 'dot'; got {similarity}.")
    self.encoder = encoder
    self.similarity = similarity
    self.item_to_text = item_to_text or str
    self.embedding_key = embedding_key
_embed_item(item)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
47
48
49
50
def _embed_item(self, item: T) -> np.ndarray:
    if self.embedding_key is not None and isinstance(item, dict) and self.embedding_key in item:
        return np.asarray(item[self.embedding_key])
    return np.asarray(self.encoder.encode(self.item_to_text(item)))
_embed_query(query)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
52
53
54
55
56
57
58
def _embed_query(self, query: Any) -> np.ndarray:
    if isinstance(query, np.ndarray):
        return query
    if isinstance(query, list) and query and isinstance(query[0], (float, int)):
        return np.asarray(query)
    text = query if isinstance(query, str) else self.item_to_text(query)
    return np.asarray(self.encoder.encode(text))
_score(q, e)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
60
61
62
63
64
def _score(self, q: np.ndarray, e: np.ndarray) -> float:
    if self.similarity == "cosine":
        denom = (np.linalg.norm(q) * np.linalg.norm(e)) or 1.0
        return float(np.dot(q, e) / denom)
    return float(np.dot(q, e))
_top_k_by_similarity(items, query_embedding, k)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
66
67
68
69
70
71
72
73
74
75
76
77
def _top_k_by_similarity(
    self,
    items: Sequence[T],
    query_embedding: np.ndarray,
    k: int,
) -> list[T]:
    items_list = list(items)
    if not items_list:
        return []
    scores = [self._score(query_embedding, self._embed_item(item)) for item in items_list]
    ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
    return [item for item, _ in ranked[:k]]
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
79
80
81
82
83
84
85
86
87
88
89
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    if query is None:
        raise ValueError("DenseRetrievalSelector requires a `query`.")
    query_emb = self._embed_query(query)
    return self._top_k_by_similarity(items, query_emb, k)
_Encoder

Bases: Protocol

Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
13
14
15
16
@runtime_checkable
class _Encoder(Protocol):
    def encode(self, text: str) -> Any:
        ...
encode(text)
Source code in steerability/algorithms/input_control/common/selectors/dense_retrieval.py 
15
16
def encode(self, text: str) -> Any:
    ...

mmr

Maximum marginal relevance for diversity.

T = TypeVar('T') module-attribute
MMRSelector

Bases: BaseSelector[T]

Maximum-marginal-relevance selection.

At each step, picks the item that maximizes λ · sim(item, query) − (1 − λ) · max_j sim(item, picked_j). Used when callers want a diverse subset of candidates instead of the top-k by raw similarity.

Parameters:

Name Type Description Default
encoder _Encoder

Object exposing encode(text) -> embedding (a 1-D array-like).

 required
lambda_param float

Trade-off in [0, 1]. 1.0 = pure relevance, 0.0 = pure diversity. Defaults to 0.5.

 0.5
item_to_text Callable[[T], str] | None

Mapping (item) -> str. Default is str().

 None
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
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
77
78
79
80
81
82
class MMRSelector(BaseSelector[T]):
    """Maximum-marginal-relevance selection.

    At each step, picks the item that maximizes `λ · sim(item, query) − (1 − λ) · max_j sim(item, picked_j)`.
    Used when callers want a diverse subset of candidates instead of the top-k by raw similarity.

    Args:
        encoder: Object exposing `encode(text) -> embedding` (a 1-D array-like).
        lambda_param: Trade-off in [0, 1]. 1.0 = pure relevance, 0.0 = pure diversity. Defaults to 0.5.
        item_to_text: Mapping `(item) -> str`. Default is `str()`.
    """

    def __init__(
        self,
        encoder: _Encoder,
        lambda_param: float = 0.5,
        item_to_text: Callable[[T], str] | None = None,
    ) -> None:
        if not 0.0 <= lambda_param <= 1.0:
            raise ValueError(f"lambda_param must be in [0, 1]; got {lambda_param}.")
        self.encoder = encoder
        self.lambda_param = lambda_param
        self.item_to_text = item_to_text or str

    @staticmethod
    def _cosine(a: np.ndarray, b: np.ndarray) -> float:
        denom = (np.linalg.norm(a) * np.linalg.norm(b)) or 1.0
        return float(np.dot(a, b) / denom)

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        items_list = list(items)
        if not items_list:
            return []

        item_embs = [np.asarray(self.encoder.encode(self.item_to_text(item))) for item in items_list]
        if query is None:
            query_emb = np.zeros_like(item_embs[0])
        else:
            query_text = query if isinstance(query, str) else self.item_to_text(query)
            query_emb = np.asarray(self.encoder.encode(query_text))

        relevance = np.array([self._cosine(query_emb, emb) for emb in item_embs])
        picked: list[int] = []
        remaining = set(range(len(items_list)))

        while remaining and len(picked) < k:
            if not picked:
                best = max(remaining, key=lambda i: relevance[i])
            else:
                def mmr_score(i: int) -> float:
                    diversity_penalty = max(self._cosine(item_embs[i], item_embs[j]) for j in picked)
                    return self.lambda_param * relevance[i] - (1 - self.lambda_param) * diversity_penalty

                best = max(remaining, key=mmr_score)
            picked.append(best)
            remaining.discard(best)

        return [items_list[i] for i in picked]
encoder = encoder instance-attribute
item_to_text = item_to_text or str instance-attribute
lambda_param = lambda_param instance-attribute
__init__(encoder, lambda_param=0.5, item_to_text=None)
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
31
32
33
34
35
36
37
38
39
40
41
def __init__(
    self,
    encoder: _Encoder,
    lambda_param: float = 0.5,
    item_to_text: Callable[[T], str] | None = None,
) -> None:
    if not 0.0 <= lambda_param <= 1.0:
        raise ValueError(f"lambda_param must be in [0, 1]; got {lambda_param}.")
    self.encoder = encoder
    self.lambda_param = lambda_param
    self.item_to_text = item_to_text or str
_cosine(a, b) staticmethod
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
43
44
45
46
@staticmethod
def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    denom = (np.linalg.norm(a) * np.linalg.norm(b)) or 1.0
    return float(np.dot(a, b) / denom)
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/mmr.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
78
79
80
81
82
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    items_list = list(items)
    if not items_list:
        return []

    item_embs = [np.asarray(self.encoder.encode(self.item_to_text(item))) for item in items_list]
    if query is None:
        query_emb = np.zeros_like(item_embs[0])
    else:
        query_text = query if isinstance(query, str) else self.item_to_text(query)
        query_emb = np.asarray(self.encoder.encode(query_text))

    relevance = np.array([self._cosine(query_emb, emb) for emb in item_embs])
    picked: list[int] = []
    remaining = set(range(len(items_list)))

    while remaining and len(picked) < k:
        if not picked:
            best = max(remaining, key=lambda i: relevance[i])
        else:
            def mmr_score(i: int) -> float:
                diversity_penalty = max(self._cosine(item_embs[i], item_embs[j]) for j in picked)
                return self.lambda_param * relevance[i] - (1 - self.lambda_param) * diversity_penalty

            best = max(remaining, key=mmr_score)
        picked.append(best)
        remaining.discard(best)

    return [items_list[i] for i in picked]
_Encoder

Bases: Protocol

Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
13
14
15
16
@runtime_checkable
class _Encoder(Protocol):
    def encode(self, text: str) -> Any:
        ...
encode(text)
Source code in steerability/algorithms/input_control/common/selectors/mmr.py 
15
16
def encode(self, text: str) -> Any:
    ...

random

Sample k items uniformly without replacement.

T = TypeVar('T') module-attribute
RandomSelector

Bases: BaseSelector[T]

Sample k items uniformly without replacement.

If k >= len(items), returns a copy of all items in random order. The selector owns its RNG, so reseed() gives per-call determinism without disturbing other consumers of the process-global generator.

Source code in steerability/algorithms/input_control/common/selectors/random.py 
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
class RandomSelector(BaseSelector[T]):
    """Sample `k` items uniformly without replacement.

    If `k >= len(items)`, returns a copy of all items in random order. The selector owns its
    RNG, so `reseed()` gives per-call determinism without disturbing other consumers of the
    process-global generator.
    """

    def __init__(self, seed: int | None = None) -> None:
        self._rng = random.Random(seed)

    def reseed(self, seed: int) -> None:
        """Re-seed the owned RNG for one call's deterministic selection."""
        self._rng.seed(seed)

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        items_list = list(items)
        if k >= len(items_list):
            shuffled = items_list[:]
            self._rng.shuffle(shuffled)
            return shuffled
        return self._rng.sample(items_list, k)
_rng = random.Random(seed) instance-attribute
__init__(seed=None)
Source code in steerability/algorithms/input_control/common/selectors/random.py 
20
21
def __init__(self, seed: int | None = None) -> None:
    self._rng = random.Random(seed)
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
reseed(seed)

Re-seed the owned RNG for one call's deterministic selection.

Source code in steerability/algorithms/input_control/common/selectors/random.py 
23
24
25
def reseed(self, seed: int) -> None:
    """Re-seed the owned RNG for one call's deterministic selection."""
    self._rng.seed(seed)
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/random.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    items_list = list(items)
    if k >= len(items_list):
        shuffled = items_list[:]
        self._rng.shuffle(shuffled)
        return shuffled
    return self._rng.sample(items_list, k)

top_k

Score items via a BaseScorer and return the top-k.

T = TypeVar('T') module-attribute
TopKSelector

Bases: BaseSelector[T]

Pick the k items with the highest scores, where scores come from a BaseScorer.

The scorer's score() is called once with the full item list. Items must be str-castable so they can be passed as prompts to the scorer.

Parameters:

Name Type Description Default
scorer BaseScorer

A BaseScorer instance.

 required
item_to_prompt Callable[[T], str] | None

Optional mapping (item) -> str used when items are not strings (e.g. dicts). Defaults to str().

 None
Source code in steerability/algorithms/input_control/common/selectors/top_k.py 
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
43
44
class TopKSelector(BaseSelector[T]):
    """Pick the `k` items with the highest scores, where scores come from a `BaseScorer`.

    The scorer's `score()` is called once with the full item list. Items must be `str`-castable so they
    can be passed as `prompts` to the scorer.

    Args:
        scorer: A `BaseScorer` instance.
        item_to_prompt: Optional mapping `(item) -> str` used when items are not strings (e.g. dicts).
            Defaults to `str()`.
    """

    def __init__(self, scorer: BaseScorer, item_to_prompt: Callable[[T], str] | None = None) -> None:
        self.scorer = scorer
        self.item_to_prompt = item_to_prompt or str

    def select(
        self,
        items: Sequence[T],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[T]:
        items_list = list(items)
        if not items_list:
            return []
        prompts = [self.item_to_prompt(item) for item in items_list]
        queries = None
        if query is not None:
            queries = [query if isinstance(query, dict) else {"query": query}] * len(prompts)
        scores = self.scorer.score(prompts, queries=queries)
        ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
        return [item for item, _ in ranked[:k]]
item_to_prompt = item_to_prompt or str instance-attribute
scorer = scorer instance-attribute
__init__(scorer, item_to_prompt=None)
Source code in steerability/algorithms/input_control/common/selectors/top_k.py 
24
25
26
def __init__(self, scorer: BaseScorer, item_to_prompt: Callable[[T], str] | None = None) -> None:
    self.scorer = scorer
    self.item_to_prompt = item_to_prompt or str
prepare(model=None, tokenizer=None, data=None, **kwargs)

Optional offline setup (e.g. train an encoder, precompute embeddings).

Default no-op. Parent controls invoke selector.prepare(...) during their own steer() step when a selector implements it.

Source code in steerability/algorithms/input_control/common/selectors/base.py 
23
24
25
26
27
28
29
30
31
32
33
34
def prepare(
    self,
    model=None,
    tokenizer=None,
    data=None,
    **kwargs,
) -> None:
    """Optional offline setup (e.g. train an encoder, precompute embeddings).

    Default no-op. Parent controls invoke `selector.prepare(...)` during their own `steer()` step
    when a selector implements it.
    """
select(items, query=None, k=1, context=None)
Source code in steerability/algorithms/input_control/common/selectors/top_k.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def select(
    self,
    items: Sequence[T],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[T]:
    items_list = list(items)
    if not items_list:
        return []
    prompts = [self.item_to_prompt(item) for item in items_list]
    queries = None
    if query is not None:
        queries = [query if isinstance(query, dict) else {"query": query}] * len(prompts)
    scores = self.scorer.score(prompts, queries=queries)
    ranked = sorted(zip(items_list, scores), key=lambda pair: pair[1], reverse=True)
    return [item for item, _ in ranked[:k]]