PRewrite

steerability.algorithms.input_control.prewrite

args

Arguments for the PRewrite input control.

control

PRewrite, an RL-trained instruction rewriter.

Reference:

• "PRewrite: Prompt Rewriting with Reinforcement Learning" Weize Kong, Spurthi Amba Hombaiah, Mingyang Zhang, Qiaozhu Mei, Michael Bendersky https://arxiv.org/abs/2401.08189

logger = logging.getLogger(__name__) module-attribute

PRewrite

Bases: InputControl

Train (or load) an LLM rewriter that rewrites a seed instruction; optionally select the best rewrite by dev-set evaluation, then apply it as the system prompt at inference time.

When train_rewriter=True, the rewriter is trained with GRPO (group-relative policy optimization). The reward is downstream task performance, computed by applying each rewrite with the frozen task model over a dev set and scoring with a per-row SampleScorer, or a user-supplied reward_fn.

Memory shape: TextMemory(slots={"instruction": str}).

Reference:

• "PRewrite: Prompt Rewriting with Reinforcement Learning" Weize Kong, Spurthi Amba Hombaiah, Mingyang Zhang, Qiaozhu Mei, Michael Bendersky https://arxiv.org/abs/2401.08189

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

class PRewrite(InputControl):
    """Train (or load) an LLM rewriter that rewrites a seed instruction; optionally select the best
    rewrite by dev-set evaluation, then apply it as the system prompt at inference time.

    When `train_rewriter=True`, the rewriter is trained with GRPO (group-relative policy optimization).
    The reward is downstream task performance, computed by applying each rewrite with the frozen task
    model over a dev set and scoring with a per-row `SampleScorer`, or a user-supplied `reward_fn`.

    Memory shape: `TextMemory(slots={"instruction": str})`.

    Reference:

      - "PRewrite: Prompt Rewriting with Reinforcement Learning"
        Weize Kong, Spurthi Amba Hombaiah, Mingyang Zhang, Qiaozhu Mei, Michael Bendersky
        [https://arxiv.org/abs/2401.08189](https://arxiv.org/abs/2401.08189)
    """

    Args = PRewriteArgs
    supports_batching: bool = True

    # placeholder fields; injected from PRewriteArgs at __init__
    initial_instruction: str = ""
    rewriter_model_name_or_path: str | None = None
    rewriter_model: Any = None
    rewriter_tokenizer: Any = None
    train_rewriter: bool = False
    meta_prompt: str | None = None
    strategy: str = "search"
    k_candidates: int = 10
    dev_set: list[dict] | None = None
    row_scorer: Any = None
    training_seeds: list[str] | None = None
    reward_fn: Any = None
    grpo_config: dict | None = None
    reward_dev_size: int | None = None
    trust_remote_code: bool = False
    rewriter_gen_kwargs: dict | None = None
    eval_gen_kwargs: dict | None = None

    # method-owned state
    memory: TextMemory | None = None
    tokenizer: Any = None
    _formatter: SystemPromptFormatter | None = None

    def steer_access(self) -> ModelAccess:
        """`ModelAccess.ROLLOUTS`; rewriting and dev-set scoring generate through the
        session, and adaptation is a formatter. With a precomputed `memory`, steering is
        model-free (`ModelAccess.FACTS`)."""
        if getattr(getattr(self, "args", None), "memory", None) is not None:
            return ModelAccess.FACTS
        return ModelAccess.ROLLOUTS

    def export_state(self) -> dict:
        """The optimized instruction 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 optimized
        memory (the search-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/prewrite", kwargs

    def fit_identity(self):
        """The optimizer-relevant args (everything except `memory`), 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.args.memory is not None:
            self.memory = self.args.memory
            self._formatter = SystemPromptFormatter()
            return

        task_lm = SessionLM(session) if session is not None else model
        rewriter_lm, rewriter_tok = self._resolve_rewriter(task_lm, tokenizer)
        meta_prompt = self.meta_prompt or meta_prompts.DEFAULT

        if self.train_rewriter:
            reward_fn = self._build_reward_fn(task_lm=task_lm, task_tok=tokenizer)
            rewriter_lm = self._grpo_train_rewriter(rewriter_lm, rewriter_tok, meta_prompt, reward_fn)

        proposer = LLMMetaPromptProposer(
            llm=rewriter_lm,
            tokenizer=rewriter_tok,
            meta_prompt_template=meta_prompt,
            gen_kwargs=self.rewriter_gen_kwargs,
            use_chat_template=False if self.train_rewriter else None,
            parse_fn=parse_concise_instruction,
        )

        if self.strategy == "inference":
            candidates = proposer.propose(seed=self.initial_instruction, n=1)
            best = candidates[0] if candidates else self.initial_instruction
        else:
            n = max(self.k_candidates, 1)
            candidates = proposer.propose(seed=self.initial_instruction, n=n)
            if not candidates:
                logger.warning("PRewrite proposer returned no candidates; falling back to initial instruction.")
                best = self.initial_instruction
            else:
                scorer = TaskEvaluationScorer(
                    task_lm=task_lm,
                    tokenizer=tokenizer,
                    dev_set=self.dev_set,
                    row_scorer=self.row_scorer,
                    gen_kwargs=self.eval_gen_kwargs or {"max_new_tokens": 32, "do_sample": False},
                )
                selector = TopKSelector(scorer=scorer)
                best = selector.select(candidates, k=1)[0]

        stripped = (best or "").strip()
        if not stripped:
            logger.info(
                "PRewrite: chosen rewrite was empty/whitespace; falling back to initial_instruction."
            )
            best = self.initial_instruction
        else:
            best = stripped

        self.memory = TextMemory(slots={"instruction": best})
        self._formatter = SystemPromptFormatter()

    def _resolve_rewriter(self, task_lm, tokenizer) -> tuple[Any, Any]:
        """Pick the rewriter LLM.

        Resolution order:

          1. Pre-loaded `rewriter_model` (+ `rewriter_tokenizer`) if supplied.
          2. Load from `rewriter_model_name_or_path` if supplied.
          3. Default: reuse the task model through the session (forbidden under
             `train_rewriter=True`; rejected at args time).
        """
        if self.rewriter_model is not None:
            rewriter_tok = self.rewriter_tokenizer
            if rewriter_tok is None:
                source = getattr(self.rewriter_model, "name_or_path", None) or getattr(
                    getattr(self.rewriter_model, "config", None), "_name_or_path", None
                )
                if source is None:
                    raise ValueError(
                        "PRewrite: `rewriter_model` was supplied without `rewriter_tokenizer` "
                        "and no name_or_path could be inferred from the model."
                    )
                rewriter_tok = AutoTokenizer.from_pretrained(source, trust_remote_code=self.trust_remote_code)
            return self.rewriter_model, rewriter_tok
        if self.rewriter_model_name_or_path is None:
            return task_lm, tokenizer
        rewriter_lm = AutoModelForCausalLM.from_pretrained(
            self.rewriter_model_name_or_path,
            device_map="auto",
            trust_remote_code=self.trust_remote_code,
        )
        rewriter_tok = AutoTokenizer.from_pretrained(
            self.rewriter_model_name_or_path, trust_remote_code=self.trust_remote_code
        )
        if rewriter_tok.pad_token_id is None and rewriter_tok.eos_token_id is not None:
            rewriter_tok.pad_token = rewriter_tok.eos_token
        return rewriter_lm, rewriter_tok

    def _build_reward_fn(self, task_lm, task_tok):
        """Build the GRPO reward callable for rewriter training.

        Uses a user-supplied `reward_fn` if present. Otherwise builds a `TaskEvaluationScorer` that
        applies each rewrite with the frozen task model over `dev_set` and aggregates `row_scorer`
        to a scalar. The reward's `task_lm` generates through the steering session and stays frozen;
        only the rewriter is trained.
        """
        if self.reward_fn is not None:
            return self.reward_fn
        if self.row_scorer is not None and self.dev_set:
            scorer = TaskEvaluationScorer(
                task_lm=task_lm,
                tokenizer=task_tok,
                dev_set=self.dev_set,
                row_scorer=self.row_scorer,
                gen_kwargs=self.eval_gen_kwargs or {"max_new_tokens": 32, "do_sample": False},
                max_dev_size=self.reward_dev_size,
            )
            return make_scorer_reward_func(scorer, parse_fn=parse_concise_instruction)
        raise ValueError(
            "PRewrite.train_rewriter requires a GRPO reward source: a callable `reward_fn` or "
            "(`row_scorer` and `dev_set`)."
        )

    def _grpo_train_rewriter(self, rewriter_lm, rewriter_tok, meta_prompt: str, reward_fn):
        """GRPO-train the rewriter on a pool of seed instructions, using a callable/scorer reward.

        Returns the trained rewriter.
        """
        from datasets import Dataset

        from steerability.algorithms.structural_control.wrappers.trl.grpotrainer import GRPO, GRPOArgs

        seeds = self.training_seeds or [self.initial_instruction]
        train_dataset = Dataset.from_dict(
            {"prompt": [meta_prompt.format(seed=s) for s in seeds]}
        )

        grpo_args = GRPOArgs(
            train_dataset=train_dataset,
            reward_funcs=[reward_fn],
            **(self.grpo_config or {}),
        )
        grpo = GRPO(grpo_args)
        return grpo.steer(rewriter_lm, rewriter_tok)

    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("PRewrite.adapt_messages called before .steer().")
        return self._formatter.apply_to_messages(messages, self.memory, runtime_kwargs)

    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("PRewrite.adapt called before .steer().")
        if not isinstance(input_ids, torch.Tensor):
            input_ids = torch.tensor(input_ids, dtype=torch.long)
        return self._formatter.apply_to_ids(input_ids, self.memory, self.tokenizer, runtime_kwargs)

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

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

dev_set = None class-attribute instance-attribute

enabled = True class-attribute instance-attribute

eval_gen_kwargs = None class-attribute instance-attribute

grpo_config = None class-attribute instance-attribute

initial_instruction = '' class-attribute instance-attribute

k_candidates = 10 class-attribute instance-attribute

memory = None class-attribute instance-attribute

meta_prompt = None class-attribute instance-attribute

reward_dev_size = None class-attribute instance-attribute

reward_fn = None class-attribute instance-attribute

rewriter_gen_kwargs = None class-attribute instance-attribute

rewriter_model = None class-attribute instance-attribute

rewriter_model_name_or_path = None class-attribute instance-attribute

rewriter_tokenizer = None class-attribute instance-attribute

row_scorer = None class-attribute instance-attribute

strategy = 'search' class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

tokenizer = None class-attribute instance-attribute

train_rewriter = False class-attribute instance-attribute

training_seeds = None class-attribute instance-attribute

trust_remote_code = False 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/prewrite/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("PRewrite.adapt called before .steer().")
    if not isinstance(input_ids, torch.Tensor):
        input_ids = torch.tensor(input_ids, dtype=torch.long)
    return self._formatter.apply_to_ids(input_ids, self.memory, self.tokenizer, runtime_kwargs)

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/prewrite/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("PRewrite.adapt_messages called before .steer().")
    return self._formatter.apply_to_messages(messages, self.memory, runtime_kwargs)

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 optimized instruction memory under the "memory" key (after steer()).

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

def export_state(self) -> dict:
    """The optimized instruction 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), or None when the recipe already carries a memory.

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

def fit_identity(self):
    """The optimizer-relevant args (everything except `memory`), 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 optimized memory (the search-only args stay inert).

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

def frozen_form(self, state: dict) -> tuple[str, dict]:
    """A same-class frozen form: the recipe args with `memory=` set to the optimized
    memory (the search-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/prewrite", kwargs

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, 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/prewrite/control.py

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

    if self.args.memory is not None:
        self.memory = self.args.memory
        self._formatter = SystemPromptFormatter()
        return

    task_lm = SessionLM(session) if session is not None else model
    rewriter_lm, rewriter_tok = self._resolve_rewriter(task_lm, tokenizer)
    meta_prompt = self.meta_prompt or meta_prompts.DEFAULT

    if self.train_rewriter:
        reward_fn = self._build_reward_fn(task_lm=task_lm, task_tok=tokenizer)
        rewriter_lm = self._grpo_train_rewriter(rewriter_lm, rewriter_tok, meta_prompt, reward_fn)

    proposer = LLMMetaPromptProposer(
        llm=rewriter_lm,
        tokenizer=rewriter_tok,
        meta_prompt_template=meta_prompt,
        gen_kwargs=self.rewriter_gen_kwargs,
        use_chat_template=False if self.train_rewriter else None,
        parse_fn=parse_concise_instruction,
    )

    if self.strategy == "inference":
        candidates = proposer.propose(seed=self.initial_instruction, n=1)
        best = candidates[0] if candidates else self.initial_instruction
    else:
        n = max(self.k_candidates, 1)
        candidates = proposer.propose(seed=self.initial_instruction, n=n)
        if not candidates:
            logger.warning("PRewrite proposer returned no candidates; falling back to initial instruction.")
            best = self.initial_instruction
        else:
            scorer = TaskEvaluationScorer(
                task_lm=task_lm,
                tokenizer=tokenizer,
                dev_set=self.dev_set,
                row_scorer=self.row_scorer,
                gen_kwargs=self.eval_gen_kwargs or {"max_new_tokens": 32, "do_sample": False},
            )
            selector = TopKSelector(scorer=scorer)
            best = selector.select(candidates, k=1)[0]

    stripped = (best or "").strip()
    if not stripped:
        logger.info(
            "PRewrite: chosen rewrite was empty/whitespace; falling back to initial_instruction."
        )
        best = self.initial_instruction
    else:
        best = stripped

    self.memory = TextMemory(slots={"instruction": best})
    self._formatter = SystemPromptFormatter()

steer_access()

ModelAccess.ROLLOUTS; rewriting and dev-set scoring generate through the session, and adaptation is a formatter. With a precomputed memory, steering is model-free (ModelAccess.FACTS).

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

def steer_access(self) -> ModelAccess:
    """`ModelAccess.ROLLOUTS`; rewriting and dev-set scoring generate through the
    session, and adaptation is a formatter. With a precomputed `memory`, steering is
    model-free (`ModelAccess.FACTS`)."""
    if getattr(getattr(self, "args", None), "memory", None) is not None:
        return ModelAccess.FACTS
    return ModelAccess.ROLLOUTS

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

utils

meta_prompts

Meta-prompt templates for PRewrite.

Wording matches Appendix B of Kong et al., 2024 (arXiv:2401.08189) up to format-string adaptation for LLMMetaPromptProposer. The {seed} placeholder receives the initial instruction.

DEFAULT = 'You are a prompt engineer. Rewrite the following instruction to make it clearer, more specific, and more effective for a downstream language model, while preserving its intent.\n\nOriginal instruction:\n{seed}\n\nRewritten instruction:' module-attribute

reward

Reward-function adapters for GRPO-trained PRewrite rewriters.

Wraps a TaskEvaluationScorer into the callable shape TRL's GRPOTrainer expects, reward_func(prompts, completions, **kwargs), returning one float per completion. The reward applies the rewritten instruction with the frozen task model over a dev set and scores the answers with a per-row SampleScorer (Kong et al., 2024).

For a callable reward on trl==0.16.1, GRPOTrainer calls reward_func(prompts=prompts, completions=completions, **reward_kwargs). Completions are plain strings when the dataset "prompt" column is text (the PRewrite case) and conversational message lists when it is a chat. Both shapes are handled here.

make_scorer_reward_func(scorer, parse_fn=None)

Wrap a TaskEvaluationScorer as a GRPO reward function.

Each completion is a candidate rewritten instruction; its reward is the scorer's dev-set score for that instruction. Identical rewrites are scored once and the result reused. PRewrite rewrites are input-agnostic (the rewriter sees only the seed meta-prompt, not the task inputs), so a generation group often contains duplicates.

Parameters:

Name	Type	Description	Default
scorer	TaskEvaluationScorer	Configured TaskEvaluationScorer. scorer.score(list_of_rewrites) returns one float per rewrite (a full dev-set pass per rewrite under the frozen task model).	required
parse_fn	Callable[[str], list[str]] | None	Optional parser applied to each completion before scoring (e.g. parse_concise_instruction), keeping the reward consistent with how the proposer parses candidates at inference time. If it returns an empty list for a completion, the raw stripped completion text is scored instead.	None

Returns:

Type	Description
Callable	A callable reward_func(prompts, completions, **kwargs) -> list[float] aligned to completions.

Source code in steerability/algorithms/input_control/prewrite/utils/reward.py

def make_scorer_reward_func(
    scorer: TaskEvaluationScorer,
    parse_fn: Callable[[str], list[str]] | None = None,
) -> Callable:
    """Wrap a `TaskEvaluationScorer` as a GRPO reward function.

    Each completion is a candidate rewritten instruction; its reward is the scorer's dev-set score for
    that instruction. Identical rewrites are scored once and the result reused. PRewrite rewrites are
    input-agnostic (the rewriter sees only the seed meta-prompt, not the task inputs), so a generation
    group often contains duplicates.

    Args:
        scorer: Configured `TaskEvaluationScorer`. `scorer.score(list_of_rewrites)` returns one float per
            rewrite (a full dev-set pass per rewrite under the frozen task model).
        parse_fn: Optional parser applied to each completion before scoring (e.g.
            `parse_concise_instruction`), keeping the reward consistent with how the proposer parses
            candidates at inference time. If it returns an empty list for a completion, the raw stripped
            completion text is scored instead.

    Returns:
        A callable `reward_func(prompts, completions, **kwargs) -> list[float]` aligned to `completions`.
    """

    def _extract(completion: Any) -> str:
        text = _completion_text(completion).strip()
        if parse_fn is not None:
            parsed = parse_fn(text)
            if parsed:
                return parsed[0].strip()
        return text

    def reward_func(prompts, completions, **kwargs) -> list[float]:
        rewrites = [_extract(completion) for completion in completions]
        unique = list(dict.fromkeys(rewrites))  # order-preserving dedup
        if not unique:
            return []
        scores = scorer.score(unique)
        lookup = dict(zip(unique, scores))
        return [lookup[rewrite] for rewrite in rewrites]

    return reward_func