CPO

steerability.algorithms.input_control.cpo

args

Arguments for the CPO input control.

control

CPO, causal prompt optimization.

Reference:

• "Causal Prompt Optimization" Chen et al., 2026 https://arxiv.org/abs/2602.01711

logger = logging.getLogger(__name__) module-attribute

CPO

Bases: InputControl

Causal prompt optimization.

Stage 1 (steer): build a ⟨query, prompt, score⟩ dataset (either supplied or generated from train_dataset), then train a CausalRewardScorer over PCA-reduced embeddings.

Stage 2 (adapt_messages): per-query, run a B-wide K-retained R-round tree search using the prompt proposer to expand parents and the causal scorer to retain survivors. Returns the best survivor as the system prompt.

The query cache keys on a sha256 hash of the query text. By default the raw query is hashed, so whitespace, casing, and punctuation differences yield distinct cache entries. Supply cache_key_normalizer (a str -> str callable) to normalize the query before hashing so near-duplicate queries share a cached prompt.

Reference:

• "Causal Prompt Optimization" Chen et al., 2026 https://arxiv.org/abs/2602.01711

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

class CPO(InputControl):
    """Causal prompt optimization.

    Stage 1 (`steer`): build a `⟨query, prompt, score⟩` dataset (either supplied or generated from
    `train_dataset`), then train a `CausalRewardScorer` over PCA-reduced embeddings.

    Stage 2 (`adapt_messages`): per-query, run a B-wide K-retained R-round tree search using the prompt
    proposer to expand parents and the causal scorer to retain survivors. Returns the best survivor as
    the system prompt.

    The query cache keys on a sha256 hash of the query text. By default the raw query is hashed, so
    whitespace, casing, and punctuation differences yield distinct cache entries. Supply
    `cache_key_normalizer` (a `str -> str` callable) to normalize the query before hashing so
    near-duplicate queries share a cached prompt.

    Reference:

      - "Causal Prompt Optimization"
        Chen et al., 2026
        [https://arxiv.org/abs/2602.01711](https://arxiv.org/abs/2602.01711)
    """

    Args = CPOArgs
    supports_batching: bool = True

    # placeholder fields; injected from CPOArgs at __init__
    seed_prompt: str = ""
    train_dataset: list[dict] = []
    offline_data: list[dict] | None = None
    n_prompts_per_query: int = 200
    embedding_model: str = "nomic-ai/nomic-embed-text-v1.5"
    trust_remote_code: bool = True
    pca_query_dim: int = 40
    pca_prompt_dim: int = 15
    prompt_lm: Any = None
    rounds: int = 3
    candidates_per_parent: int = 5
    retained_per_round: int = 3
    row_scorer: Any = None
    cache_queries: bool = True
    cache_key_normalizer: Any = None
    use_dml: bool | None = None
    refinement_meta_prompt: str | None = None
    proposer_gen_kwargs: dict | None = None
    eval_gen_kwargs: dict | None = None
    memory: CPOMemory | None = None

    # method-owned state
    tokenizer: Any = None
    _formatter: SystemPromptFormatter | None = None
    _proposer: LLMMetaPromptProposer | None = None
    _encoder: TextEncoder | None = None

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

        With `prompt_lm` supplied the proposer is a control-owned auxiliary and every phase is
        prompt-only. With `prompt_lm` unset the live pipeline model is bound as the proposer at
        steer and consulted per query at adapt, so the generate phase requires
        `Capability.IN_PROCESS_TORCH`."""
        if self.prompt_lm is not None:
            return Requirements()
        return Requirements(generate=needs(
            Capability.IN_PROCESS_TORCH,
            hint="set prompt_lm to run CPO's per-query search off the pipeline model",
        ))

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.ROLLOUTS` with `prompt_lm` supplied (offline data generation rides the
        session); `ModelAccess.MODULE` with `prompt_lm` unset (the live model is bound as the
        proposer). A precomputed `memory` with `prompt_lm` supplied needs only structural
        facts (`ModelAccess.FACTS`)."""
        if self.prompt_lm is not None:
            has_memory = getattr(getattr(self, "args", None), "memory", None) is not None
            return ModelAccess.FACTS if has_memory else ModelAccess.ROLLOUTS
        return ModelAccess.MODULE

    def export_state(self) -> dict:
        """The trained scorer memory under the `"memory"` key (after `steer()`)."""
        return {"memory": self.memory} if self.memory is not None else {}

    def frozen_form(self, state: dict) -> tuple[str, dict]:
        """A same-class frozen form: the recipe args with `memory=` set to the trained memory
        (the training-only args stay inert)."""
        from dataclasses import fields

        kwargs = {f.name: getattr(self.args, f.name) for f in fields(self.args) if f.init}
        kwargs["memory"] = state["memory"]
        return "input_control/cpo", kwargs

    def fit_identity(self):
        """The optimizer-relevant args (everything except `memory` and the per-query search
        knobs), or None when the recipe already carries a memory."""
        from dataclasses import fields

        if self.args.memory is not None:
            return None
        return {
            f.name: getattr(self.args, f.name)
            for f in fields(self.args) if f.init and f.name != "memory"
        }

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

        if self.prompt_lm is not None:
            proposer_lm = self.prompt_lm
            encoder_device = None
        else:
            proposer_lm = model
            encoder_device = next(model.parameters()).device if model is not None else None
        self._encoder = TextEncoder(
            self.embedding_model,
            device=encoder_device,
            trust_remote_code=self.trust_remote_code,
        )

        self._proposer = LLMMetaPromptProposer(
            llm=proposer_lm,
            tokenizer=tokenizer,
            meta_prompt_template=self.refinement_meta_prompt or refinement_meta_prompt.CPO_DEFAULT,
            gen_kwargs=self.proposer_gen_kwargs,
            parse_fn=parse_concise_instruction,
        )

        if self.args.memory is not None:
            memory = self.args.memory
            if isinstance(memory, (str, Path)):
                memory = CPOMemory.load(Path(memory), encoder=self._encoder)
            self.memory = memory
            self._formatter = SystemPromptFormatter()
            return

        task_lm = model if model is not None else (SessionLM(session) if session is not None else None)
        offline_data = self.offline_data or self._generate_offline_data(task_lm, tokenizer)
        scorer = causal_reward.train(
            offline_data=offline_data,
            embedding_model=self.embedding_model,
            pca_query_dim=self.pca_query_dim,
            pca_prompt_dim=self.pca_prompt_dim,
            seed_prompt=self.seed_prompt,
            use_dml=self.use_dml,
            encoder=self._encoder,
            trust_remote_code=self.trust_remote_code,
        )

        self.memory = CPOMemory(causal_scorer=scorer)
        self._formatter = SystemPromptFormatter()

    def _generate_offline_data(self, task_lm, tokenizer) -> list[dict]:
        """Build ⟨query, prompt, score⟩ rows from `train_dataset` × proposer × row scorer.

        Each training row contributes `n_prompts_per_query` ⟨q, p, s⟩ triples (including the seed
        prompt as one of them, so the seed is always in-distribution for the reward model). Candidate
        prompts are scored with `TaskEvaluationScorer` using a one-row dev set per training query.
        """
        if not self.train_dataset:
            raise RuntimeError("offline_data is None and train_dataset is empty; nothing to fit on.")

        gen_kwargs = self.eval_gen_kwargs or {"max_new_tokens": 32, "do_sample": False}
        rows: list[dict] = []

        for example in self.train_dataset:
            query = example.get("input", example.get("query", ""))
            dev_row: dict[str, Any] = {"input": query}
            if "reference" in example:
                dev_row["reference"] = example["reference"]

            candidate_prompts = [self.seed_prompt]
            if self.n_prompts_per_query > 1:
                candidate_prompts.extend(
                    self._proposer.propose(seed=self.seed_prompt, n=self.n_prompts_per_query - 1)
                )

            scorer = TaskEvaluationScorer(
                task_lm=task_lm,
                tokenizer=tokenizer,
                dev_set=[dev_row],
                row_scorer=self.row_scorer,
                gen_kwargs=gen_kwargs,
            )
            scores = scorer.score(candidate_prompts)
            rows.extend(
                {"query": query, "prompt": prompt, "score": score}
                for prompt, score in zip(candidate_prompts, scores)
            )
        return rows

    def adapt_messages(
        self,
        messages: list[list[dict]],
        runtime_kwargs: dict | None = None,
    ) -> list[list[dict]] | None:
        if self.memory is None or self._formatter is None:
            raise RuntimeError("CPO.adapt_messages called before .steer().")
        out: list[list[dict]] = []
        for chat in messages:
            query_text = self._extract_query(chat)
            best = self._best_prompt_for_query(query_text)
            slot_memory = TextMemory(slots={"instruction": best})
            adapted = self._formatter.apply_to_messages([chat], slot_memory, runtime_kwargs)
            out.append(adapted[0])
        return out

    def adapt(
        self,
        input_ids: list[int] | torch.Tensor,
        runtime_kwargs: dict | None = None,
    ) -> list[int] | torch.Tensor:
        if self.memory is None or self._formatter is None or self.tokenizer is None:
            raise RuntimeError("CPO.adapt called before .steer().")
        if not isinstance(input_ids, torch.Tensor):
            input_ids = torch.tensor(input_ids, dtype=torch.long)
        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)

        adapted_rows: list[torch.Tensor] = []
        for row in input_ids:
            query_text = self.tokenizer.decode(row, skip_special_tokens=True)
            best = self._best_prompt_for_query(query_text)
            slot_memory = TextMemory(slots={"instruction": best})
            adapted = self._formatter.apply_to_ids(
                row.unsqueeze(0), slot_memory, self.tokenizer, runtime_kwargs
            )
            adapted_rows.append(adapted[0])

        # left-pad to max length (causal LM convention) and stack
        max_len = max(r.size(0) for r in adapted_rows)
        pad_id = self.tokenizer.pad_token_id
        if pad_id is None:
            pad_id = self.tokenizer.eos_token_id
        if pad_id is None:
            raise RuntimeError(
                "CPO.adapt: tokenizer has neither pad_token_id nor eos_token_id; cannot pad batch."
            )
        padded = torch.stack([
            torch.cat([torch.full((max_len - r.size(0),), pad_id, dtype=r.dtype, device=r.device), r])
            for r in adapted_rows
        ])
        return padded

    @staticmethod
    def _extract_query(chat: list[dict]) -> str:
        for msg in reversed(chat):
            if msg.get("role") == "user":
                return str(msg.get("content", ""))
        return str(chat[-1].get("content", "")) if chat else ""

    def _best_prompt_for_query(self, query: str) -> str:
        if self.memory is None:
            raise RuntimeError("CPO not steered; call .steer() first.")
        key_source = self.cache_key_normalizer(query) if self.cache_key_normalizer else query
        cache_key = hashlib.sha256(key_source.encode("utf-8")).hexdigest()
        if self.cache_queries and cache_key in self.memory.query_cache:
            return self.memory.query_cache[cache_key]
        best = self._tree_search(query)
        if self.cache_queries:
            self.memory.query_cache[cache_key] = best
        return best

    def _tree_search(self, query: str) -> str:
        """B-wide, K-retained, R-round search from `seed_prompt`.

        Survivors of each round are carried into the next round's candidate pool (elitism), so the
        search never returns a prompt that the trained reward model scores below the seed. When every
        proposal scores worse, the seed survives every round and is returned.
        """
        survivors: list[str] = [self.seed_prompt]
        for _ in range(self.rounds):
            candidates: list[str] = list(survivors)  # elitism: parents compete with their children
            for parent in survivors:
                spawned = self._proposer.propose(seed=parent, n=self.candidates_per_parent)
                candidates.extend([c.strip() for c in spawned if isinstance(c, str) and c.strip()])
            candidates = list(dict.fromkeys(candidates))  # dedupe, preserve order
            scores = self.memory.causal_scorer.score(
                candidates,
                queries=[{"text": query}] * len(candidates),
            )
            order = np.argsort(scores)[::-1]
            survivors = [candidates[i] for i in order[: self.retained_per_round]]
        if not survivors:
            return self.seed_prompt
        final_scores = self.memory.causal_scorer.score(
            survivors,
            queries=[{"text": query}] * len(survivors),
        )
        return survivors[int(np.argmax(final_scores))]

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

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

cache_key_normalizer = None class-attribute instance-attribute

cache_queries = True class-attribute instance-attribute

candidates_per_parent = 5 class-attribute instance-attribute

embedding_model = 'nomic-ai/nomic-embed-text-v1.5' class-attribute instance-attribute

enabled = True class-attribute instance-attribute

eval_gen_kwargs = None class-attribute instance-attribute

memory = None class-attribute instance-attribute

n_prompts_per_query = 200 class-attribute instance-attribute

offline_data = None class-attribute instance-attribute

pca_prompt_dim = 15 class-attribute instance-attribute

pca_query_dim = 40 class-attribute instance-attribute

prompt_lm = None class-attribute instance-attribute

proposer_gen_kwargs = None class-attribute instance-attribute

refinement_meta_prompt = None class-attribute instance-attribute

retained_per_round = 3 class-attribute instance-attribute

rounds = 3 class-attribute instance-attribute

row_scorer = None class-attribute instance-attribute

seed_prompt = '' class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

train_dataset = [] class-attribute instance-attribute

trust_remote_code = True class-attribute instance-attribute

use_dml = None class-attribute instance-attribute

adapt(input_ids, runtime_kwargs=None)

Transform input_ids into a steered prompt.

May read instance state (e.g. self.memory) that was populated by steer().

Parameters:

Name	Type	Description	Default
input_ids	list[int] | Tensor	The user's prompt token IDs.	required
runtime_kwargs	dict | None	Per-call parameters.	None

Returns:

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

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

def adapt(
    self,
    input_ids: list[int] | torch.Tensor,
    runtime_kwargs: dict | None = None,
) -> list[int] | torch.Tensor:
    if self.memory is None or self._formatter is None or self.tokenizer is None:
        raise RuntimeError("CPO.adapt called before .steer().")
    if not isinstance(input_ids, torch.Tensor):
        input_ids = torch.tensor(input_ids, dtype=torch.long)
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)

    adapted_rows: list[torch.Tensor] = []
    for row in input_ids:
        query_text = self.tokenizer.decode(row, skip_special_tokens=True)
        best = self._best_prompt_for_query(query_text)
        slot_memory = TextMemory(slots={"instruction": best})
        adapted = self._formatter.apply_to_ids(
            row.unsqueeze(0), slot_memory, self.tokenizer, runtime_kwargs
        )
        adapted_rows.append(adapted[0])

    # left-pad to max length (causal LM convention) and stack
    max_len = max(r.size(0) for r in adapted_rows)
    pad_id = self.tokenizer.pad_token_id
    if pad_id is None:
        pad_id = self.tokenizer.eos_token_id
    if pad_id is None:
        raise RuntimeError(
            "CPO.adapt: tokenizer has neither pad_token_id nor eos_token_id; cannot pad batch."
        )
    padded = torch.stack([
        torch.cat([torch.full((max_len - r.size(0),), pad_id, dtype=r.dtype, device=r.device), r])
        for r in adapted_rows
    ])
    return padded

adapt_messages(messages, runtime_kwargs=None)

Optional message-level adaptation, called BEFORE chat-template tokenization.

Default returns None (no message-level changes). Subclasses that modify chat structure (set/replace system prompt, insert example turns) override this. When this method returns a non-None result for a chat-input generation, the pipeline does NOT additionally call this control's adapt() for that call; a control may therefore implement both entry points (message-level for chat input, token-level for raw text/tensor input) without being applied twice. Other input controls in the pipeline still run per their own phase.

Parameters:

Name	Type	Description	Default
messages	list[list[dict]]	A batch of chats; outer list is the batch, inner list is the message sequence for one chat.	required
runtime_kwargs	dict | None	Per-call parameters.	None

Returns:

Type	Description
list[list[dict]] | None	The transformed messages, or None to indicate no change.

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

def adapt_messages(
    self,
    messages: list[list[dict]],
    runtime_kwargs: dict | None = None,
) -> list[list[dict]] | None:
    if self.memory is None or self._formatter is None:
        raise RuntimeError("CPO.adapt_messages called before .steer().")
    out: list[list[dict]] = []
    for chat in messages:
        query_text = self._extract_query(chat)
        best = self._best_prompt_for_query(query_text)
        slot_memory = TextMemory(slots={"instruction": best})
        adapted = self._formatter.apply_to_messages([chat], slot_memory, runtime_kwargs)
        out.append(adapted[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()

The trained scorer memory under the "memory" key (after steer()).

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

def export_state(self) -> dict:
    """The trained scorer memory under the `"memory"` key (after `steer()`)."""
    return {"memory": self.memory} if self.memory is not None else {}

fit_identity()

The optimizer-relevant args (everything except memory and the per-query search knobs), or None when the recipe already carries a memory.

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

def fit_identity(self):
    """The optimizer-relevant args (everything except `memory` and the per-query search
    knobs), or None when the recipe already carries a memory."""
    from dataclasses import fields

    if self.args.memory is not None:
        return None
    return {
        f.name: getattr(self.args, f.name)
        for f in fields(self.args) if f.init and f.name != "memory"
    }

frozen_form(state)

A same-class frozen form: the recipe args with memory= set to the trained memory (the training-only args stay inert).

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

def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class frozen form: the recipe args with `memory=` set to the trained memory
    (the training-only args stay inert)."""
    from dataclasses import fields

    kwargs = {f.name: getattr(self.args, f.name) for f in fields(self.args) if f.init}
    kwargs["memory"] = state["memory"]
    return "input_control/cpo", kwargs

requirements()

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

With prompt_lm supplied the proposer is a control-owned auxiliary and every phase is prompt-only. With prompt_lm unset the live pipeline model is bound as the proposer at steer and consulted per query at adapt, so the generate phase requires Capability.IN_PROCESS_TORCH.

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

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

    With `prompt_lm` supplied the proposer is a control-owned auxiliary and every phase is
    prompt-only. With `prompt_lm` unset the live pipeline model is bound as the proposer at
    steer and consulted per query at adapt, so the generate phase requires
    `Capability.IN_PROCESS_TORCH`."""
    if self.prompt_lm is not None:
        return Requirements()
    return Requirements(generate=needs(
        Capability.IN_PROCESS_TORCH,
        hint="set prompt_lm to run CPO's per-query search off the pipeline model",
    ))

steer(model=None, tokenizer=None, session=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/cpo/control.py

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

    if self.prompt_lm is not None:
        proposer_lm = self.prompt_lm
        encoder_device = None
    else:
        proposer_lm = model
        encoder_device = next(model.parameters()).device if model is not None else None
    self._encoder = TextEncoder(
        self.embedding_model,
        device=encoder_device,
        trust_remote_code=self.trust_remote_code,
    )

    self._proposer = LLMMetaPromptProposer(
        llm=proposer_lm,
        tokenizer=tokenizer,
        meta_prompt_template=self.refinement_meta_prompt or refinement_meta_prompt.CPO_DEFAULT,
        gen_kwargs=self.proposer_gen_kwargs,
        parse_fn=parse_concise_instruction,
    )

    if self.args.memory is not None:
        memory = self.args.memory
        if isinstance(memory, (str, Path)):
            memory = CPOMemory.load(Path(memory), encoder=self._encoder)
        self.memory = memory
        self._formatter = SystemPromptFormatter()
        return

    task_lm = model if model is not None else (SessionLM(session) if session is not None else None)
    offline_data = self.offline_data or self._generate_offline_data(task_lm, tokenizer)
    scorer = causal_reward.train(
        offline_data=offline_data,
        embedding_model=self.embedding_model,
        pca_query_dim=self.pca_query_dim,
        pca_prompt_dim=self.pca_prompt_dim,
        seed_prompt=self.seed_prompt,
        use_dml=self.use_dml,
        encoder=self._encoder,
        trust_remote_code=self.trust_remote_code,
    )

    self.memory = CPOMemory(causal_scorer=scorer)
    self._formatter = SystemPromptFormatter()

steer_access()

ModelAccess.ROLLOUTS with prompt_lm supplied (offline data generation rides the session); ModelAccess.MODULE with prompt_lm unset (the live model is bound as the proposer). A precomputed memory with prompt_lm supplied needs only structural facts (ModelAccess.FACTS).

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

def steer_access(self) -> ModelAccess:
    """`ModelAccess.ROLLOUTS` with `prompt_lm` supplied (offline data generation rides the
    session); `ModelAccess.MODULE` with `prompt_lm` unset (the live model is bound as the
    proposer). A precomputed `memory` with `prompt_lm` supplied needs only structural
    facts (`ModelAccess.FACTS`)."""
    if self.prompt_lm is not None:
        has_memory = getattr(getattr(self, "args", None), "memory", None) is not None
        return ModelAccess.FACTS if has_memory else ModelAccess.ROLLOUTS
    return ModelAccess.MODULE

steer_fits()

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

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

Returns:

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

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

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 ()

CPOMemory dataclass

Custom Memory for CPO; holds the trained scorer and a query→best-prompt cache.

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

@dataclass
class CPOMemory:
    """Custom Memory for CPO; holds the trained scorer and a query→best-prompt cache."""
    causal_scorer: CausalRewardScorer
    query_cache: dict[str, str] = field(default_factory=dict)

    def save(self, path: Path) -> None:
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        self.causal_scorer.save(path / "scorer.pkl")
        (path / "cache.json").write_text(json.dumps(self.query_cache))

    @classmethod
    def load(cls, path: Path, encoder: TextEncoder) -> "CPOMemory":
        path = Path(path)
        scorer = CausalRewardScorer.load(path / "scorer.pkl", encoder=encoder)
        cache = json.loads((path / "cache.json").read_text())
        return cls(causal_scorer=scorer, query_cache=cache)

causal_scorer instance-attribute

query_cache = field(default_factory=dict) class-attribute instance-attribute

load(path, encoder) classmethod

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

@classmethod
def load(cls, path: Path, encoder: TextEncoder) -> "CPOMemory":
    path = Path(path)
    scorer = CausalRewardScorer.load(path / "scorer.pkl", encoder=encoder)
    cache = json.loads((path / "cache.json").read_text())
    return cls(causal_scorer=scorer, query_cache=cache)

save(path)

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

def save(self, path: Path) -> None:
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    self.causal_scorer.save(path / "scorer.pkl")
    (path / "cache.json").write_text(json.dumps(self.query_cache))

utils

causal_reward

Causal reward model training (CPO Stage 1).

The CPO paper (Chen et al. 2026, arXiv:2602.01711) trains the reward model with Double Machine Learning (DML) on PCA-reduced query and prompt embeddings, treating the prompt as a continuous treatment relative to a fixed seed prompt. We expose a single CausalRewardScorer that conforms to BaseScorer and a train(...) factory.

When econml is installed we use CausalForestDML. Otherwise we fall back to a plain GradientBoostingRegressor over the concatenated (query, prompt) features, regressing on score directly. The fallback drops the causal interpretation. econml is not a toolkit dependency; users who want the paper's estimator install it separately.

logger = logging.getLogger(__name__) module-attribute

CausalRewardScorer dataclass

Bases: BaseScorer

Trained CPO reward model. Implements BaseScorer.score(prompts, queries).

Internally:

• embeds prompts and queries with the supplied encoder
• applies separately-fitted PCAs to reduce dimensionality
• holds either an econml CausalForestDML (mode="dml") or a GradientBoostingRegressor (mode="gbr") over the concatenated reduced features

Source code in steerability/algorithms/input_control/cpo/utils/causal_reward.py

@dataclass
class CausalRewardScorer(BaseScorer):
    """Trained CPO reward model. Implements `BaseScorer.score(prompts, queries)`.

    Internally:

      - embeds prompts and queries with the supplied encoder
      - applies separately-fitted PCAs to reduce dimensionality
      - holds either an econml `CausalForestDML` (`mode="dml"`) or a `GradientBoostingRegressor`
        (`mode="gbr"`) over the concatenated reduced features
    """

    encoder: TextEncoder
    query_pca: PCA
    prompt_pca: PCA
    seed_prompt_emb: np.ndarray  # the seed prompt's reduced embedding (treatment baseline for DML)
    estimator: object  # CausalForestDML or GradientBoostingRegressor
    mode: str  # "dml" or "gbr"

    def _featurize(self, prompts: Sequence[str], queries: Sequence[dict] | None) -> tuple[np.ndarray, np.ndarray]:
        prompt_full = self.encoder.encode_batch(list(prompts))
        prompt_red = self.prompt_pca.transform(prompt_full)

        if queries is not None:
            query_texts = [q.get("text", str(q)) for q in queries]
        else:
            query_texts = [""] * len(prompts)
        query_full = self.encoder.encode_batch(query_texts)
        query_red = self.query_pca.transform(query_full)
        return query_red, prompt_red

    def score(
        self,
        prompts: Sequence[str],
        queries: Sequence[dict] | None = None,
    ) -> list[float]:
        if not prompts:
            return []
        if queries is not None and len(queries) != len(prompts):
            raise ValueError(
                f"`queries` (len={len(queries)}) must align with `prompts` (len={len(prompts)})."
            )

        query_red, prompt_red = self._featurize(prompts, queries)

        if self.mode == "dml":
            # DML treatment effect of (prompt - seed_prompt) given query covariate
            treatment = prompt_red - self.seed_prompt_emb[None, :]
            # CausalForestDML.effect(X, T0, T1) is signed; we want effect of moving from seed to candidate
            t0 = np.zeros_like(treatment)
            effects = self.estimator.effect(query_red, T0=t0, T1=treatment)  # type: ignore[attr-defined]
            return [float(e) for e in np.atleast_1d(effects)]

        # GBR fallback: regress score on concatenated features
        x = np.concatenate([query_red, prompt_red], axis=1)
        preds = self.estimator.predict(x)  # type: ignore[attr-defined]
        return [float(p) for p in np.atleast_1d(preds)]

    def save(self, path: Path) -> None:
        path = Path(path)
        with open(path, "wb") as f:
            pickle.dump(
                {
                    "query_pca": self.query_pca,
                    "prompt_pca": self.prompt_pca,
                    "seed_prompt_emb": self.seed_prompt_emb,
                    "estimator": self.estimator,
                    "mode": self.mode,
                    "encoder_name_or_path": getattr(self.encoder, "_name_or_path", None),
                },
                f,
            )

    @classmethod
    def load(cls, path: Path, encoder: TextEncoder) -> "CausalRewardScorer":
        with open(path, "rb") as f:
            payload = pickle.load(f)
        return cls(
            encoder=encoder,
            query_pca=payload["query_pca"],
            prompt_pca=payload["prompt_pca"],
            seed_prompt_emb=payload["seed_prompt_emb"],
            estimator=payload["estimator"],
            mode=payload["mode"],
        )

encoder instance-attribute

estimator instance-attribute

mode instance-attribute

prompt_pca instance-attribute

query_pca instance-attribute

seed_prompt_emb instance-attribute

load(path, encoder) classmethod

Source code in steerability/algorithms/input_control/cpo/utils/causal_reward.py

@classmethod
def load(cls, path: Path, encoder: TextEncoder) -> "CausalRewardScorer":
    with open(path, "rb") as f:
        payload = pickle.load(f)
    return cls(
        encoder=encoder,
        query_pca=payload["query_pca"],
        prompt_pca=payload["prompt_pca"],
        seed_prompt_emb=payload["seed_prompt_emb"],
        estimator=payload["estimator"],
        mode=payload["mode"],
    )

save(path)

Source code in steerability/algorithms/input_control/cpo/utils/causal_reward.py

def save(self, path: Path) -> None:
    path = Path(path)
    with open(path, "wb") as f:
        pickle.dump(
            {
                "query_pca": self.query_pca,
                "prompt_pca": self.prompt_pca,
                "seed_prompt_emb": self.seed_prompt_emb,
                "estimator": self.estimator,
                "mode": self.mode,
                "encoder_name_or_path": getattr(self.encoder, "_name_or_path", None),
            },
            f,
        )

score(prompts, queries=None)

Source code in steerability/algorithms/input_control/cpo/utils/causal_reward.py

def score(
    self,
    prompts: Sequence[str],
    queries: Sequence[dict] | None = None,
) -> list[float]:
    if not prompts:
        return []
    if queries is not None and len(queries) != len(prompts):
        raise ValueError(
            f"`queries` (len={len(queries)}) must align with `prompts` (len={len(prompts)})."
        )

    query_red, prompt_red = self._featurize(prompts, queries)

    if self.mode == "dml":
        # DML treatment effect of (prompt - seed_prompt) given query covariate
        treatment = prompt_red - self.seed_prompt_emb[None, :]
        # CausalForestDML.effect(X, T0, T1) is signed; we want effect of moving from seed to candidate
        t0 = np.zeros_like(treatment)
        effects = self.estimator.effect(query_red, T0=t0, T1=treatment)  # type: ignore[attr-defined]
        return [float(e) for e in np.atleast_1d(effects)]

    # GBR fallback: regress score on concatenated features
    x = np.concatenate([query_red, prompt_red], axis=1)
    preds = self.estimator.predict(x)  # type: ignore[attr-defined]
    return [float(p) for p in np.atleast_1d(preds)]

train(offline_data, embedding_model, pca_query_dim, pca_prompt_dim, seed_prompt, use_dml=None, encoder=None, device=None, trust_remote_code=False)

Fit the CPO reward model on a list of {query, prompt, score} rows.

Parameters:

Name	Type	Description	Default
offline_data	Sequence[dict]	list of dicts with keys "query", "prompt", "score".	required
embedding_model	str	HF model id for the encoder. Ignored if encoder is supplied.	required
pca_query_dim	int	PCA target dimension for queries.	required
pca_prompt_dim	int	PCA target dimension for prompts.	required
seed_prompt	str	The base prompt used as the treatment baseline.	required
use_dml	bool | None	If True, require econml; if False, force the GBR fallback. If None (default), use DML when econml is importable, otherwise fall back.	None
encoder	TextEncoder | None	Optional pre-built TextEncoder. Useful for tests to avoid reloading.	None
device	str | None	Optional device for the encoder.	None
trust_remote_code	bool	Trust remote code when loading embedding_model. Ignored if encoder is supplied.	False

Source code in steerability/algorithms/input_control/cpo/utils/causal_reward.py

def train(
    offline_data: Sequence[dict],
    embedding_model: str,
    pca_query_dim: int,
    pca_prompt_dim: int,
    seed_prompt: str,
    use_dml: bool | None = None,
    encoder: TextEncoder | None = None,
    device: str | None = None,
    trust_remote_code: bool = False,
) -> CausalRewardScorer:
    """Fit the CPO reward model on a list of `{query, prompt, score}` rows.

    Args:
        offline_data: list of dicts with keys `"query"`, `"prompt"`, `"score"`.
        embedding_model: HF model id for the encoder. Ignored if `encoder` is supplied.
        pca_query_dim: PCA target dimension for queries.
        pca_prompt_dim: PCA target dimension for prompts.
        seed_prompt: The base prompt used as the treatment baseline.
        use_dml: If True, require econml; if False, force the GBR fallback. If None (default), use DML
            when econml is importable, otherwise fall back.
        encoder: Optional pre-built `TextEncoder`. Useful for tests to avoid reloading.
        device: Optional device for the encoder.
        trust_remote_code: Trust remote code when loading `embedding_model`. Ignored if `encoder` is supplied.
    """
    if not offline_data:
        raise ValueError("offline_data is empty.")
    enc = (
        encoder
        if encoder is not None
        else TextEncoder(embedding_model, device=device, trust_remote_code=trust_remote_code)
    )

    queries = [row["query"] for row in offline_data]
    prompts = [row["prompt"] for row in offline_data]
    scores = np.array([row["score"] for row in offline_data], dtype=float)

    if scores.size:
        values, counts = np.unique(scores, return_counts=True)
        top_value = float(values[int(np.argmax(counts))])
        top_frac = float(counts.max()) / float(scores.size)
        if values.size == 1:
            warnings.warn(
                f"CPO offline_data scores are constant (value={top_value:.4f} over {scores.size} rows). "
                "A reward model fit on constant targets predicts that constant for every prompt, so the "
                "tree search degenerates to arbitrary tie-breaking. Use a metric/dev set that produces "
                "score variance.",
                UserWarning,
                stacklevel=2,
            )
        elif top_frac >= 0.95:
            warnings.warn(
                f"CPO offline_data scores are nearly saturated: value {top_value:.4f} covers "
                f"{100.0 * top_frac:.0f}% of {scores.size} rows. The reward model will have almost no "
                "signal to rank prompts; consider harder queries, stricter or graded scoring, or more "
                "diverse candidate prompts.",
                UserWarning,
                stacklevel=2,
            )

    query_full = enc.encode_batch(queries)
    prompt_full = enc.encode_batch(prompts + [seed_prompt])

    query_pca = fit_pca(query_full, pca_query_dim)
    prompt_pca = fit_pca(prompt_full, pca_prompt_dim)

    query_red = query_pca.transform(query_full)
    prompt_red = prompt_pca.transform(prompt_full[:-1])
    seed_prompt_red = prompt_pca.transform(prompt_full[-1:])[0]

    DML = _try_import_dml() if use_dml is None else (_try_import_dml() if use_dml else None)
    if use_dml is True and DML is None:
        raise ModuleNotFoundError(
            "CPO with use_dml=True requires econml's CausalForestDML, which is not installed. "
            'Install it separately with `pip install "econml>=0.16,<0.17"` (the range the toolkit '
            "has been run with), or leave use_dml=None to fall back to GradientBoostingRegressor.",
            name="econml",
        )

    if DML is not None:
        treatment = prompt_red - seed_prompt_red[None, :]
        try:
            estimator = DML(
                model_y=GradientBoostingRegressor(),
                model_t=GradientBoostingRegressor(),
                discrete_treatment=False,
                random_state=0,
            )
            estimator.fit(Y=scores, T=treatment, X=query_red)
            mode = "dml"
        except Exception as exc:  # econml occasionally fails on tiny / degenerate data
            logger.warning("CausalForestDML.fit failed (%s); falling back to GBR.", exc)
            estimator = GradientBoostingRegressor()
            estimator.fit(np.concatenate([query_red, prompt_red], axis=1), scores)
            mode = "gbr"
    else:
        estimator = GradientBoostingRegressor()
        estimator.fit(np.concatenate([query_red, prompt_red], axis=1), scores)
        mode = "gbr"

    return CausalRewardScorer(
        encoder=enc,
        query_pca=query_pca,
        prompt_pca=prompt_pca,
        seed_prompt_emb=seed_prompt_red,
        estimator=estimator,
        mode=mode,
    )

embeddings

Sentence embeddings + PCA reduction for CPO.

Uses a single-tower transformer with mean-pooled last_hidden_state. The CPO paper specifies nomic-embed-text-v1.5 as the default, but any HF text encoder works.

logger = logging.getLogger(__name__) module-attribute

TextEncoder

Mean-pooled HF encoder. Not a BaseScorer; this is purely a featurizer for CPO's reward model.

Source code in steerability/algorithms/input_control/cpo/utils/embeddings.py

class TextEncoder:
    """Mean-pooled HF encoder. Not a `BaseScorer`; this is purely a featurizer for CPO's reward model."""

    def __init__(
        self,
        model_name_or_path: str,
        device: str | torch.device | None = None,
        trust_remote_code: bool = True,
    ) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code)
        self.model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code)
        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_batch(self, texts: Sequence[str]) -> np.ndarray:
        if not texts:
            return np.zeros((0, self.model.config.hidden_size), dtype=np.float32)
        inputs = self.tokenizer(
            list(texts), return_tensors="pt", truncation=True, padding=True
        ).to(self.device)
        outputs = self.model(**inputs)
        hidden = outputs.last_hidden_state  # [B, T, H]
        mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
        pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return pooled.cpu().numpy()

    def encode(self, text: str) -> np.ndarray:
        return self.encode_batch([text])[0]

device property

model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code) instance-attribute

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code) instance-attribute

encode(text)

Source code in steerability/algorithms/input_control/cpo/utils/embeddings.py

def encode(self, text: str) -> np.ndarray:
    return self.encode_batch([text])[0]

encode_batch(texts)

Source code in steerability/algorithms/input_control/cpo/utils/embeddings.py

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

fit_pca(matrix, n_components)

Fit a PCA whose number of components is clipped to the available rank.

Source code in steerability/algorithms/input_control/cpo/utils/embeddings.py

def fit_pca(matrix: np.ndarray, n_components: int) -> PCA:
    """Fit a PCA whose number of components is clipped to the available rank."""
    n_samples, n_features = matrix.shape
    max_components = max(1, min(n_components, n_samples, n_features))
    pca = PCA(n_components=max_components)
    pca.fit(matrix)
    return pca

refinement_meta_prompt

Self-refinement meta-prompt for CPO.

Follows Appendix EC.3 of Chen et al., 2026 (arXiv:2602.01711), with a single {seed} placeholder so it composes with the common LLMMetaPromptProposer.

CPO_DEFAULT = 'You are an expert prompt engineer. Refine the prompt below to make it more effective at the task it describes. Preserve intent; vary wording, structure, and specificity. Write a single, complete, ready-to-use instruction for the specific task -- not a generic template. Never include bracketed placeholders (e.g., [topic]); do not explain your changes. Reply with the refined prompt only, in one or two sentences.\n\nCurrent prompt:\n{seed}\n\nRefined prompt:' module-attribute