FewShot

steerability.algorithms.input_control.few_shot

args

control

Few-shot learning control for prompt adaptation.

FewShot

Bases: InputControl

Implementation of few-shot learning control for prompt adaptation.

FewShot enables selective behavioral steering by prepending specific examples to user prompts, guiding model responses through demonstration.

The method operates in two modes:

1. Pool-based sampling: Maintains pools of positive and negative examples from which k examples are dynamically selected using configurable sampling strategies (random, semantic similarity, etc.).

2. Runtime injection: Accepts examples directly at inference time through runtime_kwargs, enabling context-specific demonstrations without predefined pools. Useful for dynamic or user-provided examples.

The selected examples are formatted into a system prompt with clear positive/negative labels and prepended to the user query using the model's chat template, allowing the model to learn the desired behavior pattern from the demonstrations.

Runtime keyword arguments:

• positive_examples (list[dict], optional): Positive examples to use for this specific query (overrides pool-based selection).
• negative_examples (list[dict], optional): Negative examples to use for this specific query (overrides pool-based selection).

Notes:

• Requires a tokenizer with chat_template support for optimal formatting
• Examples are automatically labeled as "### Positive example" or "### Negative example"
• When both pools and runtime examples are available, runtime examples take precedence
• If no examples are provided, the original input is returned unchanged
• Keys with a leading underscore in example dicts are reserved by the framework (e.g. _polarity); user fields should not start with _.

Source code in steerability/algorithms/input_control/few_shot/control.py

class FewShot(InputControl):
    """
    Implementation of few-shot learning control for prompt adaptation.

    FewShot enables selective behavioral steering by prepending specific examples to user prompts, guiding model
    responses through demonstration.

    The method operates in two modes:

    1. **Pool-based sampling**: Maintains pools of positive and negative examples from which k examples are dynamically
        selected using configurable sampling strategies (random, semantic similarity, etc.).

    2. **Runtime injection**: Accepts examples directly at inference time through runtime_kwargs, enabling
        context-specific demonstrations without predefined pools. Useful for dynamic or user-provided examples.

    The selected examples are formatted into a system prompt with clear positive/negative labels and prepended to the
    user query using the model's chat template, allowing the model to learn the desired behavior pattern from the
    demonstrations.

    Runtime keyword arguments:

    - `positive_examples` (`list[dict]`, `optional`): Positive examples to use for this specific query (overrides pool-based selection).
    - `negative_examples` (`list[dict]`, `optional`): Negative examples to use for this specific query (overrides pool-based selection).

    Notes:

    - Requires a tokenizer with chat_template support for optimal formatting
    - Examples are automatically labeled as "### Positive example" or "### Negative example"
    - When both pools and runtime examples are available, runtime examples take precedence
    - If no examples are provided, the original input is returned unchanged
    - Keys with a leading underscore in example dicts are reserved by the framework (e.g. `_polarity`);
        user fields should not start with `_`.
    """

    Args = FewShotArgs

    RUNTIME_KWARGS_SCHEMA = [
        {
            "name": "positive_examples",
            "type": "list[dict]",
            "scope": "call",
            "help": "Positive examples for this call, applied to every prompt row; overrides pool-based selection.",
        },
        {
            "name": "negative_examples",
            "type": "list[dict]",
            "scope": "call",
            "help": "Negative examples for this call, applied to every prompt row; overrides pool-based selection.",
        },
    ]

    supports_batching: bool = True

    # placeholders (dataclass attrs from FewShotArgs override these at __init__ time)
    tokenizer: PreTrainedTokenizerBase | None = None
    directive: str | None = None
    positive_example_pool: Sequence[dict] | None = None
    negative_example_pool: Sequence[dict] | None = None
    k_positive: int | None = None
    k_negative: int | None = None
    selector: Any = None  # str | BaseSelector | None — resolved in steer()
    formatter: Any = None  # BaseFormatter | None — resolved in steer()

    # method-owned state populated in steer()
    pool: PoolMemory[dict] | None = None
    _selector: BaseSelector | None = None
    _formatter: FewShotBlockFormatter | None = None

    def steer(
            self,
            model=None,
            tokenizer: PreTrainedTokenizerBase | None = None,
            **kwargs,
    ) -> None:
        self.tokenizer = tokenizer

        # build the example pool
        self.pool = PoolMemory[dict]()
        for example in self.positive_example_pool or []:
            self.pool.add(example, polarity="pos")
        for example in self.negative_example_pool or []:
            self.pool.add(example, polarity="neg")

        # resolve selector argument (instance | name | None) into a BaseSelector[dict]
        self._selector = selector_from_arg(self.selector)

        # selectors that need offline preparation (e.g. EPR) get a chance here
        prepare = getattr(self._selector, "prepare", None)
        if callable(prepare):
            prepare(model=model, tokenizer=tokenizer, data=self.pool)

        # formatter is shared between adapt and adapt_messages
        self._formatter = self.formatter or FewShotBlockFormatter()

    def adapt(
        self,
        input_ids: list[int] | torch.Tensor,
        runtime_kwargs: dict | None = None,
    ) -> list[int] | torch.Tensor:
        """Add few-shot examples to the model's prompt and return adapted token ids.

        Both the chat-template path and the no-template fallback route through the resolved
        `BaseFormatter`, which renders the example block from the pool.

        Assumes `input_ids` represents the user's prompt before any chat templating. Pre-templated
        input will be re-templated and produce malformed output; use `adapt_messages` for chat input.

        Args:
            input_ids: The user's prompt token IDs.
            runtime_kwargs: May contain `positive_examples` / `negative_examples` to override pool sampling.

        Returns:
            The transformed token IDs.

        Raises:
            RuntimeError: If tokenizer is not set (requires calling `steer()` first)

        Warnings:
            UserWarning: Issued when no examples are configured or none remain after selection.
        """
        if self.tokenizer is None:
            raise RuntimeError("FewShot needs a tokenizer; call .steer() first.")

        # infer mode from arguments
        using_runtime_examples = (
            runtime_kwargs
            and ("positive_examples" in runtime_kwargs or "negative_examples" in runtime_kwargs)
        )
        using_pool_mode = self.positive_example_pool is not None or self.negative_example_pool is not None

        using_directive = bool(self.directive)
        if not (using_runtime_examples or using_pool_mode or using_directive):
            warnings.warn(
                "FewShot: nothing to inject (no examples, no directive). Returning input unchanged.",
                UserWarning,
            )
            return input_ids

        # determine input format
        is_tensor = isinstance(input_ids, torch.Tensor)
        original_device = input_ids.device if is_tensor else None
        original_dtype = input_ids.dtype if is_tensor else None

        # normalize to 2D list format [batch_size, seq_len]
        if is_tensor:
            if input_ids.ndim == 1:
                batch_input_ids = [input_ids.tolist()]
                single_sequence = True
            else:
                batch_input_ids = input_ids.tolist()
                single_sequence = False
        else:
            if isinstance(input_ids[0], int):
                batch_input_ids = [input_ids]
                single_sequence = True
            else:
                batch_input_ids = input_ids
                single_sequence = False

        use_chat_template = has_chat_template(self.tokenizer)

        adapted_batch: list[list[int]] = []
        for input_ids_single in batch_input_ids:
            original_text = self.tokenizer.decode(input_ids_single, skip_special_tokens=True)

            # sample or gather examples independently per item
            if using_runtime_examples:
                examples = self._gather_runtime_examples(runtime_kwargs)
            else:
                examples = self._sample_from_pools(query=original_text)

            if not examples and not self.directive:
                warnings.warn(
                    "FewShot: nothing to inject for this item. Returning input unchanged.",
                    UserWarning,
                )
                adapted_batch.append(list(input_ids_single))
                continue

            slot_memory = TextMemory(slots={
                "examples": examples,
                "directive": self.directive or "",
            })

            if use_chat_template:
                chat = [{"role": "user", "content": original_text}]
                adapted_chat = self._formatter.apply_to_messages([chat], slot_memory)[0]
                rendered = render_messages(self.tokenizer, adapted_chat, add_generation_prompt=True)
                adapted_tokens = self.tokenizer(rendered, add_special_tokens=False)["input_ids"]
            else:
                input_tensor = torch.tensor(input_ids_single, dtype=torch.long).unsqueeze(0)
                adapted_tensor = self._formatter.apply_to_ids(input_tensor, slot_memory, self.tokenizer)
                adapted_tokens = adapted_tensor[0].tolist()

            adapted_batch.append(adapted_tokens)

        # pad to uniform length for batched output
        max_len = max(len(seq) for seq in adapted_batch)
        if self.tokenizer.pad_token_id is None:
            raise RuntimeError(
                "FewShot: tokenizer has no pad_token_id; cannot pad batch sequences. "
                "Set a pad token before using FewShot with batched inputs."
            )
        pad_id = self.tokenizer.pad_token_id

        padded_batch = [seq + [pad_id] * (max_len - len(seq)) for seq in adapted_batch]

        # convert back to original format
        if is_tensor:
            result = torch.tensor(padded_batch, dtype=original_dtype, device=original_device)
            if single_sequence:
                result = result.squeeze(0)
            return result
        else:
            if single_sequence:
                return padded_batch[0]
            return padded_batch

    def _select_with_polarity(self, polarity: str, k: int, query: Any = None) -> list[dict]:
        """Run the resolved selector against the pool subset matching `polarity`."""
        if self.pool is None or self._selector is None:
            return []
        polarities = self.pool.metadata.get("polarity", [])
        items = [item for item, pol in zip(self.pool.items, polarities) if pol == polarity]
        if not items or k <= 0:
            return []
        return self._selector.select(items, query=query, k=k)

    def _sample_from_pools(self, query: Any = None) -> list[dict[str, Any]]:
        """Sample examples from the pools, attaching polarity labels for downstream formatting."""
        all_examples: list[dict[str, Any]] = []
        if self.positive_example_pool and self.k_positive and self.k_positive > 0:
            for example in self._select_with_polarity("pos", self.k_positive, query=query):
                all_examples.append({**example, "_polarity": "positive"})
        if self.negative_example_pool and self.k_negative and self.k_negative > 0:
            for example in self._select_with_polarity("neg", self.k_negative, query=query):
                all_examples.append({**example, "_polarity": "negative"})
        return all_examples

    def adapt_messages(
        self,
        messages: list[list[dict]],
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]] | None:
        """Insert a single system message containing the directive and labeled example blocks.

        Runtime examples (`positive_examples` / `negative_examples` in `runtime_kwargs`) take precedence
        over pool-based selection. If there are no examples and no directive, returns None (no change).
        """
        runtime_kwargs = runtime_kwargs or {}
        using_runtime = "positive_examples" in runtime_kwargs or "negative_examples" in runtime_kwargs
        using_pools = self.positive_example_pool is not None or self.negative_example_pool is not None
        if not (using_runtime or using_pools or self.directive):
            return None

        out: list[list[dict]] = []
        for chat in messages:
            if using_runtime:
                examples = self._gather_runtime_examples(runtime_kwargs)
            else:
                examples = self._sample_from_pools(query=chat)
            if not examples and not self.directive:
                out.append(list(chat))
                continue
            slot_memory = TextMemory(slots={
                "examples": examples,
                "directive": self.directive or "",
            })
            adapted_batch = self._formatter.apply_to_messages([chat], slot_memory)
            out.append(adapted_batch[0])
        return out

    @staticmethod
    def _gather_runtime_examples(runtime_kwargs: dict[str, Any]) -> list[dict[str, Any]]:
        """Gather examples from runtime_kwargs."""
        examples = []
        if "positive_examples" in runtime_kwargs:
            for example in runtime_kwargs["positive_examples"]:
                examples.append({**example, "_polarity": "positive"})
        if "negative_examples" in runtime_kwargs:
            for example in runtime_kwargs["negative_examples"]:
                examples.append({**example, "_polarity": "negative"})
        return examples

RUNTIME_KWARGS_SCHEMA = [{'name': 'positive_examples', 'type': 'list[dict]', 'scope': 'call', 'help': 'Positive examples for this call, applied to every prompt row; overrides pool-based selection.'}, {'name': 'negative_examples', 'type': 'list[dict]', 'scope': 'call', 'help': 'Negative examples for this call, applied to every prompt row; overrides pool-based selection.'}] class-attribute instance-attribute

args = self.Args.validate(*args, **kwargs) instance-attribute

directive = None class-attribute instance-attribute

enabled = True class-attribute instance-attribute

formatter = None class-attribute instance-attribute

k_negative = None class-attribute instance-attribute

k_positive = None class-attribute instance-attribute

memory = None class-attribute instance-attribute

negative_example_pool = None class-attribute instance-attribute

pool = None class-attribute instance-attribute

positive_example_pool = None class-attribute instance-attribute

selector = None class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

adapt(input_ids, runtime_kwargs=None)

Add few-shot examples to the model's prompt and return adapted token ids.

Both the chat-template path and the no-template fallback route through the resolved BaseFormatter, which renders the example block from the pool.

Assumes input_ids represents the user's prompt before any chat templating. Pre-templated input will be re-templated and produce malformed output; use adapt_messages for chat input.

Parameters:

Name	Type	Description	Default
input_ids	list[int] | Tensor	The user's prompt token IDs.	required
runtime_kwargs	dict | None	May contain positive_examples / negative_examples to override pool sampling.	None

Returns:

Type	Description
list[int] | Tensor	The transformed token IDs.

Raises:

Type	Description
RuntimeError	If tokenizer is not set (requires calling steer() first)

Warns:

Type	Description
UserWarning	Issued when no examples are configured or none remain after selection.

Source code in steerability/algorithms/input_control/few_shot/control.py

def adapt(
    self,
    input_ids: list[int] | torch.Tensor,
    runtime_kwargs: dict | None = None,
) -> list[int] | torch.Tensor:
    """Add few-shot examples to the model's prompt and return adapted token ids.

    Both the chat-template path and the no-template fallback route through the resolved
    `BaseFormatter`, which renders the example block from the pool.

    Assumes `input_ids` represents the user's prompt before any chat templating. Pre-templated
    input will be re-templated and produce malformed output; use `adapt_messages` for chat input.

    Args:
        input_ids: The user's prompt token IDs.
        runtime_kwargs: May contain `positive_examples` / `negative_examples` to override pool sampling.

    Returns:
        The transformed token IDs.

    Raises:
        RuntimeError: If tokenizer is not set (requires calling `steer()` first)

    Warnings:
        UserWarning: Issued when no examples are configured or none remain after selection.
    """
    if self.tokenizer is None:
        raise RuntimeError("FewShot needs a tokenizer; call .steer() first.")

    # infer mode from arguments
    using_runtime_examples = (
        runtime_kwargs
        and ("positive_examples" in runtime_kwargs or "negative_examples" in runtime_kwargs)
    )
    using_pool_mode = self.positive_example_pool is not None or self.negative_example_pool is not None

    using_directive = bool(self.directive)
    if not (using_runtime_examples or using_pool_mode or using_directive):
        warnings.warn(
            "FewShot: nothing to inject (no examples, no directive). Returning input unchanged.",
            UserWarning,
        )
        return input_ids

    # determine input format
    is_tensor = isinstance(input_ids, torch.Tensor)
    original_device = input_ids.device if is_tensor else None
    original_dtype = input_ids.dtype if is_tensor else None

    # normalize to 2D list format [batch_size, seq_len]
    if is_tensor:
        if input_ids.ndim == 1:
            batch_input_ids = [input_ids.tolist()]
            single_sequence = True
        else:
            batch_input_ids = input_ids.tolist()
            single_sequence = False
    else:
        if isinstance(input_ids[0], int):
            batch_input_ids = [input_ids]
            single_sequence = True
        else:
            batch_input_ids = input_ids
            single_sequence = False

    use_chat_template = has_chat_template(self.tokenizer)

    adapted_batch: list[list[int]] = []
    for input_ids_single in batch_input_ids:
        original_text = self.tokenizer.decode(input_ids_single, skip_special_tokens=True)

        # sample or gather examples independently per item
        if using_runtime_examples:
            examples = self._gather_runtime_examples(runtime_kwargs)
        else:
            examples = self._sample_from_pools(query=original_text)

        if not examples and not self.directive:
            warnings.warn(
                "FewShot: nothing to inject for this item. Returning input unchanged.",
                UserWarning,
            )
            adapted_batch.append(list(input_ids_single))
            continue

        slot_memory = TextMemory(slots={
            "examples": examples,
            "directive": self.directive or "",
        })

        if use_chat_template:
            chat = [{"role": "user", "content": original_text}]
            adapted_chat = self._formatter.apply_to_messages([chat], slot_memory)[0]
            rendered = render_messages(self.tokenizer, adapted_chat, add_generation_prompt=True)
            adapted_tokens = self.tokenizer(rendered, add_special_tokens=False)["input_ids"]
        else:
            input_tensor = torch.tensor(input_ids_single, dtype=torch.long).unsqueeze(0)
            adapted_tensor = self._formatter.apply_to_ids(input_tensor, slot_memory, self.tokenizer)
            adapted_tokens = adapted_tensor[0].tolist()

        adapted_batch.append(adapted_tokens)

    # pad to uniform length for batched output
    max_len = max(len(seq) for seq in adapted_batch)
    if self.tokenizer.pad_token_id is None:
        raise RuntimeError(
            "FewShot: tokenizer has no pad_token_id; cannot pad batch sequences. "
            "Set a pad token before using FewShot with batched inputs."
        )
    pad_id = self.tokenizer.pad_token_id

    padded_batch = [seq + [pad_id] * (max_len - len(seq)) for seq in adapted_batch]

    # convert back to original format
    if is_tensor:
        result = torch.tensor(padded_batch, dtype=original_dtype, device=original_device)
        if single_sequence:
            result = result.squeeze(0)
        return result
    else:
        if single_sequence:
            return padded_batch[0]
        return padded_batch

adapt_messages(messages, runtime_kwargs=None)

Insert a single system message containing the directive and labeled example blocks.

Runtime examples (positive_examples / negative_examples in runtime_kwargs) take precedence over pool-based selection. If there are no examples and no directive, returns None (no change).

Source code in steerability/algorithms/input_control/few_shot/control.py

def adapt_messages(
    self,
    messages: list[list[dict]],
    runtime_kwargs: dict | None = None,
) -> list[list[dict]] | None:
    """Insert a single system message containing the directive and labeled example blocks.

    Runtime examples (`positive_examples` / `negative_examples` in `runtime_kwargs`) take precedence
    over pool-based selection. If there are no examples and no directive, returns None (no change).
    """
    runtime_kwargs = runtime_kwargs or {}
    using_runtime = "positive_examples" in runtime_kwargs or "negative_examples" in runtime_kwargs
    using_pools = self.positive_example_pool is not None or self.negative_example_pool is not None
    if not (using_runtime or using_pools or self.directive):
        return None

    out: list[list[dict]] = []
    for chat in messages:
        if using_runtime:
            examples = self._gather_runtime_examples(runtime_kwargs)
        else:
            examples = self._sample_from_pools(query=chat)
        if not examples and not self.directive:
            out.append(list(chat))
            continue
        slot_memory = TextMemory(slots={
            "examples": examples,
            "directive": self.directive or "",
        })
        adapted_batch = self._formatter.apply_to_messages([chat], slot_memory)
        out.append(adapted_batch[0])
    return out

cleanup()

Release resources allocated during steer().

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

Source code in steerability/algorithms/core/base_control.py

def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

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

clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

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

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed forwarded to the clone's reseed().	None

Returns:

Type	Description
BaseControl	The clone.

Source code in steerability/algorithms/core/base_control.py

def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

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

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

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone

export_state()

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

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

Returns:

Type	Description
dict[str, Any]	Mapping from logical name to artifact value.

Source code in steerability/algorithms/core/base_control.py

def export_state(self) -> dict[str, Any]:
    """Steer-time state to persist when freezing, keyed by logical name.

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

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

fit_identity()

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

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

Returns:

Type	Description
Any | None	The fit-identity object, or None.

Source code in steerability/algorithms/core/base_control.py

def fit_identity(self) -> Any | None:
    """The object whose canonical form defines this control's fit digest, or None.

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

    Returns:
        The fit-identity object, or None.
    """
    return None

frozen_form(state)

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

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

Parameters:

Name	Type	Description	Default
state	dict[str, Any]	The output of export_state().	required

Returns:

Type	Description
tuple[str, dict[str, Any]]	The frozen method key and its constructor kwargs.

Raises:

Type	Description
NotFreezableError	If the control produces fits or state but declares no frozen form.

Source code in steerability/algorithms/core/base_control.py

def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The `(registry method key, constructor kwargs)` of this control's frozen form.

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

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

    Returns:
        The frozen method key and its constructor kwargs.

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

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

requirements()

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

Input controls transform the prompt client-side, so the generate phase requires nothing beyond the session contract on any backend. A control whose steer() reads the live pipeline model (e.g. for rollouts or scoring) overrides this with a steer-phase Capability.IN_PROCESS_TORCH requirement.

Returns:

Type	Description
Requirements	The control's phase-keyed requirements.

Source code in steerability/algorithms/input_control/base.py

def requirements(self) -> Requirements:
    """Backend requirements computed from this instance's configuration, per phase.

    Input controls transform the prompt client-side, so the generate phase requires nothing
    beyond the session contract on any backend. A control whose `steer()` reads the live
    pipeline model (e.g. for rollouts or scoring) overrides this with a steer-phase
    `Capability.IN_PROCESS_TORCH` requirement.

    Returns:
        The control's phase-keyed requirements.
    """
    return Requirements()

steer(model=None, tokenizer=None, **kwargs)

Optional offline preparation. Default is no-op.

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

Source code in steerability/algorithms/input_control/few_shot/control.py

def steer(
        self,
        model=None,
        tokenizer: PreTrainedTokenizerBase | None = None,
        **kwargs,
) -> None:
    self.tokenizer = tokenizer

    # build the example pool
    self.pool = PoolMemory[dict]()
    for example in self.positive_example_pool or []:
        self.pool.add(example, polarity="pos")
    for example in self.negative_example_pool or []:
        self.pool.add(example, polarity="neg")

    # resolve selector argument (instance | name | None) into a BaseSelector[dict]
    self._selector = selector_from_arg(self.selector)

    # selectors that need offline preparation (e.g. EPR) get a chance here
    prepare = getattr(self._selector, "prepare", None)
    if callable(prepare):
        prepare(model=model, tokenizer=tokenizer, data=self.pool)

    # formatter is shared between adapt and adapt_messages
    self._formatter = self.formatter or FewShotBlockFormatter()

steer_access()

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

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

Returns:

Type	Description
ModelAccess	The declared access rung.

Source code in steerability/algorithms/core/base_control.py

def steer_access(self) -> ModelAccess:
    """The model access this instance's steer step requires, on the `ModelAccess` ladder.

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

    Returns:
        The declared access rung.
    """
    return ModelAccess.FACTS

steer_fits()

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

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

Returns:

Type	Description
tuple[tuple[str, str], ...]	The declared fit artifacts, in declaration order.

Source code in steerability/algorithms/core/base_control.py

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

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

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

selectors

Example selectors for few-shot learning prompt adaptation.

Selectors determine which examples are picked from the pool when constructing the few-shot prompt.

Available selectors:

• RandomSelector (re-exported from common.selectors): uniform random sampling. Suitable for homogeneous pools where any example is roughly as informative as any other.
• EPRSelector: learned dense retriever (Rubin et al. 2021). Constructed by the caller (it requires a scoring LM) and passed to FewShot via selector=.

EPRSelector

Bases: DenseRetrievalSelector

Learned dense retriever for few-shot example selection (EPR).

At runtime, select() encodes the query and returns the nearest items in the example pool by cosine similarity over precomputed embeddings. Its prepare() runs the EPR training procedure:

1. For each training pair (x, y), use BM25 over y to retrieve a candidate set of size candidate_set_size.
2. Use the scoring LM to label k_pos positives and k_neg negatives in the candidate set.
3. Train a single-tower BERT-base encoder via contrastive loss with in-batch and hard negatives.

The trained encoder is then used at inference time to embed the query and rank pool items.

Parameters:

Name	Type	Description	Default
scoring_lm	PreTrainedModel	Causal LM used to score (anchor, candidate) pairs during offline labeling.	required
scoring_tokenizer	PreTrainedTokenizerBase	Tokenizer paired with scoring_lm.	required
base_encoder	str	Hugging Face model id of the encoder backbone.	'bert-base-uncased'
candidate_set_size	int	Number of BM25-nearest candidates considered per anchor (paper's L).	50
k_pos	int	Number of positives extracted per anchor.	5
k_neg	int	Number of hard negatives extracted per anchor.	5
train_epochs	int	Epochs of contrastive training over the labeled pairs.	3
batch_size	int	Mini-batch size during contrastive training.	8
max_anchors	int | None	Optional cap on anchors used during labeling (useful for fast CI).	None
input_field	str	Item key holding the input text.	'input'
output_field	str	Item key holding the output / target text.	'output'
device	Any	Optional device override for the encoder.	None

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/selector.py

class EPRSelector(DenseRetrievalSelector):
    """Learned dense retriever for few-shot example selection (EPR).

    At runtime, `select()` encodes the query and returns the nearest items in the example pool by
    cosine similarity over precomputed embeddings. Its `prepare()` runs the EPR training procedure:

      1. For each training pair `(x, y)`, use BM25 over `y` to retrieve a candidate set of size
         `candidate_set_size`.
      2. Use the scoring LM to label `k_pos` positives and `k_neg` negatives in the candidate set.
      3. Train a single-tower BERT-base encoder via contrastive loss with in-batch and hard negatives.

    The trained encoder is then used at inference time to embed the query and rank pool items.

    Args:
        scoring_lm: Causal LM used to score (anchor, candidate) pairs during offline labeling.
        scoring_tokenizer: Tokenizer paired with `scoring_lm`.
        base_encoder: Hugging Face model id of the encoder backbone.
        candidate_set_size: Number of BM25-nearest candidates considered per anchor (paper's `L`).
        k_pos: Number of positives extracted per anchor.
        k_neg: Number of hard negatives extracted per anchor.
        train_epochs: Epochs of contrastive training over the labeled pairs.
        batch_size: Mini-batch size during contrastive training.
        max_anchors: Optional cap on anchors used during labeling (useful for fast CI).
        input_field: Item key holding the input text.
        output_field: Item key holding the output / target text.
        device: Optional device override for the encoder.
    """

    def __init__(
        self,
        scoring_lm: PreTrainedModel,
        scoring_tokenizer: PreTrainedTokenizerBase,
        base_encoder: str = "bert-base-uncased",
        candidate_set_size: int = 50,
        k_pos: int = 5,
        k_neg: int = 5,
        train_epochs: int = 3,
        batch_size: int = 8,
        max_anchors: int | None = None,
        input_field: str = "input",
        output_field: str = "output",
        device: Any = None,
    ) -> None:
        # skip DenseRetrievalSelector's __init__; we own the encoder and create it later in prepare()
        self.encoder = None  # filled in prepare()
        self.similarity = "cosine"
        self.item_to_text = self._format_item
        self.embedding_key = "_epr_embedding"

        self.scoring_lm = scoring_lm
        self.scoring_tokenizer = scoring_tokenizer
        self.base_encoder_name = base_encoder
        self.candidate_set_size = candidate_set_size
        self.k_pos = k_pos
        self.k_neg = k_neg
        self.train_epochs = train_epochs
        self.batch_size = batch_size
        self.max_anchors = max_anchors
        self.input_field = input_field
        self.output_field = output_field
        self.device = device

    def _format_item(self, item: Any) -> str:
        if isinstance(item, dict):
            return f"Input: {item.get(self.input_field, '')}\nOutput: {item.get(self.output_field, '')}"
        return str(item)

    def _extract_query_text(self, query: Any) -> str:
        # when called from FewShot.adapt_messages, `query` is the chat (a list of dicts).
        if isinstance(query, list) and query and isinstance(query[0], dict):
            for msg in reversed(query):
                if msg.get("role") == "user":
                    return str(msg.get("content", ""))
            return str(query[-1].get("content", ""))
        if isinstance(query, dict):
            return str(query.get(self.input_field, "") or query.get("query", ""))
        if isinstance(query, str):
            return query
        return str(query)

    def prepare(
        self,
        model=None,
        tokenizer=None,
        data: PoolMemory[dict] | Sequence[dict] | None = None,
        **kwargs,
    ) -> None:
        if data is None:
            logger.warning("EPRSelector.prepare called with no data; nothing to train on.")
            return

        if isinstance(data, PoolMemory):
            items = list(data.items)
        else:
            items = list(data)

        if not items:
            logger.warning("EPRSelector.prepare: empty pool; skipping training.")
            return

        candidate_sets = bm25_index.build_and_query(
            items,
            query_field=self.output_field,
            candidate_set_size=self.candidate_set_size,
        )
        labeled = lm_labeling.label_pairs(
            items=items,
            candidate_sets=candidate_sets,
            scoring_lm=self.scoring_lm,
            scoring_tokenizer=self.scoring_tokenizer,
            k_pos=self.k_pos,
            k_neg=self.k_neg,
            input_field=self.input_field,
            output_field=self.output_field,
            max_anchors=self.max_anchors,
        )
        self.encoder = train_encoder.train(
            labeled=labeled,
            items=items,
            base_encoder=self.base_encoder_name,
            epochs=self.train_epochs,
            batch_size=self.batch_size,
            input_field=self.input_field,
            output_field=self.output_field,
            device=self.device,
        )

        # pre-compute pool embeddings into PoolMemory metadata for inference-time lookup.
        embeddings = [self.encoder.encode(self._format_item(item)) for item in items]
        if isinstance(data, PoolMemory):
            data.metadata[self.embedding_key] = embeddings
            for item, emb in zip(data.items, embeddings):
                if isinstance(item, dict):
                    item[self.embedding_key] = emb
        else:
            for item, emb in zip(items, embeddings):
                if isinstance(item, dict):
                    item[self.embedding_key] = emb

    def select(
        self,
        items: Sequence[dict],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[dict]:
        if self.encoder is None:
            raise RuntimeError("EPRSelector.select called before prepare(); the encoder is not trained.")
        query_text = self._extract_query_text(query)
        query_emb = np.asarray(self.encoder.encode(query_text))
        return self._top_k_by_similarity(items, query_emb, k)

base_encoder_name = base_encoder instance-attribute

batch_size = batch_size instance-attribute

candidate_set_size = candidate_set_size instance-attribute

device = device instance-attribute

embedding_key = '_epr_embedding' instance-attribute

encoder = None instance-attribute

input_field = input_field instance-attribute

item_to_text = self._format_item instance-attribute

k_neg = k_neg instance-attribute

k_pos = k_pos instance-attribute

max_anchors = max_anchors instance-attribute

output_field = output_field instance-attribute

scoring_lm = scoring_lm instance-attribute

scoring_tokenizer = scoring_tokenizer instance-attribute

similarity = 'cosine' instance-attribute

train_epochs = train_epochs instance-attribute

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/few_shot/selectors/epr/selector.py

def prepare(
    self,
    model=None,
    tokenizer=None,
    data: PoolMemory[dict] | Sequence[dict] | None = None,
    **kwargs,
) -> None:
    if data is None:
        logger.warning("EPRSelector.prepare called with no data; nothing to train on.")
        return

    if isinstance(data, PoolMemory):
        items = list(data.items)
    else:
        items = list(data)

    if not items:
        logger.warning("EPRSelector.prepare: empty pool; skipping training.")
        return

    candidate_sets = bm25_index.build_and_query(
        items,
        query_field=self.output_field,
        candidate_set_size=self.candidate_set_size,
    )
    labeled = lm_labeling.label_pairs(
        items=items,
        candidate_sets=candidate_sets,
        scoring_lm=self.scoring_lm,
        scoring_tokenizer=self.scoring_tokenizer,
        k_pos=self.k_pos,
        k_neg=self.k_neg,
        input_field=self.input_field,
        output_field=self.output_field,
        max_anchors=self.max_anchors,
    )
    self.encoder = train_encoder.train(
        labeled=labeled,
        items=items,
        base_encoder=self.base_encoder_name,
        epochs=self.train_epochs,
        batch_size=self.batch_size,
        input_field=self.input_field,
        output_field=self.output_field,
        device=self.device,
    )

    # pre-compute pool embeddings into PoolMemory metadata for inference-time lookup.
    embeddings = [self.encoder.encode(self._format_item(item)) for item in items]
    if isinstance(data, PoolMemory):
        data.metadata[self.embedding_key] = embeddings
        for item, emb in zip(data.items, embeddings):
            if isinstance(item, dict):
                item[self.embedding_key] = emb
    else:
        for item, emb in zip(items, embeddings):
            if isinstance(item, dict):
                item[self.embedding_key] = emb

select(items, query=None, k=1, context=None)

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/selector.py

def select(
    self,
    items: Sequence[dict],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[dict]:
    if self.encoder is None:
        raise RuntimeError("EPRSelector.select called before prepare(); the encoder is not trained.")
    query_text = self._extract_query_text(query)
    query_emb = np.asarray(self.encoder.encode(query_text))
    return self._top_k_by_similarity(items, query_emb, k)

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

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)

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

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

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

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)

selector_from_arg(value)

Resolve a selector argument value to a BaseSelector[dict] instance.

Accepts an instance directly, a string name (looked up in the registry), or None (defaults to RandomSelector).

Source code in steerability/algorithms/input_control/few_shot/selectors/__init__.py

def selector_from_arg(value: BaseSelector | str | None) -> BaseSelector:
    """Resolve a `selector` argument value to a `BaseSelector[dict]` instance.

    Accepts an instance directly, a string name (looked up in the registry), or None (defaults to
    `RandomSelector`).
    """
    if value is None:
        return RandomSelector()
    if isinstance(value, BaseSelector):
        return value
    if isinstance(value, str):
        if value not in SELECTOR_REGISTRY:
            raise ValueError(f"Unknown selector name: {value!r}. Known: {sorted(SELECTOR_REGISTRY)}.")
        return SELECTOR_REGISTRY[value]()
    raise TypeError(f"selector must be BaseSelector, str, or None; got {type(value).__name__}.")

epr

EPR, Efficient Prompt Retrieval (Rubin, Herzig, Berant 2021).

Reference:

• "Learning To Retrieve Prompts for In-Context Learning" Ohad Rubin, Jonathan Herzig, Jonathan Berant https://arxiv.org/abs/2112.08633

EPRSelector

Bases: DenseRetrievalSelector

Learned dense retriever for few-shot example selection (EPR).

At runtime, select() encodes the query and returns the nearest items in the example pool by cosine similarity over precomputed embeddings. Its prepare() runs the EPR training procedure:

1. For each training pair (x, y), use BM25 over y to retrieve a candidate set of size candidate_set_size.
2. Use the scoring LM to label k_pos positives and k_neg negatives in the candidate set.
3. Train a single-tower BERT-base encoder via contrastive loss with in-batch and hard negatives.

The trained encoder is then used at inference time to embed the query and rank pool items.

Parameters:

Name	Type	Description	Default
scoring_lm	PreTrainedModel	Causal LM used to score (anchor, candidate) pairs during offline labeling.	required
scoring_tokenizer	PreTrainedTokenizerBase	Tokenizer paired with scoring_lm.	required
base_encoder	str	Hugging Face model id of the encoder backbone.	'bert-base-uncased'
candidate_set_size	int	Number of BM25-nearest candidates considered per anchor (paper's L).	50
k_pos	int	Number of positives extracted per anchor.	5
k_neg	int	Number of hard negatives extracted per anchor.	5
train_epochs	int	Epochs of contrastive training over the labeled pairs.	3
batch_size	int	Mini-batch size during contrastive training.	8
max_anchors	int | None	Optional cap on anchors used during labeling (useful for fast CI).	None
input_field	str	Item key holding the input text.	'input'
output_field	str	Item key holding the output / target text.	'output'
device	Any	Optional device override for the encoder.	None

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/selector.py

class EPRSelector(DenseRetrievalSelector):
    """Learned dense retriever for few-shot example selection (EPR).

    At runtime, `select()` encodes the query and returns the nearest items in the example pool by
    cosine similarity over precomputed embeddings. Its `prepare()` runs the EPR training procedure:

      1. For each training pair `(x, y)`, use BM25 over `y` to retrieve a candidate set of size
         `candidate_set_size`.
      2. Use the scoring LM to label `k_pos` positives and `k_neg` negatives in the candidate set.
      3. Train a single-tower BERT-base encoder via contrastive loss with in-batch and hard negatives.

    The trained encoder is then used at inference time to embed the query and rank pool items.

    Args:
        scoring_lm: Causal LM used to score (anchor, candidate) pairs during offline labeling.
        scoring_tokenizer: Tokenizer paired with `scoring_lm`.
        base_encoder: Hugging Face model id of the encoder backbone.
        candidate_set_size: Number of BM25-nearest candidates considered per anchor (paper's `L`).
        k_pos: Number of positives extracted per anchor.
        k_neg: Number of hard negatives extracted per anchor.
        train_epochs: Epochs of contrastive training over the labeled pairs.
        batch_size: Mini-batch size during contrastive training.
        max_anchors: Optional cap on anchors used during labeling (useful for fast CI).
        input_field: Item key holding the input text.
        output_field: Item key holding the output / target text.
        device: Optional device override for the encoder.
    """

    def __init__(
        self,
        scoring_lm: PreTrainedModel,
        scoring_tokenizer: PreTrainedTokenizerBase,
        base_encoder: str = "bert-base-uncased",
        candidate_set_size: int = 50,
        k_pos: int = 5,
        k_neg: int = 5,
        train_epochs: int = 3,
        batch_size: int = 8,
        max_anchors: int | None = None,
        input_field: str = "input",
        output_field: str = "output",
        device: Any = None,
    ) -> None:
        # skip DenseRetrievalSelector's __init__; we own the encoder and create it later in prepare()
        self.encoder = None  # filled in prepare()
        self.similarity = "cosine"
        self.item_to_text = self._format_item
        self.embedding_key = "_epr_embedding"

        self.scoring_lm = scoring_lm
        self.scoring_tokenizer = scoring_tokenizer
        self.base_encoder_name = base_encoder
        self.candidate_set_size = candidate_set_size
        self.k_pos = k_pos
        self.k_neg = k_neg
        self.train_epochs = train_epochs
        self.batch_size = batch_size
        self.max_anchors = max_anchors
        self.input_field = input_field
        self.output_field = output_field
        self.device = device

    def _format_item(self, item: Any) -> str:
        if isinstance(item, dict):
            return f"Input: {item.get(self.input_field, '')}\nOutput: {item.get(self.output_field, '')}"
        return str(item)

    def _extract_query_text(self, query: Any) -> str:
        # when called from FewShot.adapt_messages, `query` is the chat (a list of dicts).
        if isinstance(query, list) and query and isinstance(query[0], dict):
            for msg in reversed(query):
                if msg.get("role") == "user":
                    return str(msg.get("content", ""))
            return str(query[-1].get("content", ""))
        if isinstance(query, dict):
            return str(query.get(self.input_field, "") or query.get("query", ""))
        if isinstance(query, str):
            return query
        return str(query)

    def prepare(
        self,
        model=None,
        tokenizer=None,
        data: PoolMemory[dict] | Sequence[dict] | None = None,
        **kwargs,
    ) -> None:
        if data is None:
            logger.warning("EPRSelector.prepare called with no data; nothing to train on.")
            return

        if isinstance(data, PoolMemory):
            items = list(data.items)
        else:
            items = list(data)

        if not items:
            logger.warning("EPRSelector.prepare: empty pool; skipping training.")
            return

        candidate_sets = bm25_index.build_and_query(
            items,
            query_field=self.output_field,
            candidate_set_size=self.candidate_set_size,
        )
        labeled = lm_labeling.label_pairs(
            items=items,
            candidate_sets=candidate_sets,
            scoring_lm=self.scoring_lm,
            scoring_tokenizer=self.scoring_tokenizer,
            k_pos=self.k_pos,
            k_neg=self.k_neg,
            input_field=self.input_field,
            output_field=self.output_field,
            max_anchors=self.max_anchors,
        )
        self.encoder = train_encoder.train(
            labeled=labeled,
            items=items,
            base_encoder=self.base_encoder_name,
            epochs=self.train_epochs,
            batch_size=self.batch_size,
            input_field=self.input_field,
            output_field=self.output_field,
            device=self.device,
        )

        # pre-compute pool embeddings into PoolMemory metadata for inference-time lookup.
        embeddings = [self.encoder.encode(self._format_item(item)) for item in items]
        if isinstance(data, PoolMemory):
            data.metadata[self.embedding_key] = embeddings
            for item, emb in zip(data.items, embeddings):
                if isinstance(item, dict):
                    item[self.embedding_key] = emb
        else:
            for item, emb in zip(items, embeddings):
                if isinstance(item, dict):
                    item[self.embedding_key] = emb

    def select(
        self,
        items: Sequence[dict],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[dict]:
        if self.encoder is None:
            raise RuntimeError("EPRSelector.select called before prepare(); the encoder is not trained.")
        query_text = self._extract_query_text(query)
        query_emb = np.asarray(self.encoder.encode(query_text))
        return self._top_k_by_similarity(items, query_emb, k)

base_encoder_name = base_encoder instance-attribute

batch_size = batch_size instance-attribute

candidate_set_size = candidate_set_size instance-attribute

device = device instance-attribute

embedding_key = '_epr_embedding' instance-attribute

encoder = None instance-attribute

input_field = input_field instance-attribute

item_to_text = self._format_item instance-attribute

k_neg = k_neg instance-attribute

k_pos = k_pos instance-attribute

max_anchors = max_anchors instance-attribute

output_field = output_field instance-attribute

scoring_lm = scoring_lm instance-attribute

scoring_tokenizer = scoring_tokenizer instance-attribute

similarity = 'cosine' instance-attribute

train_epochs = train_epochs instance-attribute

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/few_shot/selectors/epr/selector.py

def prepare(
    self,
    model=None,
    tokenizer=None,
    data: PoolMemory[dict] | Sequence[dict] | None = None,
    **kwargs,
) -> None:
    if data is None:
        logger.warning("EPRSelector.prepare called with no data; nothing to train on.")
        return

    if isinstance(data, PoolMemory):
        items = list(data.items)
    else:
        items = list(data)

    if not items:
        logger.warning("EPRSelector.prepare: empty pool; skipping training.")
        return

    candidate_sets = bm25_index.build_and_query(
        items,
        query_field=self.output_field,
        candidate_set_size=self.candidate_set_size,
    )
    labeled = lm_labeling.label_pairs(
        items=items,
        candidate_sets=candidate_sets,
        scoring_lm=self.scoring_lm,
        scoring_tokenizer=self.scoring_tokenizer,
        k_pos=self.k_pos,
        k_neg=self.k_neg,
        input_field=self.input_field,
        output_field=self.output_field,
        max_anchors=self.max_anchors,
    )
    self.encoder = train_encoder.train(
        labeled=labeled,
        items=items,
        base_encoder=self.base_encoder_name,
        epochs=self.train_epochs,
        batch_size=self.batch_size,
        input_field=self.input_field,
        output_field=self.output_field,
        device=self.device,
    )

    # pre-compute pool embeddings into PoolMemory metadata for inference-time lookup.
    embeddings = [self.encoder.encode(self._format_item(item)) for item in items]
    if isinstance(data, PoolMemory):
        data.metadata[self.embedding_key] = embeddings
        for item, emb in zip(data.items, embeddings):
            if isinstance(item, dict):
                item[self.embedding_key] = emb
    else:
        for item, emb in zip(items, embeddings):
            if isinstance(item, dict):
                item[self.embedding_key] = emb

select(items, query=None, k=1, context=None)

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/selector.py

def select(
    self,
    items: Sequence[dict],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[dict]:
    if self.encoder is None:
        raise RuntimeError("EPRSelector.select called before prepare(); the encoder is not trained.")
    query_text = self._extract_query_text(query)
    query_emb = np.asarray(self.encoder.encode(query_text))
    return self._top_k_by_similarity(items, query_emb, k)

selector

EPRSelector, a learned dense retriever for few-shot example selection.

Reference:

• "Learning To Retrieve Prompts for In-Context Learning" Ohad Rubin, Jonathan Herzig, Jonathan Berant https://arxiv.org/abs/2112.08633

logger = logging.getLogger(__name__) module-attribute

EPRSelector

Bases: DenseRetrievalSelector

Learned dense retriever for few-shot example selection (EPR).

At runtime, select() encodes the query and returns the nearest items in the example pool by cosine similarity over precomputed embeddings. Its prepare() runs the EPR training procedure:

1. For each training pair (x, y), use BM25 over y to retrieve a candidate set of size candidate_set_size.
2. Use the scoring LM to label k_pos positives and k_neg negatives in the candidate set.
3. Train a single-tower BERT-base encoder via contrastive loss with in-batch and hard negatives.

The trained encoder is then used at inference time to embed the query and rank pool items.

Parameters:

Name	Type	Description	Default
scoring_lm	PreTrainedModel	Causal LM used to score (anchor, candidate) pairs during offline labeling.	required
scoring_tokenizer	PreTrainedTokenizerBase	Tokenizer paired with scoring_lm.	required
base_encoder	str	Hugging Face model id of the encoder backbone.	'bert-base-uncased'
candidate_set_size	int	Number of BM25-nearest candidates considered per anchor (paper's L).	50
k_pos	int	Number of positives extracted per anchor.	5
k_neg	int	Number of hard negatives extracted per anchor.	5
train_epochs	int	Epochs of contrastive training over the labeled pairs.	3
batch_size	int	Mini-batch size during contrastive training.	8
max_anchors	int | None	Optional cap on anchors used during labeling (useful for fast CI).	None
input_field	str	Item key holding the input text.	'input'
output_field	str	Item key holding the output / target text.	'output'
device	Any	Optional device override for the encoder.	None

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/selector.py

class EPRSelector(DenseRetrievalSelector):
    """Learned dense retriever for few-shot example selection (EPR).

    At runtime, `select()` encodes the query and returns the nearest items in the example pool by
    cosine similarity over precomputed embeddings. Its `prepare()` runs the EPR training procedure:

      1. For each training pair `(x, y)`, use BM25 over `y` to retrieve a candidate set of size
         `candidate_set_size`.
      2. Use the scoring LM to label `k_pos` positives and `k_neg` negatives in the candidate set.
      3. Train a single-tower BERT-base encoder via contrastive loss with in-batch and hard negatives.

    The trained encoder is then used at inference time to embed the query and rank pool items.

    Args:
        scoring_lm: Causal LM used to score (anchor, candidate) pairs during offline labeling.
        scoring_tokenizer: Tokenizer paired with `scoring_lm`.
        base_encoder: Hugging Face model id of the encoder backbone.
        candidate_set_size: Number of BM25-nearest candidates considered per anchor (paper's `L`).
        k_pos: Number of positives extracted per anchor.
        k_neg: Number of hard negatives extracted per anchor.
        train_epochs: Epochs of contrastive training over the labeled pairs.
        batch_size: Mini-batch size during contrastive training.
        max_anchors: Optional cap on anchors used during labeling (useful for fast CI).
        input_field: Item key holding the input text.
        output_field: Item key holding the output / target text.
        device: Optional device override for the encoder.
    """

    def __init__(
        self,
        scoring_lm: PreTrainedModel,
        scoring_tokenizer: PreTrainedTokenizerBase,
        base_encoder: str = "bert-base-uncased",
        candidate_set_size: int = 50,
        k_pos: int = 5,
        k_neg: int = 5,
        train_epochs: int = 3,
        batch_size: int = 8,
        max_anchors: int | None = None,
        input_field: str = "input",
        output_field: str = "output",
        device: Any = None,
    ) -> None:
        # skip DenseRetrievalSelector's __init__; we own the encoder and create it later in prepare()
        self.encoder = None  # filled in prepare()
        self.similarity = "cosine"
        self.item_to_text = self._format_item
        self.embedding_key = "_epr_embedding"

        self.scoring_lm = scoring_lm
        self.scoring_tokenizer = scoring_tokenizer
        self.base_encoder_name = base_encoder
        self.candidate_set_size = candidate_set_size
        self.k_pos = k_pos
        self.k_neg = k_neg
        self.train_epochs = train_epochs
        self.batch_size = batch_size
        self.max_anchors = max_anchors
        self.input_field = input_field
        self.output_field = output_field
        self.device = device

    def _format_item(self, item: Any) -> str:
        if isinstance(item, dict):
            return f"Input: {item.get(self.input_field, '')}\nOutput: {item.get(self.output_field, '')}"
        return str(item)

    def _extract_query_text(self, query: Any) -> str:
        # when called from FewShot.adapt_messages, `query` is the chat (a list of dicts).
        if isinstance(query, list) and query and isinstance(query[0], dict):
            for msg in reversed(query):
                if msg.get("role") == "user":
                    return str(msg.get("content", ""))
            return str(query[-1].get("content", ""))
        if isinstance(query, dict):
            return str(query.get(self.input_field, "") or query.get("query", ""))
        if isinstance(query, str):
            return query
        return str(query)

    def prepare(
        self,
        model=None,
        tokenizer=None,
        data: PoolMemory[dict] | Sequence[dict] | None = None,
        **kwargs,
    ) -> None:
        if data is None:
            logger.warning("EPRSelector.prepare called with no data; nothing to train on.")
            return

        if isinstance(data, PoolMemory):
            items = list(data.items)
        else:
            items = list(data)

        if not items:
            logger.warning("EPRSelector.prepare: empty pool; skipping training.")
            return

        candidate_sets = bm25_index.build_and_query(
            items,
            query_field=self.output_field,
            candidate_set_size=self.candidate_set_size,
        )
        labeled = lm_labeling.label_pairs(
            items=items,
            candidate_sets=candidate_sets,
            scoring_lm=self.scoring_lm,
            scoring_tokenizer=self.scoring_tokenizer,
            k_pos=self.k_pos,
            k_neg=self.k_neg,
            input_field=self.input_field,
            output_field=self.output_field,
            max_anchors=self.max_anchors,
        )
        self.encoder = train_encoder.train(
            labeled=labeled,
            items=items,
            base_encoder=self.base_encoder_name,
            epochs=self.train_epochs,
            batch_size=self.batch_size,
            input_field=self.input_field,
            output_field=self.output_field,
            device=self.device,
        )

        # pre-compute pool embeddings into PoolMemory metadata for inference-time lookup.
        embeddings = [self.encoder.encode(self._format_item(item)) for item in items]
        if isinstance(data, PoolMemory):
            data.metadata[self.embedding_key] = embeddings
            for item, emb in zip(data.items, embeddings):
                if isinstance(item, dict):
                    item[self.embedding_key] = emb
        else:
            for item, emb in zip(items, embeddings):
                if isinstance(item, dict):
                    item[self.embedding_key] = emb

    def select(
        self,
        items: Sequence[dict],
        query: Any = None,
        k: int = 1,
        context: dict | None = None,
    ) -> list[dict]:
        if self.encoder is None:
            raise RuntimeError("EPRSelector.select called before prepare(); the encoder is not trained.")
        query_text = self._extract_query_text(query)
        query_emb = np.asarray(self.encoder.encode(query_text))
        return self._top_k_by_similarity(items, query_emb, k)

base_encoder_name = base_encoder instance-attribute

batch_size = batch_size instance-attribute

candidate_set_size = candidate_set_size instance-attribute

device = device instance-attribute

embedding_key = '_epr_embedding' instance-attribute

encoder = None instance-attribute

input_field = input_field instance-attribute

item_to_text = self._format_item instance-attribute

k_neg = k_neg instance-attribute

k_pos = k_pos instance-attribute

max_anchors = max_anchors instance-attribute

output_field = output_field instance-attribute

scoring_lm = scoring_lm instance-attribute

scoring_tokenizer = scoring_tokenizer instance-attribute

similarity = 'cosine' instance-attribute

train_epochs = train_epochs instance-attribute

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/few_shot/selectors/epr/selector.py

def prepare(
    self,
    model=None,
    tokenizer=None,
    data: PoolMemory[dict] | Sequence[dict] | None = None,
    **kwargs,
) -> None:
    if data is None:
        logger.warning("EPRSelector.prepare called with no data; nothing to train on.")
        return

    if isinstance(data, PoolMemory):
        items = list(data.items)
    else:
        items = list(data)

    if not items:
        logger.warning("EPRSelector.prepare: empty pool; skipping training.")
        return

    candidate_sets = bm25_index.build_and_query(
        items,
        query_field=self.output_field,
        candidate_set_size=self.candidate_set_size,
    )
    labeled = lm_labeling.label_pairs(
        items=items,
        candidate_sets=candidate_sets,
        scoring_lm=self.scoring_lm,
        scoring_tokenizer=self.scoring_tokenizer,
        k_pos=self.k_pos,
        k_neg=self.k_neg,
        input_field=self.input_field,
        output_field=self.output_field,
        max_anchors=self.max_anchors,
    )
    self.encoder = train_encoder.train(
        labeled=labeled,
        items=items,
        base_encoder=self.base_encoder_name,
        epochs=self.train_epochs,
        batch_size=self.batch_size,
        input_field=self.input_field,
        output_field=self.output_field,
        device=self.device,
    )

    # pre-compute pool embeddings into PoolMemory metadata for inference-time lookup.
    embeddings = [self.encoder.encode(self._format_item(item)) for item in items]
    if isinstance(data, PoolMemory):
        data.metadata[self.embedding_key] = embeddings
        for item, emb in zip(data.items, embeddings):
            if isinstance(item, dict):
                item[self.embedding_key] = emb
    else:
        for item, emb in zip(items, embeddings):
            if isinstance(item, dict):
                item[self.embedding_key] = emb

select(items, query=None, k=1, context=None)

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/selector.py

def select(
    self,
    items: Sequence[dict],
    query: Any = None,
    k: int = 1,
    context: dict | None = None,
) -> list[dict]:
    if self.encoder is None:
        raise RuntimeError("EPRSelector.select called before prepare(); the encoder is not trained.")
    query_text = self._extract_query_text(query)
    query_emb = np.asarray(self.encoder.encode(query_text))
    return self._top_k_by_similarity(items, query_emb, k)

utils

bm25_index

BM25 candidate retrieval for EPR labeling.

Prefers rank_bm25 when installed; falls back to a minimal pure-NumPy implementation otherwise so smoke tests can run on the [epr] extra-free baseline.

build_and_query(items, query_field='output', candidate_set_size=50)

For each item, return BM25-nearest candidate indices in the same pool (excluding self).

Used during EPR's offline labeling: for each training pair (x, y), retrieve a candidate set Ē of size L=candidate_set_size from the rest of the pool, ranking by BM25 over the target sequence y (the query_field).

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/bm25_index.py

def build_and_query(
    items: Sequence[dict],
    query_field: str = "output",
    candidate_set_size: int = 50,
) -> dict[int, list[int]]:
    """For each item, return BM25-nearest candidate indices in the same pool (excluding self).

    Used during EPR's offline labeling: for each training pair (x, y), retrieve a candidate set Ē of
    size L=`candidate_set_size` from the rest of the pool, ranking by BM25 over the target sequence y
    (the `query_field`).
    """
    docs = [str(item.get(query_field, "")) for item in items]
    index = build_index(docs)
    out: dict[int, list[int]] = {}
    for i, doc in enumerate(docs):
        ranked = query(index, doc, top_k=candidate_set_size + 1)
        # exclude the item itself from its own candidate set
        out[i] = [j for j in ranked if j != i][:candidate_set_size]
    return out

build_index(corpus)

Return an opaque index object exposing get_scores(query_tokens) -> list[float].

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/bm25_index.py

def build_index(corpus: Sequence[str]):
    """Return an opaque index object exposing `get_scores(query_tokens) -> list[float]`."""
    tokenized = [_tokenize(doc) for doc in corpus]
    try:
        from rank_bm25 import BM25Okapi  # type: ignore
        return BM25Okapi(tokenized)
    except ImportError:
        return _NaiveBM25(tokenized)

query(index, query_text, top_k)

Return the indices of the top-k most BM25-relevant docs in the indexed corpus.

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/bm25_index.py

def query(index, query_text: str, top_k: int) -> list[int]:
    """Return the indices of the top-k most BM25-relevant docs in the indexed corpus."""
    tokens = _tokenize(query_text)
    scores = list(index.get_scores(tokens))
    ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    return ranked[:top_k]

lm_labeling

LM-based labeling of (anchor, candidate) pairs for EPR.

For each training row (x_i, y_i), compute, for each candidate row (x_j, y_j) from the BM25 set, the log-probability log p_LM(y_i | (x_j, y_j), x_i). The top-k_pos highest-scoring candidates become positives, the bottom-k_neg become hard negatives.

logger = logging.getLogger(__name__) module-attribute

LabeledPair dataclass

One labeled training pair: an anchor row + a list of (candidate_index, polarity) labels.

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/lm_labeling.py

@dataclass
class LabeledPair:
    """One labeled training pair: an anchor row + a list of (candidate_index, polarity) labels."""
    anchor_index: int
    positives: list[int]
    negatives: list[int]

anchor_index instance-attribute

negatives instance-attribute

positives instance-attribute

label_pairs(items, candidate_sets, scoring_lm, scoring_tokenizer, k_pos=5, k_neg=5, input_field='input', output_field='output', max_anchors=None)

Label each anchor with its k_pos best and k_neg worst candidates by LM scoring.

max_anchors caps the number of anchors processed (useful for fast CI); when None, scores every anchor in candidate_sets.

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/lm_labeling.py

def label_pairs(
    items: Sequence[dict],
    candidate_sets: dict[int, list[int]],
    scoring_lm,
    scoring_tokenizer,
    k_pos: int = 5,
    k_neg: int = 5,
    input_field: str = "input",
    output_field: str = "output",
    max_anchors: int | None = None,
) -> list[LabeledPair]:
    """Label each anchor with its `k_pos` best and `k_neg` worst candidates by LM scoring.

    `max_anchors` caps the number of anchors processed (useful for fast CI); when None, scores every
    anchor in `candidate_sets`.
    """
    items = list(items)
    anchor_indices = list(candidate_sets.keys())
    if max_anchors is not None:
        anchor_indices = anchor_indices[:max_anchors]

    out: list[LabeledPair] = []
    for anchor_idx in anchor_indices:
        candidates = candidate_sets.get(anchor_idx, [])
        if not candidates:
            continue

        scored: list[tuple[int, float]] = []
        anchor = items[anchor_idx]
        for cand_idx in candidates:
            score = _score_candidate_for_anchor(
                scoring_lm,
                scoring_tokenizer,
                anchor,
                items[cand_idx],
                input_field,
                output_field,
            )
            scored.append((cand_idx, score))
        scored.sort(key=lambda pair: pair[1], reverse=True)
        positives = [idx for idx, _ in scored[:k_pos]]
        negatives = [idx for idx, _ in scored[-k_neg:]]
        out.append(LabeledPair(anchor_index=anchor_idx, positives=positives, negatives=negatives))
    return out

train_encoder

Contrastive (DPR-style) training of the EPR dual encoder.

The encoder is a single transformer (we use the same model for both query and document encodings, sharing weights). Training uses an InfoNCE-style objective with in-batch negatives and the LM-labeled hard negatives.

logger = logging.getLogger(__name__) module-attribute

EPREncoder

A thin encoder wrapper exposing encode(text) -> np.ndarray.

Uses mean-pooled last_hidden_state with attention masking. Backbone is shared between queries and pool items (a single-tower retriever, like DPR's text-only variant).

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/train_encoder.py

class EPREncoder:
    """A thin encoder wrapper exposing `encode(text) -> np.ndarray`.

    Uses mean-pooled last_hidden_state with attention masking. Backbone is shared between queries and
    pool items (a single-tower retriever, like DPR's text-only variant).
    """

    def __init__(self, model_name_or_path: str = "bert-base-uncased", device: str | torch.device | None = None):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        self.model = AutoModel.from_pretrained(model_name_or_path)
        if device is not None:
            self.model = self.model.to(device)
        self.model.eval()

    @property
    def device(self) -> torch.device:
        return next(self.model.parameters()).device

    @torch.no_grad()
    def encode(self, text: str) -> np.ndarray:
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(self.device)
        outputs = self.model(**inputs)
        hidden = outputs.last_hidden_state  # [1, T, H]
        mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)  # [1, T, 1]
        pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return pooled[0].cpu().numpy()

device property

model = AutoModel.from_pretrained(model_name_or_path) instance-attribute

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) instance-attribute

encode(text)

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/train_encoder.py

@torch.no_grad()
def encode(self, text: str) -> np.ndarray:
    inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(self.device)
    outputs = self.model(**inputs)
    hidden = outputs.last_hidden_state  # [1, T, H]
    mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)  # [1, T, 1]
    pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
    return pooled[0].cpu().numpy()

train(labeled, items, base_encoder='bert-base-uncased', epochs=3, batch_size=8, **kwargs)

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/train_encoder.py

def train(
    labeled,
    items,
    base_encoder: str = "bert-base-uncased",
    epochs: int = 3,
    batch_size: int = 8,
    **kwargs,
) -> EPREncoder:
    return train_encoder(
        items=items,
        labeled=labeled,
        base_encoder=base_encoder,
        epochs=epochs,
        batch_size=batch_size,
        **kwargs,
    )

train_encoder(items, labeled, base_encoder='bert-base-uncased', epochs=3, batch_size=8, learning_rate=2e-05, input_field='input', output_field='output', device=None)

Train (or fine-tune) a dual encoder for EPR.

Returns:

Type	Description
EPREncoder	An EPREncoder ready for inference-time encode().

Source code in steerability/algorithms/input_control/few_shot/selectors/epr/utils/train_encoder.py

def train_encoder(
    items: Sequence[dict],
    labeled: Sequence[LabeledPair],
    base_encoder: str = "bert-base-uncased",
    epochs: int = 3,
    batch_size: int = 8,
    learning_rate: float = 2e-5,
    input_field: str = "input",
    output_field: str = "output",
    device: str | torch.device | None = None,
) -> EPREncoder:
    """Train (or fine-tune) a dual encoder for EPR.

    Returns:
        An `EPREncoder` ready for inference-time `encode()`.
    """
    encoder = EPREncoder(model_name_or_path=base_encoder, device=device)
    if not labeled:
        logger.warning("EPR train_encoder: no labeled pairs supplied; returning untrained encoder.")
        return encoder

    encoder.model.train()
    optimizer = torch.optim.AdamW(encoder.model.parameters(), lr=learning_rate)

    dataset = _ContrastiveDataset(items, labeled, input_field, output_field)

    def collate(batch):
        out = _collate(batch, encoder.tokenizer)
        for key, tensors in out.items():
            out[key] = {k: v.to(encoder.device) for k, v in tensors.items()}
        return out

    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate)

    for epoch in range(epochs):
        epoch_loss = 0.0
        steps = 0
        for batch in loader:
            optimizer.zero_grad()
            q_emb = _mean_pool(encoder.model, batch["queries"])  # [B, H]
            p_emb = _mean_pool(encoder.model, batch["positives"])  # [B, H]

            # in-batch negatives
            logits = q_emb @ p_emb.T  # [B, B]; row i should pick column i (the positive)
            if "negatives" in batch:
                n_emb = _mean_pool(encoder.model, batch["negatives"])  # [Bn, H]
                neg_logits = q_emb @ n_emb.T  # [B, Bn]
                logits = torch.cat([logits, neg_logits], dim=1)

            targets = torch.arange(q_emb.size(0), device=q_emb.device)
            loss = nn.functional.cross_entropy(logits, targets)
            loss.backward()
            optimizer.step()
            epoch_loss += float(loss.item())
            steps += 1
        logger.debug("EPR epoch %d loss=%.4f", epoch, epoch_loss / max(steps, 1))

    encoder.model.eval()
    return encoder